diff --git a/AconfigFlags.bp b/AconfigFlags.bp
index 3c5686b..c768121 100644
--- a/AconfigFlags.bp
+++ b/AconfigFlags.bp
@@ -437,10 +437,23 @@
     name: "android.companion.virtualdevice.flags-aconfig",
     package: "android.companion.virtualdevice.flags",
     container: "system",
+    exportable: true,
     srcs: ["core/java/android/companion/virtual/flags/*.aconfig"],
 }
 
 java_aconfig_library {
+    name: "android.companion.virtualdevice.flags-aconfig-java-export",
+    aconfig_declarations: "android.companion.virtualdevice.flags-aconfig",
+    defaults: ["framework-minus-apex-aconfig-java-defaults"],
+    mode: "exported",
+    min_sdk_version: "30",
+    apex_available: [
+        "//apex_available:platform",
+        "com.android.permission",
+    ],
+}
+
+java_aconfig_library {
     name: "android.companion.virtual.flags-aconfig-java",
     aconfig_declarations: "android.companion.virtual.flags-aconfig",
     defaults: ["framework-minus-apex-aconfig-java-defaults"],
diff --git a/Ravenwood.bp b/Ravenwood.bp
index 5f32ba0..ec58210 100644
--- a/Ravenwood.bp
+++ b/Ravenwood.bp
@@ -50,7 +50,7 @@
 framework_minus_apex_cmd = "$(location hoststubgen) " +
     "@$(location :ravenwood-standard-options) " +
     "--debug-log $(location hoststubgen_framework-minus-apex.log) " +
-    "--out-impl-jar $(location ravenwood.jar) " +
+    "--out-jar $(location ravenwood.jar) " +
     "--in-jar $(location :framework-minus-apex-for-hoststubgen) " +
     "--policy-override-file $(location :ravenwood-framework-policies) " +
     "--annotation-allowed-classes-file $(location :ravenwood-annotation-allowed-classes) "
@@ -183,7 +183,7 @@
         "--stats-file $(location hoststubgen_services.core_stats.csv) " +
         "--supported-api-list-file $(location hoststubgen_services.core_apis.csv) " +
 
-        "--out-impl-jar $(location ravenwood.jar) " +
+        "--out-jar $(location ravenwood.jar) " +
 
         "--gen-keep-all-file $(location hoststubgen_services.core_keep_all.txt) " +
         "--gen-input-dump-file $(location hoststubgen_services.core_dump.txt) " +
diff --git a/apct-tests/perftests/protolog/Android.bp b/apct-tests/perftests/tracing/Android.bp
similarity index 100%
rename from apct-tests/perftests/protolog/Android.bp
rename to apct-tests/perftests/tracing/Android.bp
diff --git a/apct-tests/perftests/protolog/AndroidManifest.xml b/apct-tests/perftests/tracing/AndroidManifest.xml
similarity index 100%
rename from apct-tests/perftests/protolog/AndroidManifest.xml
rename to apct-tests/perftests/tracing/AndroidManifest.xml
diff --git a/apct-tests/perftests/protolog/AndroidTest.xml b/apct-tests/perftests/tracing/AndroidTest.xml
similarity index 100%
rename from apct-tests/perftests/protolog/AndroidTest.xml
rename to apct-tests/perftests/tracing/AndroidTest.xml
diff --git a/apct-tests/perftests/protolog/OWNERS b/apct-tests/perftests/tracing/OWNERS
similarity index 100%
rename from apct-tests/perftests/protolog/OWNERS
rename to apct-tests/perftests/tracing/OWNERS
diff --git a/apct-tests/perftests/protolog/src/com/android/internal/protolog/ProtologPerfTest.java b/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
similarity index 60%
rename from apct-tests/perftests/protolog/src/com/android/internal/protolog/ProtologPerfTest.java
rename to apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
index e1edb37..f4759b8 100644
--- a/apct-tests/perftests/protolog/src/com/android/internal/protolog/ProtologPerfTest.java
+++ b/apct-tests/perftests/tracing/src/com/android/internal/protolog/ProtoLogPerfTest.java
@@ -15,43 +15,54 @@
  */
 package com.android.internal.protolog;
 
-import android.perftests.utils.BenchmarkState;
-import android.perftests.utils.PerfStatusReporter;
+import android.app.Activity;
+import android.os.Bundle;
+import android.perftests.utils.Stats;
+
+import androidx.test.InstrumentationRegistry;
 
 import com.android.internal.protolog.common.IProtoLogGroup;
 import com.android.internal.protolog.common.LogLevel;
 
 import org.junit.Before;
 import org.junit.BeforeClass;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
 import org.junit.runners.Parameterized.Parameters;
 
-import java.util.Arrays;
+import java.util.ArrayList;
 import java.util.Collection;
 
 @RunWith(Parameterized.class)
-public class ProtologPerfTest {
-    @Rule public PerfStatusReporter mPerfStatusReporter = new PerfStatusReporter();
-
-    @Parameters(name="logToProto_{0}_logToLogcat_{1}")
-    public static Collection<Object[]> params() {
-        return Arrays.asList(new Object[][] {
-                { true, true },
-                { true, false },
-                { false, true },
-                { false, false }
-        });
-    }
-
+public class ProtoLogPerfTest {
     private final boolean mLogToProto;
     private final boolean mLogToLogcat;
 
-    public ProtologPerfTest(boolean logToProto, boolean logToLogcat) {
-        mLogToProto = logToProto;
-        mLogToLogcat = logToLogcat;
+    /**
+     * Generates the parameters used for this test class
+     */
+    @Parameters(name = "LOG_TO_{0}")
+    public static Collection<Object[]> params() {
+        var params = new ArrayList<Object[]>();
+
+        for (LogTo logTo : LogTo.values()) {
+            params.add(new Object[] { logTo });
+        }
+
+        return params;
+    }
+
+    public ProtoLogPerfTest(LogTo logTo) {
+        mLogToProto = switch (logTo) {
+            case ALL, PROTO_ONLY -> true;
+            case LOGCAT_ONLY, NONE -> false;
+        };
+
+        mLogToLogcat = switch (logTo) {
+            case ALL, LOGCAT_ONLY -> true;
+            case PROTO_ONLY, NONE -> false;
+        };
     }
 
     @BeforeClass
@@ -66,11 +77,11 @@
     }
 
     @Test
-    public void logProcessedProtoLogMessageWithoutArgs() {
+    public void log_Processed_NoArgs() {
         final var protoLog = ProtoLog.getSingleInstance();
+        final var perfMonitor = new PerfMonitor();
 
-        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
-        while (state.keepRunning()) {
+        while (perfMonitor.keepRunning()) {
             protoLog.log(
                     LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 123,
                     0, (Object[]) null);
@@ -78,11 +89,11 @@
     }
 
     @Test
-    public void logProcessedProtoLogMessageWithArgs() {
+    public void log_Processed_WithArgs() {
         final var protoLog = ProtoLog.getSingleInstance();
+        final var perfMonitor = new PerfMonitor();
 
-        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
-        while (state.keepRunning()) {
+        while (perfMonitor.keepRunning()) {
             protoLog.log(
                     LogLevel.INFO, TestProtoLogGroup.TEST_GROUP, 123,
                     0b1110101001010100,
@@ -91,18 +102,58 @@
     }
 
     @Test
-    public void logNonProcessedProtoLogMessageWithNoArgs() {
-        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
-        while (state.keepRunning()) {
+    public void log_Unprocessed_NoArgs() {
+        final var perfMonitor = new PerfMonitor();
+
+        while (perfMonitor.keepRunning()) {
             ProtoLog.d(TestProtoLogGroup.TEST_GROUP, "Test message");
         }
     }
 
     @Test
-    public void logNonProcessedProtoLogMessageWithArgs() {
-        BenchmarkState state = mPerfStatusReporter.getBenchmarkState();
-        while (state.keepRunning()) {
-            ProtoLog.d(TestProtoLogGroup.TEST_GROUP, "Test messag %s, %d, %b", "arg1", 2, true);
+    public void log_Unprocessed_WithArgs() {
+        final var perfMonitor = new PerfMonitor();
+
+        while (perfMonitor.keepRunning()) {
+            ProtoLog.d(TestProtoLogGroup.TEST_GROUP, "Test message %s, %d, %b", "arg1", 2, true);
+        }
+    }
+
+    private class PerfMonitor {
+        private int mIteration = 0;
+
+        private static final int WARM_UP_ITERATIONS = 10;
+        private static final int ITERATIONS = 1000;
+
+        private final ArrayList<Long> mResults = new ArrayList<>();
+
+        private long mStartTimeNs;
+
+        public boolean keepRunning() {
+            final long currentTime = System.nanoTime();
+
+            mIteration++;
+
+            if (mIteration > ITERATIONS) {
+                reportResults();
+                return false;
+            }
+
+            if (mIteration > WARM_UP_ITERATIONS) {
+                mResults.add(currentTime - mStartTimeNs);
+            }
+
+            mStartTimeNs = System.nanoTime();
+            return true;
+        }
+
+        public void reportResults() {
+            final var stats = new Stats(mResults);
+
+            Bundle status = new Bundle();
+            status.putLong("protologging_time_mean_ns", (long) stats.getMean());
+            status.putLong("protologging_time_median_ns", (long) stats.getMedian());
+            InstrumentationRegistry.getInstrumentation().sendStatus(Activity.RESULT_OK, status);
         }
     }
 
@@ -168,4 +219,11 @@
             return ordinal();
         }
     }
+
+    private enum LogTo {
+        PROTO_ONLY,
+        LOGCAT_ONLY,
+        ALL,
+        NONE,
+    }
 }
diff --git a/apex/jobscheduler/service/aconfig/job.aconfig b/apex/jobscheduler/service/aconfig/job.aconfig
index 613b784..e5389b4 100644
--- a/apex/jobscheduler/service/aconfig/job.aconfig
+++ b/apex/jobscheduler/service/aconfig/job.aconfig
@@ -65,3 +65,13 @@
    description: "Remove started user if user will be stopped due to user switch"
    bug: "321598070"
 }
+
+flag {
+   name: "use_correct_process_state_for_logging"
+   namespace: "backstage_power"
+   description: "Use correct process state for statsd logging"
+   bug: "361308212"
+   metadata {
+       purpose: PURPOSE_BUGFIX
+   }
+}
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java b/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
index d65a66c..be8e304 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobServiceContext.java
@@ -473,6 +473,14 @@
             mInitialDownloadedBytesFromCalling = TrafficStats.getUidRxBytes(job.getUid());
             mInitialUploadedBytesFromCalling = TrafficStats.getUidTxBytes(job.getUid());
 
+            int procState = mService.getUidProcState(job.getUid());
+            if (Flags.useCorrectProcessStateForLogging()
+                    && procState > ActivityManager.PROCESS_STATE_TRANSIENT_BACKGROUND) {
+                // Try to get the latest proc state from AMS, there might be some delay
+                // for the proc states worse than TRANSIENT_BACKGROUND.
+                procState = mActivityManagerInternal.getUidProcessState(job.getUid());
+            }
+
             FrameworkStatsLog.write(FrameworkStatsLog.SCHEDULED_JOB_STATE_CHANGED,
                     job.isProxyJob() ? new int[]{sourceUid, job.getUid()} : new int[]{sourceUid},
                     // Given that the source tag is set by the calling app, it should be connected
@@ -517,7 +525,7 @@
                     job.getEstimatedNetworkDownloadBytes(),
                     job.getEstimatedNetworkUploadBytes(),
                     job.getWorkCount(),
-                    ActivityManager.processStateAmToProto(mService.getUidProcState(job.getUid())),
+                    ActivityManager.processStateAmToProto(procState),
                     job.getNamespaceHash(),
                     /* system_measured_source_download_bytes */ 0,
                     /* system_measured_source_upload_bytes */ 0,
@@ -1528,6 +1536,13 @@
         mJobPackageTracker.noteInactive(completedJob,
                 loggingInternalStopReason, loggingDebugReason);
         final int sourceUid = completedJob.getSourceUid();
+        int procState = mService.getUidProcState(completedJob.getUid());
+        if (Flags.useCorrectProcessStateForLogging()
+                && procState > ActivityManager.PROCESS_STATE_TRANSIENT_BACKGROUND) {
+            // Try to get the latest proc state from AMS, there might be some delay
+            // for the proc states worse than TRANSIENT_BACKGROUND.
+            procState = mActivityManagerInternal.getUidProcessState(completedJob.getUid());
+        }
         FrameworkStatsLog.write(FrameworkStatsLog.SCHEDULED_JOB_STATE_CHANGED,
                 completedJob.isProxyJob()
                         ? new int[]{sourceUid, completedJob.getUid()} : new int[]{sourceUid},
@@ -1573,7 +1588,7 @@
                 completedJob.getEstimatedNetworkUploadBytes(),
                 completedJob.getWorkCount(),
                 ActivityManager
-                        .processStateAmToProto(mService.getUidProcState(completedJob.getUid())),
+                        .processStateAmToProto(procState),
                 completedJob.getNamespaceHash(),
                 TrafficStats.getUidRxBytes(completedJob.getSourceUid())
                         - mInitialDownloadedBytesFromSource,
diff --git a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
index d92351d..c9d3407 100644
--- a/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
+++ b/apex/jobscheduler/service/java/com/android/server/usage/AppStandbyController.java
@@ -1989,10 +1989,8 @@
                 mAdminProtectedPackages.put(userId, packageNames);
             }
         }
-        if (android.app.admin.flags.Flags.disallowUserControlBgUsageFix()) {
-            if (!Flags.avoidIdleCheck() || mInjector.getBootPhase() >= PHASE_BOOT_COMPLETED) {
-                postCheckIdleStates(userId);
-            }
+        if (!Flags.avoidIdleCheck() || mInjector.getBootPhase() >= PHASE_BOOT_COMPLETED) {
+            postCheckIdleStates(userId);
         }
     }
 
diff --git a/api/Android.bp b/api/Android.bp
index 341be3d53..533f9f6 100644
--- a/api/Android.bp
+++ b/api/Android.bp
@@ -87,6 +87,7 @@
         "framework-permission",
         "framework-permission-s",
         "framework-profiling",
+        "framework-photopicker",
         "framework-scheduling",
         "framework-sdkextensions",
         "framework-statsd",
diff --git a/core/api/current.txt b/core/api/current.txt
index 861be40..ddfd364 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -141,7 +141,7 @@
     field public static final String MANAGE_DEVICE_POLICY_APPS_CONTROL = "android.permission.MANAGE_DEVICE_POLICY_APPS_CONTROL";
     field public static final String MANAGE_DEVICE_POLICY_APP_RESTRICTIONS = "android.permission.MANAGE_DEVICE_POLICY_APP_RESTRICTIONS";
     field public static final String MANAGE_DEVICE_POLICY_APP_USER_DATA = "android.permission.MANAGE_DEVICE_POLICY_APP_USER_DATA";
-    field @FlaggedApi("android.app.admin.flags.assist_content_user_restriction_enabled") public static final String MANAGE_DEVICE_POLICY_ASSIST_CONTENT = "android.permission.MANAGE_DEVICE_POLICY_ASSIST_CONTENT";
+    field public static final String MANAGE_DEVICE_POLICY_ASSIST_CONTENT = "android.permission.MANAGE_DEVICE_POLICY_ASSIST_CONTENT";
     field public static final String MANAGE_DEVICE_POLICY_AUDIO_OUTPUT = "android.permission.MANAGE_DEVICE_POLICY_AUDIO_OUTPUT";
     field public static final String MANAGE_DEVICE_POLICY_AUTOFILL = "android.permission.MANAGE_DEVICE_POLICY_AUTOFILL";
     field public static final String MANAGE_DEVICE_POLICY_BACKUP_SERVICE = "android.permission.MANAGE_DEVICE_POLICY_BACKUP_SERVICE";
@@ -169,7 +169,7 @@
     field public static final String MANAGE_DEVICE_POLICY_LOCK = "android.permission.MANAGE_DEVICE_POLICY_LOCK";
     field public static final String MANAGE_DEVICE_POLICY_LOCK_CREDENTIALS = "android.permission.MANAGE_DEVICE_POLICY_LOCK_CREDENTIALS";
     field public static final String MANAGE_DEVICE_POLICY_LOCK_TASK = "android.permission.MANAGE_DEVICE_POLICY_LOCK_TASK";
-    field @FlaggedApi("android.app.admin.flags.esim_management_enabled") public static final String MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS = "android.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS";
+    field public static final String MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS = "android.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS";
     field public static final String MANAGE_DEVICE_POLICY_METERED_DATA = "android.permission.MANAGE_DEVICE_POLICY_METERED_DATA";
     field public static final String MANAGE_DEVICE_POLICY_MICROPHONE = "android.permission.MANAGE_DEVICE_POLICY_MICROPHONE";
     field @FlaggedApi("android.app.admin.flags.dedicated_device_control_api_enabled") public static final String MANAGE_DEVICE_POLICY_MICROPHONE_TOGGLE = "android.permission.MANAGE_DEVICE_POLICY_MICROPHONE_TOGGLE";
@@ -7875,7 +7875,7 @@
     method public void writeToParcel(android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.app.admin.DeviceAdminInfo> CREATOR;
     field public static final int HEADLESS_DEVICE_OWNER_MODE_AFFILIATED = 1; // 0x1
-    field @FlaggedApi("android.app.admin.flags.headless_device_owner_single_user_enabled") public static final int HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER = 2; // 0x2
+    field public static final int HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER = 2; // 0x2
     field public static final int HEADLESS_DEVICE_OWNER_MODE_UNSUPPORTED = 0; // 0x0
     field public static final int USES_ENCRYPTED_STORAGE = 7; // 0x7
     field public static final int USES_POLICY_DISABLE_CAMERA = 8; // 0x8
@@ -7964,13 +7964,13 @@
     field public static final String LOCK_TASK_POLICY = "lockTask";
     field public static final String PACKAGES_SUSPENDED_POLICY = "packagesSuspended";
     field public static final String PACKAGE_UNINSTALL_BLOCKED_POLICY = "packageUninstallBlocked";
-    field @FlaggedApi("android.app.admin.flags.policy_engine_migration_v2_enabled") public static final String PASSWORD_COMPLEXITY_POLICY = "passwordComplexity";
+    field public static final String PASSWORD_COMPLEXITY_POLICY = "passwordComplexity";
     field public static final String PERMISSION_GRANT_POLICY = "permissionGrant";
     field public static final String PERSISTENT_PREFERRED_ACTIVITY_POLICY = "persistentPreferredActivity";
     field public static final String RESET_PASSWORD_TOKEN_POLICY = "resetPasswordToken";
     field public static final String SECURITY_LOGGING_POLICY = "securityLogging";
     field public static final String STATUS_BAR_DISABLED_POLICY = "statusBarDisabled";
-    field @FlaggedApi("android.app.admin.flags.policy_engine_migration_v2_enabled") public static final String USB_DATA_SIGNALING_POLICY = "usbDataSignaling";
+    field public static final String USB_DATA_SIGNALING_POLICY = "usbDataSignaling";
     field public static final String USER_CONTROL_DISABLED_PACKAGES_POLICY = "userControlDisabledPackages";
   }
 
@@ -8081,7 +8081,7 @@
     method public CharSequence getStartUserSessionMessage(@NonNull android.content.ComponentName);
     method @Deprecated public boolean getStorageEncryption(@Nullable android.content.ComponentName);
     method public int getStorageEncryptionStatus();
-    method @FlaggedApi("android.app.admin.flags.esim_management_enabled") @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS) public java.util.Set<java.lang.Integer> getSubscriptionIds();
+    method @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS) public java.util.Set<java.lang.Integer> getSubscriptionIds();
     method @Nullable public android.app.admin.SystemUpdatePolicy getSystemUpdatePolicy();
     method @Nullable public android.os.PersistableBundle getTransferOwnershipBundle();
     method @Nullable public java.util.List<android.os.PersistableBundle> getTrustAgentConfiguration(@Nullable android.content.ComponentName, @NonNull android.content.ComponentName);
@@ -8756,6 +8756,8 @@
     method public int getResultCode();
     method @NonNull public android.app.appsearch.GenericDocument getResultDocument();
     method public boolean isSuccess();
+    method @FlaggedApi("android.app.appfunctions.flags.enable_app_function_manager") @NonNull public static android.app.appfunctions.ExecuteAppFunctionResponse newFailure(int, @Nullable String, @Nullable android.os.Bundle);
+    method @FlaggedApi("android.app.appfunctions.flags.enable_app_function_manager") @NonNull public static android.app.appfunctions.ExecuteAppFunctionResponse newSuccess(@NonNull android.app.appsearch.GenericDocument, @Nullable android.os.Bundle);
     method public void writeToParcel(@NonNull android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.app.appfunctions.ExecuteAppFunctionResponse> CREATOR;
     field public static final String PROPERTY_RETURN_VALUE = "returnValue";
@@ -8767,13 +8769,6 @@
     field public static final int RESULT_TIMED_OUT = 5; // 0x5
   }
 
-  @FlaggedApi("android.app.appfunctions.flags.enable_app_function_manager") public static final class ExecuteAppFunctionResponse.Builder {
-    ctor public ExecuteAppFunctionResponse.Builder(@NonNull android.app.appsearch.GenericDocument);
-    ctor public ExecuteAppFunctionResponse.Builder(int, @NonNull String);
-    method @NonNull public android.app.appfunctions.ExecuteAppFunctionResponse build();
-    method @NonNull public android.app.appfunctions.ExecuteAppFunctionResponse.Builder setExtras(@NonNull android.os.Bundle);
-  }
-
 }
 
 package android.app.assist {
@@ -34116,7 +34111,7 @@
     field public static final String DISALLOW_AIRPLANE_MODE = "no_airplane_mode";
     field public static final String DISALLOW_AMBIENT_DISPLAY = "no_ambient_display";
     field public static final String DISALLOW_APPS_CONTROL = "no_control_apps";
-    field @FlaggedApi("android.app.admin.flags.assist_content_user_restriction_enabled") public static final String DISALLOW_ASSIST_CONTENT = "no_assist_content";
+    field public static final String DISALLOW_ASSIST_CONTENT = "no_assist_content";
     field public static final String DISALLOW_AUTOFILL = "no_autofill";
     field public static final String DISALLOW_BLUETOOTH = "no_bluetooth";
     field public static final String DISALLOW_BLUETOOTH_SHARING = "no_bluetooth_sharing";
@@ -34166,7 +34161,7 @@
     field public static final String DISALLOW_SHARE_INTO_MANAGED_PROFILE = "no_sharing_into_profile";
     field public static final String DISALLOW_SHARE_LOCATION = "no_share_location";
     field public static final String DISALLOW_SHARING_ADMIN_CONFIGURED_WIFI = "no_sharing_admin_configured_wifi";
-    field @FlaggedApi("android.app.admin.flags.esim_management_enabled") public static final String DISALLOW_SIM_GLOBALLY = "no_sim_globally";
+    field public static final String DISALLOW_SIM_GLOBALLY = "no_sim_globally";
     field public static final String DISALLOW_SMS = "no_sms";
     field public static final String DISALLOW_SYSTEM_ERROR_DIALOGS = "no_system_error_dialogs";
     field @FlaggedApi("com.android.net.thread.platform.flags.thread_user_restriction_enabled") public static final String DISALLOW_THREAD_NETWORK = "no_thread_network";
@@ -43968,8 +43963,11 @@
     field @FlaggedApi("com.android.internal.telephony.flags.carrier_enabled_satellite_flag") public static final String KEY_SATELLITE_CONNECTION_HYSTERESIS_SEC_INT = "satellite_connection_hysteresis_sec_int";
     field @FlaggedApi("com.android.internal.telephony.flags.carrier_enabled_satellite_flag") public static final String KEY_SATELLITE_ENTITLEMENT_STATUS_REFRESH_DAYS_INT = "satellite_entitlement_status_refresh_days_int";
     field @FlaggedApi("com.android.internal.telephony.flags.carrier_enabled_satellite_flag") public static final String KEY_SATELLITE_ENTITLEMENT_SUPPORTED_BOOL = "satellite_entitlement_supported_bool";
+    field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_ESOS_INACTIVITY_TIMEOUT_SEC_INT = "satellite_esos_inactivity_timeout_sec_int";
     field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_ESOS_SUPPORTED_BOOL = "satellite_esos_supported_bool";
-    field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_SCREEN_OFF_INACTIVITY_TIMEOUT_SEC_INT = "satellite_screen_off_inactivity_timeout_duration_sec_int";
+    field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_P2P_SMS_INACTIVITY_TIMEOUT_SEC_INT = "satellite_p2p_sms_inactivity_timeout_sec_int";
+    field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_ROAMING_P2P_SMS_SUPPORTED_BOOL = "satellite_roaming_p2p_sms_supported_bool";
+    field @FlaggedApi("com.android.internal.telephony.flags.carrier_roaming_nb_iot_ntn") public static final String KEY_SATELLITE_SCREEN_OFF_INACTIVITY_TIMEOUT_SEC_INT = "satellite_screen_off_inactivity_timeout_sec_int";
     field public static final String KEY_SHOW_4G_FOR_3G_DATA_ICON_BOOL = "show_4g_for_3g_data_icon_bool";
     field public static final String KEY_SHOW_4G_FOR_LTE_DATA_ICON_BOOL = "show_4g_for_lte_data_icon_bool";
     field public static final String KEY_SHOW_APN_SETTING_CDMA_BOOL = "show_apn_setting_cdma_bool";
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index e87bc50..ba16037 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -61,6 +61,7 @@
     field @FlaggedApi("com.android.internal.telephony.flags.use_oem_domain_selection_service") public static final String BIND_DOMAIN_SELECTION_SERVICE = "android.permission.BIND_DOMAIN_SELECTION_SERVICE";
     field public static final String BIND_DOMAIN_VERIFICATION_AGENT = "android.permission.BIND_DOMAIN_VERIFICATION_AGENT";
     field public static final String BIND_EUICC_SERVICE = "android.permission.BIND_EUICC_SERVICE";
+    field @FlaggedApi("android.crashrecovery.flags.enable_crashrecovery") public static final String BIND_EXPLICIT_HEALTH_CHECK_SERVICE = "android.permission.BIND_EXPLICIT_HEALTH_CHECK_SERVICE";
     field public static final String BIND_EXTERNAL_STORAGE_SERVICE = "android.permission.BIND_EXTERNAL_STORAGE_SERVICE";
     field public static final String BIND_FIELD_CLASSIFICATION_SERVICE = "android.permission.BIND_FIELD_CLASSIFICATION_SERVICE";
     field public static final String BIND_GBA_SERVICE = "android.permission.BIND_GBA_SERVICE";
@@ -1422,7 +1423,7 @@
     field public static final int STATUS_DEVICE_ADMIN_NOT_SUPPORTED = 13; // 0xd
     field public static final int STATUS_HAS_DEVICE_OWNER = 1; // 0x1
     field public static final int STATUS_HAS_PAIRED = 8; // 0x8
-    field @FlaggedApi("android.app.admin.flags.headless_device_owner_single_user_enabled") public static final int STATUS_HEADLESS_ONLY_SYSTEM_USER = 17; // 0x11
+    field public static final int STATUS_HEADLESS_ONLY_SYSTEM_USER = 17; // 0x11
     field public static final int STATUS_HEADLESS_SYSTEM_USER_MODE_NOT_SUPPORTED = 16; // 0x10
     field public static final int STATUS_MANAGED_USERS_NOT_SUPPORTED = 9; // 0x9
     field public static final int STATUS_NONSYSTEM_USER_EXISTS = 5; // 0x5
@@ -3469,6 +3470,7 @@
   public static interface VirtualDeviceManager.ActivityListener {
     method @FlaggedApi("android.companion.virtualdevice.flags.activity_control_api") public default void onActivityLaunchBlocked(int, @NonNull android.content.ComponentName, @NonNull android.os.UserHandle, @Nullable android.content.IntentSender);
     method public void onDisplayEmpty(int);
+    method @FlaggedApi("android.companion.virtualdevice.flags.activity_control_api") public default void onSecureWindowShown(int, @NonNull android.content.ComponentName, @NonNull android.os.UserHandle);
     method @Deprecated public void onTopActivityChanged(int, @NonNull android.content.ComponentName);
     method public default void onTopActivityChanged(int, @NonNull android.content.ComponentName, int);
   }
@@ -4631,6 +4633,7 @@
     method public int getCommittedSessionId();
     method @NonNull public java.util.List<android.content.rollback.PackageRollbackInfo> getPackages();
     method public int getRollbackId();
+    method @FlaggedApi("android.crashrecovery.flags.enable_crashrecovery") public int getRollbackImpactLevel();
     method public boolean isStaged();
     method public void writeToParcel(android.os.Parcel, int);
     field @NonNull public static final android.os.Parcelable.Creator<android.content.rollback.RollbackInfo> CREATOR;
@@ -6746,6 +6749,14 @@
     field public static final int STATUS_OK = 0; // 0x0
   }
 
+  @FlaggedApi("android.media.soundtrigger.generic_model_api") public static final class SoundTrigger.GenericSoundModel extends android.hardware.soundtrigger.SoundTrigger.SoundModel implements android.os.Parcelable {
+    ctor public SoundTrigger.GenericSoundModel(@NonNull java.util.UUID, @NonNull java.util.UUID, @Nullable byte[], int);
+    ctor public SoundTrigger.GenericSoundModel(@NonNull java.util.UUID, @NonNull java.util.UUID, @Nullable byte[]);
+    method public int describeContents();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.hardware.soundtrigger.SoundTrigger.GenericSoundModel> CREATOR;
+  }
+
   public static final class SoundTrigger.Keyphrase implements android.os.Parcelable {
     ctor public SoundTrigger.Keyphrase(int, int, @NonNull java.util.Locale, @NonNull String, @Nullable int[]);
     method public int describeContents();
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index ec23cfe..42f7615 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -597,19 +597,19 @@
     method @RequiresPermission(android.Manifest.permission.FORCE_DEVICE_POLICY_MANAGER_LOGS) public long forceNetworkLogs();
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void forceRemoveActiveAdmin(@NonNull android.content.ComponentName, int);
     method @RequiresPermission(android.Manifest.permission.FORCE_DEVICE_POLICY_MANAGER_LOGS) public long forceSecurityLogs();
-    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public void forceSetMaxPolicyStorageLimit(int);
+    method @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public void forceSetMaxPolicyStorageLimit(int);
     method public void forceUpdateUserSetupComplete(int);
     method @NonNull public java.util.Set<java.lang.String> getDefaultCrossProfilePackages();
     method @Deprecated public int getDeviceOwnerType(@NonNull android.content.ComponentName);
     method @Nullable public String getDevicePolicyManagementRoleHolderUpdaterPackage();
     method @NonNull public java.util.Set<java.lang.String> getDisallowedSystemApps(@NonNull android.content.ComponentName, int, @NonNull String);
-    method @FlaggedApi("android.app.admin.flags.headless_device_owner_provisioning_fix_enabled") @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public int getHeadlessDeviceOwnerMode();
+    method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public int getHeadlessDeviceOwnerMode();
     method public long getLastBugReportRequestTime();
     method public long getLastNetworkLogRetrievalTime();
     method public long getLastSecurityLogRetrievalTime();
     method public java.util.List<java.lang.String> getOwnerInstalledCaCerts(@NonNull android.os.UserHandle);
     method @NonNull @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_ADMINS) public java.util.Set<java.lang.String> getPolicyExemptApps();
-    method @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public int getPolicySizeForAdmin(@NonNull android.app.admin.EnforcingAdmin);
+    method @RequiresPermission("android.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT") public int getPolicySizeForAdmin(@NonNull android.app.admin.EnforcingAdmin);
     method public boolean isCurrentInputMethodSetByOwner();
     method public boolean isFactoryResetProtectionPolicySupported();
     method @RequiresPermission(anyOf={android.Manifest.permission.MANAGE_USERS, android.Manifest.permission.INTERACT_ACROSS_USERS}) public boolean isNewUserDisclaimerAcknowledged();
@@ -680,7 +680,7 @@
   }
 
   public final class EnforcingAdmin implements android.os.Parcelable {
-    ctor @FlaggedApi("android.app.admin.flags.device_policy_size_tracking_internal_bug_fix_enabled") public EnforcingAdmin(@NonNull String, @NonNull android.app.admin.Authority, @NonNull android.os.UserHandle, @Nullable android.content.ComponentName);
+    ctor public EnforcingAdmin(@NonNull String, @NonNull android.app.admin.Authority, @NonNull android.os.UserHandle, @Nullable android.content.ComponentName);
   }
 
   public final class FlagUnion extends android.app.admin.ResolutionMechanism<java.lang.Integer> {
@@ -1269,10 +1269,6 @@
 
 package android.content.rollback {
 
-  public final class RollbackInfo implements android.os.Parcelable {
-    method @FlaggedApi("android.content.pm.recoverability_detection") public int getRollbackImpactLevel();
-  }
-
   public final class RollbackManager {
     method @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS) public void blockRollbackManager(long);
     method @RequiresPermission(android.Manifest.permission.TEST_MANAGE_ROLLBACKS) public void expireRollbackForPackage(@NonNull String);
@@ -1542,6 +1538,10 @@
     method @Deprecated public final void setPreviewSurface(android.view.Surface) throws java.io.IOException;
   }
 
+  public final class Sensor {
+    method public int getHandle();
+  }
+
   public final class SensorPrivacyManager {
     method @FlaggedApi("com.android.internal.camera.flags.camera_privacy_allowlist") @RequiresPermission(android.Manifest.permission.MANAGE_SENSOR_PRIVACY) public void setCameraPrivacyAllowlist(@NonNull java.util.List<java.lang.String>);
     method @RequiresPermission(android.Manifest.permission.MANAGE_SENSOR_PRIVACY) public void setSensorPrivacy(int, int, boolean);
@@ -1797,11 +1797,15 @@
 
   public class InputSettings {
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_bounce_keys_flag") public static int getAccessibilityBounceKeysThreshold(@NonNull android.content.Context);
+    method @FlaggedApi("com.android.hardware.input.keyboard_repeat_keys") public static int getAccessibilityRepeatKeysDelay(@NonNull android.content.Context);
+    method @FlaggedApi("com.android.hardware.input.keyboard_repeat_keys") public static int getAccessibilityRepeatKeysTimeout(@NonNull android.content.Context);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_slow_keys_flag") public static int getAccessibilitySlowKeysThreshold(@NonNull android.content.Context);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_mouse_keys") public static boolean isAccessibilityMouseKeysEnabled(@NonNull android.content.Context);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_sticky_keys_flag") public static boolean isAccessibilityStickyKeysEnabled(@NonNull android.content.Context);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_bounce_keys_flag") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilityBounceKeysThreshold(@NonNull android.content.Context, int);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_mouse_keys") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilityMouseKeysEnabled(@NonNull android.content.Context, boolean);
+    method @FlaggedApi("com.android.hardware.input.keyboard_repeat_keys") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilityRepeatKeysDelay(@NonNull android.content.Context, int);
+    method @FlaggedApi("com.android.hardware.input.keyboard_repeat_keys") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilityRepeatKeysTimeout(@NonNull android.content.Context, int);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_slow_keys_flag") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilitySlowKeysThreshold(@NonNull android.content.Context, int);
     method @FlaggedApi("com.android.hardware.input.keyboard_a11y_sticky_keys_flag") @RequiresPermission(android.Manifest.permission.WRITE_SETTINGS) public static void setAccessibilityStickyKeysEnabled(@NonNull android.content.Context, boolean);
     method @RequiresPermission(android.Manifest.permission.WRITE_SECURE_SETTINGS) public static void setMaximumObscuringOpacityForTouch(@NonNull android.content.Context, @FloatRange(from=0, to=1) float);
@@ -3134,6 +3138,7 @@
 
   public abstract class DreamOverlayService extends android.app.Service {
     ctor public DreamOverlayService();
+    method @FlaggedApi("android.service.dreams.publish_preview_state_to_overlay") public final boolean isDreamInPreviewMode();
     method @Nullable public final android.os.IBinder onBind(@NonNull android.content.Intent);
     method public void onEndDream();
     method public abstract void onStartDream(@NonNull android.view.WindowManager.LayoutParams);
@@ -3666,7 +3671,11 @@
     method @Nullable public android.view.Display.Mode getUserPreferredDisplayMode();
     method public boolean hasAccess(int);
     method @RequiresPermission(android.Manifest.permission.MODIFY_USER_PREFERRED_DISPLAY_MODE) public void setUserPreferredDisplayMode(@NonNull android.view.Display.Mode);
+    field public static final int FLAG_ALWAYS_UNLOCKED = 512; // 0x200
+    field public static final int FLAG_OWN_FOCUS = 2048; // 0x800
+    field public static final int FLAG_TOUCH_FEEDBACK_DISABLED = 1024; // 0x400
     field public static final int FLAG_TRUSTED = 128; // 0x80
+    field public static final int REMOVE_MODE_DESTROY_CONTENT = 1; // 0x1
     field public static final int TYPE_EXTERNAL = 2; // 0x2
     field public static final int TYPE_INTERNAL = 1; // 0x1
     field public static final int TYPE_OVERLAY = 4; // 0x4
diff --git a/core/java/android/app/AppOpsManager.java b/core/java/android/app/AppOpsManager.java
index 21396a1..8fd3326 100644
--- a/core/java/android/app/AppOpsManager.java
+++ b/core/java/android/app/AppOpsManager.java
@@ -7459,15 +7459,15 @@
         }
 
         /**
-         * Similar to {@link #onOpChanged(String, String, int)} but includes the device for which
-         * the op mode has changed.
+         * Similar to {@link #onOpChanged(String, String)} but includes user and the device for
+         * which the op mode has changed.
          *
          * <p> Implement this method if callbacks are required on all devices.
          * If not implemented explicitly, the default implementation will notify for op changes
-         * on the default device {@link VirtualDeviceManager#PERSISTENT_DEVICE_ID_DEFAULT} only.
+         * on the default device only.
          *
-         * <p> If implemented, {@link #onOpChanged(String, String, int)}
-         * will not be called automatically.
+         * <p> If implemented, {@link #onOpChanged(String, String)} will not be called
+         * automatically.
          *
          * @param op The Op that changed.
          * @param packageName Package of the app whose Op changed.
diff --git a/core/java/android/app/AutomaticZenRule.java b/core/java/android/app/AutomaticZenRule.java
index 62b5412..e0a9371 100644
--- a/core/java/android/app/AutomaticZenRule.java
+++ b/core/java/android/app/AutomaticZenRule.java
@@ -49,7 +49,7 @@
 
     /**
      * Rule is of an unknown type. This is the default value if not provided by the owning app,
-     * and the value returned if the true type was added in an API level lower than the calling
+     * and the value returned if the true type was added in an API level higher than the calling
      * app's targetSdk.
      */
     @FlaggedApi(Flags.FLAG_MODES_API)
@@ -315,7 +315,8 @@
     }
 
     /**
-     * Gets the zen policy.
+     * Gets the {@link ZenPolicy} applied if {@link #getInterruptionFilter()} is
+     * {@link NotificationManager#INTERRUPTION_FILTER_PRIORITY}.
      */
     @Nullable
     public ZenPolicy getZenPolicy() {
@@ -345,6 +346,17 @@
 
     /**
      * Sets the interruption filter that is applied when this rule is active.
+     *
+     * <ul>
+     *     <li>When {@link NotificationManager#INTERRUPTION_FILTER_PRIORITY}, the rule will use
+     *     the {@link ZenPolicy} supplied to {@link #setZenPolicy} (or a default one).
+     *     <li>When {@link NotificationManager#INTERRUPTION_FILTER_ALARMS} or
+     *     {@link NotificationManager#INTERRUPTION_FILTER_NONE}, the rule will use a fixed
+     *     {@link ZenPolicy} matching the filter.
+     *     <li>When {@link NotificationManager#INTERRUPTION_FILTER_ALL}, the rule will not block
+     *     notifications, but can still have {@link ZenDeviceEffects}.
+     * </ul>
+     *
      * @param interruptionFilter The do not disturb mode to enter when this rule is active.
      */
     public void setInterruptionFilter(@InterruptionFilter int interruptionFilter) {
@@ -374,7 +386,8 @@
     }
 
     /**
-     * Sets the zen policy.
+     * Sets the {@link ZenPolicy} applied if {@link #getInterruptionFilter()} is
+     * {@link NotificationManager#INTERRUPTION_FILTER_PRIORITY}.
      *
      * <p>When updating an existing rule via {@link NotificationManager#updateAutomaticZenRule},
      * a {@code null} value here means the previous policy is retained.
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index e99ba84..7a36fbb 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -2706,14 +2706,9 @@
         if (mAllowlistToken == null) {
             mAllowlistToken = processAllowlistToken;
         }
-        if (Flags.secureAllowlistToken()) {
-            // Propagate this token to all pending intents that are unmarshalled from the parcel,
-            // or keep the one we're already propagating, if that's the case.
-            if (!parcel.hasClassCookie(PendingIntent.class)) {
-                parcel.setClassCookie(PendingIntent.class, mAllowlistToken);
-            }
-        } else {
-            // Propagate this token to all pending intents that are unmarshalled from the parcel.
+        // Propagate this token to all pending intents that are unmarshalled from the parcel,
+        // or keep the one we're already propagating, if that's the case.
+        if (!parcel.hasClassCookie(PendingIntent.class)) {
             parcel.setClassCookie(PendingIntent.class, mAllowlistToken);
         }
 
@@ -3333,28 +3328,22 @@
             PendingIntent.addOnMarshaledListener(addedListener);
         }
         try {
-            if (Flags.secureAllowlistToken()) {
-                boolean mustClearCookie = false;
-                if (!parcel.hasClassCookie(Notification.class)) {
-                    // This is the "root" notification, and not an "inner" notification (including
-                    // publicVersion or anything else that might be embedded in extras). So we want
-                    // to use its token for every inner notification (might be null).
-                    parcel.setClassCookie(Notification.class, mAllowlistToken);
-                    mustClearCookie = true;
-                }
-                try {
-                    // IMPORTANT: Add marshaling code in writeToParcelImpl as we
-                    // want to intercept all pending events written to the parcel.
-                    writeToParcelImpl(parcel, flags);
-                } finally {
-                    if (mustClearCookie) {
-                        parcel.removeClassCookie(Notification.class, mAllowlistToken);
-                    }
-                }
-            } else {
+            boolean mustClearCookie = false;
+            if (!parcel.hasClassCookie(Notification.class)) {
+                // This is the "root" notification, and not an "inner" notification (including
+                // publicVersion or anything else that might be embedded in extras). So we want
+                // to use its token for every inner notification (might be null).
+                parcel.setClassCookie(Notification.class, mAllowlistToken);
+                mustClearCookie = true;
+            }
+            try {
                 // IMPORTANT: Add marshaling code in writeToParcelImpl as we
                 // want to intercept all pending events written to the parcel.
                 writeToParcelImpl(parcel, flags);
+            } finally {
+                if (mustClearCookie) {
+                    parcel.removeClassCookie(Notification.class, mAllowlistToken);
+                }
             }
 
             synchronized (this) {
@@ -3371,13 +3360,9 @@
     private void writeToParcelImpl(Parcel parcel, int flags) {
         parcel.writeInt(1);
 
-        if (Flags.secureAllowlistToken()) {
-            // Always use the same token as the root notification (might be null).
-            IBinder rootNotificationToken = (IBinder) parcel.getClassCookie(Notification.class);
-            parcel.writeStrongBinder(rootNotificationToken);
-        } else {
-            parcel.writeStrongBinder(mAllowlistToken);
-        }
+        // Always use the same token as the root notification (might be null).
+        IBinder rootNotificationToken = (IBinder) parcel.getClassCookie(Notification.class);
+        parcel.writeStrongBinder(rootNotificationToken);
 
         parcel.writeLong(when);
         parcel.writeLong(creationTime);
diff --git a/core/java/android/app/StatusBarManager.java b/core/java/android/app/StatusBarManager.java
index 14195c4..63e03914 100644
--- a/core/java/android/app/StatusBarManager.java
+++ b/core/java/android/app/StatusBarManager.java
@@ -1326,6 +1326,7 @@
         private boolean mClock;
         private boolean mNotificationIcons;
         private boolean mRotationSuggestion;
+        private boolean mQuickSettings;
 
         /** @hide */
         public DisableInfo(int flags1, int flags2) {
@@ -1338,6 +1339,7 @@
             mClock = (flags1 & DISABLE_CLOCK) != 0;
             mNotificationIcons = (flags1 & DISABLE_NOTIFICATION_ICONS) != 0;
             mRotationSuggestion = (flags2 & DISABLE2_ROTATE_SUGGESTIONS) != 0;
+            mQuickSettings = (flags2 & DISABLE2_QUICK_SETTINGS) != 0;
         }
 
         /** @hide */
@@ -1471,6 +1473,20 @@
         }
 
         /**
+         * @hide
+         */
+        public void setQuickSettingsDisabled(boolean disabled) {
+            mQuickSettings = disabled;
+        }
+
+        /**
+         * @hide
+         */
+        public boolean isQuickSettingsDisabled() {
+            return mQuickSettings;
+        }
+
+        /**
          * @return {@code true} if no components are disabled (default state)
          * @hide
          */
@@ -1478,7 +1494,7 @@
         public boolean areAllComponentsEnabled() {
             return !mStatusBarExpansion && !mNavigateHome && !mNotificationPeeking && !mRecents
                     && !mSearch && !mSystemIcons && !mClock && !mNotificationIcons
-                    && !mRotationSuggestion;
+                    && !mRotationSuggestion && !mQuickSettings;
         }
 
         /** @hide */
@@ -1492,6 +1508,7 @@
             mClock = false;
             mNotificationIcons = false;
             mRotationSuggestion = false;
+            mQuickSettings = false;
         }
 
         /**
@@ -1502,7 +1519,7 @@
         public boolean areAllComponentsDisabled() {
             return mStatusBarExpansion && mNavigateHome && mNotificationPeeking
                     && mRecents && mSearch && mSystemIcons && mClock && mNotificationIcons
-                    && mRotationSuggestion;
+                    && mRotationSuggestion && mQuickSettings;
         }
 
         /** @hide */
@@ -1516,6 +1533,7 @@
             mClock = true;
             mNotificationIcons = true;
             mRotationSuggestion = true;
+            mQuickSettings = true;
         }
 
         @NonNull
@@ -1533,6 +1551,7 @@
             sb.append(" mClock=").append(mClock ? "disabled" : "enabled");
             sb.append(" mNotificationIcons=").append(mNotificationIcons ? "disabled" : "enabled");
             sb.append(" mRotationSuggestion=").append(mRotationSuggestion ? "disabled" : "enabled");
+            sb.append(" mQuickSettings=").append(mQuickSettings ? "disabled" : "enabled");
 
             return sb.toString();
 
@@ -1557,6 +1576,7 @@
             if (mClock) disable1 |= DISABLE_CLOCK;
             if (mNotificationIcons) disable1 |= DISABLE_NOTIFICATION_ICONS;
             if (mRotationSuggestion) disable2 |= DISABLE2_ROTATE_SUGGESTIONS;
+            if (mQuickSettings) disable2 |= DISABLE2_QUICK_SETTINGS;
 
             return new Pair<Integer, Integer>(disable1, disable2);
         }
diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index 8b3ee24..e44e776 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -31,6 +31,7 @@
 import android.app.ambientcontext.AmbientContextManager;
 import android.app.ambientcontext.IAmbientContextManager;
 import android.app.appfunctions.AppFunctionManager;
+import android.app.appfunctions.AppFunctionManagerConfiguration;
 import android.app.appfunctions.IAppFunctionManager;
 import android.app.appsearch.AppSearchManagerFrameworkInitializer;
 import android.app.blob.BlobStoreManagerFrameworkInitializer;
@@ -937,8 +938,10 @@
                         @Override
                         public AppFunctionManager createService(ContextImpl ctx)
                                 throws ServiceNotFoundException {
+                            if (!AppFunctionManagerConfiguration.isSupported(ctx)) {
+                                return null;
+                            }
                             IAppFunctionManager service;
-                            //TODO(b/357551503): If the feature not present avoid look up every time
                             service = IAppFunctionManager.Stub.asInterface(
                                     ServiceManager.getServiceOrThrow(Context.APP_FUNCTION_SERVICE));
                             return new AppFunctionManager(service, ctx.getOuterContext());
diff --git a/core/java/android/app/admin/AccountTypePolicyKey.java b/core/java/android/app/admin/AccountTypePolicyKey.java
index 02e492b..515c1c6 100644
--- a/core/java/android/app/admin/AccountTypePolicyKey.java
+++ b/core/java/android/app/admin/AccountTypePolicyKey.java
@@ -24,7 +24,6 @@
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
-import android.app.admin.flags.Flags;
 import android.os.Bundle;
 import android.os.Parcel;
 
@@ -54,9 +53,7 @@
     @TestApi
     public AccountTypePolicyKey(@NonNull String key, @NonNull String accountType) {
         super(key);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxStringLength(accountType, "accountType");
-        }
+        PolicySizeVerifier.enforceMaxStringLength(accountType, "accountType");
         mAccountType = Objects.requireNonNull((accountType));
     }
 
diff --git a/core/java/android/app/admin/BundlePolicyValue.java b/core/java/android/app/admin/BundlePolicyValue.java
index c993671..00e67e6 100644
--- a/core/java/android/app/admin/BundlePolicyValue.java
+++ b/core/java/android/app/admin/BundlePolicyValue.java
@@ -18,7 +18,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.admin.flags.Flags;
 import android.os.Bundle;
 import android.os.Parcel;
 
@@ -31,9 +30,7 @@
 
     public BundlePolicyValue(Bundle value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxBundleFieldsLength(value);
-        }
+        PolicySizeVerifier.enforceMaxBundleFieldsLength(value);
     }
 
     private BundlePolicyValue(Parcel source) {
diff --git a/core/java/android/app/admin/ComponentNamePolicyValue.java b/core/java/android/app/admin/ComponentNamePolicyValue.java
index a7a2f7d..f092b7b 100644
--- a/core/java/android/app/admin/ComponentNamePolicyValue.java
+++ b/core/java/android/app/admin/ComponentNamePolicyValue.java
@@ -18,7 +18,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.admin.flags.Flags;
 import android.content.ComponentName;
 import android.os.Parcel;
 
@@ -31,9 +30,7 @@
 
     public ComponentNamePolicyValue(@NonNull ComponentName value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxComponentNameLength(value);
-        }
+        PolicySizeVerifier.enforceMaxComponentNameLength(value);
     }
 
     private ComponentNamePolicyValue(Parcel source) {
diff --git a/core/java/android/app/admin/DeviceAdminInfo.java b/core/java/android/app/admin/DeviceAdminInfo.java
index 46c9e78..4f2efa4 100644
--- a/core/java/android/app/admin/DeviceAdminInfo.java
+++ b/core/java/android/app/admin/DeviceAdminInfo.java
@@ -16,9 +16,6 @@
 
 package android.app.admin;
 
-import static android.app.admin.flags.Flags.FLAG_HEADLESS_DEVICE_OWNER_SINGLE_USER_ENABLED;
-
-import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.app.admin.flags.Flags;
@@ -195,7 +192,6 @@
      * DPCs should set the value of attribute "headless-device-owner-mode" inside the
      * "headless-system-user" tag as "single_user".
      */
-    @FlaggedApi(FLAG_HEADLESS_DEVICE_OWNER_SINGLE_USER_ENABLED)
     public static final int HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER = 2;
 
     /**
diff --git a/core/java/android/app/admin/DevicePolicyIdentifiers.java b/core/java/android/app/admin/DevicePolicyIdentifiers.java
index 156512a..c0e435c 100644
--- a/core/java/android/app/admin/DevicePolicyIdentifiers.java
+++ b/core/java/android/app/admin/DevicePolicyIdentifiers.java
@@ -16,8 +16,6 @@
 
 package android.app.admin;
 
-import static android.app.admin.flags.Flags.FLAG_POLICY_ENGINE_MIGRATION_V2_ENABLED;
-
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
@@ -185,13 +183,11 @@
     /**
      * String identifier for {@link DevicePolicyManager#setUsbDataSignalingEnabled}.
      */
-    @FlaggedApi(FLAG_POLICY_ENGINE_MIGRATION_V2_ENABLED)
     public static final String USB_DATA_SIGNALING_POLICY = "usbDataSignaling";
 
     /**
      * String identifier for {@link DevicePolicyManager#setRequiredPasswordComplexity}.
      */
-    @FlaggedApi(FLAG_POLICY_ENGINE_MIGRATION_V2_ENABLED)
     public static final String PASSWORD_COMPLEXITY_POLICY = "passwordComplexity";
 
     /**
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index 5088ea6..1ddec17 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -54,12 +54,8 @@
 import static android.Manifest.permission.SET_TIME;
 import static android.Manifest.permission.SET_TIME_ZONE;
 import static android.app.admin.DeviceAdminInfo.HEADLESS_DEVICE_OWNER_MODE_UNSUPPORTED;
-import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED;
 import static android.app.admin.flags.Flags.FLAG_DEVICE_THEFT_API_ENABLED;
-import static android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED;
 import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_ENABLED;
-import static android.app.admin.flags.Flags.FLAG_HEADLESS_DEVICE_OWNER_PROVISIONING_FIX_ENABLED;
-import static android.app.admin.flags.Flags.FLAG_HEADLESS_DEVICE_OWNER_SINGLE_USER_ENABLED;
 import static android.app.admin.flags.Flags.onboardingBugreportV2Enabled;
 import static android.app.admin.flags.Flags.onboardingConsentlessBugreports;
 import static android.app.admin.flags.Flags.FLAG_IS_MTE_POLICY_ENFORCED;
@@ -2988,7 +2984,6 @@
      * @hide
      */
     @SystemApi
-    @FlaggedApi(FLAG_HEADLESS_DEVICE_OWNER_SINGLE_USER_ENABLED)
     public static final int STATUS_HEADLESS_ONLY_SYSTEM_USER = 17;
 
     /**
@@ -17744,7 +17739,6 @@
      * @throws SecurityException if the caller is not authorized to call this method.
      * @return ids of all managed subscriptions currently downloaded by an admin on the device.
      */
-    @FlaggedApi(FLAG_ESIM_MANAGEMENT_ENABLED)
     @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS)
     @NonNull
     public Set<Integer> getSubscriptionIds() {
@@ -17813,7 +17807,6 @@
      */
     @TestApi
     @RequiresPermission(permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT)
-    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
     public void forceSetMaxPolicyStorageLimit(int storageLimit) {
         if (mService != null) {
             try {
@@ -17831,7 +17824,6 @@
      */
     @TestApi
     @RequiresPermission(permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT)
-    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
     public int getPolicySizeForAdmin(@NonNull EnforcingAdmin admin) {
         if (mService != null) {
             try {
@@ -17850,13 +17842,9 @@
      * @hide
      */
     @TestApi
-    @FlaggedApi(FLAG_HEADLESS_DEVICE_OWNER_PROVISIONING_FIX_ENABLED)
     @RequiresPermission(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS)
     @DeviceAdminInfo.HeadlessDeviceOwnerMode
     public int getHeadlessDeviceOwnerMode() {
-        if (!Flags.headlessDeviceOwnerProvisioningFixEnabled()) {
-            return HEADLESS_DEVICE_OWNER_MODE_UNSUPPORTED;
-        }
         if (mService != null) {
             try {
                 return mService.getHeadlessDeviceOwnerMode(mContext.getPackageName());
diff --git a/core/java/android/app/admin/EnforcingAdmin.java b/core/java/android/app/admin/EnforcingAdmin.java
index f70a53f..5f9bb9c 100644
--- a/core/java/android/app/admin/EnforcingAdmin.java
+++ b/core/java/android/app/admin/EnforcingAdmin.java
@@ -16,9 +16,6 @@
 
 package android.app.admin;
 
-import static android.app.admin.flags.Flags.FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED;
-
-import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
@@ -64,7 +61,6 @@
      *
      * @hide
      */
-    @FlaggedApi(FLAG_DEVICE_POLICY_SIZE_TRACKING_INTERNAL_BUG_FIX_ENABLED)
     @TestApi
     public EnforcingAdmin(
             @NonNull String packageName, @NonNull Authority authority,
diff --git a/core/java/android/app/admin/LockTaskPolicy.java b/core/java/android/app/admin/LockTaskPolicy.java
index 68b4ad8..ab32d46 100644
--- a/core/java/android/app/admin/LockTaskPolicy.java
+++ b/core/java/android/app/admin/LockTaskPolicy.java
@@ -19,7 +19,6 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
-import android.app.admin.flags.Flags;
 import android.os.Parcel;
 import android.os.Parcelable;
 
@@ -135,10 +134,8 @@
     }
 
     private void setPackagesInternal(Set<String> packages) {
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            for (String p : packages) {
-                PolicySizeVerifier.enforceMaxPackageNameLength(p);
-            }
+        for (String p : packages) {
+            PolicySizeVerifier.enforceMaxPackageNameLength(p);
         }
         mPackages = new HashSet<>(packages);
     }
diff --git a/core/java/android/app/admin/PackagePermissionPolicyKey.java b/core/java/android/app/admin/PackagePermissionPolicyKey.java
index 1a04f6c..226c576 100644
--- a/core/java/android/app/admin/PackagePermissionPolicyKey.java
+++ b/core/java/android/app/admin/PackagePermissionPolicyKey.java
@@ -25,7 +25,6 @@
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
-import android.app.admin.flags.Flags;
 import android.os.Bundle;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -59,10 +58,8 @@
     public PackagePermissionPolicyKey(@NonNull String identifier, @NonNull String packageName,
             @NonNull String permissionName) {
         super(identifier);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
-            PolicySizeVerifier.enforceMaxStringLength(permissionName, "permissionName");
-        }
+        PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
+        PolicySizeVerifier.enforceMaxStringLength(permissionName, "permissionName");
         mPackageName = Objects.requireNonNull((packageName));
         mPermissionName = Objects.requireNonNull((permissionName));
     }
diff --git a/core/java/android/app/admin/PackagePolicyKey.java b/core/java/android/app/admin/PackagePolicyKey.java
index 9e31a23..8fa21db 100644
--- a/core/java/android/app/admin/PackagePolicyKey.java
+++ b/core/java/android/app/admin/PackagePolicyKey.java
@@ -24,7 +24,6 @@
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
-import android.app.admin.flags.Flags;
 import android.os.Bundle;
 import android.os.Parcel;
 import android.os.Parcelable;
@@ -55,9 +54,7 @@
     @TestApi
     public PackagePolicyKey(@NonNull String key, @NonNull String packageName) {
         super(key);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
-        }
+        PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
         mPackageName = Objects.requireNonNull((packageName));
     }
 
diff --git a/core/java/android/app/admin/PackageSetPolicyValue.java b/core/java/android/app/admin/PackageSetPolicyValue.java
index 8b253a2..24c50b0 100644
--- a/core/java/android/app/admin/PackageSetPolicyValue.java
+++ b/core/java/android/app/admin/PackageSetPolicyValue.java
@@ -18,7 +18,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.admin.flags.Flags;
 import android.os.Parcel;
 
 import java.util.HashSet;
@@ -32,10 +31,8 @@
 
     public PackageSetPolicyValue(@NonNull Set<String> value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            for (String packageName : value) {
-                PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
-            }
+        for (String packageName : value) {
+            PolicySizeVerifier.enforceMaxPackageNameLength(packageName);
         }
     }
 
diff --git a/core/java/android/app/admin/StringPolicyValue.java b/core/java/android/app/admin/StringPolicyValue.java
index 6efe9ad..bb07c23 100644
--- a/core/java/android/app/admin/StringPolicyValue.java
+++ b/core/java/android/app/admin/StringPolicyValue.java
@@ -18,7 +18,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.admin.flags.Flags;
 import android.os.Parcel;
 
 import java.util.Objects;
@@ -30,9 +29,7 @@
 
     public StringPolicyValue(@NonNull String value) {
         super(value);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxStringLength(value, "policyValue");
-        }
+        PolicySizeVerifier.enforceMaxStringLength(value, "policyValue");
     }
 
     private StringPolicyValue(Parcel source) {
diff --git a/core/java/android/app/admin/UserRestrictionPolicyKey.java b/core/java/android/app/admin/UserRestrictionPolicyKey.java
index 9054287..16cfba4 100644
--- a/core/java/android/app/admin/UserRestrictionPolicyKey.java
+++ b/core/java/android/app/admin/UserRestrictionPolicyKey.java
@@ -21,7 +21,6 @@
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
-import android.app.admin.flags.Flags;
 import android.os.Bundle;
 import android.os.Parcel;
 
@@ -45,9 +44,7 @@
     @TestApi
     public UserRestrictionPolicyKey(@NonNull String identifier, @NonNull String restriction) {
         super(identifier);
-        if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxStringLength(restriction, "restriction");
-        }
+        PolicySizeVerifier.enforceMaxStringLength(restriction, "restriction");
         mRestriction = Objects.requireNonNull(restriction);
     }
 
diff --git a/core/java/android/app/admin/flags/flags.aconfig b/core/java/android/app/admin/flags/flags.aconfig
index 56f4792..d9bd77f 100644
--- a/core/java/android/app/admin/flags/flags.aconfig
+++ b/core/java/android/app/admin/flags/flags.aconfig
@@ -4,6 +4,7 @@
 package: "android.app.admin.flags"
 container: "system"
 
+# Fully rolled out and must not be used.
 flag {
   name: "policy_engine_migration_v2_enabled"
   is_exported: true
@@ -28,16 +29,6 @@
 }
 
 flag {
-  name: "device_policy_size_tracking_internal_bug_fix_enabled"
-  namespace: "enterprise"
-  description: "Bug fix for tracking the total policy size and have a max threshold"
-  bug: "281543351"
-  metadata {
-      purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
   name: "onboarding_bugreport_v2_enabled"
   is_exported: true
   namespace: "enterprise"
@@ -77,13 +68,6 @@
 }
 
 flag {
-  name: "permission_migration_for_zero_trust_impl_enabled"
-  namespace: "enterprise"
-  description: "(Implementation) Migrate existing APIs to permission based, and enable DMRH to call them to collect Zero Trust signals."
-  bug: "289520697"
-}
-
-flag {
   name: "device_theft_api_enabled"
   is_exported: true
   namespace: "enterprise"
@@ -105,6 +89,14 @@
   bug: "289520697"
 }
 
+flag {
+  name: "coexistence_migration_for_supervision_enabled"
+  is_exported: true
+  namespace: "enterprise"
+  description: "Migrate existing APIs that are used by supervision (Kids Module) to be coexistable."
+  bug: "356894721"
+}
+
 # Fully rolled out and must not be used.
 flag {
   name: "security_log_v2_enabled"
@@ -125,13 +117,6 @@
 }
 
 flag {
-  name: "dumpsys_policy_engine_migration_enabled"
-  namespace: "enterprise"
-  description: "Update DumpSys to include information about migrated APIs in DPE"
-  bug: "304999634"
-}
-
-flag {
     name: "allow_querying_profile_type"
     is_exported: true
     namespace: "enterprise"
@@ -180,6 +165,7 @@
   bug: "295301164"
 }
 
+# Fully rolled out and must not be used.
 flag {
   name: "headless_device_owner_single_user_enabled"
   is_exported: true
@@ -207,26 +193,6 @@
 }
 
 flag {
-  name: "power_exemption_bg_usage_fix"
-  namespace: "enterprise"
-  description: "Ensure aps with EXEMPT_FROM_POWER_RESTRICTIONS can execute in the background"
-  bug: "333379020"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
-  name: "disallow_user_control_bg_usage_fix"
-  namespace: "enterprise"
-  description: "Make DPM.setUserControlDisabledPackages() ensure background usage is allowed"
-  bug: "326031059"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
   name: "disallow_user_control_stopped_state_fix"
   namespace: "enterprise"
   description: "Ensure DPM.setUserControlDisabledPackages() clears FLAG_STOPPED for the app"
@@ -244,16 +210,6 @@
 }
 
 flag {
-  name: "headless_device_owner_provisioning_fix_enabled"
-  namespace: "enterprise"
-  description: "Fix provisioning for single-user headless DO"
-  bug: "289515470"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
   name: "dmrh_set_app_restrictions"
   namespace: "enterprise"
   description: "Allow DMRH to set application restrictions (both on the profile and the parent)"
@@ -264,13 +220,6 @@
 }
 
 flag {
-  name: "allow_screen_brightness_control_on_cope"
-  namespace: "enterprise"
-  description: "Allow COPE admin to control screen brightness and timeout."
-  bug: "323894620"
-}
-
-flag {
   name: "always_persist_do"
   namespace: "enterprise"
   description: "Always write device_owners2.xml so that migration flags aren't lost"
@@ -288,16 +237,6 @@
 }
 
 flag {
-  name: "headless_device_owner_delegate_security_logging_bug_fix"
-  namespace: "enterprise"
-  description: "Fix delegate security logging for single user headless DO."
-  bug: "289515470"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
   name: "headless_single_user_bad_device_admin_state_fix"
   namespace: "enterprise"
   description: "Fix the bad state in DPMS caused by an earlier bug related to the headless single user change"
@@ -318,16 +257,6 @@
 }
 
 flag {
-  name: "delete_private_space_under_restriction"
-  namespace: "enterprise"
-  description: "Delete private space if user restriction is set"
-  bug: "328758346"
-  metadata {
-    purpose: PURPOSE_BUGFIX
-  }
-}
-
-flag {
     name: "unmanaged_mode_migration"
     namespace: "enterprise"
     description: "Migrate APIs for unmanaged mode"
diff --git a/core/java/android/app/appfunctions/AppFunctionManager.java b/core/java/android/app/appfunctions/AppFunctionManager.java
index b6240a7..8f609de 100644
--- a/core/java/android/app/appfunctions/AppFunctionManager.java
+++ b/core/java/android/app/appfunctions/AppFunctionManager.java
@@ -49,9 +49,8 @@
     /**
      * Creates an instance.
      *
-     * @param mService An interface to the backing service.
-     * @param context A {@link Context}.
-     *
+     * @param service An interface to the backing service.
+     * @param context  A {@link Context}.
      * @hide
      */
     public AppFunctionManager(IAppFunctionManager service, Context context) {
@@ -108,8 +107,8 @@
                             } catch (RuntimeException e) {
                                 // Ideally shouldn't happen since errors are wrapped into the
                                 // response, but we catch it here for additional safety.
-                                callback.accept(new ExecuteAppFunctionResponse.Builder(
-                                        getResultCode(e), e.getMessage()).build());
+                                callback.accept(ExecuteAppFunctionResponse.newFailure(
+                                        getResultCode(e), e.getMessage(), /*extras=*/ null));
                             }
                         }
                     });
diff --git a/core/java/android/app/appfunctions/AppFunctionManagerConfiguration.java b/core/java/android/app/appfunctions/AppFunctionManagerConfiguration.java
new file mode 100644
index 0000000..e4784b4
--- /dev/null
+++ b/core/java/android/app/appfunctions/AppFunctionManagerConfiguration.java
@@ -0,0 +1,63 @@
+/*
+ * 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.app.appfunctions;
+
+import static android.app.appfunctions.flags.Flags.enableAppFunctionManager;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.content.pm.PackageManager;
+
+/**
+ * Represents the system configuration of support for the {@code AppFunctionManager} and
+ * associated systems.
+ *
+ * @hide
+ */
+public class AppFunctionManagerConfiguration {
+    private final Context mContext;
+
+    /**
+     * Constructs a new instance of {@code AppFunctionManagerConfiguration}.
+     * @param context context
+     */
+    public AppFunctionManagerConfiguration(@NonNull final Context context) {
+        mContext = context;
+    }
+
+    /**
+     * Indicates whether the current target is intended to support {@code AppFunctionManager}.
+     * @return {@code true} if supported; otherwise {@code false}
+     */
+    public boolean isSupported() {
+        return enableAppFunctionManager() && !isWatch();
+
+    }
+
+    /**
+     * Indicates whether the current target is intended to support {@code AppFunctionManager}.
+     * @param context context
+     * @return {@code true} if supported; otherwise {@code false}
+     */
+    public static boolean isSupported(@NonNull final Context context) {
+        return new AppFunctionManagerConfiguration(context).isSupported();
+    }
+
+    private boolean isWatch() {
+        return mContext.getPackageManager().hasSystemFeature(PackageManager.FEATURE_WATCH);
+    }
+}
diff --git a/core/java/android/app/appfunctions/AppFunctionManagerHelper.java b/core/java/android/app/appfunctions/AppFunctionManagerHelper.java
new file mode 100644
index 0000000..3169f0e
--- /dev/null
+++ b/core/java/android/app/appfunctions/AppFunctionManagerHelper.java
@@ -0,0 +1,210 @@
+/*
+ * 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.app.appfunctions;
+
+import static android.app.appfunctions.AppFunctionRuntimeMetadata.PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID;
+import static android.app.appfunctions.AppFunctionRuntimeMetadata.PROPERTY_ENABLED;
+import static android.app.appfunctions.AppFunctionStaticMetadataHelper.APP_FUNCTION_INDEXER_PACKAGE;
+import static android.app.appfunctions.AppFunctionStaticMetadataHelper.STATIC_PROPERTY_ENABLED_BY_DEFAULT;
+import static android.app.appfunctions.flags.Flags.FLAG_ENABLE_APP_FUNCTION_MANAGER;
+
+import android.annotation.CallbackExecutor;
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.app.appsearch.AppSearchManager;
+import android.app.appsearch.AppSearchResult;
+import android.app.appsearch.GlobalSearchSession;
+import android.app.appsearch.JoinSpec;
+import android.app.appsearch.SearchResult;
+import android.app.appsearch.SearchResults;
+import android.app.appsearch.SearchSpec;
+import android.os.OutcomeReceiver;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.Executor;
+
+/**
+ * Helper class containing utilities for {@link AppFunctionManager}.
+ *
+ * @hide
+ */
+@FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+public class AppFunctionManagerHelper {
+
+    /**
+     * Returns (through a callback) a boolean indicating whether the app function is enabled.
+     * <p>
+     * This method can only check app functions that are owned by the caller owned by packages
+     * visible to the caller.
+     * <p>
+     * If operation fails, the callback's {@link OutcomeReceiver#onError} is called with errors:
+     * <ul>
+     *     <li>{@link IllegalArgumentException}, if the function is not found</li>
+     *     <li>{@link SecurityException}, if the caller does not have permission to query the
+     *         target package
+     *     </li>
+     *  </ul>
+     *
+     * @param functionIdentifier the identifier of the app function to check (unique within the
+     *                           target package) and in most cases, these are automatically
+     *                           generated by the AppFunctions SDK
+     * @param targetPackage      the package name of the app function's owner
+     * @param appSearchExecutor  the executor to run the metadata search mechanism through AppSearch
+     * @param callbackExecutor   the executor to run the callback
+     * @param callback           the callback to receive the function enabled check result
+     * @hide
+     */
+    public static void isAppFunctionEnabled(
+            @NonNull String functionIdentifier,
+            @NonNull String targetPackage,
+            @NonNull AppSearchManager appSearchManager,
+            @NonNull Executor appSearchExecutor,
+            @NonNull @CallbackExecutor Executor callbackExecutor,
+            @NonNull OutcomeReceiver<Boolean, Exception> callback
+    ) {
+        Objects.requireNonNull(functionIdentifier);
+        Objects.requireNonNull(targetPackage);
+        Objects.requireNonNull(appSearchManager);
+        Objects.requireNonNull(appSearchExecutor);
+        Objects.requireNonNull(callbackExecutor);
+        Objects.requireNonNull(callback);
+
+        appSearchManager.createGlobalSearchSession(appSearchExecutor,
+                (searchSessionResult) -> {
+                    if (!searchSessionResult.isSuccess()) {
+                        callbackExecutor.execute(() ->
+                                callback.onError(failedResultToException(searchSessionResult)));
+                        return;
+                    }
+                    try (GlobalSearchSession searchSession = searchSessionResult.getResultValue()) {
+                        SearchResults results = searchJoinedStaticWithRuntimeAppFunctions(
+                                searchSession, targetPackage, functionIdentifier);
+                        results.getNextPage(appSearchExecutor,
+                                listAppSearchResult -> callbackExecutor.execute(() -> {
+                                    if (listAppSearchResult.isSuccess()) {
+                                        callback.onResult(getEnabledStateFromSearchResults(
+                                                Objects.requireNonNull(
+                                                        listAppSearchResult.getResultValue())));
+                                    } else {
+                                        callback.onError(
+                                                failedResultToException(listAppSearchResult));
+                                    }
+                                }));
+                    } catch (Exception e) {
+                        callbackExecutor.execute(() -> callback.onError(e));
+                    }
+                });
+    }
+
+    /**
+     * Searches joined app function static and runtime metadata using the function Id and the
+     * package.
+     *
+     * @hide
+     */
+    private static @NonNull SearchResults searchJoinedStaticWithRuntimeAppFunctions(
+            @NonNull GlobalSearchSession session,
+            @NonNull String targetPackage,
+            @NonNull String functionIdentifier) {
+        SearchSpec runtimeSearchSpec = getAppFunctionRuntimeMetadataSearchSpecByFunctionId(
+                targetPackage);
+        JoinSpec joinSpec = new JoinSpec.Builder(
+                PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID)
+                .setNestedSearch(
+                        functionIdentifier,
+                        runtimeSearchSpec).build();
+        SearchSpec joinedStaticWithRuntimeSearchSpec = new SearchSpec.Builder()
+                .setJoinSpec(joinSpec)
+                .addFilterPackageNames(APP_FUNCTION_INDEXER_PACKAGE)
+                .addFilterSchemas(
+                        AppFunctionStaticMetadataHelper.getStaticSchemaNameForPackage(
+                                targetPackage))
+                .setTermMatch(SearchSpec.TERM_MATCH_EXACT_ONLY)
+                .build();
+        return session.search(functionIdentifier, joinedStaticWithRuntimeSearchSpec);
+    }
+
+    /**
+     * Finds whether the function is enabled or not from the search results returned by
+     * {@link #searchJoinedStaticWithRuntimeAppFunctions}.
+     *
+     * @throws IllegalArgumentException if the function is not found in the results
+     * @hide
+     */
+    private static boolean getEnabledStateFromSearchResults(
+            @NonNull List<SearchResult> joinedStaticRuntimeResults) {
+        if (joinedStaticRuntimeResults.isEmpty()) {
+            // Function not found.
+            throw new IllegalArgumentException("App function not found.");
+        } else {
+            List<SearchResult> runtimeMetadataResults =
+                    joinedStaticRuntimeResults.getFirst().getJoinedResults();
+            if (!runtimeMetadataResults.isEmpty()) {
+                Boolean result = (Boolean) runtimeMetadataResults
+                        .getFirst().getGenericDocument()
+                        .getProperty(PROPERTY_ENABLED);
+                if (result != null) {
+                    return result;
+                }
+            }
+            // Runtime metadata not found. Using the default value in the static metadata.
+            return joinedStaticRuntimeResults.getFirst().getGenericDocument()
+                    .getPropertyBoolean(STATIC_PROPERTY_ENABLED_BY_DEFAULT);
+        }
+    }
+
+    /**
+     * Returns search spec that queries app function metadata for a specific package name by its
+     * function identifier.
+     *
+     * @hide
+     */
+    public static @NonNull SearchSpec getAppFunctionRuntimeMetadataSearchSpecByFunctionId(
+            @NonNull String targetPackage) {
+        return new SearchSpec.Builder()
+                .addFilterPackageNames(APP_FUNCTION_INDEXER_PACKAGE)
+                .addFilterSchemas(
+                        AppFunctionRuntimeMetadata.getRuntimeSchemaNameForPackage(
+                                targetPackage))
+                .addFilterProperties(
+                        AppFunctionRuntimeMetadata.getRuntimeSchemaNameForPackage(
+                                targetPackage),
+                        List.of(AppFunctionRuntimeMetadata.PROPERTY_FUNCTION_ID))
+                .setTermMatch(SearchSpec.TERM_MATCH_EXACT_ONLY)
+                .build();
+    }
+
+    /**
+     * Converts a failed app search result codes into an exception.
+     *
+     * @hide
+     */
+    public static @NonNull Exception failedResultToException(
+            @NonNull AppSearchResult appSearchResult) {
+        return switch (appSearchResult.getResultCode()) {
+            case AppSearchResult.RESULT_INVALID_ARGUMENT -> new IllegalArgumentException(
+                    appSearchResult.getErrorMessage());
+            case AppSearchResult.RESULT_IO_ERROR -> new IOException(
+                    appSearchResult.getErrorMessage());
+            case AppSearchResult.RESULT_SECURITY_ERROR -> new SecurityException(
+                    appSearchResult.getErrorMessage());
+            default -> new IllegalStateException(appSearchResult.getErrorMessage());
+        };
+    }
+}
diff --git a/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java b/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java
new file mode 100644
index 0000000..fdd12b0
--- /dev/null
+++ b/core/java/android/app/appfunctions/AppFunctionRuntimeMetadata.java
@@ -0,0 +1,204 @@
+/*
+ * 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.app.appfunctions;
+
+import static android.app.appfunctions.flags.Flags.FLAG_ENABLE_APP_FUNCTION_MANAGER;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.appsearch.AppSearchSchema;
+import android.app.appsearch.GenericDocument;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.Objects;
+
+/**
+ * Represents runtime function metadata of an app function.
+ *
+ * <p>This is a temporary solution for app function indexing, as later we would like to index the
+ * actual function signature entity class shape instead of just the schema info.
+ *
+ * @hide
+ */
+// TODO(b/357551503): Link to canonical docs rather than duplicating once they
+// are available.
+@FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+public class AppFunctionRuntimeMetadata extends GenericDocument {
+    public static final String RUNTIME_SCHEMA_TYPE = "AppFunctionRuntimeMetadata";
+    public static final String APP_FUNCTION_INDEXER_PACKAGE = "android";
+    public static final String APP_FUNCTION_RUNTIME_METADATA_DB = "appfunctions-db";
+    public static final String APP_FUNCTION_RUNTIME_NAMESPACE = "app_functions_runtime";
+    public static final String PROPERTY_FUNCTION_ID = "functionId";
+    public static final String PROPERTY_PACKAGE_NAME = "packageName";
+    public static final String PROPERTY_ENABLED = "enabled";
+    public static final String PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID =
+            "appFunctionStaticMetadataQualifiedId";
+    private static final String TAG = "AppSearchAppFunction";
+    private static final String RUNTIME_SCHEMA_TYPE_SEPARATOR = "-";
+
+    public AppFunctionRuntimeMetadata(@NonNull GenericDocument genericDocument) {
+        super(genericDocument);
+    }
+
+    /**
+     * Returns a per-app runtime metadata schema name, to store all functions for that package.
+     */
+    public static String getRuntimeSchemaNameForPackage(@NonNull String pkg) {
+        return RUNTIME_SCHEMA_TYPE + RUNTIME_SCHEMA_TYPE_SEPARATOR + Objects.requireNonNull(pkg);
+    }
+
+    /**
+     * Returns the document id for an app function's runtime metadata.
+     */
+    public static String getDocumentIdForAppFunction(
+            @NonNull String pkg, @NonNull String functionId) {
+        return pkg + "/" + functionId;
+    }
+
+    /**
+     * Different packages have different visibility requirements. To allow for different visibility,
+     * we need to have per-package app function schemas.
+     * <p>This schema should be set visible to callers from the package owner itself and for callers
+     * with {@link android.permission.EXECUTE_APP_FUNCTIONS_TRUSTED} or {@link
+     * android.permission.EXECUTE_APP_FUNCTIONS} permissions.
+     *
+     * @param packageName The package name to create a schema for.
+     */
+    @NonNull
+    public static AppSearchSchema createAppFunctionRuntimeSchema(@NonNull String packageName) {
+        return new AppSearchSchema.Builder(getRuntimeSchemaNameForPackage(packageName))
+                .addProperty(
+                        new AppSearchSchema.StringPropertyConfig.Builder(PROPERTY_FUNCTION_ID)
+                                .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
+                                .setIndexingType(
+                                        AppSearchSchema.StringPropertyConfig
+                                                .INDEXING_TYPE_EXACT_TERMS)
+                                .setTokenizerType(
+                                        AppSearchSchema.StringPropertyConfig
+                                                .TOKENIZER_TYPE_VERBATIM)
+                                .build())
+                .addProperty(
+                        new AppSearchSchema.StringPropertyConfig.Builder(PROPERTY_PACKAGE_NAME)
+                                .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
+                                .setIndexingType(
+                                        AppSearchSchema.StringPropertyConfig
+                                                .INDEXING_TYPE_EXACT_TERMS)
+                                .setTokenizerType(
+                                        AppSearchSchema.StringPropertyConfig
+                                                .TOKENIZER_TYPE_VERBATIM)
+                                .build())
+                .addProperty(
+                        new AppSearchSchema.BooleanPropertyConfig.Builder(PROPERTY_ENABLED)
+                                .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
+                                .build())
+                .addProperty(
+                        new AppSearchSchema.StringPropertyConfig.Builder(
+                                PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID)
+                                .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
+                                .setJoinableValueType(
+                                        AppSearchSchema.StringPropertyConfig
+                                                .JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                .build())
+                .addParentType(RUNTIME_SCHEMA_TYPE)
+                .build();
+    }
+
+    /**
+     * Returns the function id. This might look like "com.example.message#send_message".
+     */
+    @NonNull
+    public String getFunctionId() {
+        return Objects.requireNonNull(getPropertyString(PROPERTY_FUNCTION_ID));
+    }
+
+    /**
+     * Returns the package name of the package that owns this function.
+     */
+    @NonNull
+    public String getPackageName() {
+        return Objects.requireNonNull(getPropertyString(PROPERTY_PACKAGE_NAME));
+    }
+
+    /**
+     * Returns if the function is set to be enabled or not. If not set, the {@link
+     * AppFunctionStaticMetadataHelper#STATIC_PROPERTY_ENABLED_BY_DEFAULT} value would be used.
+     */
+    @Nullable
+    public Boolean getEnabled() {
+        return (Boolean) getProperty(PROPERTY_ENABLED);
+    }
+
+    /**
+     * Returns the qualified id linking to the static metadata of the app function.
+     */
+    @Nullable
+    @VisibleForTesting
+    public String getAppFunctionStaticMetadataQualifiedId() {
+        return getPropertyString(PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID);
+    }
+
+    public static final class Builder extends GenericDocument.Builder<Builder> {
+        /**
+         * Creates a Builder for a {@link AppFunctionRuntimeMetadata}.
+         *
+         * @param packageName               the name of the package that owns the function.
+         * @param functionId                the id of the function.
+         * @param staticMetadataQualifiedId the qualified static metadata id that this runtime
+         *                                  metadata refers to.
+         */
+        public Builder(
+                @NonNull String packageName,
+                @NonNull String functionId,
+                @NonNull String staticMetadataQualifiedId) {
+            super(
+                    APP_FUNCTION_RUNTIME_NAMESPACE,
+                    getDocumentIdForAppFunction(
+                            Objects.requireNonNull(packageName),
+                            Objects.requireNonNull(functionId)),
+                    getRuntimeSchemaNameForPackage(packageName));
+            setPropertyString(PROPERTY_PACKAGE_NAME, packageName);
+            setPropertyString(PROPERTY_FUNCTION_ID, functionId);
+
+            // Set qualified id automatically
+            setPropertyString(
+                    PROPERTY_APP_FUNCTION_STATIC_METADATA_QUALIFIED_ID,
+                    staticMetadataQualifiedId);
+        }
+
+
+        /**
+         * Sets an indicator specifying if the function is enabled or not. This would override the
+         * default enabled state in the static metadata ({@link
+         * AppFunctionStaticMetadataHelper#STATIC_PROPERTY_ENABLED_BY_DEFAULT}).
+         */
+        @NonNull
+        public Builder setEnabled(boolean enabled) {
+            setPropertyBoolean(PROPERTY_ENABLED, enabled);
+            return this;
+        }
+
+        /**
+         * Creates the {@link AppFunctionRuntimeMetadata} GenericDocument.
+         */
+        @NonNull
+        public AppFunctionRuntimeMetadata build() {
+            return new AppFunctionRuntimeMetadata(super.build());
+        }
+    }
+}
diff --git a/core/java/android/app/appfunctions/AppFunctionService.java b/core/java/android/app/appfunctions/AppFunctionService.java
index 6259d16..22bc908 100644
--- a/core/java/android/app/appfunctions/AppFunctionService.java
+++ b/core/java/android/app/appfunctions/AppFunctionService.java
@@ -81,8 +81,9 @@
                         // Apps should handle exceptions. But if they don't, report the error on
                         // behalf of them.
                         safeCallback.onResult(
-                                new ExecuteAppFunctionResponse.Builder(
-                                        getResultCode(ex), getExceptionMessage(ex)).build());
+                                ExecuteAppFunctionResponse.newFailure(
+                                        getResultCode(ex),
+                                        ex.getMessage(), /*extras=*/  null));
                     }
                 }
             };
@@ -117,8 +118,4 @@
     public abstract void onExecuteFunction(
             @NonNull ExecuteAppFunctionRequest request,
             @NonNull Consumer<ExecuteAppFunctionResponse> callback);
-
-    private String getExceptionMessage(Exception exception) {
-        return exception.getMessage() == null ? "" : exception.getMessage();
-    }
 }
diff --git a/core/java/android/app/appfunctions/AppFunctionStaticMetadataHelper.java b/core/java/android/app/appfunctions/AppFunctionStaticMetadataHelper.java
new file mode 100644
index 0000000..6d4172a
--- /dev/null
+++ b/core/java/android/app/appfunctions/AppFunctionStaticMetadataHelper.java
@@ -0,0 +1,72 @@
+/*
+ * 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.app.appfunctions;
+
+import static android.app.appfunctions.flags.Flags.FLAG_ENABLE_APP_FUNCTION_MANAGER;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.app.appsearch.util.DocumentIdUtil;
+
+import java.util.Objects;
+
+/**
+ * Contains constants and helper related to static metadata represented with {@code
+ * com.android.server.appsearch.appsindexer.appsearchtypes.AppFunctionStaticMetadata}.
+ * <p>
+ * The constants listed here **must not change** and be kept consistent with the canonical
+ * static metadata class.
+ *
+ * @hide
+ */
+@FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+public class AppFunctionStaticMetadataHelper {
+    public static final String STATIC_SCHEMA_TYPE = "AppFunctionStaticMetadata";
+    public static final String STATIC_PROPERTY_ENABLED_BY_DEFAULT = "enabledByDefault";
+
+    public static final String APP_FUNCTION_STATIC_NAMESPACE = "app_functions";
+
+    // These are constants that has to be kept the same with {@code
+    // com.android.server.appsearch.appsindexer.appsearchtypes.AppSearchHelper}.
+    public static final String APP_FUNCTION_STATIC_METADATA_DB = "apps-db";
+    public static final String APP_FUNCTION_INDEXER_PACKAGE = "android";
+
+    /**
+     * Returns a per-app static metadata schema name, to store all functions for that package.
+     */
+    public static String getStaticSchemaNameForPackage(@NonNull String pkg) {
+        return STATIC_SCHEMA_TYPE + "-" + Objects.requireNonNull(pkg);
+    }
+
+    /** Returns the document id for an app function's static metadata. */
+    public static String getDocumentIdForAppFunction(
+            @NonNull String pkg, @NonNull String functionId) {
+        return pkg + "/" + functionId;
+    }
+
+    /**
+     * Returns the fully qualified Id used in AppSearch for the given package and function id
+     * app function static metadata.
+     */
+    public static String getStaticMetadataQualifiedId(String packageName, String functionId) {
+        return DocumentIdUtil.createQualifiedId(
+                APP_FUNCTION_INDEXER_PACKAGE,
+                APP_FUNCTION_STATIC_METADATA_DB,
+                APP_FUNCTION_STATIC_NAMESPACE,
+                getDocumentIdForAppFunction(packageName, functionId));
+    }
+}
diff --git a/core/java/android/app/appfunctions/ExecuteAppFunctionResponse.java b/core/java/android/app/appfunctions/ExecuteAppFunctionResponse.java
index 9fb3375..58d4088 100644
--- a/core/java/android/app/appfunctions/ExecuteAppFunctionResponse.java
+++ b/core/java/android/app/appfunctions/ExecuteAppFunctionResponse.java
@@ -162,6 +162,55 @@
     }
 
     /**
+     * Returns a successful response.
+     *
+     * @param resultDocument The return value of the executed function.
+     * @param extras         The additional metadata data relevant to this function execution
+     *                       response.
+     */
+    @NonNull
+    @FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+    public static ExecuteAppFunctionResponse newSuccess(@NonNull GenericDocument resultDocument,
+                                                        @Nullable Bundle extras) {
+        Objects.requireNonNull(resultDocument);
+        Bundle actualExtras = getActualExtras(extras);
+        GenericDocumentWrapper resultDocumentWrapper = new GenericDocumentWrapper(resultDocument);
+
+        return new ExecuteAppFunctionResponse(
+                resultDocumentWrapper, actualExtras, RESULT_OK, /*errorMessage=*/null);
+    }
+
+    /**
+     * Returns a failure response.
+     *
+     * @param resultCode   The result code of the app function execution.
+     * @param extras       The additional metadata data relevant to this function execution
+     *                     response.
+     * @param errorMessage The error message associated with the result, if any.
+     */
+    @NonNull
+    @FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+    public static ExecuteAppFunctionResponse newFailure(@ResultCode int resultCode,
+                                                        @Nullable String errorMessage,
+                                                        @Nullable Bundle extras) {
+        if (resultCode == RESULT_OK) {
+            throw new IllegalArgumentException("resultCode must not be RESULT_OK");
+        }
+        Bundle actualExtras = getActualExtras(extras);
+        GenericDocumentWrapper emptyWrapper = new GenericDocumentWrapper(
+                new GenericDocument.Builder<>("", "", "").build());
+        return new ExecuteAppFunctionResponse(
+                emptyWrapper, actualExtras, resultCode, errorMessage);
+    }
+
+    private static Bundle getActualExtras(@Nullable Bundle extras) {
+        if (extras == null) {
+            return Bundle.EMPTY;
+        }
+        return extras;
+    }
+
+    /**
      * Returns a generic document containing the return value of the executed function.
      *
      * <p>The {@link #PROPERTY_RETURN_VALUE} key can be used to obtain the return value.</p>
@@ -250,64 +299,4 @@
     @Retention(RetentionPolicy.SOURCE)
     public @interface ResultCode {
     }
-
-    /**
-     * The builder for creating {@link ExecuteAppFunctionResponse} instances.
-     */
-    @FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
-    public static final class Builder {
-        @NonNull
-        private GenericDocument mResultDocument =
-                new GenericDocument.Builder<>("", "", "").build();
-        @NonNull
-        private Bundle mExtras = Bundle.EMPTY;
-        private int mResultCode;
-        @Nullable
-        private String mErrorMessage;
-
-        /**
-         * Creates a new builder for {@link ExecuteAppFunctionResponse}.
-         */
-        private Builder() {
-        }
-
-        /**
-         * Creates a new builder for {@link ExecuteAppFunctionResponse} to build a success response
-         * with a result code of {@link #RESULT_OK} and a resultDocument.
-         */
-        public Builder(@NonNull GenericDocument resultDocument) {
-            Objects.requireNonNull(resultDocument);
-            mResultDocument = resultDocument;
-            mResultCode = RESULT_OK;
-        }
-
-        /**
-         * Creates a new builder for {@link ExecuteAppFunctionResponse} to build an error response
-         * with a result code and an error message.
-         */
-        public Builder(@ResultCode int resultCode,
-                       @NonNull String errorMessage) {
-            mResultCode = resultCode;
-            mErrorMessage = Objects.requireNonNull(errorMessage);
-        }
-
-        /**
-         * Sets the extras of the app function execution.
-         */
-        @NonNull
-        public Builder setExtras(@NonNull Bundle extras) {
-            mExtras = Objects.requireNonNull(extras);
-            return this;
-        }
-
-        /**
-         * Builds the {@link ExecuteAppFunctionResponse} instance.
-         */
-        @NonNull
-        public ExecuteAppFunctionResponse build() {
-            return new ExecuteAppFunctionResponse(
-                    new GenericDocumentWrapper(mResultDocument),
-                    mExtras, mResultCode, mErrorMessage);
-        }
-    }
 }
diff --git a/core/java/android/app/supervision/OWNERS b/core/java/android/app/supervision/OWNERS
new file mode 100644
index 0000000..afc5495
--- /dev/null
+++ b/core/java/android/app/supervision/OWNERS
@@ -0,0 +1,2 @@
+jparks@google.com
+romkal@google.com
diff --git a/core/java/android/companion/flags.aconfig b/core/java/android/companion/flags.aconfig
index ee9114f..93d62cf 100644
--- a/core/java/android/companion/flags.aconfig
+++ b/core/java/android/companion/flags.aconfig
@@ -10,13 +10,6 @@
 }
 
 flag {
-    name: "companion_transport_apis"
-    namespace: "companion"
-    description: "Grants access to the companion transport apis."
-    bug: "288297505"
-}
-
-flag {
     name: "association_tag"
     is_exported: true
     namespace: "companion"
diff --git a/core/java/android/companion/virtual/IVirtualDeviceActivityListener.aidl b/core/java/android/companion/virtual/IVirtualDeviceActivityListener.aidl
index 7c674f9..767f52a 100644
--- a/core/java/android/companion/virtual/IVirtualDeviceActivityListener.aidl
+++ b/core/java/android/companion/virtual/IVirtualDeviceActivityListener.aidl
@@ -54,4 +54,13 @@
      */
     void onActivityLaunchBlocked(int displayId, in ComponentName componentName, in UserHandle user,
             in IntentSender intentSender);
+
+    /**
+     * Called when a secure surface is shown on the device.
+     *
+     * @param displayId The display ID on which the secure surface was shown.
+     * @param componentName The component name of the activity that showed the secure surface.
+     * @param user The user associated with the activity.
+     */
+    void onSecureWindowShown(int displayId, in ComponentName componentName, in UserHandle user);
 }
diff --git a/core/java/android/companion/virtual/VirtualDeviceInternal.java b/core/java/android/companion/virtual/VirtualDeviceInternal.java
index b7bf2d1..de20a68 100644
--- a/core/java/android/companion/virtual/VirtualDeviceInternal.java
+++ b/core/java/android/companion/virtual/VirtualDeviceInternal.java
@@ -151,7 +151,24 @@
                         Binder.restoreCallingIdentity(token);
                     }
                 }
+
+                @Override
+                public void onSecureWindowShown(int displayId, ComponentName componentName,
+                        UserHandle user) {
+                    final long token = Binder.clearCallingIdentity();
+                    try {
+                        synchronized (mActivityListenersLock) {
+                            for (int i = 0; i < mActivityListeners.size(); i++) {
+                                mActivityListeners.valueAt(i)
+                                        .onSecureWindowShown(displayId, componentName, user);
+                            }
+                        }
+                    } finally {
+                        Binder.restoreCallingIdentity(token);
+                    }
+                }
             };
+
     private final IVirtualDeviceSoundEffectListener mSoundEffectListener =
             new IVirtualDeviceSoundEffectListener.Stub() {
                 @Override
@@ -584,6 +601,12 @@
                     mActivityListener.onActivityLaunchBlocked(
                             displayId, componentName, user, intentSender));
         }
+
+        public void onSecureWindowShown(int displayId, ComponentName componentName,
+                UserHandle user) {
+            mExecutor.execute(() ->
+                    mActivityListener.onSecureWindowShown(displayId, componentName, user));
+        }
     }
 
     /**
diff --git a/core/java/android/companion/virtual/VirtualDeviceManager.java b/core/java/android/companion/virtual/VirtualDeviceManager.java
index 40aa6837..cf34452 100644
--- a/core/java/android/companion/virtual/VirtualDeviceManager.java
+++ b/core/java/android/companion/virtual/VirtualDeviceManager.java
@@ -1255,6 +1255,20 @@
         @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_ACTIVITY_CONTROL_API)
         default void onActivityLaunchBlocked(int displayId, @NonNull ComponentName componentName,
                 @NonNull UserHandle user, @Nullable IntentSender intentSender) {}
+
+        /**
+         * Called when a window with a secure surface is shown on the device.
+         *
+         * @param displayId The display ID on which the window was shown.
+         * @param componentName The component name of the activity that showed the window.
+         * @param user The user associated with the activity.
+         *
+         * @see Display#FLAG_SECURE
+         * @see WindowManager.LayoutParams#FLAG_SECURE
+         */
+        @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_ACTIVITY_CONTROL_API)
+        default void onSecureWindowShown(int displayId, @NonNull ComponentName componentName,
+                @NonNull UserHandle user) {}
     }
 
     /**
diff --git a/core/java/android/companion/virtual/flags/flags.aconfig b/core/java/android/companion/virtual/flags/flags.aconfig
index 260fcfe..297fe8a 100644
--- a/core/java/android/companion/virtual/flags/flags.aconfig
+++ b/core/java/android/companion/virtual/flags/flags.aconfig
@@ -50,6 +50,7 @@
      name: "activity_control_api"
      description: "Enable APIs for fine grained activity policy, fallback and callbacks"
      bug: "333443509"
+     is_exported: true
 }
 
 flag {
diff --git a/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java b/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
index 82f183f..68bc9bc 100644
--- a/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
+++ b/core/java/android/companion/virtual/sensor/VirtualSensorConfig.java
@@ -17,7 +17,13 @@
 package android.companion.virtual.sensor;
 
 
+import static android.hardware.Sensor.REPORTING_MODE_CONTINUOUS;
+import static android.hardware.Sensor.REPORTING_MODE_ONE_SHOT;
+import static android.hardware.Sensor.REPORTING_MODE_ON_CHANGE;
+import static android.hardware.Sensor.REPORTING_MODE_SPECIAL_TRIGGER;
+
 import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -30,6 +36,8 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 import java.util.Objects;
 
 
@@ -71,6 +79,17 @@
 
     private final int mFlags;
 
+    /** @hide */
+    @IntDef(prefix = "REPORTING_MODE_", value = {
+            REPORTING_MODE_CONTINUOUS,
+            REPORTING_MODE_ON_CHANGE,
+            REPORTING_MODE_ONE_SHOT,
+            REPORTING_MODE_SPECIAL_TRIGGER
+    })
+
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface ReportingMode {}
+
     private VirtualSensorConfig(int type, @NonNull String name, @Nullable String vendor,
             float maximumRange, float resolution, float power, int minDelay, int maxDelay,
             int flags) {
@@ -240,7 +259,7 @@
      * @see Sensor#getReportingMode()
      */
     @FlaggedApi(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
-    public int getReportingMode() {
+    public @ReportingMode int getReportingMode() {
         return ((mFlags & REPORTING_MODE_MASK) >> REPORTING_MODE_SHIFT);
     }
 
@@ -442,11 +461,11 @@
          */
         @FlaggedApi(Flags.FLAG_DEVICE_AWARE_DISPLAY_POWER)
         @NonNull
-        public VirtualSensorConfig.Builder setReportingMode(int reportingMode) {
-            if (reportingMode != Sensor.REPORTING_MODE_CONTINUOUS
-                    && reportingMode != Sensor.REPORTING_MODE_ON_CHANGE
-                    && reportingMode != Sensor.REPORTING_MODE_ONE_SHOT
-                    && reportingMode != Sensor.REPORTING_MODE_SPECIAL_TRIGGER) {
+        public VirtualSensorConfig.Builder setReportingMode(@ReportingMode int reportingMode) {
+            if (reportingMode != REPORTING_MODE_CONTINUOUS
+                    && reportingMode != REPORTING_MODE_ON_CHANGE
+                    && reportingMode != REPORTING_MODE_ONE_SHOT
+                    && reportingMode != REPORTING_MODE_SPECIAL_TRIGGER) {
                 throw new IllegalArgumentException("Invalid reporting mode: " + reportingMode);
             }
             mFlags |= reportingMode << REPORTING_MODE_SHIFT;
diff --git a/core/java/android/content/AttributionSource.java b/core/java/android/content/AttributionSource.java
index 7fcfbbc..ffed536 100644
--- a/core/java/android/content/AttributionSource.java
+++ b/core/java/android/content/AttributionSource.java
@@ -629,11 +629,10 @@
      * A builder for {@link AttributionSource}
      */
     public static final class Builder {
+        private boolean mHasBeenUsed;
         private @NonNull final AttributionSourceState mAttributionSourceState =
                 new AttributionSourceState();
 
-        private long mBuilderFieldsSet = 0L;
-
         /**
          * Creates a new Builder.
          *
@@ -642,8 +641,17 @@
          */
         public Builder(int uid) {
             mAttributionSourceState.uid = uid;
+            mAttributionSourceState.pid = Process.INVALID_PID;
+            mAttributionSourceState.deviceId = Context.DEVICE_ID_DEFAULT;
+            mAttributionSourceState.token = sDefaultToken;
         }
 
+        /**
+         * Creates a builder that is ready to build a new {@link AttributionSource} where
+         * all fields (primitive, immutable data, pointers) are copied from the given
+         * {@link AttributionSource}. Builder methods can still be used to mutate fields further.
+         * @param current The source to copy fields from.
+         */
         public Builder(@NonNull AttributionSource current) {
             if (current == null) {
                 throw new IllegalArgumentException("current AttributionSource can not be null");
@@ -652,9 +660,11 @@
             mAttributionSourceState.pid = current.getPid();
             mAttributionSourceState.packageName = current.getPackageName();
             mAttributionSourceState.attributionTag = current.getAttributionTag();
-            mAttributionSourceState.token = current.getToken();
             mAttributionSourceState.renouncedPermissions =
                 current.mAttributionSourceState.renouncedPermissions;
+            mAttributionSourceState.deviceId = current.getDeviceId();
+            mAttributionSourceState.next = current.mAttributionSourceState.next;
+            mAttributionSourceState.token = current.getToken();
         }
 
         /**
@@ -666,7 +676,6 @@
          */
         public @NonNull Builder setPid(int value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x2;
             mAttributionSourceState.pid = value;
             return this;
         }
@@ -676,7 +685,6 @@
          */
         public @NonNull Builder setPackageName(@Nullable String value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x4;
             mAttributionSourceState.packageName = value;
             return this;
         }
@@ -686,7 +694,6 @@
          */
         public @NonNull Builder setAttributionTag(@Nullable String value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x8;
             mAttributionSourceState.attributionTag = value;
             return this;
         }
@@ -717,11 +724,11 @@
          */
         @SystemApi
         @RequiresPermission(android.Manifest.permission.RENOUNCE_PERMISSIONS)
-        public @NonNull Builder setRenouncedPermissions(@Nullable Set<String> value) {
+        public @NonNull Builder setRenouncedPermissions(
+                @Nullable Set<String> renouncedPermissions) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x10;
-            mAttributionSourceState.renouncedPermissions = (value != null)
-                    ? value.toArray(new String[0]) : null;
+            mAttributionSourceState.renouncedPermissions = (renouncedPermissions != null)
+                    ? renouncedPermissions.toArray(new String[0]) : null;
             return this;
         }
 
@@ -734,7 +741,6 @@
         @FlaggedApi(Flags.FLAG_DEVICE_AWARE_PERMISSION_APIS_ENABLED)
         public @NonNull Builder setDeviceId(int deviceId) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x20;
             mAttributionSourceState.deviceId = deviceId;
             return this;
         }
@@ -744,7 +750,6 @@
          */
         public @NonNull Builder setNext(@Nullable AttributionSource value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x40;
             mAttributionSourceState.next = (value != null) ? new AttributionSourceState[]
                     {value.mAttributionSourceState} : mAttributionSourceState.next;
             return this;
@@ -759,7 +764,6 @@
             if (value == null) {
                 throw new IllegalArgumentException("Null AttributionSource not permitted.");
             }
-            mBuilderFieldsSet |= 0x40;
             mAttributionSourceState.next =
                     new AttributionSourceState[]{value.mAttributionSourceState};
             return this;
@@ -768,28 +772,7 @@
         /** Builds the instance. This builder should not be touched after calling this! */
         public @NonNull AttributionSource build() {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x80; // Mark builder used
-
-            if ((mBuilderFieldsSet & 0x2) == 0) {
-                mAttributionSourceState.pid = Process.INVALID_PID;
-            }
-            if ((mBuilderFieldsSet & 0x4) == 0) {
-                mAttributionSourceState.packageName = null;
-            }
-            if ((mBuilderFieldsSet & 0x8) == 0) {
-                mAttributionSourceState.attributionTag = null;
-            }
-            if ((mBuilderFieldsSet & 0x10) == 0) {
-                mAttributionSourceState.renouncedPermissions = null;
-            }
-            if ((mBuilderFieldsSet & 0x20) == 0) {
-                mAttributionSourceState.deviceId = Context.DEVICE_ID_DEFAULT;
-            }
-            if ((mBuilderFieldsSet & 0x40) == 0) {
-                mAttributionSourceState.next = null;
-            }
-
-            mAttributionSourceState.token = sDefaultToken;
+            mHasBeenUsed = true;
 
             if (mAttributionSourceState.next == null) {
                 // The NDK aidl backend doesn't support null parcelable arrays.
@@ -799,7 +782,7 @@
         }
 
         private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x80) != 0) {
+            if (mHasBeenUsed) {
                 throw new IllegalStateException(
                         "This Builder should not be reused. Use a new Builder instance instead");
             }
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index 3bf0f032..12c5d07 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -33,6 +33,7 @@
 import android.annotation.Nullable;
 import android.annotation.PermissionMethod;
 import android.annotation.PermissionName;
+import android.annotation.RequiresFeature;
 import android.annotation.RequiresPermission;
 import android.annotation.StringDef;
 import android.annotation.StringRes;
@@ -6678,6 +6679,8 @@
      * Use with {@link #getSystemService(String)} to retrieve a {@link
      * android.webkit.WebViewUpdateManager} for accessing the WebView update service.
      *
+     * <p>This can only be used on devices with {@link PackageManager#FEATURE_WEBVIEW}.
+     *
      * @see #getSystemService(String)
      * @see android.webkit.WebViewUpdateManager
      * @hide
@@ -6685,6 +6688,7 @@
     @FlaggedApi(android.webkit.Flags.FLAG_UPDATE_SERVICE_IPC_WRAPPER)
     @SystemApi(client = SystemApi.Client.MODULE_LIBRARIES)
     @SuppressLint("ServiceName")
+    @RequiresFeature(PackageManager.FEATURE_WEBVIEW)
     public static final String WEBVIEW_UPDATE_SERVICE = "webviewupdate";
 
     /**
diff --git a/core/java/android/content/Intent.java b/core/java/android/content/Intent.java
index abb0d8d..da3cc1b 100644
--- a/core/java/android/content/Intent.java
+++ b/core/java/android/content/Intent.java
@@ -7650,13 +7650,6 @@
             | FLAG_GRANT_PREFIX_URI_PERMISSION;
 
     /**
-     * Flags that are not normally set by application code, but set for you by the system.
-     */
-    private static final int SYSTEM_ONLY_FLAGS = FLAG_ACTIVITY_LAUNCHED_FROM_HISTORY
-            | FLAG_ACTIVITY_BROUGHT_TO_FRONT
-            | FLAG_RECEIVER_FROM_SHELL;
-
-    /**
      * Local flag indicating this instance was created by copy constructor.
      */
     private static final int LOCAL_FLAG_FROM_COPY = 1 << 0;
@@ -7709,11 +7702,6 @@
     @TestApi
     public static final int EXTENDED_FLAG_FILTER_MISMATCH = 1 << 0;
 
-    /**
-     * Extended flags that are not normally set by application code, but set for you by the system.
-     */
-    private static final int SYSTEM_ONLY_EXTENDED_FLAGS = EXTENDED_FLAG_FILTER_MISMATCH;
-
     // ---------------------------------------------------------------------
     // ---------------------------------------------------------------------
     // toUri() and parseUri() options.
@@ -12657,28 +12645,6 @@
         }
     }
 
-    /**
-     * Prepare this {@link Intent} to enter system_server.
-     *
-     * @hide
-     */
-    public void prepareToEnterSystemServer() {
-        // Refuse possible leaked file descriptors
-        if (hasFileDescriptors()) {
-            throw new IllegalArgumentException("File descriptors passed in Intent");
-        }
-        // These flags are set only by the system, and should be stripped out as soon as the intent
-        // is received by system_server from the caller so it can be properly updated later.
-        removeFlags(SYSTEM_ONLY_FLAGS);
-        removeExtendedFlags(SYSTEM_ONLY_EXTENDED_FLAGS);
-        if (mOriginalIntent != null) {
-            mOriginalIntent.prepareToEnterSystemServer();
-        }
-        if (mSelector != null) {
-            mSelector.prepareToEnterSystemServer();
-        }
-    }
-
     /** @hide */
     public boolean hasWebURI() {
         if (getData() == null) {
diff --git a/core/java/android/content/rollback/RollbackInfo.java b/core/java/android/content/rollback/RollbackInfo.java
index d128055..a20159d 100644
--- a/core/java/android/content/rollback/RollbackInfo.java
+++ b/core/java/android/content/rollback/RollbackInfo.java
@@ -19,8 +19,6 @@
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.SystemApi;
-import android.annotation.TestApi;
-import android.content.pm.Flags;
 import android.content.pm.PackageManager;
 import android.content.pm.VersionedPackage;
 import android.os.Parcel;
@@ -136,11 +134,8 @@
      * Get rollback impact level. Refer {@link
      * android.content.pm.PackageInstaller.SessionParams#setRollbackImpactLevel(int)} for more info
      * on impact level.
-     *
-     * @hide
      */
-    @TestApi
-    @FlaggedApi(Flags.FLAG_RECOVERABILITY_DETECTION)
+    @FlaggedApi(android.crashrecovery.flags.Flags.FLAG_ENABLE_CRASHRECOVERY)
     public @PackageManager.RollbackImpactLevel int getRollbackImpactLevel() {
         return mRollbackImpactLevel;
     }
diff --git a/core/java/android/hardware/Sensor.java b/core/java/android/hardware/Sensor.java
index 10c3730..e0b9f60 100644
--- a/core/java/android/hardware/Sensor.java
+++ b/core/java/android/hardware/Sensor.java
@@ -17,7 +17,9 @@
 
 package android.hardware;
 
+import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.annotation.TestApi;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.hardware.input.InputSensorInfo;
 import android.os.Build;
@@ -1182,6 +1184,8 @@
 
     /** @hide */
     @UnsupportedAppUsage
+    @SuppressLint("UnflaggedApi") // Promotion to TestApi
+    @TestApi
     public int getHandle() {
         return mHandle;
     }
diff --git a/core/java/android/hardware/biometrics/AuthenticateOptions.java b/core/java/android/hardware/biometrics/AuthenticateOptions.java
index 7766071..4dc6ea19 100644
--- a/core/java/android/hardware/biometrics/AuthenticateOptions.java
+++ b/core/java/android/hardware/biometrics/AuthenticateOptions.java
@@ -74,4 +74,7 @@
 
     /** The attribution tag, if any. */
     @Nullable String getAttributionTag();
+
+    /** If the authentication is requested due to mandatory biometrics being active. */
+    boolean isMandatoryBiometrics();
 }
diff --git a/core/java/android/hardware/biometrics/IBiometricAuthenticator.aidl b/core/java/android/hardware/biometrics/IBiometricAuthenticator.aidl
index 17cd18c..b195225 100644
--- a/core/java/android/hardware/biometrics/IBiometricAuthenticator.aidl
+++ b/core/java/android/hardware/biometrics/IBiometricAuthenticator.aidl
@@ -49,7 +49,7 @@
     void prepareForAuthentication(boolean requireConfirmation, IBinder token, long operationId,
             int userId, IBiometricSensorReceiver sensorReceiver, String opPackageName,
             long requestId, int cookie, boolean allowBackgroundAuthentication,
-            boolean isForLegacyFingerprintManager);
+            boolean isForLegacyFingerprintManager, boolean isMandatoryBiometrics);
 
     // Starts authentication with the previously prepared client.
     void startPreparedClient(int cookie);
diff --git a/core/java/android/hardware/camera2/CaptureRequest.java b/core/java/android/hardware/camera2/CaptureRequest.java
index fbed50a..ac72116 100644
--- a/core/java/android/hardware/camera2/CaptureRequest.java
+++ b/core/java/android/hardware/camera2/CaptureRequest.java
@@ -1537,9 +1537,14 @@
      * be made, and for firing pre-capture flash pulses to estimate
      * scene brightness and required final capture flash power, when
      * the flash is enabled.</p>
-     * <p>Normally, this entry should be set to START for only a
-     * single request, and the application should wait until the
-     * sequence completes before starting a new one.</p>
+     * <p>Flash is enabled during precapture sequence when:</p>
+     * <ul>
+     * <li>AE mode is ON_ALWAYS_FLASH</li>
+     * <li>AE mode is ON_AUTO_FLASH and the scene is deemed too dark without flash, or</li>
+     * <li>AE mode is ON and flash mode is TORCH or SINGLE</li>
+     * </ul>
+     * <p>Normally, this entry should be set to START for only single request, and the
+     * application should wait until the sequence completes before starting a new one.</p>
      * <p>When a precapture metering sequence is finished, the camera device
      * may lock the auto-exposure routine internally to be able to accurately expose the
      * subsequent still capture image (<code>{@link CaptureRequest#CONTROL_CAPTURE_INTENT android.control.captureIntent} == STILL_CAPTURE</code>).
@@ -2705,6 +2710,13 @@
      * in {@link CameraCharacteristics#FLASH_SINGLE_STRENGTH_DEFAULT_LEVEL android.flash.singleStrengthDefaultLevel}.
      * If {@link CaptureRequest#CONTROL_AE_MODE android.control.aeMode} is set to any of <code>ON_AUTO_FLASH</code>, <code>ON_ALWAYS_FLASH</code>,
      * <code>ON_AUTO_FLASH_REDEYE</code>, <code>ON_EXTERNAL_FLASH</code> values, then the strengthLevel will be ignored.</p>
+     * <p>When AE mode is ON and flash mode is TORCH or SINGLE, the application should make sure
+     * the AE mode, flash mode, and flash strength level remain the same between precapture
+     * trigger request and final capture request. The flash strength level being set during
+     * precapture sequence is used by the camera device as a reference. The actual strength
+     * may be less, and the auto-exposure routine makes sure proper conversions of sensor
+     * exposure time and sensitivities between precapture and final capture for the specified
+     * strength level.</p>
      * <p><b>Range of valid values:</b><br>
      * <code>[1-{@link CameraCharacteristics#FLASH_TORCH_STRENGTH_MAX_LEVEL android.flash.torchStrengthMaxLevel}]</code> when the {@link CaptureRequest#FLASH_MODE android.flash.mode} is
      * set to TORCH;
diff --git a/core/java/android/hardware/camera2/CaptureResult.java b/core/java/android/hardware/camera2/CaptureResult.java
index d652b4c..34ce92c 100644
--- a/core/java/android/hardware/camera2/CaptureResult.java
+++ b/core/java/android/hardware/camera2/CaptureResult.java
@@ -935,9 +935,14 @@
      * be made, and for firing pre-capture flash pulses to estimate
      * scene brightness and required final capture flash power, when
      * the flash is enabled.</p>
-     * <p>Normally, this entry should be set to START for only a
-     * single request, and the application should wait until the
-     * sequence completes before starting a new one.</p>
+     * <p>Flash is enabled during precapture sequence when:</p>
+     * <ul>
+     * <li>AE mode is ON_ALWAYS_FLASH</li>
+     * <li>AE mode is ON_AUTO_FLASH and the scene is deemed too dark without flash, or</li>
+     * <li>AE mode is ON and flash mode is TORCH or SINGLE</li>
+     * </ul>
+     * <p>Normally, this entry should be set to START for only single request, and the
+     * application should wait until the sequence completes before starting a new one.</p>
      * <p>When a precapture metering sequence is finished, the camera device
      * may lock the auto-exposure routine internally to be able to accurately expose the
      * subsequent still capture image (<code>{@link CaptureRequest#CONTROL_CAPTURE_INTENT android.control.captureIntent} == STILL_CAPTURE</code>).
@@ -2821,8 +2826,6 @@
      * boost when the light level threshold is exceeded.</p>
      * <p>This state indicates when low light boost is 'ACTIVE' and applied. Similarly, it can
      * indicate when it is not being applied by returning 'INACTIVE'.</p>
-     * <p>This key will be absent from the CaptureResult if AE mode is not set to
-     * 'ON_LOW_LIGHT_BOOST_BRIGHTNESS_PRIORITY.</p>
      * <p>The default value will always be 'INACTIVE'.</p>
      * <p><b>Possible values:</b></p>
      * <ul>
@@ -2996,6 +2999,13 @@
      * in {@link CameraCharacteristics#FLASH_SINGLE_STRENGTH_DEFAULT_LEVEL android.flash.singleStrengthDefaultLevel}.
      * If {@link CaptureRequest#CONTROL_AE_MODE android.control.aeMode} is set to any of <code>ON_AUTO_FLASH</code>, <code>ON_ALWAYS_FLASH</code>,
      * <code>ON_AUTO_FLASH_REDEYE</code>, <code>ON_EXTERNAL_FLASH</code> values, then the strengthLevel will be ignored.</p>
+     * <p>When AE mode is ON and flash mode is TORCH or SINGLE, the application should make sure
+     * the AE mode, flash mode, and flash strength level remain the same between precapture
+     * trigger request and final capture request. The flash strength level being set during
+     * precapture sequence is used by the camera device as a reference. The actual strength
+     * may be less, and the auto-exposure routine makes sure proper conversions of sensor
+     * exposure time and sensitivities between precapture and final capture for the specified
+     * strength level.</p>
      * <p><b>Range of valid values:</b><br>
      * <code>[1-{@link CameraCharacteristics#FLASH_TORCH_STRENGTH_MAX_LEVEL android.flash.torchStrengthMaxLevel}]</code> when the {@link CaptureRequest#FLASH_MODE android.flash.mode} is
      * set to TORCH;
diff --git a/core/java/android/hardware/camera2/impl/CameraDeviceImpl.java b/core/java/android/hardware/camera2/impl/CameraDeviceImpl.java
index a60c48e..c7dba6c 100644
--- a/core/java/android/hardware/camera2/impl/CameraDeviceImpl.java
+++ b/core/java/android/hardware/camera2/impl/CameraDeviceImpl.java
@@ -80,6 +80,7 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -90,6 +91,17 @@
     private final String TAG;
     private final boolean DEBUG = false;
 
+    private static final ThreadFactory sThreadFactory = new ThreadFactory() {
+        private static final ThreadFactory mFactory = Executors.defaultThreadFactory();
+
+        @Override
+        public Thread newThread(Runnable r) {
+            Thread thread = mFactory.newThread(r);
+            thread.setName("CameraDeviceExecutor");
+            return thread;
+        }
+    };
+
     private static final int REQUEST_ID_NONE = -1;
 
     /**
@@ -354,7 +366,11 @@
         mCameraId = cameraId;
         if (Flags.singleThreadExecutor()) {
             mDeviceCallback = new ClientStateCallback(executor, callback);
-            mDeviceExecutor = Executors.newSingleThreadExecutor();
+            if (Flags.singleThreadExecutorNaming()) {
+                mDeviceExecutor = Executors.newSingleThreadExecutor(sThreadFactory);
+            } else {
+                mDeviceExecutor = Executors.newSingleThreadExecutor();
+            }
         } else {
             mDeviceCallback = callback;
             mDeviceExecutor = executor;
diff --git a/core/java/android/hardware/display/BrightnessInfo.java b/core/java/android/hardware/display/BrightnessInfo.java
index 109b0a8..6a96a54 100644
--- a/core/java/android/hardware/display/BrightnessInfo.java
+++ b/core/java/android/hardware/display/BrightnessInfo.java
@@ -158,6 +158,8 @@
                 return "thermal";
             case BRIGHTNESS_MAX_REASON_POWER_IC:
                 return "power IC";
+            case BRIGHTNESS_MAX_REASON_WEAR_BEDTIME_MODE:
+                return "wear bedtime";
         }
         return "invalid";
     }
diff --git a/core/java/android/hardware/face/FaceAuthenticateOptions.java b/core/java/android/hardware/face/FaceAuthenticateOptions.java
index 518f902a..8babbfa 100644
--- a/core/java/android/hardware/face/FaceAuthenticateOptions.java
+++ b/core/java/android/hardware/face/FaceAuthenticateOptions.java
@@ -120,6 +120,8 @@
     }
 
 
+    /** If the authentication is requested due to mandatory biometrics being active. */
+    private boolean mIsMandatoryBiometrics;
 
     // Code below generated by codegen v1.0.23.
     //
@@ -188,7 +190,8 @@
             @AuthenticateReason int authenticateReason,
             @PowerManager.WakeReason int wakeReason,
             @NonNull String opPackageName,
-            @Nullable String attributionTag) {
+            @Nullable String attributionTag,
+            boolean isMandatoryBiometrics) {
         this.mUserId = userId;
         this.mSensorId = sensorId;
         this.mDisplayState = displayState;
@@ -229,6 +232,7 @@
         com.android.internal.util.AnnotationValidations.validate(
                 NonNull.class, null, mOpPackageName);
         this.mAttributionTag = attributionTag;
+        this.mIsMandatoryBiometrics = isMandatoryBiometrics;
 
         // onConstructed(); // You can define this method to get a callback
     }
@@ -261,7 +265,7 @@
      * The reason for this operation when requested by the system (sysui),
      * otherwise AUTHENTICATE_REASON_UNKNOWN.
      *
-     * See packages/SystemUI/src/com/android/systemui/deviceentry/shared/FaceAuthReason.kt
+     * See frameworks/base/packages/SystemUI/src/com/android/keyguard/FaceAuthReason.kt
      * for more details about each reason.
      */
     @DataClass.Generated.Member
@@ -299,6 +303,14 @@
     }
 
     /**
+     * If the authentication is requested due to mandatory biometrics being active.
+     */
+    @DataClass.Generated.Member
+    public boolean isMandatoryBiometrics() {
+        return mIsMandatoryBiometrics;
+    }
+
+    /**
      * The sensor id for this operation.
      */
     @DataClass.Generated.Member
@@ -332,6 +344,15 @@
         return this;
     }
 
+    /**
+     * If the authentication is requested due to mandatory biometrics being active.
+     */
+    @DataClass.Generated.Member
+    public @NonNull FaceAuthenticateOptions setIsMandatoryBiometrics( boolean value) {
+        mIsMandatoryBiometrics = value;
+        return this;
+    }
+
     @Override
     @DataClass.Generated.Member
     public boolean equals(@Nullable Object o) {
@@ -351,7 +372,8 @@
                 && mAuthenticateReason == that.mAuthenticateReason
                 && mWakeReason == that.mWakeReason
                 && java.util.Objects.equals(mOpPackageName, that.mOpPackageName)
-                && java.util.Objects.equals(mAttributionTag, that.mAttributionTag);
+                && java.util.Objects.equals(mAttributionTag, that.mAttributionTag)
+                && mIsMandatoryBiometrics == that.mIsMandatoryBiometrics;
     }
 
     @Override
@@ -368,6 +390,7 @@
         _hash = 31 * _hash + mWakeReason;
         _hash = 31 * _hash + java.util.Objects.hashCode(mOpPackageName);
         _hash = 31 * _hash + java.util.Objects.hashCode(mAttributionTag);
+        _hash = 31 * _hash + Boolean.hashCode(mIsMandatoryBiometrics);
         return _hash;
     }
 
@@ -377,9 +400,10 @@
         // You can override field parcelling by defining methods like:
         // void parcelFieldName(Parcel dest, int flags) { ... }
 
-        byte flg = 0;
+        int flg = 0;
+        if (mIsMandatoryBiometrics) flg |= 0x80;
         if (mAttributionTag != null) flg |= 0x40;
-        dest.writeByte(flg);
+        dest.writeInt(flg);
         dest.writeInt(mUserId);
         dest.writeInt(mSensorId);
         dest.writeInt(mDisplayState);
@@ -400,7 +424,8 @@
         // You can override field unparcelling by defining methods like:
         // static FieldType unparcelFieldName(Parcel in) { ... }
 
-        byte flg = in.readByte();
+        int flg = in.readInt();
+        boolean isMandatoryBiometrics = (flg & 0x80) != 0;
         int userId = in.readInt();
         int sensorId = in.readInt();
         int displayState = in.readInt();
@@ -449,6 +474,7 @@
         com.android.internal.util.AnnotationValidations.validate(
                 NonNull.class, null, mOpPackageName);
         this.mAttributionTag = attributionTag;
+        this.mIsMandatoryBiometrics = isMandatoryBiometrics;
 
         // onConstructed(); // You can define this method to get a callback
     }
@@ -481,6 +507,7 @@
         private @PowerManager.WakeReason int mWakeReason;
         private @NonNull String mOpPackageName;
         private @Nullable String mAttributionTag;
+        private boolean mIsMandatoryBiometrics;
 
         private long mBuilderFieldsSet = 0L;
 
@@ -524,7 +551,7 @@
          * The reason for this operation when requested by the system (sysui),
          * otherwise AUTHENTICATE_REASON_UNKNOWN.
          *
-         * See packages/SystemUI/src/com/android/systemui/deviceentry/shared/FaceAuthReason.kt
+         * See frameworks/base/packages/SystemUI/src/com/android/keyguard/FaceAuthReason.kt
          * for more details about each reason.
          */
         @DataClass.Generated.Member
@@ -573,10 +600,21 @@
             return this;
         }
 
+        /**
+         * If the authentication is requested due to mandatory biometrics being active.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setIsMandatoryBiometrics(boolean value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x80;
+            mIsMandatoryBiometrics = value;
+            return this;
+        }
+
         /** Builds the instance. This builder should not be touched after calling this! */
         public @NonNull FaceAuthenticateOptions build() {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x80; // Mark builder used
+            mBuilderFieldsSet |= 0x100; // Mark builder used
 
             if ((mBuilderFieldsSet & 0x1) == 0) {
                 mUserId = defaultUserId();
@@ -606,12 +644,13 @@
                     mAuthenticateReason,
                     mWakeReason,
                     mOpPackageName,
-                    mAttributionTag);
+                    mAttributionTag,
+                    mIsMandatoryBiometrics);
             return o;
         }
 
         private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x80) != 0) {
+            if ((mBuilderFieldsSet & 0x100) != 0) {
                 throw new IllegalStateException(
                         "This Builder should not be reused. Use a new Builder instance instead");
             }
@@ -619,10 +658,10 @@
     }
 
     @DataClass.Generated(
-            time = 1677119626034L,
+            time = 1723436679828L,
             codegenVersion = "1.0.23",
             sourceFile = "frameworks/base/core/java/android/hardware/face/FaceAuthenticateOptions.java",
-            inputSignatures = "private final  int mUserId\nprivate  int mSensorId\nprivate final @android.hardware.biometrics.AuthenticateOptions.DisplayState int mDisplayState\npublic static final  int AUTHENTICATE_REASON_UNKNOWN\npublic static final  int AUTHENTICATE_REASON_STARTED_WAKING_UP\npublic static final  int AUTHENTICATE_REASON_PRIMARY_BOUNCER_SHOWN\npublic static final  int AUTHENTICATE_REASON_ASSISTANT_VISIBLE\npublic static final  int AUTHENTICATE_REASON_ALTERNATE_BIOMETRIC_BOUNCER_SHOWN\npublic static final  int AUTHENTICATE_REASON_NOTIFICATION_PANEL_CLICKED\npublic static final  int AUTHENTICATE_REASON_OCCLUDING_APP_REQUESTED\npublic static final  int AUTHENTICATE_REASON_PICK_UP_GESTURE_TRIGGERED\npublic static final  int AUTHENTICATE_REASON_QS_EXPANDED\npublic static final  int AUTHENTICATE_REASON_SWIPE_UP_ON_BOUNCER\npublic static final  int AUTHENTICATE_REASON_UDFPS_POINTER_DOWN\nprivate final @android.hardware.face.FaceAuthenticateOptions.AuthenticateReason int mAuthenticateReason\nprivate final @android.os.PowerManager.WakeReason int mWakeReason\nprivate @android.annotation.NonNull java.lang.String mOpPackageName\nprivate @android.annotation.Nullable java.lang.String mAttributionTag\nprivate static  int defaultUserId()\nprivate static  int defaultSensorId()\nprivate static  int defaultDisplayState()\nprivate static  int defaultAuthenticateReason()\nprivate static  int defaultWakeReason()\nprivate static  java.lang.String defaultOpPackageName()\nprivate static  java.lang.String defaultAttributionTag()\nclass FaceAuthenticateOptions extends java.lang.Object implements [android.hardware.biometrics.AuthenticateOptions, android.os.Parcelable]\n@com.android.internal.util.DataClass(genParcelable=true, genAidl=true, genBuilder=true, genSetters=true, genEqualsHashCode=true)")
+            inputSignatures = "private final  int mUserId\nprivate  int mSensorId\nprivate final @android.hardware.biometrics.AuthenticateOptions.DisplayState int mDisplayState\npublic static final  int AUTHENTICATE_REASON_UNKNOWN\npublic static final  int AUTHENTICATE_REASON_STARTED_WAKING_UP\npublic static final  int AUTHENTICATE_REASON_PRIMARY_BOUNCER_SHOWN\npublic static final  int AUTHENTICATE_REASON_ASSISTANT_VISIBLE\npublic static final  int AUTHENTICATE_REASON_ALTERNATE_BIOMETRIC_BOUNCER_SHOWN\npublic static final  int AUTHENTICATE_REASON_NOTIFICATION_PANEL_CLICKED\npublic static final  int AUTHENTICATE_REASON_OCCLUDING_APP_REQUESTED\npublic static final  int AUTHENTICATE_REASON_PICK_UP_GESTURE_TRIGGERED\npublic static final  int AUTHENTICATE_REASON_QS_EXPANDED\npublic static final  int AUTHENTICATE_REASON_SWIPE_UP_ON_BOUNCER\npublic static final  int AUTHENTICATE_REASON_UDFPS_POINTER_DOWN\nprivate final @android.hardware.face.FaceAuthenticateOptions.AuthenticateReason int mAuthenticateReason\nprivate final @android.os.PowerManager.WakeReason int mWakeReason\nprivate @android.annotation.NonNull java.lang.String mOpPackageName\nprivate @android.annotation.Nullable java.lang.String mAttributionTag\nprivate  boolean mIsMandatoryBiometrics\nprivate static  int defaultUserId()\nprivate static  int defaultSensorId()\nprivate static  int defaultDisplayState()\nprivate static  int defaultAuthenticateReason()\nprivate static  int defaultWakeReason()\nprivate static  java.lang.String defaultOpPackageName()\nprivate static  java.lang.String defaultAttributionTag()\nclass FaceAuthenticateOptions extends java.lang.Object implements [android.hardware.biometrics.AuthenticateOptions, android.os.Parcelable]\n@com.android.internal.util.DataClass(genParcelable=true, genAidl=true, genBuilder=true, genSetters=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/core/java/android/hardware/fingerprint/FingerprintAuthenticateOptions.java b/core/java/android/hardware/fingerprint/FingerprintAuthenticateOptions.java
index dc66542..ddf1e5b 100644
--- a/core/java/android/hardware/fingerprint/FingerprintAuthenticateOptions.java
+++ b/core/java/android/hardware/fingerprint/FingerprintAuthenticateOptions.java
@@ -97,6 +97,11 @@
         return null;
     }
 
+    /**
+     * If the authentication is requested due to mandatory biometrics being active.
+     */
+    private boolean mIsMandatoryBiometrics;
+
     // Code below generated by codegen v1.0.23.
     //
     // DO NOT MODIFY!
@@ -118,7 +123,8 @@
             @AuthenticateOptions.DisplayState int displayState,
             @NonNull String opPackageName,
             @Nullable String attributionTag,
-            @Nullable AuthenticateReason.Vendor vendorReason) {
+            @Nullable AuthenticateReason.Vendor vendorReason,
+            boolean isMandatoryBiometrics) {
         this.mUserId = userId;
         this.mSensorId = sensorId;
         this.mIgnoreEnrollmentState = ignoreEnrollmentState;
@@ -130,6 +136,7 @@
                 NonNull.class, null, mOpPackageName);
         this.mAttributionTag = attributionTag;
         this.mVendorReason = vendorReason;
+        this.mIsMandatoryBiometrics = isMandatoryBiometrics;
 
         // onConstructed(); // You can define this method to get a callback
     }
@@ -199,6 +206,14 @@
     }
 
     /**
+     * If the authentication is requested due to mandatory biometrics being active.
+     */
+    @DataClass.Generated.Member
+    public boolean isMandatoryBiometrics() {
+        return mIsMandatoryBiometrics;
+    }
+
+    /**
      * The sensor id for this operation.
      */
     @DataClass.Generated.Member
@@ -244,6 +259,15 @@
         return this;
     }
 
+    /**
+     * If the authentication is requested due to mandatory biometrics being active.
+     */
+    @DataClass.Generated.Member
+    public @NonNull FingerprintAuthenticateOptions setIsMandatoryBiometrics( boolean value) {
+        mIsMandatoryBiometrics = value;
+        return this;
+    }
+
     @Override
     @DataClass.Generated.Member
     public boolean equals(@Nullable Object o) {
@@ -263,7 +287,8 @@
                 && mDisplayState == that.mDisplayState
                 && java.util.Objects.equals(mOpPackageName, that.mOpPackageName)
                 && java.util.Objects.equals(mAttributionTag, that.mAttributionTag)
-                && java.util.Objects.equals(mVendorReason, that.mVendorReason);
+                && java.util.Objects.equals(mVendorReason, that.mVendorReason)
+                && mIsMandatoryBiometrics == that.mIsMandatoryBiometrics;
     }
 
     @Override
@@ -280,6 +305,7 @@
         _hash = 31 * _hash + java.util.Objects.hashCode(mOpPackageName);
         _hash = 31 * _hash + java.util.Objects.hashCode(mAttributionTag);
         _hash = 31 * _hash + java.util.Objects.hashCode(mVendorReason);
+        _hash = 31 * _hash + Boolean.hashCode(mIsMandatoryBiometrics);
         return _hash;
     }
 
@@ -289,11 +315,12 @@
         // You can override field parcelling by defining methods like:
         // void parcelFieldName(Parcel dest, int flags) { ... }
 
-        byte flg = 0;
+        int flg = 0;
         if (mIgnoreEnrollmentState) flg |= 0x4;
+        if (mIsMandatoryBiometrics) flg |= 0x80;
         if (mAttributionTag != null) flg |= 0x20;
         if (mVendorReason != null) flg |= 0x40;
-        dest.writeByte(flg);
+        dest.writeInt(flg);
         dest.writeInt(mUserId);
         dest.writeInt(mSensorId);
         dest.writeInt(mDisplayState);
@@ -313,8 +340,9 @@
         // You can override field unparcelling by defining methods like:
         // static FieldType unparcelFieldName(Parcel in) { ... }
 
-        byte flg = in.readByte();
+        int flg = in.readInt();
         boolean ignoreEnrollmentState = (flg & 0x4) != 0;
+        boolean isMandatoryBiometrics = (flg & 0x80) != 0;
         int userId = in.readInt();
         int sensorId = in.readInt();
         int displayState = in.readInt();
@@ -333,6 +361,7 @@
                 NonNull.class, null, mOpPackageName);
         this.mAttributionTag = attributionTag;
         this.mVendorReason = vendorReason;
+        this.mIsMandatoryBiometrics = isMandatoryBiometrics;
 
         // onConstructed(); // You can define this method to get a callback
     }
@@ -365,6 +394,7 @@
         private @NonNull String mOpPackageName;
         private @Nullable String mAttributionTag;
         private @Nullable AuthenticateReason.Vendor mVendorReason;
+        private boolean mIsMandatoryBiometrics;
 
         private long mBuilderFieldsSet = 0L;
 
@@ -456,10 +486,21 @@
             return this;
         }
 
+        /**
+         * If the authentication is requested due to mandatory biometrics being active.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setIsMandatoryBiometrics(boolean value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x80;
+            mIsMandatoryBiometrics = value;
+            return this;
+        }
+
         /** Builds the instance. This builder should not be touched after calling this! */
         public @NonNull FingerprintAuthenticateOptions build() {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x80; // Mark builder used
+            mBuilderFieldsSet |= 0x100; // Mark builder used
 
             if ((mBuilderFieldsSet & 0x1) == 0) {
                 mUserId = defaultUserId();
@@ -489,12 +530,13 @@
                     mDisplayState,
                     mOpPackageName,
                     mAttributionTag,
-                    mVendorReason);
+                    mVendorReason,
+                    mIsMandatoryBiometrics);
             return o;
         }
 
         private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x80) != 0) {
+            if ((mBuilderFieldsSet & 0x100) != 0) {
                 throw new IllegalStateException(
                         "This Builder should not be reused. Use a new Builder instance instead");
             }
@@ -502,10 +544,10 @@
     }
 
     @DataClass.Generated(
-            time = 1689703591032L,
+            time = 1723436831455L,
             codegenVersion = "1.0.23",
             sourceFile = "frameworks/base/core/java/android/hardware/fingerprint/FingerprintAuthenticateOptions.java",
-            inputSignatures = "private final  int mUserId\nprivate  int mSensorId\nprivate final  boolean mIgnoreEnrollmentState\nprivate final @android.hardware.biometrics.AuthenticateOptions.DisplayState int mDisplayState\nprivate @android.annotation.NonNull java.lang.String mOpPackageName\nprivate @android.annotation.Nullable java.lang.String mAttributionTag\nprivate @android.annotation.Nullable android.hardware.biometrics.common.AuthenticateReason.Vendor mVendorReason\nprivate static  int defaultUserId()\nprivate static  int defaultSensorId()\nprivate static  boolean defaultIgnoreEnrollmentState()\nprivate static  int defaultDisplayState()\nprivate static  java.lang.String defaultOpPackageName()\nprivate static  java.lang.String defaultAttributionTag()\nprivate static  android.hardware.biometrics.common.AuthenticateReason.Vendor defaultVendorReason()\nclass FingerprintAuthenticateOptions extends java.lang.Object implements [android.hardware.biometrics.AuthenticateOptions, android.os.Parcelable]\n@com.android.internal.util.DataClass(genParcelable=true, genAidl=true, genBuilder=true, genSetters=true, genEqualsHashCode=true)")
+            inputSignatures = "private final  int mUserId\nprivate  int mSensorId\nprivate final  boolean mIgnoreEnrollmentState\nprivate final @android.hardware.biometrics.AuthenticateOptions.DisplayState int mDisplayState\nprivate @android.annotation.NonNull java.lang.String mOpPackageName\nprivate @android.annotation.Nullable java.lang.String mAttributionTag\nprivate @android.annotation.Nullable android.hardware.biometrics.common.AuthenticateReason.Vendor mVendorReason\nprivate  boolean mIsMandatoryBiometrics\nprivate static  int defaultUserId()\nprivate static  int defaultSensorId()\nprivate static  boolean defaultIgnoreEnrollmentState()\nprivate static  int defaultDisplayState()\nprivate static  java.lang.String defaultOpPackageName()\nprivate static  java.lang.String defaultAttributionTag()\nprivate static  android.hardware.biometrics.common.AuthenticateReason.Vendor defaultVendorReason()\nclass FingerprintAuthenticateOptions extends java.lang.Object implements [android.hardware.biometrics.AuthenticateOptions, android.os.Parcelable]\n@com.android.internal.util.DataClass(genParcelable=true, genAidl=true, genBuilder=true, genSetters=true, genEqualsHashCode=true)")
     @Deprecated
     private void __metadata() {}
 
diff --git a/core/java/android/hardware/input/InputManagerGlobal.java b/core/java/android/hardware/input/InputManagerGlobal.java
index 03cf7c5..2a36238 100644
--- a/core/java/android/hardware/input/InputManagerGlobal.java
+++ b/core/java/android/hardware/input/InputManagerGlobal.java
@@ -567,7 +567,7 @@
         Objects.requireNonNull(listener, "listener must not be null");
 
         synchronized (mOnTabletModeChangedListeners) {
-            if (mOnTabletModeChangedListeners == null) {
+            if (mOnTabletModeChangedListeners.isEmpty()) {
                 initializeTabletModeListenerLocked();
             }
             int idx = findOnTabletModeChangedListenerLocked(listener);
diff --git a/core/java/android/hardware/input/InputSettings.java b/core/java/android/hardware/input/InputSettings.java
index c5d0caf22..8592ded 100644
--- a/core/java/android/hardware/input/InputSettings.java
+++ b/core/java/android/hardware/input/InputSettings.java
@@ -20,10 +20,12 @@
 import static com.android.hardware.input.Flags.FLAG_KEYBOARD_A11Y_MOUSE_KEYS;
 import static com.android.hardware.input.Flags.FLAG_KEYBOARD_A11Y_SLOW_KEYS_FLAG;
 import static com.android.hardware.input.Flags.FLAG_KEYBOARD_A11Y_STICKY_KEYS_FLAG;
+import static com.android.hardware.input.Flags.FLAG_KEYBOARD_REPEAT_KEYS;
 import static com.android.hardware.input.Flags.keyboardA11yBounceKeysFlag;
 import static com.android.hardware.input.Flags.keyboardA11ySlowKeysFlag;
 import static com.android.hardware.input.Flags.keyboardA11yStickyKeysFlag;
 import static com.android.hardware.input.Flags.keyboardA11yMouseKeys;
+import static com.android.hardware.input.Flags.keyboardRepeatKeys;
 import static com.android.hardware.input.Flags.touchpadTapDragging;
 import static com.android.hardware.input.Flags.touchpadVisualizer;
 import static com.android.input.flags.Flags.enableInputFilterRustImpl;
@@ -40,6 +42,7 @@
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.sysprop.InputProperties;
+import android.view.ViewConfiguration;
 
 /**
  * InputSettings encapsulates reading and writing settings related to input
@@ -90,6 +93,30 @@
      */
     public static final int DEFAULT_STYLUS_POINTER_ICON_ENABLED = 1;
 
+    /**
+     * The minimum allowed repeat keys timeout before starting key repeats.
+     * @hide
+     */
+    public static final int MIN_KEY_REPEAT_TIMEOUT_MILLIS = 150;
+
+    /**
+     * The maximum allowed repeat keys timeout before starting key repeats.
+     * @hide
+     */
+    public static final int MAX_KEY_REPEAT_TIMEOUT_MILLIS = 2000;
+
+    /**
+     * The minimum allowed repeat keys delay between successive key repeats.
+     * @hide
+     */
+    public static final int MIN_KEY_REPEAT_DELAY_MILLIS = 20;
+
+    /**
+     * The maximum allowed repeat keys delay between successive key repeats.
+     * @hide
+     */
+    public static final int MAX_KEY_REPEAT_DELAY_MILLIS = 2000;
+
     private InputSettings() {
     }
 
@@ -767,4 +794,141 @@
                 Settings.Secure.ACCESSIBILITY_MOUSE_KEYS_ENABLED, enabled ? 1 : 0,
                 UserHandle.USER_CURRENT);
     }
+
+    /**
+     * Whether "Repeat keys" feature flag is enabled.
+     *
+     * <p>
+     * ‘Repeat keys’ is a feature which allows users to generate key repeats when a particular
+     * key on the physical keyboard is held down. This accessibility feature allows the user
+     * to configure the timeout before the key repeats begin as well as the delay
+     * between successive key repeats.
+     * </p>
+     *
+     * @hide
+     */
+    public static boolean isRepeatKeysFeatureFlagEnabled() {
+        return keyboardRepeatKeys();
+    }
+
+    /**
+     * Get Accessibility repeat keys timeout duration in milliseconds.
+     * The default key repeat timeout is {@link ViewConfiguration#DEFAULT_KEY_REPEAT_TIMEOUT_MS}.
+     *
+     * @param context The application context
+     * @return The time duration for which a key should be pressed after
+     *         which the pressed key will be repeated. The timeout must be between
+     *         {@link #MIN_KEY_REPEAT_TIMEOUT_MILLIS} and
+     *         {@link #MAX_KEY_REPEAT_TIMEOUT_MILLIS}
+     *
+     * <p>
+     * ‘Repeat keys’ is a feature which allows users to generate key repeats when a particular
+     * key on the physical keyboard is held down. This accessibility feature allows the user
+     * to configure the timeout before the key repeats begin as well as the delay
+     * between successive key repeats.
+     * </p>
+     *
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_KEYBOARD_REPEAT_KEYS)
+    public static int getAccessibilityRepeatKeysTimeout(@NonNull Context context) {
+        return Settings.Secure.getIntForUser(context.getContentResolver(),
+                Settings.Secure.KEY_REPEAT_TIMEOUT_MS, ViewConfiguration.getKeyRepeatTimeout(),
+                UserHandle.USER_CURRENT);
+    }
+
+    /**
+     * Get Accessibility repeat keys delay rate in milliseconds.
+     * The default key repeat delay is {@link ViewConfiguration#DEFAULT_KEY_REPEAT_DELAY_MS}.
+     *
+     * @param context The application context
+     * @return Time duration between successive key repeats when a key is
+     *         pressed down. The delay duration must be between
+     *         {@link #MIN_KEY_REPEAT_DELAY_MILLIS} and
+     *         {@link #MAX_KEY_REPEAT_DELAY_MILLIS}
+     *
+     * <p>
+     * ‘Repeat keys’ is a feature which allows users to generate key repeats when a particular
+     * key on the physical keyboard is held down. This accessibility feature allows the user
+     * to configure the timeout before the key repeats begin as well as the delay
+     * between successive key repeats.
+     * </p>
+     *
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_KEYBOARD_REPEAT_KEYS)
+    public static int getAccessibilityRepeatKeysDelay(@NonNull Context context) {
+        return Settings.Secure.getIntForUser(context.getContentResolver(),
+                Settings.Secure.KEY_REPEAT_DELAY_MS, ViewConfiguration.getKeyRepeatDelay(),
+                UserHandle.USER_CURRENT);
+    }
+
+    /**
+     * Set Accessibility repeat keys timeout duration in milliseconds.
+     *
+     * @param timeoutTimeMillis time duration for which a key should be pressed after which the
+     *                          pressed key will be repeated. The timeout must be between
+     *                          {@link #MIN_KEY_REPEAT_TIMEOUT_MILLIS} and
+     *                          {@link #MAX_KEY_REPEAT_TIMEOUT_MILLIS}
+     *
+     *  <p>
+     * ‘Repeat keys’ is a feature which allows users to generate key repeats when a particular
+     * key on the physical keyboard is held down. This accessibility feature allows the user
+     * to configure the timeout before the key repeats begin as well as the delay
+     *  between successive key repeats.
+     * </p>
+     *
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_KEYBOARD_REPEAT_KEYS)
+    @RequiresPermission(Manifest.permission.WRITE_SETTINGS)
+    public static void setAccessibilityRepeatKeysTimeout(@NonNull Context context,
+            int timeoutTimeMillis) {
+        if (timeoutTimeMillis < MIN_KEY_REPEAT_TIMEOUT_MILLIS
+                || timeoutTimeMillis > MAX_KEY_REPEAT_TIMEOUT_MILLIS) {
+            throw new IllegalArgumentException(
+                    "Provided repeat keys timeout should be in range ("
+                            + MIN_KEY_REPEAT_TIMEOUT_MILLIS + ","
+                            + MAX_KEY_REPEAT_TIMEOUT_MILLIS + ")");
+        }
+        Settings.Secure.putIntForUser(context.getContentResolver(),
+                Settings.Secure.KEY_REPEAT_TIMEOUT_MS, timeoutTimeMillis,
+                UserHandle.USER_CURRENT);
+    }
+
+    /**
+     * Set Accessibility repeat key delay duration in milliseconds.
+     *
+     * @param delayTimeMillis Time duration between successive key repeats when a key is
+     *                        pressed down. The delay duration must be between
+     *                        {@link #MIN_KEY_REPEAT_DELAY_MILLIS} and
+     *                        {@link #MAX_KEY_REPEAT_DELAY_MILLIS}
+     * <p>
+     * ‘Repeat keys’ is a feature which allows users to generate key repeats when a particular
+     * key on the physical keyboard is held down. This accessibility feature allows the user
+     * to configure the timeout before the key repeats begin as well as the delay
+     * between successive key repeats.
+     * </p>
+     *
+     * @hide
+     */
+    @TestApi
+    @FlaggedApi(FLAG_KEYBOARD_REPEAT_KEYS)
+    @RequiresPermission(Manifest.permission.WRITE_SETTINGS)
+    public static void setAccessibilityRepeatKeysDelay(@NonNull Context context,
+            int delayTimeMillis) {
+        if (delayTimeMillis < MIN_KEY_REPEAT_DELAY_MILLIS
+                || delayTimeMillis > MAX_KEY_REPEAT_DELAY_MILLIS) {
+            throw new IllegalArgumentException(
+                    "Provided repeat keys delay should be in range ("
+                            + MIN_KEY_REPEAT_DELAY_MILLIS + ","
+                            + MAX_KEY_REPEAT_DELAY_MILLIS + ")");
+        }
+        Settings.Secure.putIntForUser(context.getContentResolver(),
+                Settings.Secure.KEY_REPEAT_DELAY_MS, delayTimeMillis,
+                UserHandle.USER_CURRENT);
+    }
 }
diff --git a/core/java/android/hardware/input/input_framework.aconfig b/core/java/android/hardware/input/input_framework.aconfig
index 83c4de3..077bd82 100644
--- a/core/java/android/hardware/input/input_framework.aconfig
+++ b/core/java/android/hardware/input/input_framework.aconfig
@@ -99,3 +99,10 @@
   description: "Refactor ModifierShortcutManager internal representation of shortcuts."
   bug: "358603902"
 }
+
+flag {
+  name: "keyboard_repeat_keys"
+  namespace: "input"
+  description: "Allow configurable timeout before key repeat and repeat delay rate for key repeats"
+  bug: "336585002"
+}
diff --git a/core/java/android/hardware/soundtrigger/SoundTrigger.java b/core/java/android/hardware/soundtrigger/SoundTrigger.java
index bfff4db..9f3e3ad 100644
--- a/core/java/android/hardware/soundtrigger/SoundTrigger.java
+++ b/core/java/android/hardware/soundtrigger/SoundTrigger.java
@@ -29,6 +29,7 @@
 import static java.util.Objects.requireNonNull;
 
 import android.annotation.ElapsedRealtimeLong;
+import android.annotation.FlaggedApi;
 import android.annotation.IntDef;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
@@ -874,10 +875,9 @@
     /*****************************************************************************
      * A GenericSoundModel is a specialized {@link SoundModel} for non-voice sound
      * patterns.
-     *
-     * @hide
      ****************************************************************************/
-    public static class GenericSoundModel extends SoundModel implements Parcelable {
+    @FlaggedApi(android.media.soundtrigger.Flags.FLAG_GENERIC_MODEL_API)
+    public static final class GenericSoundModel extends SoundModel implements Parcelable {
 
         public static final @android.annotation.NonNull Parcelable.Creator<GenericSoundModel> CREATOR
                 = new Parcelable.Creator<GenericSoundModel>() {
@@ -890,11 +890,26 @@
             }
         };
 
+        /**
+         * Constructor for {@link GenericSoundModel} with version.
+         *
+         * @param uuid Unique identifier for this sound model.
+         * @param vendorUuid Unique vendor identifier for this sound model.
+         * @param data Opaque data for this sound model.
+         * @param version Vendor-specific version number of this sound model.
+         */
         public GenericSoundModel(@NonNull UUID uuid, @NonNull UUID vendorUuid,
                 @Nullable byte[] data, int version) {
             super(uuid, vendorUuid, TYPE_GENERIC_SOUND, data, version);
         }
 
+        /**
+         * Constructor for {@link GenericSoundModel} without version. The version is set to -1.
+         *
+         * @param uuid Unique identifier for this sound model.
+         * @param vendorUuid Unique vendor identifier for this sound model.
+         * @param data Opaque data for this sound model.
+         */
         @UnsupportedAppUsage
         public GenericSoundModel(@NonNull UUID uuid, @NonNull UUID vendorUuid,
                 @Nullable byte[] data) {
@@ -919,7 +934,7 @@
         }
 
         @Override
-        public void writeToParcel(Parcel dest, int flags) {
+        public void writeToParcel(@NonNull Parcel dest, int flags) {
             dest.writeString(getUuid().toString());
             if (getVendorUuid() == null) {
                 dest.writeInt(-1);
diff --git a/core/java/android/os/Handler.java b/core/java/android/os/Handler.java
index 92b630f..80f39bf 100644
--- a/core/java/android/os/Handler.java
+++ b/core/java/android/os/Handler.java
@@ -46,13 +46,13 @@
  * a {@link Message} object containing a bundle of data that will be
  * processed by the Handler's {@link #handleMessage} method (requiring that
  * you implement a subclass of Handler).
- * 
+ *
  * <p>When posting or sending to a Handler, you can either
  * allow the item to be processed as soon as the message queue is ready
  * to do so, or specify a delay before it gets processed or absolute time for
  * it to be processed.  The latter two allow you to implement timeouts,
  * ticks, and other timing-based behavior.
- * 
+ *
  * <p>When a
  * process is created for your application, its main thread is dedicated to
  * running a message queue that takes care of managing the top-level
@@ -85,13 +85,13 @@
          */
         boolean handleMessage(@NonNull Message msg);
     }
-    
+
     /**
      * Subclasses must implement this to receive messages.
      */
     public void handleMessage(@NonNull Message msg) {
     }
-    
+
     /**
      * Handle system messages here.
      */
@@ -343,8 +343,8 @@
      * The default implementation will either return the class name of the
      * message callback if any, or the hexadecimal representation of the
      * message "what" field.
-     *  
-     * @param message The message whose name is being queried 
+     *
+     * @param message The message whose name is being queried
      */
     @NonNull
     public String getMessageName(@NonNull Message message) {
@@ -367,7 +367,7 @@
 
     /**
      * Same as {@link #obtainMessage()}, except that it also sets the what member of the returned Message.
-     * 
+     *
      * @param what Value to assign to the returned Message.what field.
      * @return A Message from the global message pool.
      */
@@ -376,12 +376,12 @@
     {
         return Message.obtain(this, what);
     }
-    
+
     /**
-     * 
-     * Same as {@link #obtainMessage()}, except that it also sets the what and obj members 
+     *
+     * Same as {@link #obtainMessage()}, except that it also sets the what and obj members
      * of the returned Message.
-     * 
+     *
      * @param what Value to assign to the returned Message.what field.
      * @param obj Value to assign to the returned Message.obj field.
      * @return A Message from the global message pool.
@@ -392,7 +392,7 @@
     }
 
     /**
-     * 
+     *
      * Same as {@link #obtainMessage()}, except that it also sets the what, arg1 and arg2 members of the returned
      * Message.
      * @param what Value to assign to the returned Message.what field.
@@ -405,10 +405,10 @@
     {
         return Message.obtain(this, what, arg1, arg2);
     }
-    
+
     /**
-     * 
-     * Same as {@link #obtainMessage()}, except that it also sets the what, obj, arg1,and arg2 values on the 
+     *
+     * Same as {@link #obtainMessage()}, except that it also sets the what, obj, arg1,and arg2 values on the
      * returned Message.
      * @param what Value to assign to the returned Message.what field.
      * @param arg1 Value to assign to the returned Message.arg1 field.
@@ -423,19 +423,19 @@
 
     /**
      * Causes the Runnable r to be added to the message queue.
-     * The runnable will be run on the thread to which this handler is 
-     * attached. 
-     *  
+     * The runnable will be run on the thread to which this handler is
+     * attached.
+     *
      * @param r The Runnable that will be executed.
-     * 
-     * @return Returns true if the Runnable was successfully placed in to the 
+     *
+     * @return Returns true if the Runnable was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
     public final boolean post(@NonNull Runnable r) {
        return  sendMessageDelayed(getPostMessage(r), 0);
     }
-    
+
     /**
      * Causes the Runnable r to be added to the message queue, to be run
      * at a specific time given by <var>uptimeMillis</var>.
@@ -446,8 +446,8 @@
      * @param r The Runnable that will be executed.
      * @param uptimeMillis The absolute time at which the callback should run,
      *         using the {@link android.os.SystemClock#uptimeMillis} time-base.
-     *  
-     * @return Returns true if the Runnable was successfully placed in to the 
+     *
+     * @return Returns true if the Runnable was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.  Note that a
      *         result of true does not mean the Runnable will be processed -- if
@@ -457,7 +457,7 @@
     public final boolean postAtTime(@NonNull Runnable r, long uptimeMillis) {
         return sendMessageAtTime(getPostMessage(r), uptimeMillis);
     }
-    
+
     /**
      * Causes the Runnable r to be added to the message queue, to be run
      * at a specific time given by <var>uptimeMillis</var>.
@@ -470,21 +470,21 @@
      *         {@link #removeCallbacksAndMessages}.
      * @param uptimeMillis The absolute time at which the callback should run,
      *         using the {@link android.os.SystemClock#uptimeMillis} time-base.
-     * 
-     * @return Returns true if the Runnable was successfully placed in to the 
+     *
+     * @return Returns true if the Runnable was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.  Note that a
      *         result of true does not mean the Runnable will be processed -- if
      *         the looper is quit before the delivery time of the message
      *         occurs then the message will be dropped.
-     *         
+     *
      * @see android.os.SystemClock#uptimeMillis
      */
     public final boolean postAtTime(
             @NonNull Runnable r, @Nullable Object token, long uptimeMillis) {
         return sendMessageAtTime(getPostMessage(r, token), uptimeMillis);
     }
-    
+
     /**
      * Causes the Runnable r to be added to the message queue, to be run
      * after the specified amount of time elapses.
@@ -492,12 +492,12 @@
      * is attached.
      * <b>The time-base is {@link android.os.SystemClock#uptimeMillis}.</b>
      * Time spent in deep sleep will add an additional delay to execution.
-     *  
+     *
      * @param r The Runnable that will be executed.
      * @param delayMillis The delay (in milliseconds) until the Runnable
      *        will be executed.
-     *        
-     * @return Returns true if the Runnable was successfully placed in to the 
+     *
+     * @return Returns true if the Runnable was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.  Note that a
      *         result of true does not mean the Runnable will be processed --
@@ -507,7 +507,7 @@
     public final boolean postDelayed(@NonNull Runnable r, long delayMillis) {
         return sendMessageDelayed(getPostMessage(r), delayMillis);
     }
-    
+
     /** @hide */
     public final boolean postDelayed(Runnable r, int what, long delayMillis) {
         return sendMessageDelayed(getPostMessage(r).setWhat(what), delayMillis);
@@ -547,10 +547,10 @@
      * <b>This method is only for use in very special circumstances -- it
      * can easily starve the message queue, cause ordering problems, or have
      * other unexpected side-effects.</b>
-     *  
+     *
      * @param r The Runnable that will be executed.
-     * 
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -635,8 +635,8 @@
      * Pushes a message onto the end of the message queue after all pending messages
      * before the current time. It will be received in {@link #handleMessage},
      * in the thread attached to this handler.
-     *  
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -646,8 +646,8 @@
 
     /**
      * Sends a Message containing only the what value.
-     *  
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -659,9 +659,9 @@
     /**
      * Sends a Message containing only the what value, to be delivered
      * after the specified amount of time elapses.
-     * @see #sendMessageDelayed(android.os.Message, long) 
-     * 
-     * @return Returns true if the message was successfully placed in to the 
+     * @see #sendMessageDelayed(android.os.Message, long)
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -672,11 +672,11 @@
     }
 
     /**
-     * Sends a Message containing only the what value, to be delivered 
+     * Sends a Message containing only the what value, to be delivered
      * at a specific time.
      * @see #sendMessageAtTime(android.os.Message, long)
-     *  
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -691,8 +691,8 @@
      * Enqueue a message into the message queue after all pending messages
      * before (current time + delayMillis). You will receive it in
      * {@link #handleMessage}, in the thread attached to this handler.
-     *  
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.  Note that a
      *         result of true does not mean the message will be processed -- if
@@ -713,12 +713,12 @@
      * Time spent in deep sleep will add an additional delay to execution.
      * You will receive it in {@link #handleMessage}, in the thread attached
      * to this handler.
-     * 
+     *
      * @param uptimeMillis The absolute time at which the message should be
      *         delivered, using the
      *         {@link android.os.SystemClock#uptimeMillis} time-base.
-     *         
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.  Note that a
      *         result of true does not mean the message will be processed -- if
@@ -743,8 +743,8 @@
      * <b>This method is only for use in very special circumstances -- it
      * can easily starve the message queue, cause ordering problems, or have
      * other unexpected side-effects.</b>
-     *  
-     * @return Returns true if the message was successfully placed in to the 
+     *
+     * @return Returns true if the message was successfully placed in to the
      *         message queue.  Returns false on failure, usually because the
      *         looper processing the message queue is exiting.
      */
@@ -798,6 +798,12 @@
     /**
      * Remove any pending posts of messages with code 'what' that are in the
      * message queue.
+     *
+     * Note that `Message#what` is 0 unless otherwise set.
+     * When calling `postMessage(Runnable)` or `postAtTime(Runnable, long)`,
+     * the `Runnable` is internally wrapped with a `Message` whose `what` is 0.
+     * Calling `removeMessages(0)` will remove all messages without a `what`,
+     * including posted `Runnable`s.
      */
     public final void removeMessages(int what) {
         mQueue.removeMessages(this, what, null);
@@ -889,7 +895,7 @@
     }
 
     // if we can get rid of this method, the handler need not remember its loop
-    // we could instead export a getMessageQueue() method... 
+    // we could instead export a getMessageQueue() method...
     @NonNull
     public final Looper getLooper() {
         return mLooper;
diff --git a/core/java/android/os/Message.java b/core/java/android/os/Message.java
index a1db9be..702fdc2 100644
--- a/core/java/android/os/Message.java
+++ b/core/java/android/os/Message.java
@@ -41,6 +41,9 @@
      * what this message is about. Each {@link Handler} has its own name-space
      * for message codes, so you do not need to worry about yours conflicting
      * with other handlers.
+     *
+     * If not specified, this value is 0.
+     * Use values other than 0 to indicate custom message codes.
      */
     public int what;
 
diff --git a/core/java/android/os/SharedMemory.java b/core/java/android/os/SharedMemory.java
index c801fabf..46ae9d8 100644
--- a/core/java/android/os/SharedMemory.java
+++ b/core/java/android/os/SharedMemory.java
@@ -379,6 +379,14 @@
         private int mReferenceCount;
 
         private MemoryRegistration(int size) {
+            // Round up to the nearest page size
+            final int PAGE_SIZE = OsConstants._SC_PAGE_SIZE;
+            if (PAGE_SIZE > 0) {
+                final int remainder = size % PAGE_SIZE;
+                if (remainder != 0) {
+                    size += PAGE_SIZE - remainder;
+                }
+            }
             mSize = size;
             mReferenceCount = 1;
             VMRuntime.getRuntime().registerNativeAllocation(mSize);
diff --git a/core/java/android/os/SystemClock.java b/core/java/android/os/SystemClock.java
index 23bd30a..0ed1ab6 100644
--- a/core/java/android/os/SystemClock.java
+++ b/core/java/android/os/SystemClock.java
@@ -314,6 +314,15 @@
     }
 
     /**
+     * @see #currentNetworkTimeMillis(ITimeDetectorService)
+     * @hide
+     */
+    public static long currentNetworkTimeMillis() {
+        return currentNetworkTimeMillis(ITimeDetectorService.Stub
+                .asInterface(ServiceManager.getService(Context.TIME_DETECTOR_SERVICE)));
+    }
+
+    /**
      * Returns milliseconds since January 1, 1970 00:00:00.0 UTC, synchronized
      * using a remote network source outside the device.
      * <p>
@@ -331,14 +340,14 @@
      * at any time. Due to network delays, variations between servers, or local
      * (client side) clock drift, the accuracy of the returned times cannot be
      * guaranteed. In extreme cases, consecutive calls to {@link
-     * #currentNetworkTimeMillis()} could return times that are out of order.
+     * #currentNetworkTimeMillis(ITimeDetectorService)} could return times that
+     * are out of order.
      *
      * @throws DateTimeException when no network time can be provided.
      * @hide
      */
-    public static long currentNetworkTimeMillis() {
-        ITimeDetectorService timeDetectorService = ITimeDetectorService.Stub
-                .asInterface(ServiceManager.getService(Context.TIME_DETECTOR_SERVICE));
+    public static long currentNetworkTimeMillis(
+            ITimeDetectorService timeDetectorService) {
         if (timeDetectorService != null) {
             UnixEpochTime time;
             try {
@@ -380,16 +389,21 @@
      * at any time. Due to network delays, variations between servers, or local
      * (client side) clock drift, the accuracy of the returned times cannot be
      * guaranteed. In extreme cases, consecutive calls to {@link
-     * Clock#millis()} on the returned {@link Clock}could return times that are
+     * Clock#millis()} on the returned {@link Clock} could return times that are
      * out of order.
      *
      * @throws DateTimeException when no network time can be provided.
      */
     public static @NonNull Clock currentNetworkTimeClock() {
         return new SimpleClock(ZoneOffset.UTC) {
+            private ITimeDetectorService mSvc;
             @Override
             public long millis() {
-                return SystemClock.currentNetworkTimeMillis();
+                if (mSvc == null) {
+                    mSvc = ITimeDetectorService.Stub
+                            .asInterface(ServiceManager.getService(Context.TIME_DETECTOR_SERVICE));
+                }
+                return SystemClock.currentNetworkTimeMillis(mSvc);
             }
         };
     }
diff --git a/core/java/android/os/TransactionTooLargeException.java b/core/java/android/os/TransactionTooLargeException.java
index 79892e0..6b7cb33 100644
--- a/core/java/android/os/TransactionTooLargeException.java
+++ b/core/java/android/os/TransactionTooLargeException.java
@@ -47,7 +47,7 @@
  * If possible, try to break up big requests into smaller pieces.
  * </p><p>
  * If you are implementing a service, it may help to impose size or complexity
- * contraints on the queries that clients can perform.  For example, if the result set
+ * constraints on the queries that clients can perform.  For example, if the result set
  * could become large, then don't allow the client to request more than a few records
  * at a time.  Alternately, instead of returning all of the available data all at once,
  * return the essential information first and make the client ask for additional information
diff --git a/core/java/android/os/UserManager.java b/core/java/android/os/UserManager.java
index 536eca6..f1ec0e4e 100644
--- a/core/java/android/os/UserManager.java
+++ b/core/java/android/os/UserManager.java
@@ -1978,7 +1978,6 @@
      * @see DevicePolicyManager#clearUserRestriction(ComponentName, String)
      * @see #getUserRestrictions()
      */
-    @FlaggedApi(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED)
     public static final String DISALLOW_SIM_GLOBALLY =
             "no_sim_globally";
 
@@ -1999,7 +1998,6 @@
      * @see DevicePolicyManager#clearUserRestriction(ComponentName, String)
      * @see #getUserRestrictions()
      */
-    @FlaggedApi(android.app.admin.flags.Flags.FLAG_ASSIST_CONTENT_USER_RESTRICTION_ENABLED)
     public static final String DISALLOW_ASSIST_CONTENT = "no_assist_content";
 
     /**
diff --git a/core/java/android/permission/TEST_MAPPING b/core/java/android/permission/TEST_MAPPING
index a15d9bc..b317b80 100644
--- a/core/java/android/permission/TEST_MAPPING
+++ b/core/java/android/permission/TEST_MAPPING
@@ -1,15 +1,7 @@
 {
     "presubmit": [
         {
-            "name": "CtsPermissionTestCases",
-            "options": [
-                {
-                    "include-filter": "android.permission.cts.PermissionControllerTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.RuntimePermissionPresentationInfoTest"
-                }
-            ]
+            "name": "CtsPermissionTestCases_Platform"
         }
     ],
     "postsubmit": [
@@ -36,4 +28,4 @@
             ]
         }
     ]
-}
\ No newline at end of file
+}
diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index 85d2325..98904fe 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -1974,7 +1974,7 @@
     /**
      * Activity Action: Show Notification Policy access settings.
      * <p>
-     * Users can grant and deny access to Notification Policy (DND / Priority Modes) configuration
+     * Users can grant and deny access to Notification Policy (DND / Modes) configuration
      * from here. Managed profiles cannot grant Notification Policy access.
      * See {@link android.app.NotificationManager#isNotificationPolicyAccessGranted()} for more
      * details.
@@ -19666,6 +19666,8 @@
             public static final int PAIRED_DEVICE_OS_TYPE_ANDROID = 1;
             /** @hide */
             public static final int PAIRED_DEVICE_OS_TYPE_IOS = 2;
+            /** @hide */
+            public static final int PAIRED_DEVICE_OS_TYPE_NONE = 3;
 
             /**
              * The bluetooth settings selected BLE role for the companion.
diff --git a/core/java/android/service/chooser/flags.aconfig b/core/java/android/service/chooser/flags.aconfig
index 1f1a351..2b75493 100644
--- a/core/java/android/service/chooser/flags.aconfig
+++ b/core/java/android/service/chooser/flags.aconfig
@@ -2,6 +2,22 @@
 container: "system"
 
 flag {
+  name: "chooser_album_text"
+  is_exported: true
+  namespace: "intentresolver"
+  description: "Flag controlling the album text subtype hint for sharesheet"
+  bug: "323380224"
+}
+
+flag {
+  name: "enable_sharesheet_metadata_extra"
+  is_exported: true
+  namespace: "intentresolver"
+  description: "This flag enables sharesheet metadata to be displayed to users."
+  bug: "318942069"
+}
+
+flag {
   name: "chooser_payload_toggling"
   is_exported: true
   namespace: "intentresolver"
diff --git a/core/java/android/service/dreams/DreamOverlayService.java b/core/java/android/service/dreams/DreamOverlayService.java
index 013ec5f..711c414 100644
--- a/core/java/android/service/dreams/DreamOverlayService.java
+++ b/core/java/android/service/dreams/DreamOverlayService.java
@@ -51,11 +51,14 @@
      */
     private Executor mExecutor;
 
+    private Boolean mCurrentRedirectToWake;
+
     // An {@link IDreamOverlayClient} implementation that identifies itself when forwarding
     // requests to the {@link DreamOverlayService}
     private static class OverlayClient extends IDreamOverlayClient.Stub {
         private final WeakReference<DreamOverlayService> mService;
         private boolean mShowComplications;
+        private boolean mIsPreview;
         private ComponentName mDreamComponent;
         IDreamOverlayCallback mDreamOverlayCallback;
 
@@ -73,9 +76,11 @@
 
         @Override
         public void startDream(WindowManager.LayoutParams params, IDreamOverlayCallback callback,
-                String dreamComponent, boolean shouldShowComplications) throws RemoteException {
+                String dreamComponent, boolean isPreview, boolean shouldShowComplications)
+                throws RemoteException {
             mDreamComponent = ComponentName.unflattenFromString(dreamComponent);
             mShowComplications = shouldShowComplications;
+            mIsPreview = isPreview;
             mDreamOverlayCallback = callback;
             applyToDream(dreamOverlayService -> dreamOverlayService.startDream(this, params));
         }
@@ -122,6 +127,10 @@
             return mShowComplications;
         }
 
+        private boolean isDreamInPreviewMode() {
+            return mIsPreview;
+        }
+
         private ComponentName getComponent() {
             return mDreamComponent;
         }
@@ -132,6 +141,10 @@
         mExecutor.execute(() -> {
             endDreamInternal(mCurrentClient);
             mCurrentClient = client;
+            if (Flags.dreamWakeRedirect() && mCurrentRedirectToWake != null) {
+                mCurrentClient.redirectWake(mCurrentRedirectToWake);
+            }
+
             onStartDream(params);
         });
     }
@@ -282,8 +295,10 @@
             return;
         }
 
+        mCurrentRedirectToWake = redirect;
+
         if (mCurrentClient == null) {
-            throw new IllegalStateException("redirected wake with no dream present");
+            return;
         }
 
         mCurrentClient.redirectWake(redirect);
@@ -295,7 +310,6 @@
      *
      * @hide
      */
-    @FlaggedApi(Flags.FLAG_DREAM_WAKE_REDIRECT)
     public void onWakeRequested() {
     }
 
@@ -312,6 +326,19 @@
     }
 
     /**
+     * Returns whether dream is in preview mode.
+     */
+    @FlaggedApi(Flags.FLAG_PUBLISH_PREVIEW_STATE_TO_OVERLAY)
+    public final boolean isDreamInPreviewMode() {
+        if (mCurrentClient == null) {
+            throw new IllegalStateException(
+                    "requested if preview when no dream active");
+        }
+
+        return mCurrentClient.isDreamInPreviewMode();
+    }
+
+    /**
      * Returns the active dream component.
      * @hide
      */
diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index c3585e3..ce31e1e 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -1701,6 +1701,7 @@
                                 try {
                                     overlay.startDream(mWindow.getAttributes(), mOverlayCallback,
                                             mDreamComponent.flattenToString(),
+                                            mPreviewMode,
                                             mShouldShowComplications);
                                 } catch (RemoteException e) {
                                     Log.e(mTag, "could not send window attributes:" + e);
diff --git a/core/java/android/service/dreams/IDreamOverlayClient.aidl b/core/java/android/service/dreams/IDreamOverlayClient.aidl
index 0eb15a0..e9df402 100644
--- a/core/java/android/service/dreams/IDreamOverlayClient.aidl
+++ b/core/java/android/service/dreams/IDreamOverlayClient.aidl
@@ -31,11 +31,12 @@
     * @param callback The {@link IDreamOverlayCallback} for requesting actions such as exiting the
     *                dream.
     * @param dreamComponent The component name of the dream service requesting overlay.
+    * @param isPreview Whether the dream is in preview mode.
     * @param shouldShowComplications Whether the dream overlay should show complications, e.g. clock
     *                and weather.
     */
     void startDream(in LayoutParams params, in IDreamOverlayCallback callback,
-        in String dreamComponent, in boolean shouldShowComplications);
+        in String dreamComponent, in boolean isPreview, in boolean shouldShowComplications);
 
     /** Called when the dream is waking, to do any exit animations */
     void wakeUp();
diff --git a/core/java/android/service/dreams/flags.aconfig b/core/java/android/service/dreams/flags.aconfig
index 83e0adf..72f2de8 100644
--- a/core/java/android/service/dreams/flags.aconfig
+++ b/core/java/android/service/dreams/flags.aconfig
@@ -57,3 +57,13 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "publish_preview_state_to_overlay"
+    namespace: "systemui"
+    description: "send preview information from dream to overlay"
+    bug: "333734282"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/core/java/android/service/notification/ZenModeConfig.java b/core/java/android/service/notification/ZenModeConfig.java
index 918e591..3d8d933 100644
--- a/core/java/android/service/notification/ZenModeConfig.java
+++ b/core/java/android/service/notification/ZenModeConfig.java
@@ -1044,23 +1044,19 @@
                     rt.suppressedVisualEffects = safeInt(parser, DISALLOW_ATT_VISUAL_EFFECTS,
                             DEFAULT_SUPPRESSED_VISUAL_EFFECTS);
                 } else if (MANUAL_TAG.equals(tag)) {
-                    ZenRule manualRule = readRuleXml(parser);
-                    if (manualRule != null) {
-                        rt.manualRule = manualRule;
-
-                        // Manual rule may be present prior to modes_ui if it were on, but in that
-                        // case it would not have a set policy, so make note of the need to set
-                        // it up later.
-                        readManualRule = true;
-                        if (rt.manualRule.zenPolicy == null) {
-                            readManualRuleWithoutPolicy = true;
-                        }
+                    rt.manualRule = readRuleXml(parser);
+                    // Manual rule may be present prior to modes_ui if it were on, but in that
+                    // case it would not have a set policy, so make note of the need to set
+                    // it up later.
+                    readManualRule = true;
+                    if (rt.manualRule.zenPolicy == null) {
+                        readManualRuleWithoutPolicy = true;
                     }
                 } else if (AUTOMATIC_TAG.equals(tag)
                         || (Flags.modesApi() && AUTOMATIC_DELETED_TAG.equals(tag))) {
                     final String id = parser.getAttributeValue(null, RULE_ATT_ID);
-                    final ZenRule automaticRule = readRuleXml(parser);
-                    if (id != null && automaticRule != null) {
+                    if (id != null) {
+                        final ZenRule automaticRule = readRuleXml(parser);
                         automaticRule.id = id;
                         if (Flags.modesApi() && AUTOMATIC_DELETED_TAG.equals(tag)) {
                             String deletedRuleKey = deletedRuleKey(automaticRule);
@@ -1177,16 +1173,13 @@
         out.endTag(null, ZEN_TAG);
     }
 
+    @NonNull
     public static ZenRule readRuleXml(TypedXmlPullParser parser) {
         final ZenRule rt = new ZenRule();
         rt.enabled = safeBoolean(parser, RULE_ATT_ENABLED, true);
         rt.name = parser.getAttributeValue(null, RULE_ATT_NAME);
         final String zen = parser.getAttributeValue(null, RULE_ATT_ZEN);
-        rt.zenMode = tryParseZenMode(zen, -1);
-        if (rt.zenMode == -1) {
-            Slog.w(TAG, "Bad zen mode in rule xml:" + zen);
-            return null;
-        }
+        rt.zenMode = tryParseZenMode(zen, ZEN_MODE_IMPORTANT_INTERRUPTIONS);
         rt.conditionId = safeUri(parser, RULE_ATT_CONDITION_ID);
         rt.component = safeComponentName(parser, RULE_ATT_COMPONENT);
         rt.configurationActivity = safeComponentName(parser, RULE_ATT_CONFIG_ACTIVITY);
@@ -2939,7 +2932,7 @@
             }
         }
 
-        // TODO: b/333527800 - Rename to isActive()
+        // TODO: b/363193376 - Rename to isActive()
         public boolean isAutomaticActive() {
             if (Flags.modesApi() && Flags.modesUi()) {
                 if (!enabled || getPkg() == null) {
diff --git a/core/java/android/service/notification/ZenPolicy.java b/core/java/android/service/notification/ZenPolicy.java
index 2669391..be0d7b3 100644
--- a/core/java/android/service/notification/ZenPolicy.java
+++ b/core/java/android/service/notification/ZenPolicy.java
@@ -393,6 +393,46 @@
     }
 
     /**
+     * Base Zen Policy used when {@link android.app.NotificationManager#setInterruptionFilter} is
+     * called with {@link android.app.NotificationManager#INTERRUPTION_FILTER_ALARMS} or an
+     * {@link android.app.AutomaticZenRule} specifies said filter.
+     *
+     * <p>Note that <em>visual effects</em> for filtered notifications are unset in this base
+     * policy, so should be merged on top of the default policy's visual effects (see
+     * {@link #overwrittenWith(ZenPolicy)}).
+     *
+     * @hide
+     */
+    @FlaggedApi(Flags.FLAG_MODES_API)
+    public static ZenPolicy getBasePolicyInterruptionFilterAlarms() {
+        return new ZenPolicy.Builder()
+                .disallowAllSounds()
+                .allowAlarms(true)
+                .allowMedia(true)
+                .allowPriorityChannels(false)
+                .build();
+    }
+
+    /**
+     * Base Zen Policy used when {@link android.app.NotificationManager#setInterruptionFilter} is
+     * called with {@link android.app.NotificationManager#INTERRUPTION_FILTER_NONE} or an
+     * {@link android.app.AutomaticZenRule} specifies said filter.
+     *
+     * <p>Note that <em>visual effects</em> for filtered notifications are unset in this base
+     * policy, so it should be merged on top of the device default policy's visual effects (see
+     * {@link #overwrittenWith(ZenPolicy)}).
+     *
+     * @hide
+     */
+    @FlaggedApi(Flags.FLAG_MODES_API)
+    public static ZenPolicy getBasePolicyInterruptionFilterNone() {
+        return new ZenPolicy.Builder()
+                .disallowAllSounds()
+                .allowPriorityChannels(false)
+                .build();
+    }
+
+    /**
      * Conversation type that can bypass DND.
      * @return {@link #CONVERSATION_SENDERS_UNSET}, {@link #CONVERSATION_SENDERS_ANYONE},
      * {@link #CONVERSATION_SENDERS_IMPORTANT}, {@link #CONVERSATION_SENDERS_NONE}.
diff --git a/core/java/android/util/Half.java b/core/java/android/util/Half.java
index fe536a6..22583ac 100644
--- a/core/java/android/util/Half.java
+++ b/core/java/android/util/Half.java
@@ -19,6 +19,7 @@
 import android.annotation.HalfFloat;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.ravenwood.annotation.RavenwoodKeepWholeClass;
 
 import libcore.util.FP16;
 
@@ -92,6 +93,7 @@
  * <p>This table shows that numbers higher than 1024 lose all fractional precision.</p>
  */
 @SuppressWarnings("SimplifiableIfStatement")
+@RavenwoodKeepWholeClass
 public final class Half extends Number implements Comparable<Half> {
     /**
      * The number of bits used to represent a half-precision float value.
diff --git a/core/java/android/view/Display.java b/core/java/android/view/Display.java
index 1f7ed8b..82c52a6 100644
--- a/core/java/android/view/Display.java
+++ b/core/java/android/view/Display.java
@@ -314,6 +314,8 @@
      * @hide
      * @see #getFlags()
      */
+    @SuppressLint("UnflaggedApi") // Promotion to TestApi
+    @TestApi
     public static final int FLAG_ALWAYS_UNLOCKED = 1 << 9;
 
     /**
@@ -323,6 +325,8 @@
      * @hide
      * @see #getFlags()
      */
+    @SuppressLint("UnflaggedApi") // Promotion to TestApi
+    @TestApi
     public static final int FLAG_TOUCH_FEEDBACK_DISABLED = 1 << 10;
 
     /**
@@ -336,6 +340,8 @@
      * @see #FLAG_TRUSTED
      * @hide
      */
+    @SuppressLint("UnflaggedApi") // Promotion to TestApi
+    @TestApi
     public static final int FLAG_OWN_FOCUS = 1 << 11;
 
     /**
@@ -642,6 +648,8 @@
      * @hide
      */
     // TODO (b/114338689): Remove the flag and use WindowManager#REMOVE_CONTENT_MODE_DESTROY
+    @SuppressLint("UnflaggedApi") // Promotion to TestApi
+    @TestApi
     public static final int REMOVE_MODE_DESTROY_CONTENT = 1;
 
     /** @hide */
diff --git a/core/java/android/view/ImeBackAnimationController.java b/core/java/android/view/ImeBackAnimationController.java
index c7e93c1..b801465 100644
--- a/core/java/android/view/ImeBackAnimationController.java
+++ b/core/java/android/view/ImeBackAnimationController.java
@@ -149,15 +149,17 @@
 
     private void setPreCommitProgress(float progress) {
         if (isHideAnimationInProgress()) return;
+        setInterpolatedProgress(BACK_GESTURE.getInterpolation(progress) * PEEK_FRACTION);
+    }
+
+    private void setInterpolatedProgress(float progress) {
         if (mWindowInsetsAnimationController != null) {
             float hiddenY = mWindowInsetsAnimationController.getHiddenStateInsets().bottom;
             float shownY = mWindowInsetsAnimationController.getShownStateInsets().bottom;
             float imeHeight = shownY - hiddenY;
-            float interpolatedProgress = BACK_GESTURE.getInterpolation(progress);
-            int newY = (int) (imeHeight - interpolatedProgress * (imeHeight * PEEK_FRACTION));
+            int newY = (int) (imeHeight - progress * imeHeight);
             if (mStartRootScrollY != 0) {
-                mViewRoot.setScrollY(
-                        (int) (mStartRootScrollY * (1 - interpolatedProgress * PEEK_FRACTION)));
+                mViewRoot.setScrollY((int) (mStartRootScrollY * (1 - progress)));
             }
             mWindowInsetsAnimationController.setInsetsAndAlpha(Insets.of(0, 0, 0, newY), 1f,
                     progress);
@@ -171,21 +173,14 @@
             return;
         }
         mTriggerBack = triggerBack;
-        int currentBottomInset = mWindowInsetsAnimationController.getCurrentInsets().bottom;
-        int targetBottomInset;
-        if (triggerBack) {
-            targetBottomInset = mWindowInsetsAnimationController.getHiddenStateInsets().bottom;
-        } else {
-            targetBottomInset = mWindowInsetsAnimationController.getShownStateInsets().bottom;
-        }
-        mPostCommitAnimator = ValueAnimator.ofFloat(currentBottomInset, targetBottomInset);
+        float targetProgress = triggerBack ? 1f : 0f;
+        mPostCommitAnimator = ValueAnimator.ofFloat(
+                BACK_GESTURE.getInterpolation(mLastProgress) * PEEK_FRACTION, targetProgress);
         mPostCommitAnimator.setInterpolator(
                 triggerBack ? STANDARD_ACCELERATE : EMPHASIZED_DECELERATE);
         mPostCommitAnimator.addUpdateListener(animation -> {
-            int bottomInset = (int) ((float) animation.getAnimatedValue());
             if (mWindowInsetsAnimationController != null) {
-                mWindowInsetsAnimationController.setInsetsAndAlpha(Insets.of(0, 0, 0, bottomInset),
-                        1f, animation.getAnimatedFraction());
+                setInterpolatedProgress((float) animation.getAnimatedValue());
             } else {
                 reset();
             }
@@ -213,14 +208,8 @@
             notifyHideIme();
             // requesting IME as invisible during post-commit
             mInsetsController.setRequestedVisibleTypes(0, ime());
-            // Changes the animation state. This also notifies RootView of changed insets, which
-            // causes it to reset its scrollY to 0f (animated) if it was panned
             mInsetsController.onAnimationStateChanged(ime(), /*running*/ true);
         }
-        if (mStartRootScrollY != 0 && !triggerBack) {
-            // This causes RootView to update its scroll back to the panned position
-            mInsetsController.getHost().notifyInsetsChanged();
-        }
     }
 
     private void notifyHideIme() {
@@ -282,6 +271,10 @@
         return mPostCommitAnimator != null && mTriggerBack;
     }
 
+    boolean isAnimationInProgress() {
+        return mIsPreCommitAnimationInProgress || mWindowInsetsAnimationController != null;
+    }
+
     /**
      * Dump information about this ImeBackAnimationController
      *
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 6343313..e90b1c0 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -70,7 +70,14 @@
                         "ImeInsetsSourceConsumer#onAnimationFinished",
                         mController.getHost().getInputMethodManager(), null /* icProto */);
             }
-            boolean insetsChanged = super.onAnimationStateChanged(running);
+            boolean insetsChanged = false;
+            if (Flags.predictiveBackIme() && !running && isShowRequested()
+                    && mAnimationState == ANIMATION_STATE_HIDE) {
+                // A user controlled hide animation may have ended in the shown state (e.g.
+                // cancelled predictive back animation) -> Insets need to be reset to shown.
+                insetsChanged |= applyLocalVisibilityOverride();
+            }
+            insetsChanged |= super.onAnimationStateChanged(running);
             if (running && !isShowRequested()
                     && mController.isPredictiveBackImeHideAnimInProgress()) {
                 // IME predictive back animation switched from pre-commit to post-commit.
diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index b1df51f..8fdf91a 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -488,7 +488,7 @@
         @Override
         public Interpolator getInsetsInterpolator(boolean hasZeroInsetsIme) {
             if ((mRequestedTypes & ime()) != 0) {
-                if (mHasAnimationCallbacks && hasZeroInsetsIme) {
+                if (mHasAnimationCallbacks && !hasZeroInsetsIme) {
                     return SYNC_IME_INTERPOLATOR;
                 } else if (mShow) {
                     return LINEAR_OUT_SLOW_IN_INTERPOLATOR;
@@ -536,7 +536,7 @@
         @Override
         public long getDurationMs(boolean hasZeroInsetsIme) {
             if ((mRequestedTypes & ime()) != 0) {
-                if (mHasAnimationCallbacks && hasZeroInsetsIme) {
+                if (mHasAnimationCallbacks && !hasZeroInsetsIme) {
                     return ANIMATION_DURATION_SYNC_IME_MS;
                 } else {
                     return ANIMATION_DURATION_UNSYNC_IME_MS;
@@ -1197,7 +1197,8 @@
                 pendingRequest.listener, null /* frame */, true /* fromIme */,
                 pendingRequest.mInsetsAnimationSpec,
                 pendingRequest.animationType, pendingRequest.layoutInsetsDuringAnimation,
-                pendingRequest.useInsetsAnimationThread, statsToken);
+                pendingRequest.useInsetsAnimationThread, statsToken,
+                false /* fromPredictiveBack */);
     }
 
     @Override
@@ -1307,7 +1308,10 @@
             WindowInsetsAnimationControlListener listener,
             boolean fromIme, long durationMs, @Nullable Interpolator interpolator,
             @AnimationType int animationType, boolean fromPredictiveBack) {
-        if ((mState.calculateUncontrollableInsetsFromFrame(mFrame) & types) != 0) {
+        if ((mState.calculateUncontrollableInsetsFromFrame(mFrame) & types) != 0
+                || (fromPredictiveBack && ((mRequestedVisibleTypes & ime()) == 0))) {
+            // abort if insets are uncontrollable or if control request is from predictive back but
+            // there is already a hide anim in progress
             listener.onCancelled(null);
             return;
         }
@@ -1330,7 +1334,7 @@
         // TODO(b/342111149): Create statsToken here once ImeTracker#onStart becomes async.
         controlAnimationUnchecked(types, cancellationSignal, listener, mFrame, fromIme, spec,
                 animationType, getLayoutInsetsDuringAnimationMode(types, fromPredictiveBack),
-                false /* useInsetsAnimationThread */, null);
+                false /* useInsetsAnimationThread */, null, fromPredictiveBack);
     }
 
     private void controlAnimationUnchecked(@InsetsType int types,
@@ -1338,7 +1342,8 @@
             WindowInsetsAnimationControlListener listener, @Nullable Rect frame, boolean fromIme,
             InsetsAnimationSpec insetsAnimationSpec, @AnimationType int animationType,
             @LayoutInsetsDuringAnimation int layoutInsetsDuringAnimation,
-            boolean useInsetsAnimationThread, @Nullable ImeTracker.Token statsToken) {
+            boolean useInsetsAnimationThread, @Nullable ImeTracker.Token statsToken,
+            boolean fromPredictiveBack) {
         final boolean visible = layoutInsetsDuringAnimation == LAYOUT_INSETS_DURING_ANIMATION_SHOWN;
 
         // Basically, we accept the requested visibilities from the upstream callers...
@@ -1348,7 +1353,7 @@
         // rejecting showing IME.
         controlAnimationUncheckedInner(types, cancellationSignal, listener, frame, fromIme,
                 insetsAnimationSpec, animationType, layoutInsetsDuringAnimation,
-                useInsetsAnimationThread, statsToken);
+                useInsetsAnimationThread, statsToken, fromPredictiveBack);
 
         // We are finishing setting the requested visible types. Report them to the server
         // and/or the app.
@@ -1360,7 +1365,8 @@
             WindowInsetsAnimationControlListener listener, @Nullable Rect frame, boolean fromIme,
             InsetsAnimationSpec insetsAnimationSpec, @AnimationType int animationType,
             @LayoutInsetsDuringAnimation int layoutInsetsDuringAnimation,
-            boolean useInsetsAnimationThread, @Nullable ImeTracker.Token statsToken) {
+            boolean useInsetsAnimationThread, @Nullable ImeTracker.Token statsToken,
+            boolean fromPredictiveBack) {
         if ((types & mTypesBeingCancelled) != 0) {
             final boolean monitoredAnimation =
                     animationType == ANIMATION_TYPE_SHOW || animationType == ANIMATION_TYPE_HIDE;
@@ -1446,7 +1452,7 @@
             }
         } else {
             Pair<Integer, Boolean> typesReadyPair = collectSourceControls(
-                    fromIme, types, controls, animationType, statsToken);
+                    fromIme, types, controls, animationType, statsToken, fromPredictiveBack);
             typesReady = typesReadyPair.first;
             boolean imeReady = typesReadyPair.second;
             if (DEBUG) {
@@ -1582,7 +1588,7 @@
      */
     private Pair<Integer, Boolean> collectSourceControls(boolean fromIme, @InsetsType int types,
             SparseArray<InsetsSourceControl> controls, @AnimationType int animationType,
-            @Nullable ImeTracker.Token statsToken) {
+            @Nullable ImeTracker.Token statsToken, boolean fromPredictiveBack) {
         ImeTracker.forLogging().onProgress(statsToken,
                 ImeTracker.PHASE_CLIENT_COLLECT_SOURCE_CONTROLS);
 
@@ -1594,7 +1600,8 @@
                 continue;
             }
             boolean show = animationType == ANIMATION_TYPE_SHOW
-                    || animationType == ANIMATION_TYPE_USER;
+                    || (animationType == ANIMATION_TYPE_USER
+                            && (!fromPredictiveBack || !mHost.hasAnimationCallbacks()));
             boolean canRun = true;
             if (show) {
                 // Show request
@@ -1617,7 +1624,8 @@
                         break;
                 }
             } else {
-                consumer.requestHide(fromIme, statsToken);
+                consumer.requestHide(fromIme
+                        || (fromPredictiveBack && mHost.hasAnimationCallbacks()), statsToken);
             }
             if (!canRun) {
                 if (WARN) Log.w(TAG, String.format(
@@ -1672,9 +1680,10 @@
 
     private @LayoutInsetsDuringAnimation int getLayoutInsetsDuringAnimationMode(
             @InsetsType int types, boolean fromPredictiveBack) {
-        if (fromPredictiveBack) {
-            // When insets are animated by predictive back, we want insets to be shown to prevent a
-            // jump cut from shown to hidden at the start of the predictive back animation
+        if (fromPredictiveBack && !mHost.hasAnimationCallbacks()) {
+            // When insets are animated by predictive back and the app does not have an animation
+            // callback, we want insets to be shown to prevent a jump cut from shown to hidden at
+            // the start of the predictive back animation
             return LAYOUT_INSETS_DURING_ANIMATION_SHOWN;
         }
         // Generally, we want to layout the opposite of the current state. This is to make animation
@@ -2021,7 +2030,8 @@
                 listener /* insetsAnimationSpec */,
                 show ? ANIMATION_TYPE_SHOW : ANIMATION_TYPE_HIDE,
                 show ? LAYOUT_INSETS_DURING_ANIMATION_SHOWN : LAYOUT_INSETS_DURING_ANIMATION_HIDDEN,
-                !hasAnimationCallbacks /* useInsetsAnimationThread */, statsToken);
+                !hasAnimationCallbacks /* useInsetsAnimationThread */, statsToken,
+                false /* fromPredictiveBack */);
     }
 
     /**
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 477e35b..391d757 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -310,21 +310,22 @@
         }
         final boolean requestedVisible = (mController.getRequestedVisibleTypes() & mType) != 0;
 
+        // If we don't have control or the leash (in case of the IME), we enforce the
+        // visibility to be hidden, as otherwise we would let the app know too early.
+        if (mSourceControl == null) {
+            if (DEBUG) {
+                Log.d(TAG, TextUtils.formatSimple(
+                        "applyLocalVisibilityOverride: No control in %s for type %s, "
+                                + "requestedVisible=%s",
+                        mController.getHost().getRootViewTitle(),
+                        WindowInsets.Type.toString(mType), requestedVisible));
+            }
+            return false;
+        }
         if (Flags.refactorInsetsController()) {
-            // If we don't have control or the leash (in case of the IME), we enforce the
-            // visibility to be hidden, as otherwise we would let the app know too early.
-            if (mSourceControl == null) {
-                if (DEBUG) {
-                    Log.d(TAG, TextUtils.formatSimple(
-                            "applyLocalVisibilityOverride: No control in %s for type %s, "
-                                    + "requestedVisible=%s",
-                            mController.getHost().getRootViewTitle(),
-                            WindowInsets.Type.toString(mType), requestedVisible));
-                }
-                return false;
-                // TODO(b/323136120) add a flag to the control, to define whether a leash is needed
-            } else if (mId != InsetsSource.ID_IME_CAPTION_BAR
-                    && mSourceControl.getLeash() == null) {
+            // TODO(b/323136120) add a flag to the control, to define whether a leash is
+            //  needed and make it generic for all types
+            if (mId == InsetsSource.ID_IME && mSourceControl.getLeash() == null) {
                 if (DEBUG) {
                     Log.d(TAG, TextUtils.formatSimple(
                             "applyLocalVisibilityOverride: Set the source visibility to false, as"
@@ -338,16 +339,6 @@
                 // changed state
                 return wasVisible;
             }
-        } else {
-            // If we don't have control, we are not able to change the visibility.
-            if (mSourceControl == null) {
-                if (DEBUG) {
-                    Log.d(TAG, "applyLocalVisibilityOverride: No control in "
-                            + mController.getHost().getRootViewTitle()
-                            + " requestedVisible=" + requestedVisible);
-                }
-                return false;
-            }
         }
         if (source.isVisible() == requestedVisible) {
             return false;
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index e81f32e..7b4ea41 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -141,7 +141,6 @@
 import android.os.RemoteException;
 import android.os.SystemClock;
 import android.os.Trace;
-import android.os.Vibrator;
 import android.service.credentials.CredentialProviderService;
 import android.sysprop.DisplayProperties;
 import android.text.InputType;
@@ -27378,6 +27377,29 @@
     }
 
     /**
+     * Modifiers the input matrix such that it maps root view's coordinates to view-local
+     * coordinates.
+     *
+     * @param matrix input matrix to modify
+     * @hide
+     */
+    public void transformMatrixRootToLocal(@NonNull Matrix matrix) {
+        final ViewParent parent = mParent;
+        if (parent instanceof final View vp) {
+            vp.transformMatrixRootToLocal(matrix);
+            matrix.postTranslate(vp.mScrollX, vp.mScrollY);
+        }
+        // This method is different from transformMatrixToLocal that it doesn't perform any
+        // transformation for ViewRootImpl
+
+        matrix.postTranslate(-mLeft, -mTop);
+
+        if (!hasIdentityMatrix()) {
+            matrix.postConcat(getInverseMatrix());
+        }
+    }
+
+    /**
      * @hide
      */
     @ViewDebug.ExportedProperty(category = "layout", indexMapping = {
@@ -34156,7 +34178,8 @@
      * REQUESTED_FRAME_RATE_CATEGORY_NORMAL, REQUESTED_FRAME_RATE_CATEGORY_HIGH.
      * Keep in mind that the preferred frame rate affects the frame rate for the next frame,
      * so use this method carefully. It's important to note that the preference is valid as
-     * long as the View is invalidated.
+     * long as the View is invalidated. Please also be aware that the requested frame rate
+     * will not propagate to child views when this API is used on a ViewGroup.
      *
      * @param frameRate the preferred frame rate of the view.
      */
@@ -34175,6 +34198,8 @@
      * REQUESTED_FRAME_RATE_CATEGORY_NO_PREFERENCE, REQUESTED_FRAME_RATE_CATEGORY_LOW,
      * REQUESTED_FRAME_RATE_CATEGORY_NORMAL, and REQUESTED_FRAME_RATE_CATEGORY_HIGH.
      * Note that the frame rate value is valid as long as the View is invalidated.
+     * Please also be aware that the requested frame rate will not propagate to
+     * child views when this API is used on a ViewGroup.
      *
      * @return REQUESTED_FRAME_RATE_CATEGORY_DEFAULT by default,
      * or value passed to {@link #setRequestedFrameRate(float)}.
diff --git a/core/java/android/view/ViewConfiguration.java b/core/java/android/view/ViewConfiguration.java
index fb2b8b9..63bf392 100644
--- a/core/java/android/view/ViewConfiguration.java
+++ b/core/java/android/view/ViewConfiguration.java
@@ -373,6 +373,7 @@
     private final int mMaximumDrawingCacheSize;
     private final int mOverscrollDistance;
     private final int mOverflingDistance;
+    private final boolean mViewTouchScreenHapticScrollFeedbackEnabled;
     @UnsupportedAppUsage
     private final boolean mFadingMarqueeEnabled;
     private final long mGlobalActionsKeyTimeout;
@@ -437,6 +438,7 @@
         mSmartSelectionInitializedTimeout = SMART_SELECTION_INITIALIZED_TIMEOUT_IN_MILLISECOND;
         mSmartSelectionInitializingTimeout = SMART_SELECTION_INITIALIZING_TIMEOUT_IN_MILLISECOND;
         mPreferKeepClearForFocusEnabled = false;
+        mViewTouchScreenHapticScrollFeedbackEnabled = false;
     }
 
     /**
@@ -588,6 +590,12 @@
         mViewBasedRotaryEncoderScrollHapticsEnabledConfig =
                 res.getBoolean(
                         com.android.internal.R.bool.config_viewBasedRotaryEncoderHapticsEnabled);
+        mViewTouchScreenHapticScrollFeedbackEnabled =
+                Flags.enableTouchScrollFeedback()
+                        ? res.getBoolean(
+                        com.android.internal.R.bool
+                                .config_viewTouchScreenHapticScrollFeedbackEnabled)
+                        : false;
     }
 
     /**
@@ -1285,6 +1293,10 @@
             return mRotaryEncoderHapticScrollFeedbackEnabled;
         }
 
+        if ((source & InputDevice.SOURCE_TOUCHSCREEN) != 0) {
+            return mViewTouchScreenHapticScrollFeedbackEnabled;
+        }
+
         return false;
     }
 
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 0e1625a..f021bdf 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -6094,6 +6094,12 @@
     }
 
     boolean scrollToRectOrFocus(Rect rectangle, boolean immediate) {
+        if (mImeBackAnimationController.isAnimationInProgress()) {
+            // IME predictive back animation is currently in progress which means that scrollY is
+            // currently controlled by ImeBackAnimationController.
+            return false;
+        }
+
         final Rect ci = mAttachInfo.mContentInsets;
         final Rect vi = mAttachInfo.mVisibleInsets;
         int scrollY = 0;
diff --git a/core/java/android/view/accessibility/AccessibilityManager.java b/core/java/android/view/accessibility/AccessibilityManager.java
index a4cea33..a87e5c8 100644
--- a/core/java/android/view/accessibility/AccessibilityManager.java
+++ b/core/java/android/view/accessibility/AccessibilityManager.java
@@ -1051,6 +1051,52 @@
     }
 
     /**
+     * Registers callback for when user initialization has completed.
+     * Does nothing if the same callback is already registered.
+     *
+     * @param callback The callback to be registered
+     * @hide
+     */
+    public void registerUserInitializationCompleteCallback(
+            @NonNull IUserInitializationCompleteCallback callback) {
+        IAccessibilityManager service;
+        synchronized (mLock) {
+            service = getServiceLocked();
+            if (service == null) {
+                return;
+            }
+        }
+        try {
+            service.registerUserInitializationCompleteCallback(callback);
+        } catch (RemoteException re) {
+            Log.e(LOG_TAG, "Error while registering userInitializationCompleteCallback. ", re);
+        }
+    }
+
+    /**
+     * Unregisters callback for when user initialization has completed.
+     *
+     * @param callback The callback to be unregistered
+     * @hide
+     */
+    public void unregisterUserInitializationCompleteCallback(
+            @NonNull IUserInitializationCompleteCallback callback) {
+        IAccessibilityManager service;
+        synchronized (mLock) {
+            service = getServiceLocked();
+            if (service == null) {
+                return;
+            }
+        }
+        try {
+            service.unregisterUserInitializationCompleteCallback(callback);
+        } catch (RemoteException re) {
+            Log.e(LOG_TAG,
+                    "Error while unregistering userInitializationCompleteCallback. ", re);
+        }
+    }
+
+    /**
      * Whether the current accessibility request comes from an
      * {@link AccessibilityService} with the {@link AccessibilityServiceInfo#isAccessibilityTool}
      * property set to true.
@@ -2049,9 +2095,7 @@
      * {@link android.view.Display#DEFAULT_DISPLAY}, is or lower than
      * {@link android.view.Display#INVALID_DISPLAY}, or is already being proxy-ed.
      *
-     * @throws SecurityException if the app does not hold the
-     * {@link Manifest.permission#MANAGE_ACCESSIBILITY} permission or the
-     * {@link Manifest.permission#CREATE_VIRTUAL_DEVICE} permission.
+     * @throws SecurityException if the app does not hold the required permissions.
      *
      * @hide
      */
@@ -2079,9 +2123,7 @@
      *
      * @return {@code true} if the proxy is successfully unregistered.
      *
-     * @throws SecurityException if the app does not hold the
-     * {@link Manifest.permission#MANAGE_ACCESSIBILITY} permission or the
-     * {@link Manifest.permission#CREATE_VIRTUAL_DEVICE} permission.
+     * @throws SecurityException if the app does not hold the required permissions.
      *
      * @hide
      */
@@ -2134,8 +2176,8 @@
         try {
             return service.startFlashNotificationSequence(context.getOpPackageName(),
                     reason, mBinder);
-        } catch (RemoteException re) {
-            Log.e(LOG_TAG, "Error while start flash notification sequence", re);
+        } catch (RemoteException | SecurityException e) {
+            Log.e(LOG_TAG, "Error while start flash notification sequence", e);
             return false;
         }
     }
@@ -2164,8 +2206,8 @@
 
         try {
             return service.stopFlashNotificationSequence(context.getOpPackageName());
-        } catch (RemoteException re) {
-            Log.e(LOG_TAG, "Error while stop flash notification sequence", re);
+        } catch (RemoteException | SecurityException e) {
+            Log.e(LOG_TAG, "Error while stop flash notification sequence", e);
             return false;
         }
     }
@@ -2192,8 +2234,8 @@
         try {
             return service.startFlashNotificationEvent(context.getOpPackageName(),
                     reason, reasonPkg);
-        } catch (RemoteException re) {
-            Log.e(LOG_TAG, "Error while start flash notification event", re);
+        } catch (RemoteException | SecurityException e) {
+            Log.e(LOG_TAG, "Error while start flash notification event", e);
             return false;
         }
     }
diff --git a/core/java/android/view/accessibility/IAccessibilityManager.aidl b/core/java/android/view/accessibility/IAccessibilityManager.aidl
index 72a1fe4..2de3ce8 100644
--- a/core/java/android/view/accessibility/IAccessibilityManager.aidl
+++ b/core/java/android/view/accessibility/IAccessibilityManager.aidl
@@ -29,6 +29,7 @@
 import android.view.accessibility.IAccessibilityManagerClient;
 import android.view.accessibility.AccessibilityWindowAttributes;
 import android.view.accessibility.IMagnificationConnection;
+import android.view.accessibility.IUserInitializationCompleteCallback;
 import android.view.InputEvent;
 import android.view.IWindow;
 import android.view.MagnificationSpec;
@@ -156,13 +157,13 @@
     @EnforcePermission("INJECT_EVENTS")
     void injectInputEventToInputFilter(in InputEvent event);
 
-    @RequiresNoPermission
+    @EnforcePermission("MANAGE_ACCESSIBILITY")
     boolean startFlashNotificationSequence(String opPkg, int reason, IBinder token);
 
-    @RequiresNoPermission
+    @EnforcePermission("MANAGE_ACCESSIBILITY")
     boolean stopFlashNotificationSequence(String opPkg);
 
-    @RequiresNoPermission
+    @EnforcePermission("MANAGE_ACCESSIBILITY")
     boolean startFlashNotificationEvent(String opPkg, int reason, String reasonPkg);
 
     @RequiresNoPermission
@@ -192,4 +193,10 @@
 
     @EnforcePermission("MANAGE_ACCESSIBILITY")
     Bundle getA11yFeatureToTileMap(int userId);
+
+    @RequiresNoPermission
+    void registerUserInitializationCompleteCallback(IUserInitializationCompleteCallback callback);
+
+    @RequiresNoPermission
+    void unregisterUserInitializationCompleteCallback(IUserInitializationCompleteCallback callback);
 }
diff --git a/core/java/android/view/accessibility/IUserInitializationCompleteCallback.aidl b/core/java/android/view/accessibility/IUserInitializationCompleteCallback.aidl
new file mode 100644
index 0000000..fe6c8e2
--- /dev/null
+++ b/core/java/android/view/accessibility/IUserInitializationCompleteCallback.aidl
@@ -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 android.view.accessibility;
+
+/**
+ * A callback for when a new user finishes initializing
+ * NOTE: Must remain a oneway interface, as it is called from system_server while holding a lock.
+ * oneway allows it to return immediately and not hold the lock for longer than is necessary.
+ * @hide
+ */
+
+oneway interface IUserInitializationCompleteCallback {
+
+    /**
+     * Called when a user initialization completes.
+     *
+     * @param userId the id of the initialized user
+     */
+    @RequiresNoPermission
+    void onUserInitializationComplete(int userId);
+}
diff --git a/core/java/android/view/flags/scroll_feedback_flags.aconfig b/core/java/android/view/flags/scroll_feedback_flags.aconfig
index 338037f..e9c85684 100644
--- a/core/java/android/view/flags/scroll_feedback_flags.aconfig
+++ b/core/java/android/view/flags/scroll_feedback_flags.aconfig
@@ -14,4 +14,11 @@
     name: "use_view_based_rotary_encoder_scroll_haptics"
     description: "If enabled, the rotary encoder scroll haptic implementation in the View class will be used, and the HapticScrollFeedbackProvider logic for rotary encoder haptic will be muted."
     bug: "299587011"
-}
\ No newline at end of file
+}
+
+flag {
+    namespace: "toolkit"
+    name: "enable_touch_scroll_feedback"
+    description: "Enables touchscreen haptic scroll feedback"
+    bug: "331830899"
+}
diff --git a/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java b/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
index 2f515fe..3a008aa 100644
--- a/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
+++ b/core/java/android/view/inputmethod/IInputMethodManagerGlobalInvoker.java
@@ -323,14 +323,14 @@
     static boolean showSoftInput(@NonNull IInputMethodClient client, @Nullable IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             int lastClickToolType, @Nullable ResultReceiver resultReceiver,
-            @SoftInputShowHideReason int reason) {
+            @SoftInputShowHideReason int reason, boolean async) {
         final IInputMethodManager service = getService();
         if (service == null) {
             return false;
         }
         try {
             return service.showSoftInput(client, windowToken, statsToken, flags, lastClickToolType,
-                    resultReceiver, reason);
+                    resultReceiver, reason, async);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -339,14 +339,15 @@
     @AnyThread
     static boolean hideSoftInput(@NonNull IInputMethodClient client, @Nullable IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-            @Nullable ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
+            @Nullable ResultReceiver resultReceiver, @SoftInputShowHideReason int reason,
+            boolean async) {
         final IInputMethodManager service = getService();
         if (service == null) {
             return false;
         }
         try {
             return service.hideSoftInput(client, windowToken, statsToken, flags, resultReceiver,
-                    reason);
+                    reason, async);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
@@ -407,7 +408,7 @@
             @Nullable IRemoteInputConnection remoteInputConnection,
             @Nullable IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher) {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, boolean useAsyncShowHideMethod) {
         final IInputMethodManager service = getService();
         if (service == null) {
             return -1;
@@ -416,7 +417,7 @@
             service.startInputOrWindowGainedFocusAsync(startInputReason, client, windowToken,
                     startInputFlags, softInputMode, windowFlags, editorInfo, remoteInputConnection,
                     remoteAccessibilityInputConnection, unverifiedTargetSdkVersion, userId,
-                    imeDispatcher, advanceAngGetStartInputSequenceNumber());
+                    imeDispatcher, advanceAngGetStartInputSequenceNumber(), useAsyncShowHideMethod);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index 23d7732..2f649c2 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -56,6 +56,7 @@
 import android.annotation.UserIdInt;
 import android.app.ActivityThread;
 import android.app.PropertyInvalidatedCache;
+import android.app.compat.CompatChanges;
 import android.compat.annotation.ChangeId;
 import android.compat.annotation.EnabledSince;
 import android.compat.annotation.UnsupportedAppUsage;
@@ -441,6 +442,36 @@
     public static final long CLEAR_SHOW_FORCED_FLAG_WHEN_LEAVING = 214016041L; // This is a bug id.
 
     /**
+     * Use async method for {@link InputMethodManager#showSoftInput(View, int, ResultReceiver)},
+     * {@link InputMethodManager#showSoftInput(View, int)} and
+     * {@link InputMethodManager#hideSoftInputFromWindow(IBinder, int, ResultReceiver)},
+     * {@link InputMethodManager#hideSoftInputFromWindow(IBinder, int)} for apps targeting V+.
+     * <p>
+     * Apps can incorrectly rely on {@link InputMethodManager#showSoftInput(View, int)} and
+     * {@link InputMethodManager#hideSoftInputFromWindow(IBinder, int)} method return type
+     * to interpret result of a request rather than relying on {@link ResultReceiver}. The return
+     * type of the method was never documented to have accurate info of visibility but few apps
+     * incorrectly rely on it.
+     * <p>
+     * Starting Android V, we use async calls into system_server which returns {@code true} if
+     * method call was made but return type doesn't guarantee execution.
+     * Apps targeting older versions will fallback to existing behavior of calling synchronous
+     * methods which had undocumented result in return type.
+     *
+     * @hide
+     */
+    @ChangeId
+    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    private static final long USE_ASYNC_SHOW_HIDE_METHOD = 352594277L; // This is a bug id.
+
+    /**
+     * Version-gating is guarded by bug-fix flag.
+     */
+    private static final boolean ASYNC_SHOW_HIDE_METHOD_ENABLED =
+            !Flags.compatchangeForZerojankproxy()
+                || CompatChanges.isChangeEnabled(USE_ASYNC_SHOW_HIDE_METHOD);
+
+    /**
      * If {@code true}, avoid calling the
      * {@link com.android.server.inputmethod.InputMethodManagerService InputMethodManagerService}
      * by skipping the call to {@link IInputMethodManager#startInputOrWindowGainedFocus}
@@ -2246,6 +2277,8 @@
      *             {@link View#isFocused view focus}, and its containing window has
      *             {@link View#hasWindowFocus window focus}. Otherwise the call fails and
      *             returns {@code false}.
+     * @return {@code true} if a request was sent to system_server, {@code false} otherwise. Note:
+     * this does not return result of the request. For result use {@param resultReceiver} instead.
      */
     public boolean showSoftInput(View view, @ShowFlags int flags) {
         // Re-dispatch if there is a context mismatch.
@@ -2315,6 +2348,8 @@
      * code you receive may be either {@link #RESULT_UNCHANGED_SHOWN},
      * {@link #RESULT_UNCHANGED_HIDDEN}, {@link #RESULT_SHOWN}, or
      * {@link #RESULT_HIDDEN}.
+     * @return {@code true} if a request was sent to system_server, {@code false} otherwise. Note:
+     * this does not return result of the request. For result use {@param resultReceiver} instead.
      */
     public boolean showSoftInput(View view, @ShowFlags int flags, ResultReceiver resultReceiver) {
         return showSoftInput(view, flags, resultReceiver, SoftInputShowHideReason.SHOW_SOFT_INPUT);
@@ -2383,7 +2418,8 @@
                         flags,
                         mCurRootView.getLastClickToolType(),
                         resultReceiver,
-                        reason);
+                        reason,
+                        ASYNC_SHOW_HIDE_METHOD_ENABLED);
             }
         }
     }
@@ -2426,7 +2462,8 @@
                     flags,
                     mCurRootView.getLastClickToolType(),
                     resultReceiver,
-                    reason);
+                    reason,
+                    ASYNC_SHOW_HIDE_METHOD_ENABLED);
         }
     }
 
@@ -2459,6 +2496,9 @@
      *
      * @param windowToken The token of the window that is making the request,
      * as returned by {@link View#getWindowToken() View.getWindowToken()}.
+     * @return {@code true} if a request was sent to system_server, {@code false} otherwise. Note:
+     * this does not return result of the request. For result use {@link ResultReceiver} in
+     * {@link #hideSoftInputFromWindow(IBinder, int, ResultReceiver)} instead.
      */
     public boolean hideSoftInputFromWindow(IBinder windowToken, @HideFlags int flags) {
         return hideSoftInputFromWindow(windowToken, flags, null);
@@ -2487,6 +2527,8 @@
      * code you receive may be either {@link #RESULT_UNCHANGED_SHOWN},
      * {@link #RESULT_UNCHANGED_HIDDEN}, {@link #RESULT_SHOWN}, or
      * {@link #RESULT_HIDDEN}.
+     * @return {@code true} if a request was sent to system_server, {@code false} otherwise. Note:
+     * this does not return result of the request. For result use {@param resultReceiver} instead.
      */
     public boolean hideSoftInputFromWindow(IBinder windowToken, @HideFlags int flags,
             ResultReceiver resultReceiver) {
@@ -2530,7 +2572,7 @@
                 return true;
             } else {
                 return IInputMethodManagerGlobalInvoker.hideSoftInput(mClient, windowToken,
-                        statsToken, flags, resultReceiver, reason);
+                        statsToken, flags, resultReceiver, reason, ASYNC_SHOW_HIDE_METHOD_ENABLED);
             }
         }
     }
@@ -2573,7 +2615,7 @@
             ImeTracker.forLogging().onProgress(statsToken, ImeTracker.PHASE_CLIENT_VIEW_SERVED);
 
             return IInputMethodManagerGlobalInvoker.hideSoftInput(mClient, view.getWindowToken(),
-                    statsToken, flags, null, reason);
+                    statsToken, flags, null, reason, ASYNC_SHOW_HIDE_METHOD_ENABLED);
         }
     }
 
@@ -3350,7 +3392,7 @@
                         servedInputConnection == null ? null
                                 : servedInputConnection.asIRemoteAccessibilityInputConnection(),
                         view.getContext().getApplicationInfo().targetSdkVersion, targetUserId,
-                        mImeDispatcher);
+                        mImeDispatcher, ASYNC_SHOW_HIDE_METHOD_ENABLED);
             } else {
                 res = IInputMethodManagerGlobalInvoker.startInputOrWindowGainedFocus(
                         startInputReason, mClient, windowGainingFocus, startInputFlags,
@@ -3653,7 +3695,8 @@
                     statsToken,
                     HIDE_NOT_ALWAYS,
                     null,
-                    reason);
+                    reason,
+                    true /*async */);
         }
     }
 
@@ -3745,7 +3788,7 @@
 
             IInputMethodManagerGlobalInvoker.hideSoftInput(mClient, windowToken, statsToken,
                     0 /* flags */, null /* resultReceiver */,
-                    SoftInputShowHideReason.HIDE_SOFT_INPUT_BY_INSETS_API);
+                    SoftInputShowHideReason.HIDE_SOFT_INPUT_BY_INSETS_API, true /* async */);
         }
     }
 
diff --git a/core/java/android/view/inputmethod/flags.aconfig b/core/java/android/view/inputmethod/flags.aconfig
index e294ee2..bae8aff 100644
--- a/core/java/android/view/inputmethod/flags.aconfig
+++ b/core/java/android/view/inputmethod/flags.aconfig
@@ -136,3 +136,14 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "compatchange_for_zerojankproxy"
+  namespace: "input_method"
+  description: "Version-gate the sync/async nature of IMM#show/hideSoftInput() when using zeroJankProxy."
+  bug: "352594277"
+  is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/core/java/android/webkit/WebView.java b/core/java/android/webkit/WebView.java
index f336b5d..ffe8c80 100644
--- a/core/java/android/webkit/WebView.java
+++ b/core/java/android/webkit/WebView.java
@@ -3088,11 +3088,11 @@
         }
 
         if (Flags.updateServiceIpcWrapper()) {
-            WebViewUpdateManager manager = WebViewUpdateManager.getInstance();
-            if (manager == null) {
+            if (WebViewFactory.isWebViewSupported()) {
+                return WebViewUpdateManager.getInstance().getCurrentWebViewPackage();
+            } else {
                 return null;
             }
-            return manager.getCurrentWebViewPackage();
         } else {
             IWebViewUpdateService service = WebViewFactory.getUpdateService();
             if (service == null) {
diff --git a/core/java/android/webkit/WebViewFactory.java b/core/java/android/webkit/WebViewFactory.java
index c53a0e1..e10a398 100644
--- a/core/java/android/webkit/WebViewFactory.java
+++ b/core/java/android/webkit/WebViewFactory.java
@@ -208,7 +208,7 @@
         public MissingWebViewPackageException(Exception e) { super(e); }
     }
 
-    private static boolean isWebViewSupported() {
+    static boolean isWebViewSupported() {
         // No lock; this is a benign race as Boolean's state is final and the PackageManager call
         // will always return the same value.
         if (sWebViewSupported == null) {
@@ -318,7 +318,7 @@
         String libraryFileName;
         try {
             PackageInfo packageInfo = packageManager.getPackageInfo(packageName,
-                    PackageManager.GET_META_DATA | PackageManager.MATCH_DEBUG_TRIAGED_MISSING);
+                    PackageManager.GET_META_DATA);
             libraryFileName = getWebViewLibrary(packageInfo.applicationInfo);
         } catch (PackageManager.NameNotFoundException e) {
             Log.e(LOGTAG, "Couldn't find package " + packageName);
@@ -479,7 +479,6 @@
                 newPackageInfo = pm.getPackageInfo(
                     response.packageInfo.packageName,
                     PackageManager.GET_SHARED_LIBRARY_FILES
-                    | PackageManager.MATCH_DEBUG_TRIAGED_MISSING
                     // Make sure that we fetch the current provider even if its not
                     // installed for the current user
                     | PackageManager.MATCH_UNINSTALLED_PACKAGES
diff --git a/core/java/android/webkit/WebViewUpdateManager.java b/core/java/android/webkit/WebViewUpdateManager.java
index dd48df9..0eb71001 100644
--- a/core/java/android/webkit/WebViewUpdateManager.java
+++ b/core/java/android/webkit/WebViewUpdateManager.java
@@ -19,12 +19,14 @@
 import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.annotation.RequiresFeature;
 import android.annotation.RequiresPermission;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
 import android.app.SystemServiceRegistry;
 import android.content.Context;
 import android.content.pm.PackageInfo;
+import android.content.pm.PackageManager;
 import android.os.RemoteException;
 
 /** @hide */
@@ -43,8 +45,11 @@
      *
      * This exists for the benefit of callsites without a {@link Context}; prefer
      * {@link Context#getSystemService(Class)} otherwise.
+     *
+     * This can only be used on devices with {@link PackageManager#FEATURE_WEBVIEW}.
      */
     @SuppressLint("ManagerLookup") // service opts in to getSystemServiceWithNoContext()
+    @RequiresFeature(PackageManager.FEATURE_WEBVIEW)
     public static @Nullable WebViewUpdateManager getInstance() {
         return (WebViewUpdateManager) SystemServiceRegistry.getSystemServiceWithNoContext(
                 Context.WEBVIEW_UPDATE_SERVICE);
diff --git a/core/java/android/webkit/WebViewUpdateService.java b/core/java/android/webkit/WebViewUpdateService.java
index 6f53dde..01af182 100644
--- a/core/java/android/webkit/WebViewUpdateService.java
+++ b/core/java/android/webkit/WebViewUpdateService.java
@@ -34,11 +34,11 @@
      */
     public static WebViewProviderInfo[] getAllWebViewPackages() {
         if (Flags.updateServiceIpcWrapper()) {
-            WebViewUpdateManager manager = WebViewUpdateManager.getInstance();
-            if (manager == null) {
+            if (WebViewFactory.isWebViewSupported()) {
+                return WebViewUpdateManager.getInstance().getAllWebViewPackages();
+            } else {
                 return new WebViewProviderInfo[0];
             }
-            return manager.getAllWebViewPackages();
         } else {
             IWebViewUpdateService service = getUpdateService();
             if (service == null) {
@@ -57,11 +57,11 @@
      */
     public static WebViewProviderInfo[] getValidWebViewPackages() {
         if (Flags.updateServiceIpcWrapper()) {
-            WebViewUpdateManager manager = WebViewUpdateManager.getInstance();
-            if (manager == null) {
+            if (WebViewFactory.isWebViewSupported()) {
+                return WebViewUpdateManager.getInstance().getValidWebViewPackages();
+            } else {
                 return new WebViewProviderInfo[0];
             }
-            return manager.getValidWebViewPackages();
         } else {
             IWebViewUpdateService service = getUpdateService();
             if (service == null) {
@@ -80,11 +80,11 @@
      */
     public static String getCurrentWebViewPackageName() {
         if (Flags.updateServiceIpcWrapper()) {
-            WebViewUpdateManager manager = WebViewUpdateManager.getInstance();
-            if (manager == null) {
+            if (WebViewFactory.isWebViewSupported()) {
+                return WebViewUpdateManager.getInstance().getCurrentWebViewPackageName();
+            } else {
                 return null;
             }
-            return manager.getCurrentWebViewPackageName();
         } else {
             IWebViewUpdateService service = getUpdateService();
             if (service == null) {
diff --git a/core/java/android/webkit/flags.aconfig b/core/java/android/webkit/flags.aconfig
index defe61e..b21a490 100644
--- a/core/java/android/webkit/flags.aconfig
+++ b/core/java/android/webkit/flags.aconfig
@@ -9,3 +9,12 @@
     bug: "319292658"
     is_fixed_read_only: true
 }
+
+flag {
+    name: "mainline_apis"
+    is_exported: true
+    namespace: "webview"
+    description: "New APIs required by WebViewBootstrap mainline module"
+    bug: "310653407"
+    is_fixed_read_only: true
+}
diff --git a/core/java/android/widget/AbsListView.java b/core/java/android/widget/AbsListView.java
index ab6b512..3c854ea 100644
--- a/core/java/android/widget/AbsListView.java
+++ b/core/java/android/widget/AbsListView.java
@@ -16,6 +16,8 @@
 
 package android.widget;
 
+import static android.view.flags.Flags.enableTouchScrollFeedback;
+import static android.view.flags.Flags.scrollFeedbackApi;
 import static android.view.flags.Flags.viewVelocityApi;
 
 import android.annotation.ColorInt;
@@ -82,7 +84,6 @@
 import android.view.autofill.AutofillId;
 import android.view.contentcapture.ContentCaptureManager;
 import android.view.contentcapture.ContentCaptureSession;
-import android.view.flags.Flags;
 import android.view.inputmethod.BaseInputConnection;
 import android.view.inputmethod.CompletionInfo;
 import android.view.inputmethod.CorrectionInfo;
@@ -3703,7 +3704,6 @@
                 // If it's non-null, we're already in a scroll.
                 mScrollStrictSpan = StrictMode.enterCriticalSpan("AbsListView-scroll");
             }
-
             if (y != mLastY) {
                 // We may be here after stopping a fling and continuing to scroll.
                 // If so, we haven't disallowed intercepting touch events yet.
@@ -3735,8 +3735,15 @@
                 boolean atEdge = false;
                 if (incrementalDeltaY != 0) {
                     atEdge = trackMotionScroll(deltaY, incrementalDeltaY);
-                }
 
+                    // TODO: b/360198915 - Add unit testing for using ScrollFeedbackProvider
+                    if (enableTouchScrollFeedback()) {
+                        initHapticScrollFeedbackProviderIfNotExists();
+                        mHapticScrollFeedbackProvider.onScrollProgress(
+                                vtev.getDeviceId(), vtev.getSource(), MotionEvent.AXIS_Y,
+                                incrementalDeltaY);
+                    }
+                }
                 // Check to see if we have bumped into the scroll limit
                 motionView = this.getChildAt(motionIndex);
                 if (motionView != null) {
@@ -3745,7 +3752,6 @@
                     final int motionViewRealTop = motionView.getTop();
                     if (atEdge) {
                         // Apply overscroll
-
                         int overscroll = -incrementalDeltaY -
                                 (motionViewRealTop - motionViewPrevTop);
                         if (dispatchNestedScroll(0, overscroll - incrementalDeltaY, 0, overscroll,
@@ -3772,6 +3778,15 @@
                                     mDirection = 0; // Reset when entering overscroll.
                                     mTouchMode = TOUCH_MODE_OVERSCROLL;
                                 }
+
+                                if (enableTouchScrollFeedback()) {
+                                    initHapticScrollFeedbackProviderIfNotExists();
+                                    mHapticScrollFeedbackProvider.onScrollLimit(
+                                            vtev.getDeviceId(), vtev.getSource(),
+                                            MotionEvent.AXIS_Y,
+                                            /* isStart= */ incrementalDeltaY > 0);
+                                }
+
                                 if (incrementalDeltaY > 0) {
                                     mEdgeGlowTop.onPullDistance((float) -overscroll / getHeight(),
                                             (float) x / getWidth());
@@ -3981,7 +3996,6 @@
         if (mFastScroll != null && mFastScroll.onTouchEvent(ev)) {
             return true;
         }
-
         initVelocityTrackerIfNotExists();
         final MotionEvent vtev = MotionEvent.obtain(ev);
 
@@ -4520,7 +4534,7 @@
                     final int overscrollMode = getOverScrollMode();
 
                     if (!trackMotionScroll(delta, delta)) {
-                        if (Flags.scrollFeedbackApi()) {
+                        if (scrollFeedbackApi()) {
                             initHapticScrollFeedbackProviderIfNotExists();
                             mHapticScrollFeedbackProvider.onScrollProgress(
                                     event.getDeviceId(), event.getSource(), axis, delta);
@@ -4536,7 +4550,7 @@
                         float overscroll = (delta - (motionViewRealTop - motionViewPrevTop))
                                 / ((float) getHeight());
                         boolean hitTopLimit = delta > 0;
-                        if (Flags.scrollFeedbackApi()) {
+                        if (scrollFeedbackApi()) {
                             initHapticScrollFeedbackProviderIfNotExists();
                             mHapticScrollFeedbackProvider.onScrollLimit(
                                     event.getDeviceId(), event.getSource(), axis,
diff --git a/core/java/android/widget/TextView.java b/core/java/android/widget/TextView.java
index a346a67..a4b28ad 100644
--- a/core/java/android/widget/TextView.java
+++ b/core/java/android/widget/TextView.java
@@ -14375,7 +14375,7 @@
 
         Matrix matrix = mTempMatrix;
         matrix.reset();
-        transformMatrixToLocal(matrix);
+        transformMatrixRootToLocal(matrix);
         editorBounds.set(rect);
         // When the view has transformations like scaleX/scaleY computing the global visible
         // rectangle will already apply the transformations. The getLocalVisibleRect only offsets
diff --git a/core/java/android/window/ITaskFragmentOrganizerController.aidl b/core/java/android/window/ITaskFragmentOrganizerController.aidl
index 58b5757..b8a11cf0 100644
--- a/core/java/android/window/ITaskFragmentOrganizerController.aidl
+++ b/core/java/android/window/ITaskFragmentOrganizerController.aidl
@@ -47,6 +47,19 @@
     void unregisterOrganizer(in ITaskFragmentOrganizer organizer);
 
     /**
+     * Registers remote animations per transition type for the organizer. It will override the
+     * animations if the transition only contains windows that belong to the organized
+     * TaskFragments in the given Task.
+     */
+    void registerRemoteAnimations(in ITaskFragmentOrganizer organizer,
+        in RemoteAnimationDefinition definition);
+
+    /**
+     * Unregisters remote animations per transition type for the organizer.
+     */
+    void unregisterRemoteAnimations(in ITaskFragmentOrganizer organizer);
+
+    /**
      * Saves the state in the system, where the state can be restored if the process of
      * the TaskFragmentOrganizer is restarted.
      */
diff --git a/core/java/android/window/SnapshotDrawerUtils.java b/core/java/android/window/SnapshotDrawerUtils.java
index 2f595d1..9a7bce0 100644
--- a/core/java/android/window/SnapshotDrawerUtils.java
+++ b/core/java/android/window/SnapshotDrawerUtils.java
@@ -151,9 +151,11 @@
         @VisibleForTesting
         public void setFrames(Rect frame, Rect systemBarInsets) {
             mFrame.set(frame);
-            mSystemBarInsets.set(systemBarInsets);
             mSizeMismatch = (mFrame.width() != mSnapshotW || mFrame.height() != mSnapshotH);
-            mSystemBarBackgroundPainter.setInsets(systemBarInsets);
+            if (!Flags.drawSnapshotAspectRatioMatch() && systemBarInsets != null) {
+                mSystemBarInsets.set(systemBarInsets);
+                mSystemBarBackgroundPainter.setInsets(systemBarInsets);
+            }
         }
 
         private void drawSnapshot(boolean releaseAfterDraw) {
@@ -394,9 +396,12 @@
         final ActivityManager.RunningTaskInfo runningTaskInfo = info.taskInfo;
         final ActivityManager.TaskDescription taskDescription =
                 getOrCreateTaskDescription(runningTaskInfo);
-        drawSurface.initiateSystemBarPainter(lp.flags, lp.privateFlags,
-                attrs.insetsFlags.appearance, taskDescription, info.requestedVisibleTypes);
-        final Rect systemBarInsets = getSystemBarInsets(windowBounds, topWindowInsetsState);
+        Rect systemBarInsets = null;
+        if (!Flags.drawSnapshotAspectRatioMatch()) {
+            drawSurface.initiateSystemBarPainter(lp.flags, lp.privateFlags,
+                    attrs.insetsFlags.appearance, taskDescription, info.requestedVisibleTypes);
+            systemBarInsets = getSystemBarInsets(windowBounds, topWindowInsetsState);
+        }
         drawSurface.setFrames(windowBounds, systemBarInsets);
         drawSurface.drawSnapshot(releaseAfterDraw);
     }
diff --git a/core/java/android/window/TaskFragmentOrganizer.java b/core/java/android/window/TaskFragmentOrganizer.java
index 027d323..4cc0d8a 100644
--- a/core/java/android/window/TaskFragmentOrganizer.java
+++ b/core/java/android/window/TaskFragmentOrganizer.java
@@ -34,6 +34,7 @@
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.view.RemoteAnimationDefinition;
 import android.view.WindowManager;
 
 import com.android.window.flags.Flags;
@@ -225,6 +226,34 @@
     }
 
     /**
+     * Registers remote animations per transition type for the organizer. It will override the
+     * animations if the transition only contains windows that belong to the organized
+     * TaskFragments, and at least one of the transition window is embedded (not filling the Task).
+     * @hide
+     */
+    @CallSuper
+    public void registerRemoteAnimations(@NonNull RemoteAnimationDefinition definition) {
+        try {
+            getController().registerRemoteAnimations(mInterface, definition);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Unregisters remote animations per transition type for the organizer.
+     * @hide
+     */
+    @CallSuper
+    public void unregisterRemoteAnimations() {
+        try {
+            getController().unregisterRemoteAnimations(mInterface);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Saves the state in the system, where the state can be restored if the process of
      * the TaskFragmentOrganizer is restarted.
      *
diff --git a/core/java/android/window/TaskFragmentParentInfo.java b/core/java/android/window/TaskFragmentParentInfo.java
index 1555416..efd74c3 100644
--- a/core/java/android/window/TaskFragmentParentInfo.java
+++ b/core/java/android/window/TaskFragmentParentInfo.java
@@ -41,6 +41,8 @@
 
     private final int mDisplayId;
 
+    private final int mTaskId;
+
     private final boolean mVisible;
 
     private final boolean mHasDirectActivity;
@@ -49,9 +51,11 @@
 
     /** @hide */
     public TaskFragmentParentInfo(@NonNull Configuration configuration, int displayId,
-            boolean visible, boolean hasDirectActivity, @Nullable SurfaceControl decorSurface) {
+            int taskId, boolean visible, boolean hasDirectActivity,
+            @Nullable SurfaceControl decorSurface) {
         mConfiguration.setTo(configuration);
         mDisplayId = displayId;
+        mTaskId = taskId;
         mVisible = visible;
         mHasDirectActivity = hasDirectActivity;
         mDecorSurface = decorSurface;
@@ -61,6 +65,7 @@
     public TaskFragmentParentInfo(@NonNull TaskFragmentParentInfo info) {
         mConfiguration.setTo(info.getConfiguration());
         mDisplayId = info.mDisplayId;
+        mTaskId = info.mTaskId;
         mVisible = info.mVisible;
         mHasDirectActivity = info.mHasDirectActivity;
         mDecorSurface = info.mDecorSurface;
@@ -87,6 +92,15 @@
     }
 
     /**
+     * The id of the parent Task.
+     *
+     * @hide
+     */
+    public int getTaskId() {
+        return mTaskId;
+    }
+
+    /**
      * Whether the parent Task is visible or not
      *
      * @hide
@@ -120,7 +134,8 @@
             return false;
         }
         return getWindowingMode() == that.getWindowingMode() && mDisplayId == that.mDisplayId
-                && mVisible == that.mVisible && mHasDirectActivity == that.mHasDirectActivity
+                && mTaskId == that.mTaskId && mVisible == that.mVisible
+                && mHasDirectActivity == that.mHasDirectActivity
                 && mDecorSurface == that.mDecorSurface;
     }
 
@@ -140,6 +155,7 @@
         return TaskFragmentParentInfo.class.getSimpleName() + ":{"
                 + "config=" + mConfiguration
                 + ", displayId=" + mDisplayId
+                + ", taskId=" + mTaskId
                 + ", visible=" + mVisible
                 + ", hasDirectActivity=" + mHasDirectActivity
                 + ", decorSurface=" + mDecorSurface
@@ -163,6 +179,7 @@
         final TaskFragmentParentInfo that = (TaskFragmentParentInfo) obj;
         return mConfiguration.equals(that.mConfiguration)
                 && mDisplayId == that.mDisplayId
+                && mTaskId == that.mTaskId
                 && mVisible == that.mVisible
                 && mHasDirectActivity == that.mHasDirectActivity
                 && mDecorSurface == that.mDecorSurface;
@@ -172,6 +189,7 @@
     public int hashCode() {
         int result = mConfiguration.hashCode();
         result = 31 * result + mDisplayId;
+        result = 31 * result + mTaskId;
         result = 31 * result + (mVisible ? 1 : 0);
         result = 31 * result + (mHasDirectActivity ? 1 : 0);
         result = 31 * result + Objects.hashCode(mDecorSurface);
@@ -183,6 +201,7 @@
     public void writeToParcel(@NonNull Parcel dest, int flags) {
         mConfiguration.writeToParcel(dest, flags);
         dest.writeInt(mDisplayId);
+        dest.writeInt(mTaskId);
         dest.writeBoolean(mVisible);
         dest.writeBoolean(mHasDirectActivity);
         dest.writeTypedObject(mDecorSurface, flags);
@@ -191,6 +210,7 @@
     private TaskFragmentParentInfo(Parcel in) {
         mConfiguration.readFromParcel(in);
         mDisplayId = in.readInt();
+        mTaskId = in.readInt();
         mVisible = in.readBoolean();
         mHasDirectActivity = in.readBoolean();
         mDecorSurface = in.readTypedObject(SurfaceControl.CREATOR);
diff --git a/core/java/android/window/TransitionInfo.java b/core/java/android/window/TransitionInfo.java
index 1083f64..ec79f94 100644
--- a/core/java/android/window/TransitionInfo.java
+++ b/core/java/android/window/TransitionInfo.java
@@ -1141,6 +1141,7 @@
         // Customize activity transition animation
         private CustomActivityTransition mCustomActivityOpenTransition;
         private CustomActivityTransition mCustomActivityCloseTransition;
+        private int mUserId;
 
         private AnimationOptions(int type) {
             mType = type;
@@ -1159,6 +1160,7 @@
             mAnimations = in.readInt();
             mCustomActivityOpenTransition = in.readTypedObject(CustomActivityTransition.CREATOR);
             mCustomActivityCloseTransition = in.readTypedObject(CustomActivityTransition.CREATOR);
+            mUserId = in.readInt();
         }
 
         /** Make basic customized animation for a package */
@@ -1283,6 +1285,14 @@
             return options;
         }
 
+        public void setUserId(int userId) {
+            mUserId = userId;
+        }
+
+        public int getUserId() {
+            return mUserId;
+        }
+
         public int getType() {
             return mType;
         }
@@ -1349,6 +1359,7 @@
             dest.writeInt(mAnimations);
             dest.writeTypedObject(mCustomActivityOpenTransition, flags);
             dest.writeTypedObject(mCustomActivityCloseTransition, flags);
+            dest.writeInt(mUserId);
         }
 
         @NonNull
@@ -1406,6 +1417,7 @@
             if (mExitResId != DEFAULT_ANIMATION_RESOURCES_ID) {
                 sb.append(" exitResId=").append(mExitResId);
             }
+            sb.append(" mUserId=").append(mUserId);
             sb.append('}');
             return sb.toString();
         }
diff --git a/core/java/android/window/flags/lse_desktop_experience.aconfig b/core/java/android/window/flags/lse_desktop_experience.aconfig
index 217bca7..cab6d8e 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -115,6 +115,16 @@
 }
 
 flag {
+    name: "respect_orientation_change_for_unresizeable"
+    namespace: "lse_desktop_experience"
+    description: "Whether to resize task to respect requested orientation change of unresizeable activity"
+    bug: "353338503"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "enable_camera_compat_for_desktop_windowing"
     namespace: "lse_desktop_experience"
     description: "Whether to apply Camera Compat treatment to fixed-orientation apps in desktop windowing mode"
@@ -235,3 +245,10 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enable_desktop_windowing_persistence"
+    namespace: "lse_desktop_experience"
+    description: "Persists the desktop windowing session across reboots."
+    bug: "350456942"
+}
diff --git a/core/java/android/window/flags/responsible_apis.aconfig b/core/java/android/window/flags/responsible_apis.aconfig
index 6ce9725..cd31850 100644
--- a/core/java/android/window/flags/responsible_apis.aconfig
+++ b/core/java/android/window/flags/responsible_apis.aconfig
@@ -71,3 +71,11 @@
     bug: "339720406"
 }
 
+flag {
+    name: "bal_reduce_grace_period"
+    namespace: "responsible_apis"
+    description: "Changes to reduce or ideally remove the grace period exemption."
+    bug: "362575865"
+}
+
+
diff --git a/core/java/android/window/flags/wallpaper_manager.aconfig b/core/java/android/window/flags/wallpaper_manager.aconfig
index 8c6721a..efacc34 100644
--- a/core/java/android/window/flags/wallpaper_manager.aconfig
+++ b/core/java/android/window/flags/wallpaper_manager.aconfig
@@ -49,3 +49,13 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "avoid_rebinding_intentionally_disconnected_wallpaper"
+  namespace: "systemui"
+  description: "Prevents rebinding with intentionally disconnected wallpaper services."
+  bug: "332871851"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/core/java/android/window/flags/windowing_frontend.aconfig b/core/java/android/window/flags/windowing_frontend.aconfig
index 67fc270..a786fc2 100644
--- a/core/java/android/window/flags/windowing_frontend.aconfig
+++ b/core/java/android/window/flags/windowing_frontend.aconfig
@@ -246,3 +246,14 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "always_capture_activity_snapshot"
+  namespace: "windowing_frontend"
+  description: "Always capture activity snapshot regardless predictive back status"
+  bug: "362183912"
+  is_fixed_read_only: true
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
\ No newline at end of file
diff --git a/core/java/com/android/internal/display/BrightnessSynchronizer.java b/core/java/com/android/internal/display/BrightnessSynchronizer.java
index 9f5ed65..21fbf9d 100644
--- a/core/java/com/android/internal/display/BrightnessSynchronizer.java
+++ b/core/java/com/android/internal/display/BrightnessSynchronizer.java
@@ -134,7 +134,8 @@
      * Prints data on dumpsys.
      */
     public void dump(PrintWriter pw) {
-        pw.println("BrightnessSynchronizer");
+        pw.println("BrightnessSynchronizer:");
+        pw.println("-----------------------");
         pw.println("  mLatestIntBrightness=" + mLatestIntBrightness);
         pw.println("  mLatestFloatBrightness=" + mLatestFloatBrightness);
         pw.println("  mCurrentUpdate=" + mCurrentUpdate);
diff --git a/core/java/com/android/internal/infra/TEST_MAPPING b/core/java/com/android/internal/infra/TEST_MAPPING
index e4550c0..35f0553 100644
--- a/core/java/com/android/internal/infra/TEST_MAPPING
+++ b/core/java/com/android/internal/infra/TEST_MAPPING
@@ -9,15 +9,7 @@
       ]
     },
     {
-      "name": "CtsPermissionTestCases",
-      "options": [
-          {
-            "include-filter": "android.permission.cts.PermissionControllerTest"
-          },
-          {
-            "exclude-annotation": "androidx.test.filters.FlakyTest"
-          }
-      ]
+      "name": "CtsPermissionTestCases_Platform"
     },
     {
       "name": "FrameworksCoreTests_internal_infra"
diff --git a/core/java/com/android/internal/os/ZygoteInit.java b/core/java/com/android/internal/os/ZygoteInit.java
index b9cc457..2acda8a 100644
--- a/core/java/com/android/internal/os/ZygoteInit.java
+++ b/core/java/com/android/internal/os/ZygoteInit.java
@@ -631,21 +631,20 @@
      */
     private static Runnable forkSystemServer(String abiList, String socketName,
             ZygoteServer zygoteServer) {
-        long capabilities = posixCapabilitiesAsBits(
-                OsConstants.CAP_IPC_LOCK,
-                OsConstants.CAP_KILL,
-                OsConstants.CAP_NET_ADMIN,
-                OsConstants.CAP_NET_BIND_SERVICE,
-                OsConstants.CAP_NET_BROADCAST,
-                OsConstants.CAP_NET_RAW,
-                OsConstants.CAP_SYS_MODULE,
-                OsConstants.CAP_SYS_NICE,
-                OsConstants.CAP_SYS_PTRACE,
-                OsConstants.CAP_SYS_TIME,
-                OsConstants.CAP_SYS_TTY_CONFIG,
-                OsConstants.CAP_WAKE_ALARM,
-                OsConstants.CAP_BLOCK_SUSPEND
-        );
+        long capabilities =
+                (1L << OsConstants.CAP_IPC_LOCK) |
+                (1L << OsConstants.CAP_KILL) |
+                (1L << OsConstants.CAP_NET_ADMIN) |
+                (1L << OsConstants.CAP_NET_BIND_SERVICE) |
+                (1L << OsConstants.CAP_NET_BROADCAST) |
+                (1L << OsConstants.CAP_NET_RAW) |
+                (1L << OsConstants.CAP_SYS_MODULE) |
+                (1L << OsConstants.CAP_SYS_NICE) |
+                (1L << OsConstants.CAP_SYS_PTRACE) |
+                (1L << OsConstants.CAP_SYS_TIME) |
+                (1L << OsConstants.CAP_SYS_TTY_CONFIG) |
+                (1L << OsConstants.CAP_WAKE_ALARM) |
+                (1L << OsConstants.CAP_BLOCK_SUSPEND);
         /* Containers run without some capabilities, so drop any caps that are not available. */
         StructCapUserHeader header = new StructCapUserHeader(
                 OsConstants._LINUX_CAPABILITY_VERSION_3, 0);
@@ -742,20 +741,6 @@
     }
 
     /**
-     * Gets the bit array representation of the provided list of POSIX capabilities.
-     */
-    private static long posixCapabilitiesAsBits(int... capabilities) {
-        long result = 0;
-        for (int capability : capabilities) {
-            if ((capability < 0) || (capability > OsConstants.CAP_LAST_CAP)) {
-                throw new IllegalArgumentException(String.valueOf(capability));
-            }
-            result |= (1L << capability);
-        }
-        return result;
-    }
-
-    /**
      * This is the entry point for a Zygote process.  It creates the Zygote server, loads resources,
      * and handles other tasks related to preparing the process for forking into applications.
      *
diff --git a/core/java/com/android/internal/policy/PhoneWindow.java b/core/java/com/android/internal/policy/PhoneWindow.java
index e14249c..e0529b3 100644
--- a/core/java/com/android/internal/policy/PhoneWindow.java
+++ b/core/java/com/android/internal/policy/PhoneWindow.java
@@ -3333,6 +3333,7 @@
             Bundle args = new Bundle();
             args.putInt(Intent.EXTRA_ASSIST_INPUT_DEVICE_ID, event.getDeviceId());
             args.putLong(Intent.EXTRA_TIME, event.getEventTime());
+            args.putBoolean(Intent.EXTRA_ASSIST_INPUT_HINT_KEYBOARD, true);
             ((SearchManager) getContext().getSystemService(Context.SEARCH_SERVICE))
                     .launchAssist(args);
             return true;
diff --git a/core/java/com/android/internal/policy/TransitionAnimation.java b/core/java/com/android/internal/policy/TransitionAnimation.java
index 238e6f5..201f267 100644
--- a/core/java/com/android/internal/policy/TransitionAnimation.java
+++ b/core/java/com/android/internal/policy/TransitionAnimation.java
@@ -49,7 +49,7 @@
 import android.media.Image;
 import android.media.ImageReader;
 import android.os.Handler;
-import android.os.SystemProperties;
+import android.os.UserHandle;
 import android.util.Slog;
 import android.view.InflateException;
 import android.view.SurfaceControl;
@@ -187,23 +187,44 @@
         return createHiddenByKeyguardExit(mContext, mInterpolator, onWallpaper, toShade, subtle);
     }
 
+    /** Load keyguard unocclude animation for user. */
+    @Nullable
+    public Animation loadKeyguardUnoccludeAnimation(int userId) {
+        return loadDefaultAnimationRes(com.android.internal.R.anim.wallpaper_open_exit, userId);
+    }
+
+    /** Same as {@code loadKeyguardUnoccludeAnimation} for current user. */
     @Nullable
     public Animation loadKeyguardUnoccludeAnimation() {
-        return loadDefaultAnimationRes(com.android.internal.R.anim.wallpaper_open_exit);
+        return loadKeyguardUnoccludeAnimation(UserHandle.USER_CURRENT);
     }
 
+    /** Load voice activity open animation for user. */
     @Nullable
-    public Animation loadVoiceActivityOpenAnimation(boolean enter) {
+    public Animation loadVoiceActivityOpenAnimation(boolean enter, int userId) {
         return loadDefaultAnimationRes(enter
                 ? com.android.internal.R.anim.voice_activity_open_enter
-                : com.android.internal.R.anim.voice_activity_open_exit);
+                : com.android.internal.R.anim.voice_activity_open_exit, userId);
     }
 
+    /** Same as {@code loadVoiceActivityOpenAnimation} for current user. */
     @Nullable
-    public Animation loadVoiceActivityExitAnimation(boolean enter) {
+    public Animation loadVoiceActivityOpenAnimation(boolean enter) {
+        return loadVoiceActivityOpenAnimation(enter, UserHandle.USER_CURRENT);
+    }
+
+    /** Load voice activity exit animation for user. */
+    @Nullable
+    public Animation loadVoiceActivityExitAnimation(boolean enter, int userId) {
         return loadDefaultAnimationRes(enter
                 ? com.android.internal.R.anim.voice_activity_close_enter
-                : com.android.internal.R.anim.voice_activity_close_exit);
+                : com.android.internal.R.anim.voice_activity_close_exit, userId);
+    }
+
+    /** Same as {@code loadVoiceActivityExitAnimation} for current user. */
+    @Nullable
+    public Animation loadVoiceActivityExitAnimation(boolean enter) {
+        return loadVoiceActivityExitAnimation(enter, UserHandle.USER_CURRENT);
     }
 
     @Nullable
@@ -211,10 +232,17 @@
         return loadAnimationRes(packageName, resId);
     }
 
+    /** Load cross profile app enter animation for user. */
+    @Nullable
+    public Animation loadCrossProfileAppEnterAnimation(int userId) {
+        return loadAnimationRes(DEFAULT_PACKAGE,
+                com.android.internal.R.anim.task_open_enter_cross_profile_apps, userId);
+    }
+
+    /** Same as {@code loadCrossProfileAppEnterAnimation} for current user. */
     @Nullable
     public Animation loadCrossProfileAppEnterAnimation() {
-        return loadAnimationRes(DEFAULT_PACKAGE,
-                com.android.internal.R.anim.task_open_enter_cross_profile_apps);
+        return loadCrossProfileAppEnterAnimation(UserHandle.USER_CURRENT);
     }
 
     @Nullable
@@ -230,11 +258,11 @@
                 appRect.height(), 0, null);
     }
 
-    /** Load animation by resource Id from specific package. */
+    /** Load animation by resource Id from specific package for user. */
     @Nullable
-    public Animation loadAnimationRes(String packageName, int resId) {
+    public Animation loadAnimationRes(String packageName, int resId, int userId) {
         if (ResourceId.isValid(resId)) {
-            AttributeCache.Entry ent = getCachedAnimations(packageName, resId);
+            AttributeCache.Entry ent = getCachedAnimations(packageName, resId, userId);
             if (ent != null) {
                 return loadAnimationSafely(ent.context, resId, mTag);
             }
@@ -242,10 +270,22 @@
         return null;
     }
 
-    /** Load animation by resource Id from android package. */
+    /** Same as {@code loadAnimationRes} for current user. */
+    @Nullable
+    public Animation loadAnimationRes(String packageName, int resId) {
+        return loadAnimationRes(packageName, resId, UserHandle.USER_CURRENT);
+    }
+
+    /** Load animation by resource Id from android package for user. */
+    @Nullable
+    public Animation loadDefaultAnimationRes(int resId, int userId) {
+        return loadAnimationRes(DEFAULT_PACKAGE, resId, userId);
+    }
+
+    /** Same as {@code loadDefaultAnimationRes} for current user. */
     @Nullable
     public Animation loadDefaultAnimationRes(int resId) {
-        return loadAnimationRes(DEFAULT_PACKAGE, resId);
+        return loadAnimationRes(DEFAULT_PACKAGE, resId, UserHandle.USER_CURRENT);
     }
 
     /** Load animation by attribute Id from specific LayoutParams */
@@ -378,10 +418,10 @@
     }
 
     @Nullable
-    private AttributeCache.Entry getCachedAnimations(String packageName, int resId) {
+    private AttributeCache.Entry getCachedAnimations(String packageName, int resId, int userId) {
         if (mDebug) {
-            Slog.v(mTag, "Loading animations: package="
-                    + packageName + " resId=0x" + Integer.toHexString(resId));
+            Slog.v(mTag, "Loading animations: package=" + packageName + " resId=0x"
+                    + Integer.toHexString(resId) + " for user=" + userId);
         }
         if (packageName != null) {
             if ((resId & 0xFF000000) == 0x01000000) {
@@ -392,11 +432,16 @@
                         + packageName);
             }
             return AttributeCache.instance().get(packageName, resId,
-                    com.android.internal.R.styleable.WindowAnimation);
+                    com.android.internal.R.styleable.WindowAnimation, userId);
         }
         return null;
     }
 
+    @Nullable
+    private AttributeCache.Entry getCachedAnimations(String packageName, int resId) {
+        return getCachedAnimations(packageName, resId, UserHandle.USER_CURRENT);
+    }
+
     /** Returns window animation style ID from {@link LayoutParams} or from system in some cases */
     public int getAnimationStyleResId(@NonNull LayoutParams lp) {
         int resId = lp.windowAnimations;
diff --git a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
index 2ff8c8c..4264358 100644
--- a/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
+++ b/core/java/com/android/internal/protolog/PerfettoProtoLogImpl.java
@@ -930,37 +930,47 @@
     }
 
     private static class Message {
+        @Nullable
         private final Long mMessageHash;
+        @Nullable
         private final Integer mMessageMask;
+        @Nullable
         private final String mMessageString;
 
-        private Message(Long messageHash, int messageMask) {
+        private Message(long messageHash, int messageMask) {
             this.mMessageHash = messageHash;
             this.mMessageMask = messageMask;
             this.mMessageString = null;
         }
 
-        private Message(String messageString) {
+        private Message(@NonNull String messageString) {
             this.mMessageHash = null;
             final List<Integer> argTypes = LogDataType.parseFormatString(messageString);
             this.mMessageMask = LogDataType.logDataTypesToBitMask(argTypes);
             this.mMessageString = messageString;
         }
 
-        private int getMessageMask() {
+        @Nullable
+        private Integer getMessageMask() {
             return mMessageMask;
         }
 
+        @Nullable
         private String getMessage() {
             return mMessageString;
         }
 
+        @Nullable
         private String getMessage(@NonNull ProtoLogViewerConfigReader viewerConfigReader) {
             if (mMessageString != null) {
                 return mMessageString;
             }
 
-            return viewerConfigReader.getViewerString(mMessageHash);
+            if (mMessageHash != null) {
+                return viewerConfigReader.getViewerString(mMessageHash);
+            }
+
+            throw new RuntimeException("Both mMessageString and mMessageHash should never be null");
         }
     }
 }
diff --git a/core/java/com/android/internal/protolog/ProtoLogConfigurationService.java b/core/java/com/android/internal/protolog/ProtoLogConfigurationService.java
index 1765738..eeac139 100644
--- a/core/java/com/android/internal/protolog/ProtoLogConfigurationService.java
+++ b/core/java/com/android/internal/protolog/ProtoLogConfigurationService.java
@@ -23,6 +23,7 @@
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MESSAGES;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.GROUP_ID;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.LEVEL;
+import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.LOCATION;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.MESSAGE;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.MESSAGE_ID;
 import static android.internal.perfetto.protos.TracePacketOuterClass.TracePacket.PROTOLOG_VIEWER_CONFIG;
@@ -210,8 +211,7 @@
          *                             want to write to the trace buffer.
          * @throws FileNotFoundException if the viewerConfigFilePath is invalid.
          */
-        void trace(@NonNull ProtoLogDataSource dataSource, @NonNull String viewerConfigFilePath)
-                throws FileNotFoundException;
+        void trace(@NonNull ProtoLogDataSource dataSource, @NonNull String viewerConfigFilePath);
     }
 
     @Override
@@ -351,11 +351,7 @@
 
     private void onTracingInstanceFlush() {
         for (String fileName : mConfigFileCounts.keySet()) {
-            try {
-                mViewerConfigFileTracer.trace(mDataSource, fileName);
-            } catch (FileNotFoundException e) {
-                throw new RuntimeException(e);
-            }
+            mViewerConfigFileTracer.trace(mDataSource, fileName);
         }
     }
 
@@ -364,10 +360,16 @@
     }
 
     private static void dumpTransitionTraceConfig(@NonNull ProtoLogDataSource dataSource,
-            @NonNull String viewerConfigFilePath) throws FileNotFoundException {
-        final var pis = new ProtoInputStream(new FileInputStream(viewerConfigFilePath));
-
+            @NonNull String viewerConfigFilePath) {
         dataSource.trace(ctx -> {
+            final ProtoInputStream pis;
+            try {
+                pis = new ProtoInputStream(new FileInputStream(viewerConfigFilePath));
+            } catch (FileNotFoundException e) {
+                throw new RuntimeException(
+                        "Failed to load viewer config file " + viewerConfigFilePath, e);
+            }
+
             try {
                 final ProtoOutputStream os = ctx.newTracePacket();
 
@@ -396,11 +398,7 @@
             mConfigFileCounts.put(configFile, newCount);
             boolean lastProcessWithViewerConfig = newCount == 0;
             if (lastProcessWithViewerConfig) {
-                try {
-                    mViewerConfigFileTracer.trace(mDataSource, configFile);
-                } catch (FileNotFoundException e) {
-                    throw new RuntimeException(e);
-                }
+                mViewerConfigFileTracer.trace(mDataSource, configFile);
             }
         }
     }
@@ -446,6 +444,7 @@
                 case (int) MESSAGE -> os.write(MESSAGE, pis.readString(MESSAGE));
                 case (int) LEVEL -> os.write(LEVEL, pis.readInt(LEVEL));
                 case (int) GROUP_ID -> os.write(GROUP_ID, pis.readInt(GROUP_ID));
+                case (int) LOCATION -> os.write(LOCATION, pis.readString(LOCATION));
                 default ->
                     throw new RuntimeException(
                             "Unexpected field id " + pis.getFieldNumber());
diff --git a/core/java/com/android/internal/protolog/ProtoLogDataSource.java b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
index 5c06b87..837622f 100644
--- a/core/java/com/android/internal/protolog/ProtoLogDataSource.java
+++ b/core/java/com/android/internal/protolog/ProtoLogDataSource.java
@@ -37,6 +37,7 @@
 import android.util.proto.ProtoInputStream;
 import android.util.proto.WireTypeMismatchException;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.common.LogLevel;
 
 import java.io.IOException;
@@ -48,6 +49,7 @@
 public class ProtoLogDataSource extends DataSource<ProtoLogDataSource.Instance,
         ProtoLogDataSource.TlsState,
         ProtoLogDataSource.IncrementalState> {
+    private static final String DATASOURCE_NAME = "android.protolog";
 
     private final Instance.TracingInstanceStartCallback mOnStart;
     private final Runnable mOnFlush;
@@ -55,7 +57,13 @@
 
     public ProtoLogDataSource(Instance.TracingInstanceStartCallback onStart, Runnable onFlush,
             Instance.TracingInstanceStopCallback onStop) {
-        super("android.protolog");
+        this(onStart, onFlush, onStop, DATASOURCE_NAME);
+    }
+
+    @VisibleForTesting
+    public ProtoLogDataSource(Instance.TracingInstanceStartCallback onStart, Runnable onFlush,
+            Instance.TracingInstanceStopCallback onStop, String dataSourceName) {
+        super(dataSourceName);
         this.mOnStart = onStart;
         this.mOnFlush = onFlush;
         this.mOnStop = onStop;
@@ -195,7 +203,7 @@
                     int defaultLogFromLevelInt = configStream.readInt(DEFAULT_LOG_FROM_LEVEL);
                     if (defaultLogFromLevelInt < defaultLogFromLevel.ordinal()) {
                         defaultLogFromLevel =
-                                logLevelFromInt(configStream.readInt(DEFAULT_LOG_FROM_LEVEL));
+                                logLevelFromInt(defaultLogFromLevelInt);
                     }
                     break;
                 case (int) TRACING_MODE:
diff --git a/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java b/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
index 38ca0d8..3b24f27 100644
--- a/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
+++ b/core/java/com/android/internal/protolog/ProtoLogViewerConfigReader.java
@@ -3,7 +3,6 @@
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.GROUPS;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.Group.ID;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.Group.NAME;
-
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MESSAGES;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.MESSAGE;
 import static android.internal.perfetto.protos.Protolog.ProtoLogViewerConfig.MessageData.MESSAGE_ID;
@@ -11,7 +10,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.util.Log;
 import android.util.LongSparseArray;
 import android.util.proto.ProtoInputStream;
 
@@ -38,6 +36,7 @@
      * Returns message format string for its hash or null if unavailable
      * or the viewer config is not loaded into memory.
      */
+    @Nullable
     public synchronized String getViewerString(long messageHash) {
         return mLogMessageMap.get(messageHash);
     }
diff --git a/core/java/com/android/internal/ravenwood/RavenwoodEnvironment.java b/core/java/com/android/internal/ravenwood/RavenwoodEnvironment.java
index ee3a3c2..319efe0 100644
--- a/core/java/com/android/internal/ravenwood/RavenwoodEnvironment.java
+++ b/core/java/com/android/internal/ravenwood/RavenwoodEnvironment.java
@@ -60,7 +60,9 @@
     public static void ensureRavenwoodInitialized() {
     }
 
-    private static native void ensureRavenwoodInitialized$ravenwood();
+    private static void ensureRavenwoodInitialized$ravenwood() {
+        nativeEnsureRavenwoodInitialized();
+    }
 
     /**
      * USE IT SPARINGLY! Returns true if it's running on Ravenwood, hostside test environment.
@@ -92,7 +94,9 @@
         throw notSupportedOnDevice();
     }
 
-    private native <T> T fromAddress$ravenwood(long address);
+    private <T> T fromAddress$ravenwood(long address) {
+        return nativeFromAddress(address);
+    }
 
     /**
      * See {@link Workaround}. It's only usable on Ravenwood.
@@ -114,7 +118,14 @@
         throw notSupportedOnDevice();
     }
 
-    private native String getRavenwoodRuntimePath$ravenwood();
+    private String getRavenwoodRuntimePath$ravenwood() {
+        return nativeGetRavenwoodRuntimePath();
+    }
+
+    // Private native methods that are actually substituted on Ravenwood
+    private native <T> T nativeFromAddress(long address);
+    private native String nativeGetRavenwoodRuntimePath();
+    private static native void nativeEnsureRavenwoodInitialized();
 
     /**
      * A set of APIs used to work around missing features on Ravenwood. Ideally, this class should
diff --git a/core/java/com/android/internal/statusbar/StatusBarIcon.java b/core/java/com/android/internal/statusbar/StatusBarIcon.java
index f8a1436..1938cdb 100644
--- a/core/java/com/android/internal/statusbar/StatusBarIcon.java
+++ b/core/java/com/android/internal/statusbar/StatusBarIcon.java
@@ -51,6 +51,19 @@
         ResourceIcon
     }
 
+    public enum Shape {
+        /**
+         * Icon view should use WRAP_CONTENT -- so that the horizontal space occupied depends on the
+         * icon's shape (skinny/fat icons take less/more). Most icons will want to use this option
+         * for a nicer-looking overall spacing in the status bar, as long as the icon is "known"
+         * (i.e. not coming from a 3P package).
+         */
+        WRAP_CONTENT,
+
+        /** Icon should always be displayed in a space as wide as the status bar is tall. */
+        FIXED_SPACE,
+    }
+
     public UserHandle user;
     public String pkg;
     public Icon icon;
@@ -59,6 +72,7 @@
     public int number;
     public CharSequence contentDescription;
     public Type type;
+    public Shape shape;
 
     /**
      * Optional {@link Drawable} corresponding to {@link #icon}. This field is not parcelable, so
@@ -68,7 +82,7 @@
     @Nullable public Drawable preloadedIcon;
 
     public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
-            CharSequence contentDescription, Type type) {
+            CharSequence contentDescription, Type type, Shape shape) {
         if (icon.getType() == Icon.TYPE_RESOURCE
                 && TextUtils.isEmpty(icon.getResPackage())) {
             // This is an odd situation where someone's managed to hand us an icon without a
@@ -83,6 +97,13 @@
         this.number = number;
         this.contentDescription = contentDescription;
         this.type = type;
+        this.shape = shape;
+    }
+
+    public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
+            CharSequence contentDescription, Type type) {
+        this(user, resPackage, icon, iconLevel, number, contentDescription, type,
+                Shape.WRAP_CONTENT);
     }
 
     public StatusBarIcon(String iconPackage, UserHandle user,
@@ -107,7 +128,7 @@
     @Override
     public StatusBarIcon clone() {
         StatusBarIcon that = new StatusBarIcon(this.user, this.pkg, this.icon,
-                this.iconLevel, this.number, this.contentDescription, this.type);
+                this.iconLevel, this.number, this.contentDescription, this.type, this.shape);
         that.visible = this.visible;
         that.preloadedIcon = this.preloadedIcon;
         return that;
@@ -129,6 +150,7 @@
         this.number = in.readInt();
         this.contentDescription = in.readCharSequence();
         this.type = Type.valueOf(in.readString());
+        this.shape = Shape.valueOf(in.readString());
     }
 
     public void writeToParcel(Parcel out, int flags) {
@@ -140,6 +162,7 @@
         out.writeInt(this.number);
         out.writeCharSequence(this.contentDescription);
         out.writeString(this.type.name());
+        out.writeString(this.shape.name());
     }
 
     public int describeContents() {
diff --git a/core/java/com/android/internal/view/IInputMethodManager.aidl b/core/java/com/android/internal/view/IInputMethodManager.aidl
index b51678e..efbf887 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -76,10 +76,10 @@
 
     boolean showSoftInput(in IInputMethodClient client, @nullable IBinder windowToken,
             in ImeTracker.Token statsToken, int flags, int lastClickToolType,
-            in @nullable ResultReceiver resultReceiver, int reason);
+            in @nullable ResultReceiver resultReceiver, int reason, boolean async);
     boolean hideSoftInput(in IInputMethodClient client, @nullable IBinder windowToken,
             in ImeTracker.Token statsToken, int flags,
-            in @nullable ResultReceiver resultReceiver, int reason);
+            in @nullable ResultReceiver resultReceiver, int reason, boolean async);
 
     /**
      * A test API for CTS to request hiding the current soft input window, with the request origin
@@ -120,7 +120,8 @@
             in @nullable EditorInfo editorInfo, in @nullable IRemoteInputConnection inputConnection,
             in @nullable IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, int userId,
-            in ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq);
+            in ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq,
+            boolean useAsyncShowHideMethod);
 
     void showInputMethodPickerFromClient(in IInputMethodClient client,
             int auxiliarySubtypeMode);
diff --git a/core/jni/Android.bp b/core/jni/Android.bp
index 2abdd57..90cb10a 100644
--- a/core/jni/Android.bp
+++ b/core/jni/Android.bp
@@ -108,6 +108,7 @@
         "libtracing_perfetto",
         "libharfbuzz_ng",
         "liblog",
+        "libmediautils",
         "libminikin",
         "libz",
         "server_configurable_flags",
diff --git a/core/jni/android_media_AudioSystem.cpp b/core/jni/android_media_AudioSystem.cpp
index 638591f..9bccf5a 100644
--- a/core/jni/android_media_AudioSystem.cpp
+++ b/core/jni/android_media_AudioSystem.cpp
@@ -17,6 +17,7 @@
 
 //#define LOG_NDEBUG 0
 
+#include <atomic>
 #define LOG_TAG "AudioSystem-JNI"
 #include <android/binder_ibinder_jni.h>
 #include <android/binder_libbinder.h>
@@ -34,15 +35,16 @@
 #include <media/AudioContainers.h>
 #include <media/AudioPolicy.h>
 #include <media/AudioSystem.h>
+#include <mediautils/jthread.h>
 #include <nativehelper/JNIHelp.h>
 #include <nativehelper/ScopedLocalRef.h>
 #include <nativehelper/ScopedPrimitiveArray.h>
 #include <nativehelper/jni_macros.h>
 #include <system/audio.h>
 #include <system/audio_policy.h>
+#include <sys/system_properties.h>
 #include <utils/Log.h>
 
-#include <thread>
 #include <optional>
 #include <sstream>
 #include <memory>
@@ -57,6 +59,7 @@
 #include "android_media_AudioMixerAttributes.h"
 #include "android_media_AudioProfile.h"
 #include "android_media_MicrophoneInfo.h"
+#include "android_media_JNIUtils.h"
 #include "android_util_Binder.h"
 #include "core_jni_helpers.h"
 
@@ -3375,42 +3378,48 @@
 class JavaSystemPropertyListener {
   public:
     JavaSystemPropertyListener(JNIEnv* env, jobject javaCallback, std::string sysPropName) :
-            mCallback(env->NewGlobalRef(javaCallback)),
-            mCachedProperty(android::base::CachedProperty{std::move(sysPropName)}) {
-        mListenerThread = std::thread([this]() mutable {
-            JNIEnv* threadEnv = GetOrAttachJNIEnvironment(gVm);
-            while (!mCleanupSignal.load()) {
-                using namespace std::chrono_literals;
-                // 1s timeout so this thread can read the cleanup signal to (slowly) be able to
-                // be destroyed.
-                std::string newVal = mCachedProperty.WaitForChange(1000ms) ?: "";
-                if (newVal != "" && mLastVal != newVal) {
-                    threadEnv->CallVoidMethod(mCallback, gRunnableClassInfo.run);
-                    mLastVal = std::move(newVal);
+            mCallback {javaCallback, env},
+            mSysPropName(sysPropName),
+            mCachedProperty(android::base::CachedProperty{std::move(sysPropName)}),
+            mListenerThread([this](mediautils::stop_token stok) mutable {
+                while (!stok.stop_requested()) {
+                    using namespace std::chrono_literals;
+                    // 1s timeout so this thread can eventually respond to the stop token
+                    std::string newVal = mCachedProperty.WaitForChange(1000ms) ?: "";
+                    updateValue(newVal);
                 }
-            }
-            });
-    }
+            }) {}
 
-    ~JavaSystemPropertyListener() {
-        mCleanupSignal.store(true);
-        mListenerThread.join();
-        JNIEnv* env = GetOrAttachJNIEnvironment(gVm);
-        env->DeleteGlobalRef(mCallback);
+    void triggerUpdateIfChanged() {
+        // We must check the property without using the cached property due to thread safety issues
+        std::string newVal = base::GetProperty(mSysPropName, "");
+        updateValue(newVal);
     }
 
   private:
-    jobject mCallback;
+    void updateValue(std::string newVal) {
+        if (newVal == "") return;
+        std::lock_guard l{mLock};
+        if (mLastVal == newVal) return;
+        const auto threadEnv = GetOrAttachJNIEnvironment(gVm);
+        threadEnv->CallVoidMethod(mCallback.get(), gRunnableClassInfo.run);
+        mLastVal = std::move(newVal);
+    }
+
+    // Should outlive thread object
+    const GlobalRef mCallback;
+    const std::string mSysPropName;
     android::base::CachedProperty mCachedProperty;
-    std::thread mListenerThread;
-    std::atomic<bool> mCleanupSignal{false};
     std::string mLastVal = "";
+    std::mutex mLock;
+    const mediautils::jthread mListenerThread;
 };
 
+// A logical set keyed by address
 std::vector<std::unique_ptr<JavaSystemPropertyListener>> gSystemPropertyListeners;
 std::mutex gSysPropLock{};
 
-static void android_media_AudioSystem_listenForSystemPropertyChange(JNIEnv *env,  jobject thiz,
+static jlong android_media_AudioSystem_listenForSystemPropertyChange(JNIEnv *env,  jobject thiz,
         jstring sysProp,
         jobject javaCallback) {
     ScopedUtfChars sysPropChars{env, sysProp};
@@ -3418,6 +3427,19 @@
             std::string{sysPropChars.c_str()});
     std::unique_lock _l{gSysPropLock};
     gSystemPropertyListeners.push_back(std::move(listener));
+    return reinterpret_cast<jlong>(gSystemPropertyListeners.back().get());
+}
+
+static void android_media_AudioSystem_triggerSystemPropertyUpdate(JNIEnv *env,  jobject thiz,
+        jlong nativeHandle) {
+    std::unique_lock _l{gSysPropLock};
+    const auto iter = std::find_if(gSystemPropertyListeners.begin(), gSystemPropertyListeners.end(),
+            [nativeHandle](const auto& x) { return reinterpret_cast<jlong>(x.get()) == nativeHandle; });
+    if (iter != gSystemPropertyListeners.end()) {
+        (*iter)->triggerUpdateIfChanged();
+    } else {
+        jniThrowException(env, "java/lang/IllegalArgumentException", "Invalid handle");
+    }
 }
 
 
@@ -3595,8 +3617,11 @@
          MAKE_AUDIO_SYSTEM_METHOD(setBluetoothVariableLatencyEnabled),
          MAKE_AUDIO_SYSTEM_METHOD(isBluetoothVariableLatencyEnabled),
          MAKE_JNI_NATIVE_METHOD("listenForSystemPropertyChange",
-                                "(Ljava/lang/String;Ljava/lang/Runnable;)V",
+                                "(Ljava/lang/String;Ljava/lang/Runnable;)J",
                                 android_media_AudioSystem_listenForSystemPropertyChange),
+         MAKE_JNI_NATIVE_METHOD("triggerSystemPropertyUpdate",
+                                "(J)V",
+                                android_media_AudioSystem_triggerSystemPropertyUpdate),
 
         };
 
diff --git a/core/jni/android_util_Binder.cpp b/core/jni/android_util_Binder.cpp
index 2068bd7..46b4695 100644
--- a/core/jni/android_util_Binder.cpp
+++ b/core/jni/android_util_Binder.cpp
@@ -466,10 +466,25 @@
 public:
     sp<JavaBBinder> get(JNIEnv* env, jobject obj)
     {
-        AutoMutex _l(mLock);
-        sp<JavaBBinder> b = mBinder.promote();
-        if (b == NULL) {
-            b = new JavaBBinder(env, obj);
+        sp<JavaBBinder> b;
+        {
+            AutoMutex _l(mLock);
+            // must take lock to promote because we set the same wp<>
+            // on another thread.
+            b = mBinder.promote();
+        }
+
+        if (b) return b;
+
+        // b/360067751: constructor may trigger GC, so call outside lock
+        b = new JavaBBinder(env, obj);
+
+        {
+            AutoMutex _l(mLock);
+            // if it was constructed on another thread in the meantime,
+            // return that. 'b' will just get destructed.
+            if (sp<JavaBBinder> b2 = mBinder.promote(); b2) return b2;
+
             if (mVintf) {
                 ::android::internal::Stability::markVintf(b.get());
             }
diff --git a/core/jni/com_android_internal_content_NativeLibraryHelper.cpp b/core/jni/com_android_internal_content_NativeLibraryHelper.cpp
index fba0d81..7ad18b8 100644
--- a/core/jni/com_android_internal_content_NativeLibraryHelper.cpp
+++ b/core/jni/com_android_internal_content_NativeLibraryHelper.cpp
@@ -17,6 +17,7 @@
 #define LOG_TAG "NativeLibraryHelper"
 //#define LOG_NDEBUG 0
 
+#include <android-base/properties.h>
 #include <androidfw/ApkParsing.h>
 #include <androidfw/ZipFileRO.h>
 #include <androidfw/ZipUtils.h>
@@ -36,6 +37,7 @@
 #include <zlib.h>
 
 #include <memory>
+#include <string>
 
 #include "com_android_internal_content_FileSystemUtils.h"
 #include "core_jni_helpers.h"
@@ -125,72 +127,10 @@
     return INSTALL_SUCCEEDED;
 }
 
-/*
- * Copy the native library if needed.
- *
- * This function assumes the library and path names passed in are considered safe.
- */
-static install_status_t
-copyFileIfChanged(JNIEnv *env, void* arg, ZipFileRO* zipFile, ZipEntryRO zipEntry, const char* fileName)
-{
-    static const size_t kPageSize = getpagesize();
-    void** args = reinterpret_cast<void**>(arg);
-    jstring* javaNativeLibPath = (jstring*) args[0];
-    jboolean extractNativeLibs = *(jboolean*) args[1];
-    jboolean debuggable = *(jboolean*) args[2];
-
-    ScopedUtfChars nativeLibPath(env, *javaNativeLibPath);
-
-    uint32_t uncompLen;
-    uint32_t when;
-    uint32_t crc;
-
-    uint16_t method;
-    off64_t offset;
-    uint16_t extraFieldLength;
-    if (!zipFile->getEntryInfo(zipEntry, &method, &uncompLen, nullptr, &offset, &when, &crc,
-                               &extraFieldLength)) {
-        ALOGE("Couldn't read zip entry info\n");
-        return INSTALL_FAILED_INVALID_APK;
-    }
-
-    // Always extract wrap.sh for debuggable, even if extractNativeLibs=false. This makes it
-    // easier to use wrap.sh because it only works when it is extracted, see
-    // frameworks/base/services/core/java/com/android/server/am/ProcessList.java.
-    bool forceExtractCurrentFile = debuggable && strcmp(fileName, "wrap.sh") == 0;
-
-    if (!extractNativeLibs && !forceExtractCurrentFile) {
-        // check if library is uncompressed and page-aligned
-        if (method != ZipFileRO::kCompressStored) {
-            ALOGE("Library '%s' is compressed - will not be able to open it directly from apk.\n",
-                fileName);
-            return INSTALL_FAILED_INVALID_APK;
-        }
-
-        if (offset % kPageSize != 0) {
-            ALOGE("Library '%s' is not PAGE(%zu)-aligned - will not be able to open it directly "
-                  "from apk.\n", fileName, kPageSize);
-            return INSTALL_FAILED_INVALID_APK;
-        }
-
-#ifdef ENABLE_PUNCH_HOLES
-        // if library is uncompressed, punch hole in it in place
-        if (!punchHolesInElf64(zipFile->getZipFileName(), offset)) {
-            ALOGW("Failed to punch uncompressed elf file :%s inside apk : %s at offset: "
-                  "%" PRIu64 "",
-                  fileName, zipFile->getZipFileName(), offset);
-        }
-
-        // if extra field for this zip file is present with some length, possibility is that it is
-        // padding added for zip alignment. Punch holes there too.
-        if (!punchHolesInZip(zipFile->getZipFileName(), offset, extraFieldLength)) {
-            ALOGW("Failed to punch apk : %s at extra field", zipFile->getZipFileName());
-        }
-#endif // ENABLE_PUNCH_HOLES
-
-        return INSTALL_SUCCEEDED;
-    }
-
+static install_status_t extractNativeLibFromApk(ZipFileRO* zipFile, ZipEntryRO zipEntry,
+                                                const char* fileName,
+                                                const std::string nativeLibPath, uint32_t when,
+                                                uint32_t uncompLen, uint32_t crc) {
     // Build local file path
     const size_t fileNameLen = strlen(fileName);
     char localFileName[nativeLibPath.size() + fileNameLen + 2];
@@ -313,6 +253,88 @@
 }
 
 /*
+ * Copy the native library if needed.
+ *
+ * This function assumes the library and path names passed in are considered safe.
+ */
+static install_status_t copyFileIfChanged(JNIEnv* env, void* arg, ZipFileRO* zipFile,
+                                          ZipEntryRO zipEntry, const char* fileName) {
+    static const size_t kPageSize = getpagesize();
+    void** args = reinterpret_cast<void**>(arg);
+    jstring* javaNativeLibPath = (jstring*)args[0];
+    jboolean extractNativeLibs = *(jboolean*)args[1];
+    jboolean debuggable = *(jboolean*)args[2];
+    jboolean app_compat_16kb = *(jboolean*)args[3];
+    install_status_t ret = INSTALL_SUCCEEDED;
+
+    ScopedUtfChars nativeLibPath(env, *javaNativeLibPath);
+
+    uint32_t uncompLen;
+    uint32_t when;
+    uint32_t crc;
+
+    uint16_t method;
+    off64_t offset;
+    uint16_t extraFieldLength;
+    if (!zipFile->getEntryInfo(zipEntry, &method, &uncompLen, nullptr, &offset, &when, &crc,
+                               &extraFieldLength)) {
+        ALOGE("Couldn't read zip entry info\n");
+        return INSTALL_FAILED_INVALID_APK;
+    }
+
+    // Always extract wrap.sh for debuggable, even if extractNativeLibs=false. This makes it
+    // easier to use wrap.sh because it only works when it is extracted, see
+    // frameworks/base/services/core/java/com/android/server/am/ProcessList.java.
+    bool forceExtractCurrentFile = debuggable && strcmp(fileName, "wrap.sh") == 0;
+
+    if (!extractNativeLibs && !forceExtractCurrentFile) {
+        // check if library is uncompressed and page-aligned
+        if (method != ZipFileRO::kCompressStored) {
+            ALOGE("Library '%s' is compressed - will not be able to open it directly from apk.\n",
+                  fileName);
+            return INSTALL_FAILED_INVALID_APK;
+        }
+
+        if (offset % kPageSize != 0) {
+            // If the library is zip-aligned correctly for 4kb devices and app compat is
+            // enabled, on 16kb devices fallback to extraction
+            if (offset % 0x1000 == 0 && app_compat_16kb) {
+                ALOGI("16kB AppCompat: Library '%s' is not PAGE(%zu)-aligned - falling back to "
+                      "extraction from apk\n",
+                      fileName, kPageSize);
+                return extractNativeLibFromApk(zipFile, zipEntry, fileName, nativeLibPath.c_str(),
+                                               when, uncompLen, crc);
+            }
+
+            ALOGE("Library '%s' is not PAGE(%zu)-aligned - will not be able to open it directly "
+                  "from apk.\n",
+                  fileName, kPageSize);
+            return INSTALL_FAILED_INVALID_APK;
+        }
+
+#ifdef ENABLE_PUNCH_HOLES
+        // if library is uncompressed, punch hole in it in place
+        if (!punchHolesInElf64(zipFile->getZipFileName(), offset)) {
+            ALOGW("Failed to punch uncompressed elf file :%s inside apk : %s at offset: "
+                  "%" PRIu64 "",
+                  fileName, zipFile->getZipFileName(), offset);
+        }
+
+        // if extra field for this zip file is present with some length, possibility is that it is
+        // padding added for zip alignment. Punch holes there too.
+        if (!punchHolesInZip(zipFile->getZipFileName(), offset, extraFieldLength)) {
+            ALOGW("Failed to punch apk : %s at extra field", zipFile->getZipFileName());
+        }
+#endif // ENABLE_PUNCH_HOLES
+
+        return INSTALL_SUCCEEDED;
+    }
+
+    return extractNativeLibFromApk(zipFile, zipEntry, fileName, nativeLibPath.c_str(), when,
+                                   uncompLen, crc);
+}
+
+/*
  * An iterator over all shared libraries in a zip file. An entry is
  * considered to be a shared library if all of the conditions below are
  * satisfied :
@@ -498,12 +520,24 @@
     return status;
 }
 
+static inline bool app_compat_16kb_enabled() {
+    static const size_t kPageSize = getpagesize();
+
+    // App compat is only applicable on 16kb-page-size devices.
+    if (kPageSize != 0x4000) {
+        return false;
+    }
+
+    return android::base::GetBoolProperty("bionic.linker.16kb.app_compat.enabled", false);
+}
+
 static jint
 com_android_internal_content_NativeLibraryHelper_copyNativeBinaries(JNIEnv *env, jclass clazz,
         jlong apkHandle, jstring javaNativeLibPath, jstring javaCpuAbi,
         jboolean extractNativeLibs, jboolean debuggable)
 {
-    void* args[] = { &javaNativeLibPath, &extractNativeLibs, &debuggable };
+    jboolean app_compat_16kb = app_compat_16kb_enabled();
+    void* args[] = { &javaNativeLibPath, &extractNativeLibs, &debuggable, &app_compat_16kb };
     return (jint) iterateOverNativeFiles(env, apkHandle, javaCpuAbi, debuggable,
             copyFileIfChanged, reinterpret_cast<void*>(args));
 }
diff --git a/core/jni/platform/host/HostRuntime.cpp b/core/jni/platform/host/HostRuntime.cpp
index 020b27e..19f8299 100644
--- a/core/jni/platform/host/HostRuntime.cpp
+++ b/core/jni/platform/host/HostRuntime.cpp
@@ -391,6 +391,7 @@
 
 } // namespace android
 
+#ifndef _WIN32
 using namespace android;
 
 JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void*) {
@@ -407,3 +408,4 @@
 
     return JNI_VERSION_1_6;
 }
+#endif
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index 91c3370..5decf7f 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -843,6 +843,7 @@
     <protected-broadcast android:name="android.app.action.CONSOLIDATED_NOTIFICATION_POLICY_CHANGED" />
     <protected-broadcast android:name="android.intent.action.MAIN_USER_LOCKSCREEN_KNOWLEDGE_FACTOR_CHANGED" />
     <protected-broadcast android:name="com.android.uwb.uwbcountrycode.GEOCODE_RETRY" />
+    <protected-broadcast android:name="android.telephony.action.ACTION_SATELLITE_SUBSCRIBER_ID_LIST_CHANGED" />
 
     <!-- ====================================================================== -->
     <!--                          RUNTIME PERMISSIONS                           -->
@@ -3765,7 +3766,6 @@
         privileged app such as the Assistant app.
         <p>Protection level: internal|role
         <p>Intended for use by the DEVICE_POLICY_MANAGEMENT role only.
-        @FlaggedApi("android.app.admin.flags.assist_content_user_restriction_enabled")
     -->
     <permission android:name="android.permission.MANAGE_DEVICE_POLICY_ASSIST_CONTENT"
         android:protectionLevel="internal|role" />
@@ -4026,7 +4026,6 @@
         APIs protected by this permission on users different to the calling user.
         <p>Protection level: internal|role
         <p>Intended for use by the DEVICE_POLICY_MANAGEMENT role only.
-        @FlaggedApi("android.app.admin.flags.esim_management_enabled")
     -->
     <permission android:name="android.permission.MANAGE_DEVICE_POLICY_MANAGED_SUBSCRIPTIONS"
         android:protectionLevel="internal|role" />
@@ -7649,7 +7648,8 @@
     <permission android:name="android.permission.BIND_CARRIER_MESSAGING_CLIENT_SERVICE"
         android:protectionLevel="signature" />
 
-    <!-- Must be required by an {@link android.service.watchdog.ExplicitHealthCheckService} to
+    <!-- @FlaggedApi(android.crashrecovery.flags.Flags.FLAG_ENABLE_CRASHRECOVERY) @SystemApi
+         Must be required by an {@link android.service.watchdog.ExplicitHealthCheckService} to
          ensure that only the system can bind to it.
          @hide This is not a third-party API (intended for OEMs and system apps).
     -->
diff --git a/core/res/res/drawable/ic_zen_priority_modes.xml b/core/res/res/drawable/ic_zen_priority_modes.xml
new file mode 100644
index 0000000..e8691fc
--- /dev/null
+++ b/core/res/res/drawable/ic_zen_priority_modes.xml
@@ -0,0 +1,25 @@
+<!--
+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.
+-->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?android:attr/colorControlNormal">
+  <path
+      android:fillColor="@android:color/white"
+      android:pathData="M480,720Q580,720 650,650Q720,580 720,480Q720,380 650,310Q580,240 480,240L480,480L310,650Q345,683 388.5,701.5Q432,720 480,720ZM480,880Q397,880 324,848.5Q251,817 197,763Q143,709 111.5,636Q80,563 80,480Q80,397 111.5,324Q143,251 197,197Q251,143 324,111.5Q397,80 480,80Q563,80 636,111.5Q709,143 763,197Q817,251 848.5,324Q880,397 880,480Q880,563 848.5,636Q817,709 763,763Q709,817 636,848.5Q563,880 480,880ZM480,800Q614,800 707,707Q800,614 800,480Q800,346 707,253Q614,160 480,160Q346,160 253,253Q160,346 160,480Q160,614 253,707Q346,800 480,800ZM480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Q480,480 480,480Z"/>
+</vector>
diff --git a/core/res/res/layout/autofill_dataset_picker_header_footer.xml b/core/res/res/layout/autofill_dataset_picker_header_footer.xml
index 4d5f4f0..027f530 100644
--- a/core/res/res/layout/autofill_dataset_picker_header_footer.xml
+++ b/core/res/res/layout/autofill_dataset_picker_header_footer.xml
@@ -37,6 +37,7 @@
       <ListView
           android:id="@+id/autofill_dataset_list"
           android:layout_weight="1"
+          android:fadeScrollbars="false"
           android:layout_width="fill_parent"
           android:layout_height="0dp"
           android:drawSelectorOnTop="true"
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index b6468ee..4beeb17 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -7030,6 +7030,10 @@
          {@link InputDevice#SOURCE_ROTARY_ENCODER}s. -->
     <bool name="config_viewBasedRotaryEncoderHapticsEnabled">false</bool>
 
+    <!-- Whether the scroll haptic feedback implementation is enabled for
+    {@link InputDevice#SOURCE_TOUCHSCREEN}s. -->
+    <bool name="config_viewTouchScreenHapticScrollFeedbackEnabled">false</bool>
+
     <!-- Whether the media player is shown on the quick settings -->
     <bool name="config_quickSettingsShowMediaPlayer">true</bool>
 
diff --git a/core/res/res/values/config_telephony.xml b/core/res/res/values/config_telephony.xml
index a7240ff..69437b4 100644
--- a/core/res/res/values/config_telephony.xml
+++ b/core/res/res/values/config_telephony.xml
@@ -440,10 +440,6 @@
     <string name="config_satellite_carrier_roaming_esos_provisioned_class" translatable="false"></string>
     <java-symbol type="string" name="config_satellite_carrier_roaming_esos_provisioned_class" />
 
-    <!-- Telephony satellite gateway intent for handling carrier roaming to satellite is using ESOS messaging. -->
-    <string name="config_satellite_carrier_roaming_esos_provisioned_intent_action" translatable="false"></string>
-    <java-symbol type="string" name="config_satellite_carrier_roaming_esos_provisioned_intent_action" />
-
     <!-- The time duration in minutes to wait before retry validating a possible change
          in satellite allowed region. The default value is 10 minutes. -->
     <integer name="config_satellite_delay_minutes_before_retry_validating_possible_change_in_allowed_region">10</integer>
diff --git a/core/res/res/values/styles.xml b/core/res/res/values/styles.xml
index c084b4c..dc99634 100644
--- a/core/res/res/values/styles.xml
+++ b/core/res/res/values/styles.xml
@@ -1504,26 +1504,26 @@
 
     <!-- @hide -->
     <style name="PointerIconVectorStyleFillGreen">
-        <item name="pointerIconVectorFill">#6DD58C</item>
-        <item name="pointerIconVectorFillInverse">#6DD58C</item>
+        <item name="pointerIconVectorFill">#1AA64A</item>
+        <item name="pointerIconVectorFillInverse">#1AA64A</item>
     </style>
 
     <!-- @hide -->
     <style name="PointerIconVectorStyleFillYellow">
-        <item name="pointerIconVectorFill">#FDD663</item>
-        <item name="pointerIconVectorFillInverse">#FDD663</item>
+        <item name="pointerIconVectorFill">#F55E57</item>
+        <item name="pointerIconVectorFillInverse">#F55E57</item>
     </style>
 
     <!-- @hide -->
     <style name="PointerIconVectorStyleFillPink">
-        <item name="pointerIconVectorFill">#F2B8B5</item>
-        <item name="pointerIconVectorFillInverse">#F2B8B5</item>
+        <item name="pointerIconVectorFill">#F94AAB</item>
+        <item name="pointerIconVectorFillInverse">#F94AAB</item>
     </style>
 
     <!-- @hide -->
     <style name="PointerIconVectorStyleFillBlue">
-        <item name="pointerIconVectorFill">#8AB4F8</item>
-        <item name="pointerIconVectorFillInverse">#8AB4F8</item>
+        <item name="pointerIconVectorFill">#009DC9</item>
+        <item name="pointerIconVectorFillInverse">#009DC9</item>
     </style>
 
     <!-- @hide -->
@@ -1535,7 +1535,7 @@
     <!-- @hide -->
     <style name="PointerIconVectorStyleStrokeBlack">
         <item name="pointerIconVectorStroke">@color/black</item>
-        <item name="pointerIconVectorStrokeInverse">@color/white</item>
+        <item name="pointerIconVectorStrokeInverse">@color/black</item>
     </style>
 
     <!-- @hide -->
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index bbe661e..80ec67a 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -5490,8 +5490,11 @@
   <java-symbol type="bool" name="config_tvExternalInputLoggingDisplayNameFilterEnabled" />
   <java-symbol type="array" name="config_tvExternalInputLoggingDeviceOnScreenDisplayNames" />
   <java-symbol type="array" name="config_tvExternalInputLoggingDeviceBrandNames" />
+
+  <!-- Scroll Feedback Configs -->
   <java-symbol type="bool" name="config_viewRotaryEncoderHapticScrollFedbackEnabled" />
   <java-symbol type="bool" name="config_viewBasedRotaryEncoderHapticsEnabled" />
+  <java-symbol type="bool" name="config_viewTouchScreenHapticScrollFeedbackEnabled" />
 
   <java-symbol type="bool" name="config_quickSettingsShowMediaPlayer" />
 
@@ -5566,6 +5569,7 @@
   <java-symbol type="string"  name="recs_notification_channel_label"/>
 
   <!-- Priority Modes icons -->
+  <java-symbol type="drawable" name="ic_zen_priority_modes" />
   <java-symbol type="drawable" name="ic_zen_mode_type_bedtime" />
   <java-symbol type="drawable" name="ic_zen_mode_type_driving" />
   <java-symbol type="drawable" name="ic_zen_mode_type_immersive" />
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/AndroidManifest.xml b/core/tests/batterystatstests/BatteryStatsViewer/AndroidManifest.xml
index 94bde68..127dbfd 100644
--- a/core/tests/batterystatstests/BatteryStatsViewer/AndroidManifest.xml
+++ b/core/tests/batterystatstests/BatteryStatsViewer/AndroidManifest.xml
@@ -20,6 +20,7 @@
 
     <uses-permission android:name="android.permission.ACCESS_NETWORK_STATE"/>
     <uses-permission android:name="android.permission.BATTERY_STATS"/>
+    <uses-permission android:name="android.permission.CHANGE_COMPONENT_ENABLED_STATE"/>
     <uses-permission android:name="android.permission.INTERACT_ACROSS_USERS"/>
     <uses-permission android:name="android.permission.QUERY_ALL_PACKAGES"/>
 
@@ -31,7 +32,8 @@
         <activity android:name=".BatteryConsumerPickerActivity"
                   android:label="Battery Stats"
                   android:launchMode="singleTop"
-                  android:exported="true">
+                  android:exported="true"
+                  android:enabled="false">
             <intent-filter>
                 <action android:name="android.intent.action.MAIN"/>
                 <category android:name="android.intent.category.LAUNCHER"/>
@@ -41,5 +43,25 @@
         <activity android:name=".BatteryStatsViewerActivity"
                   android:label="Battery Stats"
                   android:parentActivityName=".BatteryConsumerPickerActivity"/>
+
+        <activity android:name=".TrampolineActivity"
+            android:exported="true"
+            android:theme="@android:style/Theme.NoDisplay">
+            <intent-filter>
+                <action android:name="com.android.settings.action.IA_SETTINGS"/>
+                <category android:name="android.intent.category.DEFAULT" />
+            </intent-filter>
+
+            <meta-data android:name="com.android.settings.category"
+                android:value="com.android.settings.category.ia.development" />
+            <meta-data android:name="com.android.settings.title"
+                android:resource="@string/settings_title" />
+            <meta-data android:name="com.android.settings.summary"
+                android:resource="@string/settings_summary" />
+            <meta-data android:name="com.android.settings.group_key"
+                android:value="debug_debugging_category" />
+            <meta-data android:name="com.android.settings.order"
+                android:value="2" />
+        </activity>
     </application>
 </manifest>
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/res/values/strings.xml b/core/tests/batterystatstests/BatteryStatsViewer/res/values/strings.xml
new file mode 100644
index 0000000..c23c148
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryStatsViewer/res/values/strings.xml
@@ -0,0 +1,20 @@
+<!--
+  ~ 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 xmlns:xliff="urn:oasis:names:tc:xliff:document:1.2">
+    <string name="settings_title">Launch Battery Stats Viewer</string>
+    <string name="settings_summary">The Battery Stats Viewer will be visible in the Launcher after it is opened once.</string>
+</resources>
diff --git a/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/TrampolineActivity.java b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/TrampolineActivity.java
new file mode 100644
index 0000000..b016488
--- /dev/null
+++ b/core/tests/batterystatstests/BatteryStatsViewer/src/com/android/frameworks/core/batterystatsviewer/TrampolineActivity.java
@@ -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.frameworks.core.batterystatsviewer;
+
+import android.app.Activity;
+import android.content.ComponentName;
+import android.content.Intent;
+import android.content.pm.PackageManager;
+import android.os.Bundle;
+
+import androidx.annotation.Nullable;
+
+public class TrampolineActivity extends Activity {
+
+    @Override
+    protected void onCreate(@Nullable Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+        showLauncherIcon();
+        launchMainActivity();
+    }
+
+    private void showLauncherIcon() {
+        PackageManager pm = getPackageManager();
+        pm.setComponentEnabledSetting(new ComponentName(this, BatteryConsumerPickerActivity.class),
+                PackageManager.COMPONENT_ENABLED_STATE_ENABLED,
+                PackageManager.DONT_KILL_APP);
+    }
+
+    private void launchMainActivity() {
+        startActivity(new Intent(this, BatteryConsumerPickerActivity.class));
+    }
+}
diff --git a/core/tests/coretests/Android.bp b/core/tests/coretests/Android.bp
index edf461a..b0e48f1 100644
--- a/core/tests/coretests/Android.bp
+++ b/core/tests/coretests/Android.bp
@@ -803,3 +803,11 @@
     include_annotations: ["android.platform.test.annotations.PlatinumTest"],
     exclude_annotations: FLAKY_OR_IGNORED,
 }
+
+test_module_config {
+    name: "FrameworksCoreTests_android_tracing",
+    base: "FrameworksCoreTests",
+    team: "trendy_team_windowing_tools",
+    test_suites: ["device-tests"],
+    include_filters: ["android.tracing"],
+}
diff --git a/core/tests/coretests/src/android/os/storage/StorageManagerIntegrationTest.java b/core/tests/coretests/src/android/os/storage/StorageManagerIntegrationTest.java
index ecd2f76..b157c95 100644
--- a/core/tests/coretests/src/android/os/storage/StorageManagerIntegrationTest.java
+++ b/core/tests/coretests/src/android/os/storage/StorageManagerIntegrationTest.java
@@ -16,8 +16,11 @@
 
 package android.os.storage;
 
+import android.content.res.ObbInfo;
+import android.os.Parcel;
 import android.os.ParcelFileDescriptor;
 import android.os.ProxyFileDescriptorCallback;
+import android.os.ServiceManager;
 import android.system.ErrnoException;
 
 import androidx.test.filters.LargeTest;
@@ -104,7 +107,14 @@
     public void testMountBadPackageNameObb() throws Exception {
         final File file = createObbFile(OBB_FILE_3_BAD_PACKAGENAME, R.raw.obb_file3_bad_packagename);
         String filePath = file.getAbsolutePath();
-        mountObb(filePath, OnObbStateChangeListener.ERROR_PERMISSION_DENIED);
+        try {
+            mountObb(filePath, OnObbStateChangeListener.ERROR_PERMISSION_DENIED);
+            fail("mountObb should throw an exception as package name is incorrect");
+        } catch (Exception ex) {
+            assertEquals("Path " + filePath
+                            + " does not contain package name " + mContext.getPackageName(),
+                    ex.getMessage());
+        }
     }
 
     /**
@@ -154,6 +164,48 @@
         }
     }
 
+    @LargeTest
+    public void testObbInfo_withValidObbInfo_success() throws Exception {
+        final File file = createObbFile(OBB_FILE_1, R.raw.obb_file1);
+        String filePath = file.getAbsolutePath();
+        try {
+            mountObb(filePath);
+            unmountObb(filePath, DONT_FORCE);
+        } catch (Exception ex) {
+            fail("No exception expected, got " + ex.getMessage());
+        }
+    }
+
+    @LargeTest
+    public void testObbInfo_withInvalidObbInfo_exception() throws Exception {
+        final File file = createObbFile(OBB_FILE_1, R.raw.obb_file1);
+        String rawPath = file.getAbsolutePath();
+        String canonicalPath = file.getCanonicalPath();
+
+        ObbInfo obbInfo = ObbInfo.CREATOR.createFromParcel(Parcel.obtain());
+        obbInfo.packageName = "com.android.obbcrash";
+        obbInfo.version = 1;
+        obbInfo.filename = canonicalPath;
+
+        try {
+            IStorageManager.Stub.asInterface(ServiceManager.getServiceOrThrow("mount")).mountObb(
+                    rawPath, canonicalPath, new ObbActionListener(), 0, obbInfo);
+            fail("mountObb should throw an exception as package name is incorrect");
+        } catch (SecurityException ex) {
+            assertEquals("Path " + canonicalPath
+                            + " does not contain package name " + mContext.getPackageName(),
+                    ex.getMessage());
+        }
+    }
+
+    private static class ObbActionListener extends IObbActionListener.Stub {
+        @SuppressWarnings("hiding")
+        @Override
+        public void onObbResult(String filename, int nonce, int status) {
+
+        }
+    }
+
     private static class MyThreadFactory implements ThreadFactory {
         Thread thread = null;
 
diff --git a/core/tests/coretests/src/android/tracing/TEST_MAPPING b/core/tests/coretests/src/android/tracing/TEST_MAPPING
new file mode 100644
index 0000000..4b7adf9
--- /dev/null
+++ b/core/tests/coretests/src/android/tracing/TEST_MAPPING
@@ -0,0 +1,8 @@
+{
+  "postsubmit": [
+    {
+       "name": "FrameworksCoreTests_android_tracing",
+        "file_patterns": [".*\\.java"]
+    }
+  ]
+}
diff --git a/core/tests/coretests/src/android/view/ImeBackAnimationControllerTest.java b/core/tests/coretests/src/android/view/ImeBackAnimationControllerTest.java
index 4d9b591c..00ffda8 100644
--- a/core/tests/coretests/src/android/view/ImeBackAnimationControllerTest.java
+++ b/core/tests/coretests/src/android/view/ImeBackAnimationControllerTest.java
@@ -254,11 +254,8 @@
         float progress = 0.5f;
         mBackAnimationController.onBackProgressed(new BackEvent(100f, 0f, progress, EDGE_LEFT));
         // verify correct ime insets manipulation
-        float interpolatedProgress = BACK_GESTURE.getInterpolation(progress);
-        int expectedInset =
-                (int) (IME_HEIGHT - interpolatedProgress * PEEK_FRACTION * IME_HEIGHT);
         verify(mWindowInsetsAnimationController, times(1)).setInsetsAndAlpha(
-                eq(Insets.of(0, 0, 0, expectedInset)), eq(1f), anyFloat());
+                eq(Insets.of(0, 0, 0, getImeHeight(progress))), eq(1f), anyFloat());
     }
 
     @Test
@@ -268,12 +265,13 @@
             WindowInsetsAnimationControlListener animationControlListener = startBackGesture();
 
             // progress back gesture
-            mBackAnimationController.onBackProgressed(new BackEvent(100f, 0f, 0.5f, EDGE_LEFT));
+            float progress = 0.5f;
+            mBackAnimationController.onBackProgressed(new BackEvent(100f, 0f, progress, EDGE_LEFT));
 
             // commit back gesture
             mBackAnimationController.onBackInvoked();
 
-            // verify setInsetsAndAlpha never called due onReady delayed
+            // verify setInsetsAndAlpha never called due to onReady delayed
             verify(mWindowInsetsAnimationController, never()).setInsetsAndAlpha(any(), anyInt(),
                     anyFloat());
             verify(mInsetsController, never()).setPredictiveBackImeHideAnimInProgress(eq(true));
@@ -283,7 +281,7 @@
 
             // verify setInsetsAndAlpha immediately called
             verify(mWindowInsetsAnimationController, times(1)).setInsetsAndAlpha(
-                    eq(Insets.of(0, 0, 0, IME_HEIGHT)), eq(1f), anyFloat());
+                    eq(Insets.of(0, 0, 0, getImeHeight(progress))), eq(1f), anyFloat());
             // verify post-commit hide anim has started
             verify(mInsetsController, times(1)).setPredictiveBackImeHideAnimInProgress(eq(true));
         });
@@ -319,4 +317,9 @@
 
         return animationControlListener.getValue();
     }
+
+    private int getImeHeight(float gestureProgress) {
+        float interpolatedProgress = BACK_GESTURE.getInterpolation(gestureProgress);
+        return (int) (IME_HEIGHT - interpolatedProgress * PEEK_FRACTION * IME_HEIGHT);
+    }
 }
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index ce7e858..bec8b1f 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -1022,7 +1022,7 @@
     }
 
     @Test
-    public void testImeRequestedVisibleDuringPredictiveBackAnim() {
+    public void testImeRequestedVisibleDuringPredictiveBackAnimWithoutCallback() {
         prepareControls();
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
             // show ime as initial state
@@ -1051,6 +1051,42 @@
     }
 
     @Test
+    public void testImeRequestedInvisibleDuringPredictiveBackAnimWithCallback() {
+        prepareControls();
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            // set WindowInsetsAnimationCallback on ViewRoot
+            mViewRoot.getView().setWindowInsetsAnimationCallback(
+                    new WindowInsetsAnimation.Callback(DISPATCH_MODE_STOP) {
+                        @Override
+                        public WindowInsets onProgress(
+                                @NonNull WindowInsets insets,
+                                @NonNull List<WindowInsetsAnimation> runningAnimations) {
+                            return insets;
+                        }
+                    });
+
+            // show ime as initial state
+            mController.show(ime(), true /* fromIme */, ImeTracker.Token.empty());
+            mController.cancelExistingAnimations(); // fast forward show animation
+            assertTrue(mController.getState().peekSource(ID_IME).isVisible());
+
+            // start control request (for predictive back animation)
+            WindowInsetsAnimationControlListener listener =
+                    mock(WindowInsetsAnimationControlListener.class);
+            mController.controlWindowInsetsAnimation(ime(), /*cancellationSignal*/ null,
+                    listener, /*fromIme*/ false, /*duration*/ -1, /*interpolator*/ null,
+                    ANIMATION_TYPE_USER, /*fromPredictiveBack*/ true);
+
+            // Verify that onReady is called (after next predraw)
+            mViewRoot.getView().getViewTreeObserver().dispatchOnPreDraw();
+            verify(listener).onReady(notNull(), eq(ime()));
+
+            // verify that insets are requested invisible during animation
+            assertFalse(isRequestedVisible(mController, ime()));
+        });
+    }
+
+    @Test
     public void testImeShowRequestCancelsPredictiveBackPostCommitAnim() {
         prepareControls();
         InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
@@ -1131,6 +1167,37 @@
         });
     }
 
+    @Test
+    public void testPredictiveBackControlRequestCancelledDuringImeHideAnim() {
+        prepareControls();
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            // show ime as initial state
+            if (!Flags.refactorInsetsController()) {
+                mController.show(ime(), true /* fromIme */, ImeTracker.Token.empty());
+            } else {
+                mController.show(ime(), false /* fromIme */, ImeTracker.Token.empty());
+            }
+            mController.cancelExistingAnimations(); // fast forward show animation
+            mViewRoot.getView().getViewTreeObserver().dispatchOnPreDraw();
+            assertTrue(mController.getState().peekSource(ID_IME).isVisible());
+
+            // start IME hide animation
+            mController.hide(ime(), true /* fromIme */, null /* statsToken */);
+            assertEquals(ANIMATION_TYPE_HIDE, mController.getAnimationType(ime()));
+
+            // start control request (for predictive back animation)
+            WindowInsetsAnimationControlListener listener =
+                    mock(WindowInsetsAnimationControlListener.class);
+            mController.controlWindowInsetsAnimation(ime(), /*cancellationSignal*/ null,
+                    listener, /*fromIme*/ false, /*duration*/ -1, /*interpolator*/ null,
+                    ANIMATION_TYPE_USER, /*fromPredictiveBack*/ true);
+
+            // verify that control request is cancelled and animation type remains HIDE
+            verify(listener).onCancelled(any());
+            assertEquals(ANIMATION_TYPE_HIDE, mController.getAnimationType(ime()));
+        });
+    }
+
     private void waitUntilNextFrame() throws Exception {
         final CountDownLatch latch = new CountDownLatch(1);
         Choreographer.getMainThreadInstance().postCallback(Choreographer.CALLBACK_COMMIT,
diff --git a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
index b183ecb..149e132 100644
--- a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
+++ b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
@@ -20,6 +20,7 @@
 
 import android.graphics.Color;
 import android.graphics.drawable.ColorDrawable;
+import android.graphics.drawable.Icon;
 import android.os.Parcel;
 import android.os.UserHandle;
 
@@ -69,22 +70,22 @@
         assertThat(copy.preloadedIcon).isEqualTo(original.preloadedIcon);
     }
 
-
     private static StatusBarIcon newStatusBarIcon() {
         final UserHandle dummyUserHandle = UserHandle.of(100);
         final String dummyIconPackageName = "com.android.internal.statusbar.test";
-        final int dummyIconId = 123;
+        final Icon dummyIcon = Icon.createWithResource(dummyIconPackageName, 123);
         final int dummyIconLevel = 1;
         final int dummyIconNumber = 2;
         final CharSequence dummyIconContentDescription = "dummyIcon";
         return new StatusBarIcon(
-                dummyIconPackageName,
                 dummyUserHandle,
-                dummyIconId,
+                dummyIconPackageName,
+                dummyIcon,
                 dummyIconLevel,
                 dummyIconNumber,
                 dummyIconContentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.FIXED_SPACE);
     }
 
     private static void assertSerializableFieldsEqual(StatusBarIcon copy, StatusBarIcon original) {
@@ -96,6 +97,7 @@
         assertThat(copy.number).isEqualTo(original.number);
         assertThat(copy.contentDescription).isEqualTo(original.contentDescription);
         assertThat(copy.type).isEqualTo(original.type);
+        assertThat(copy.shape).isEqualTo(original.shape);
     }
 
     private static StatusBarIcon parcelAndUnparcel(StatusBarIcon original) {
diff --git a/data/keyboards/Vendor_0957_Product_0033.idc b/data/keyboards/Vendor_0957_Product_0033.idc
new file mode 100644
index 0000000..7dfbe2c
--- /dev/null
+++ b/data/keyboards/Vendor_0957_Product_0033.idc
@@ -0,0 +1,23 @@
+# Copyright 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.
+#
+
+# Input Device Configuration file for Google Reference RCU Remote.
+# PID 0033 is for new G20 with start button.
+
+# Basic Parameters
+keyboard.layout = Vendor_0957_Product_0031
+# The reason why we set is follow https://docs.partner.android.com/tv/build/gtv/boot-resume
+keyboard.doNotWakeByDefault = 1
+audio.mic = 1
diff --git a/graphics/java/android/graphics/Matrix44.java b/graphics/java/android/graphics/Matrix44.java
index a99e201..683f614 100644
--- a/graphics/java/android/graphics/Matrix44.java
+++ b/graphics/java/android/graphics/Matrix44.java
@@ -19,6 +19,7 @@
 import android.annotation.FlaggedApi;
 import android.annotation.IntRange;
 import android.annotation.NonNull;
+import android.ravenwood.annotation.RavenwoodKeepWholeClass;
 
 import com.android.graphics.hwui.flags.Flags;
 
@@ -30,6 +31,7 @@
  * in row-major order. The values and operations are treated as column vectors.
  */
 @FlaggedApi(Flags.FLAG_MATRIX_44)
+@RavenwoodKeepWholeClass
 public class Matrix44 {
     final float[] mBackingArray;
     /**
diff --git a/graphics/java/android/graphics/Outline.java b/graphics/java/android/graphics/Outline.java
index 618e6dc..c7b8941 100644
--- a/graphics/java/android/graphics/Outline.java
+++ b/graphics/java/android/graphics/Outline.java
@@ -21,6 +21,7 @@
 import android.annotation.NonNull;
 import android.compat.annotation.UnsupportedAppUsage;
 import android.graphics.drawable.Drawable;
+import android.ravenwood.annotation.RavenwoodKeepWholeClass;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -35,6 +36,7 @@
  * @see android.view.View#setOutlineProvider(android.view.ViewOutlineProvider)
  * @see Drawable#getOutline(Outline)
  */
+@RavenwoodKeepWholeClass
 public final class Outline {
     private static final float RADIUS_UNDEFINED = Float.NEGATIVE_INFINITY;
 
diff --git a/graphics/java/android/graphics/ParcelableColorSpace.java b/graphics/java/android/graphics/ParcelableColorSpace.java
index 748d66c..76c17154 100644
--- a/graphics/java/android/graphics/ParcelableColorSpace.java
+++ b/graphics/java/android/graphics/ParcelableColorSpace.java
@@ -20,6 +20,7 @@
 import android.annotation.Nullable;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.ravenwood.annotation.RavenwoodKeepWholeClass;
 
 /**
  * A {@link Parcelable} wrapper for a {@link ColorSpace}. In order to enable parceling, the
@@ -27,6 +28,7 @@
  * {@link ColorSpace.Rgb} instance that has an ICC parametric transfer function as returned by
  * {@link ColorSpace.Rgb#getTransferParameters()}.
  */
+@RavenwoodKeepWholeClass
 public final class ParcelableColorSpace implements Parcelable {
     private final ColorSpace mColorSpace;
 
diff --git a/graphics/java/android/graphics/PixelFormat.java b/graphics/java/android/graphics/PixelFormat.java
index 3ec5b9c..a872e03 100644
--- a/graphics/java/android/graphics/PixelFormat.java
+++ b/graphics/java/android/graphics/PixelFormat.java
@@ -17,10 +17,12 @@
 package android.graphics;
 
 import android.annotation.IntDef;
+import android.ravenwood.annotation.RavenwoodKeepWholeClass;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 
+@RavenwoodKeepWholeClass
 public class PixelFormat {
     /** @hide */
     @IntDef({UNKNOWN, TRANSLUCENT, TRANSPARENT, OPAQUE})
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
index 69a68c8..4ce2942 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/BackupHelper.java
@@ -24,6 +24,7 @@
 
 import androidx.annotation.NonNull;
 
+import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -57,7 +58,7 @@
     void scheduleBackup() {
         if (!mBackupIdlerScheduled) {
             mBackupIdlerScheduled = true;
-            Looper.myQueue().addIdleHandler(mBackupIdler);
+            Looper.getMainLooper().getQueue().addIdleHandler(mBackupIdler);
         }
     }
 
@@ -80,9 +81,14 @@
         }
 
         if (DEBUG) Log.d(TAG, "Start to back up " + taskContainers);
+        final List<ParcelableTaskContainerData> parcelableTaskContainerDataList = new ArrayList<>(
+                taskContainers.size());
+        for (TaskContainer taskContainer : taskContainers) {
+            parcelableTaskContainerDataList.add(taskContainer.getParcelableData());
+        }
         final Bundle state = new Bundle();
-        state.setClassLoader(TaskContainer.class.getClassLoader());
-        state.putParcelableList(KEY_TASK_CONTAINERS, taskContainers);
+        state.setClassLoader(ParcelableTaskContainerData.class.getClassLoader());
+        state.putParcelableList(KEY_TASK_CONTAINERS, parcelableTaskContainerDataList);
         mController.setSavedState(state);
     }
 
@@ -91,10 +97,12 @@
             return;
         }
 
-        final List<TaskContainer> taskContainers = savedState.getParcelableArrayList(
-                KEY_TASK_CONTAINERS, TaskContainer.class);
-        for (TaskContainer taskContainer : taskContainers) {
-            if (DEBUG) Log.d(TAG, "restore task " + taskContainer.getTaskId());
+        final List<ParcelableTaskContainerData> parcelableTaskContainerDataList =
+                savedState.getParcelableArrayList(KEY_TASK_CONTAINERS,
+                        ParcelableTaskContainerData.class);
+        for (ParcelableTaskContainerData data : parcelableTaskContainerDataList) {
+            final TaskContainer taskContainer = new TaskContainer(data, mController);
+            if (DEBUG) Log.d(TAG, "Restoring task " + taskContainer.getTaskId());
             // TODO(b/289875940): implement the TaskContainer restoration.
         }
     }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
index 1eb95c1..9ea2943 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizer.java
@@ -70,6 +70,10 @@
     @NonNull
     private final TaskFragmentCallback mCallback;
 
+    @VisibleForTesting
+    @Nullable
+    TaskFragmentAnimationController mAnimationController;
+
     /**
      * Callback that notifies the controller about changes to task fragments.
      */
@@ -87,6 +91,25 @@
         mCallback = callback;
     }
 
+    @Override
+    public void unregisterOrganizer() {
+        if (mAnimationController != null) {
+            mAnimationController.unregisterRemoteAnimations();
+            mAnimationController = null;
+        }
+        super.unregisterOrganizer();
+    }
+
+    /**
+     * Overrides the animation for transitions of embedded activities organized by this organizer.
+     */
+    void overrideSplitAnimation() {
+        if (mAnimationController == null) {
+            mAnimationController = new TaskFragmentAnimationController(this);
+        }
+        mAnimationController.registerRemoteAnimations();
+    }
+
     /**
      * Starts a new Activity and puts it into split with an existing Activity side-by-side.
      * @param launchingFragmentToken    token for the launching TaskFragment. If it exists, it will
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
new file mode 100644
index 0000000..817cfce
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableSplitContainerData.java
@@ -0,0 +1,116 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import android.os.IBinder;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+/**
+ * This class holds the Parcelable data of a {@link SplitContainer}.
+ */
+class ParcelableSplitContainerData implements Parcelable {
+
+    /**
+     * A reference to the target {@link SplitContainer} that owns the data. This will not be
+     * parcelled and will be {@code null} when the data is created from a parcel.
+     */
+    @Nullable
+    final SplitContainer mSplitContainer;
+
+    @NonNull
+    final IBinder mToken;
+
+    @NonNull
+    private final IBinder mPrimaryContainerToken;
+
+    @NonNull
+    private final IBinder mSecondaryContainerToken;
+
+    // TODO(b/289875940): making this as non-null once the tag can be auto-generated from the rule.
+    @Nullable
+    final String mSplitRuleTag;
+
+    /**
+     * Whether the selection of which container is primary can be changed at runtime. Runtime
+     * updates is currently possible only for {@link SplitPinContainer}
+     *
+     * @see SplitPinContainer
+     */
+    final boolean mIsPrimaryContainerMutable;
+
+    ParcelableSplitContainerData(@NonNull SplitContainer splitContainer, @NonNull IBinder token,
+            @NonNull IBinder primaryContainerToken, @NonNull IBinder secondaryContainerToken,
+            @Nullable String splitRuleTag, boolean isPrimaryContainerMutable) {
+        mSplitContainer = splitContainer;
+        mToken = token;
+        mPrimaryContainerToken = primaryContainerToken;
+        mSecondaryContainerToken = secondaryContainerToken;
+        mSplitRuleTag = splitRuleTag;
+        mIsPrimaryContainerMutable = isPrimaryContainerMutable;
+    }
+
+    private ParcelableSplitContainerData(Parcel in) {
+        mSplitContainer = null;
+        mToken = in.readStrongBinder();
+        mPrimaryContainerToken = in.readStrongBinder();
+        mSecondaryContainerToken = in.readStrongBinder();
+        mSplitRuleTag = in.readString();
+        mIsPrimaryContainerMutable = in.readBoolean();
+    }
+
+    public static final Creator<ParcelableSplitContainerData> CREATOR = new Creator<>() {
+        @Override
+        public ParcelableSplitContainerData createFromParcel(Parcel in) {
+            return new ParcelableSplitContainerData(in);
+        }
+
+        @Override
+        public ParcelableSplitContainerData[] newArray(int size) {
+            return new ParcelableSplitContainerData[size];
+        }
+    };
+
+    @NonNull
+    private IBinder getPrimaryContainerToken() {
+        return mSplitContainer != null ? mSplitContainer.getPrimaryContainer().getToken()
+                : mPrimaryContainerToken;
+    }
+
+    @NonNull
+    private IBinder getSecondaryContainerToken() {
+        return mSplitContainer != null ? mSplitContainer.getSecondaryContainer().getToken()
+                : mSecondaryContainerToken;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeStrongBinder(mToken);
+        dest.writeStrongBinder(getPrimaryContainerToken());
+        dest.writeStrongBinder(getSecondaryContainerToken());
+        dest.writeString(mSplitRuleTag);
+        dest.writeBoolean(mIsPrimaryContainerMutable);
+    }
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
new file mode 100644
index 0000000..7377d00
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskContainerData.java
@@ -0,0 +1,127 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static android.app.ActivityTaskManager.INVALID_TASK_ID;
+
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * This class holds the Parcelable data of a {@link TaskContainer}.
+ */
+class ParcelableTaskContainerData implements Parcelable {
+
+    /**
+     * A reference to the target {@link TaskContainer} that owns the data. This will not be
+     * parcelled and will be {@code null} when the data is created from a parcel.
+     */
+    @Nullable
+    final TaskContainer mTaskContainer;
+
+    /**
+     * The unique task id.
+     */
+    final int mTaskId;
+
+    /**
+     * The parcelable data of the active TaskFragmentContainers in this Task.
+     * Note that this will only be populated before parcelling, and will not be copied when
+     * making a new instance copy.
+     */
+    @NonNull
+    private final List<ParcelableTaskFragmentContainerData>
+            mParcelableTaskFragmentContainerDataList = new ArrayList<>();
+
+    /**
+     * The parcelable data of the SplitContainers in this Task.
+     * Note that this will only be populated before parcelling, and will not be copied when
+     * making a new instance copy.
+     */
+    @NonNull
+    private final List<ParcelableSplitContainerData> mParcelableSplitContainerDataList =
+            new ArrayList<>();
+
+    ParcelableTaskContainerData(int taskId, @NonNull TaskContainer taskContainer) {
+        if (taskId == INVALID_TASK_ID) {
+            throw new IllegalArgumentException("Invalid Task id");
+        }
+
+        mTaskId = taskId;
+        mTaskContainer = taskContainer;
+    }
+
+    ParcelableTaskContainerData(@NonNull ParcelableTaskContainerData data,
+            @NonNull TaskContainer taskContainer) {
+        mTaskId = data.mTaskId;
+        mTaskContainer = taskContainer;
+    }
+
+    private ParcelableTaskContainerData(Parcel in) {
+        mTaskId = in.readInt();
+        mTaskContainer = null;
+        in.readParcelableList(mParcelableTaskFragmentContainerDataList,
+                ParcelableTaskFragmentContainerData.class.getClassLoader(),
+                ParcelableTaskFragmentContainerData.class);
+        in.readParcelableList(mParcelableSplitContainerDataList,
+                ParcelableSplitContainerData.class.getClassLoader(),
+                ParcelableSplitContainerData.class);
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeInt(mTaskId);
+        dest.writeParcelableList(getParcelableTaskFragmentContainerDataList(), flags);
+        dest.writeParcelableList(getParcelableSplitContainerDataList(), flags);
+    }
+
+    @NonNull
+    List<? extends ParcelableTaskFragmentContainerData>
+            getParcelableTaskFragmentContainerDataList() {
+        return mTaskContainer != null ? mTaskContainer.getParcelableTaskFragmentContainerDataList()
+                : mParcelableTaskFragmentContainerDataList;
+    }
+
+    @NonNull
+    List<? extends ParcelableSplitContainerData> getParcelableSplitContainerDataList() {
+        return mTaskContainer != null ? mTaskContainer.getParcelableSplitContainerDataList()
+                : mParcelableSplitContainerDataList;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    public static final Creator<ParcelableTaskContainerData> CREATOR = new Creator<>() {
+        @Override
+        public ParcelableTaskContainerData createFromParcel(Parcel in) {
+            return new ParcelableTaskContainerData(in);
+        }
+
+        @Override
+        public ParcelableTaskContainerData[] newArray(int size) {
+            return new ParcelableTaskContainerData[size];
+        }
+    };
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskFragmentContainerData.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskFragmentContainerData.java
new file mode 100644
index 0000000..a79a89a
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/ParcelableTaskFragmentContainerData.java
@@ -0,0 +1,105 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import android.app.Activity;
+import android.graphics.Rect;
+import android.os.IBinder;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+/**
+ * This class holds the Parcelable data of a {@link TaskFragmentContainer}.
+ */
+class ParcelableTaskFragmentContainerData implements Parcelable {
+
+    /**
+     * Client-created token that uniquely identifies the task fragment container instance.
+     */
+    @NonNull
+    final IBinder mToken;
+
+    /**
+     * The tag specified in launch options. {@code null} if this taskFragment container is not an
+     * overlay container.
+     */
+    @Nullable
+    final String mOverlayTag;
+
+    /**
+     * The associated {@link Activity#getActivityToken()} of the overlay container.
+     * Must be {@code null} for non-overlay container.
+     * <p>
+     * If an overlay container is associated with an activity, this overlay container will be
+     * dismissed when the associated activity is destroyed. If the overlay container is visible,
+     * activity will be launched on top of the overlay container and expanded to fill the parent
+     * container.
+     */
+    @Nullable
+    final IBinder mAssociatedActivityToken;
+
+    /**
+     * Bounds that were requested last via {@link android.window.WindowContainerTransaction}.
+     */
+    @NonNull
+    final Rect mLastRequestedBounds;
+
+    ParcelableTaskFragmentContainerData(@NonNull IBinder token, @Nullable String overlayTag,
+            @Nullable IBinder associatedActivityToken) {
+        mToken = token;
+        mOverlayTag = overlayTag;
+        mAssociatedActivityToken = associatedActivityToken;
+        mLastRequestedBounds = new Rect();
+    }
+
+    private ParcelableTaskFragmentContainerData(Parcel in) {
+        mToken = in.readStrongBinder();
+        mOverlayTag = in.readString();
+        mAssociatedActivityToken = in.readStrongBinder();
+        mLastRequestedBounds = in.readTypedObject(Rect.CREATOR);
+    }
+
+    public static final Creator<ParcelableTaskFragmentContainerData> CREATOR = new Creator<>() {
+        @Override
+        public ParcelableTaskFragmentContainerData createFromParcel(Parcel in) {
+            return new ParcelableTaskFragmentContainerData(in);
+        }
+
+        @Override
+        public ParcelableTaskFragmentContainerData[] newArray(int size) {
+            return new ParcelableTaskFragmentContainerData[size];
+        }
+    };
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel dest, int flags) {
+        dest.writeStrongBinder(mToken);
+        dest.writeString(mOverlayTag);
+        dest.writeStrongBinder(mAssociatedActivityToken);
+        dest.writeTypedObject(mLastRequestedBounds, flags);
+    }
+
+}
+
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
index 39cface..6d436ec 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitContainer.java
@@ -33,6 +33,8 @@
  */
 class SplitContainer {
     @NonNull
+    private final ParcelableSplitContainerData mParcelableData;
+    @NonNull
     private TaskFragmentContainer mPrimaryContainer;
     @NonNull
     private final TaskFragmentContainer mSecondaryContainer;
@@ -44,16 +46,6 @@
     /** @see SplitContainer#getDefaultSplitAttributes() */
     @NonNull
     private SplitAttributes mDefaultSplitAttributes;
-    @NonNull
-    private final IBinder mToken;
-
-    /**
-     * Whether the selection of which container is primary can be changed at runtime. Runtime
-     * updates is currently possible only for {@link SplitPinContainer}
-     *
-     * @see SplitPinContainer
-     */
-    private final boolean mIsPrimaryContainerMutable;
 
     SplitContainer(@NonNull TaskFragmentContainer primaryContainer,
             @NonNull Activity primaryActivity,
@@ -69,13 +61,14 @@
             @NonNull TaskFragmentContainer secondaryContainer,
             @NonNull SplitRule splitRule,
             @NonNull SplitAttributes splitAttributes, boolean isPrimaryContainerMutable) {
+        mParcelableData = new ParcelableSplitContainerData(this, new Binder("SplitContainer"),
+                primaryContainer.getToken(), secondaryContainer.getToken(), splitRule.getTag(),
+                isPrimaryContainerMutable);
         mPrimaryContainer = primaryContainer;
         mSecondaryContainer = secondaryContainer;
         mSplitRule = splitRule;
         mDefaultSplitAttributes = splitRule.getDefaultSplitAttributes();
         mCurrentSplitAttributes = splitAttributes;
-        mToken = new Binder("SplitContainer");
-        mIsPrimaryContainerMutable = isPrimaryContainerMutable;
 
         if (shouldFinishPrimaryWithSecondary(splitRule)) {
             if (mPrimaryContainer.getRunningActivityCount() == 1
@@ -94,7 +87,7 @@
     }
 
     void setPrimaryContainer(@NonNull TaskFragmentContainer primaryContainer) {
-        if (!mIsPrimaryContainerMutable) {
+        if (!mParcelableData.mIsPrimaryContainerMutable) {
             throw new IllegalStateException("Cannot update primary TaskFragmentContainer");
         }
         mPrimaryContainer = primaryContainer;
@@ -150,7 +143,12 @@
 
     @NonNull
     IBinder getToken() {
-        return mToken;
+        return mParcelableData.mToken;
+    }
+
+    @NonNull
+    ParcelableSplitContainerData getParcelableData() {
+        return mParcelableData;
     }
 
     /**
@@ -201,7 +199,7 @@
             return null;
         }
         return new SplitInfo(primaryActivityStack, secondaryActivityStack,
-                mCurrentSplitAttributes, SplitInfo.Token.createFromBinder(mToken));
+                mCurrentSplitAttributes, SplitInfo.Token.createFromBinder(mParcelableData.mToken));
     }
 
     static boolean shouldFinishPrimaryWithSecondary(@NonNull SplitRule splitRule) {
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 24b56ae..f2f2b7ea 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -58,18 +58,22 @@
 import android.app.ActivityClient;
 import android.app.ActivityOptions;
 import android.app.ActivityThread;
+import android.app.AppGlobals;
 import android.app.Application;
 import android.app.Instrumentation;
 import android.app.servertransaction.ClientTransactionListenerController;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
+import android.content.pm.PackageManager;
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
+import android.os.RemoteException;
+import android.os.SystemProperties;
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Log;
@@ -116,6 +120,7 @@
 public class SplitController implements JetpackTaskFragmentOrganizer.TaskFragmentCallback,
         ActivityEmbeddingComponent, DividerPresenter.DragEventCallback {
     static final String TAG = "SplitController";
+    static final boolean ENABLE_SHELL_TRANSITIONS = getShellTransitEnabled();
 
     // TODO(b/243518738): Move to WM Extensions if we have requirement of overlay without
     //  association. It's not set in WM Extensions nor Wm Jetpack library currently.
@@ -919,7 +924,8 @@
 
         // Update all TaskFragments in the Task. Make a copy of the list since some may be
         // removed on updating.
-        final List<TaskFragmentContainer> containers = taskContainer.getTaskFragmentContainers();
+        final List<TaskFragmentContainer> containers
+                = new ArrayList<>(taskContainer.getTaskFragmentContainers());
         for (int i = containers.size() - 1; i >= 0; i--) {
             final TaskFragmentContainer container = containers.get(i);
             // Wait until onTaskFragmentAppeared to update new container.
@@ -3307,4 +3313,17 @@
             transactionRecord.apply(false /* shouldApplyIndependently */);
         }
     }
+
+    // TODO(b/207070762): cleanup with legacy app transition
+    private static boolean getShellTransitEnabled() {
+        try {
+            if (AppGlobals.getPackageManager().hasSystemFeature(
+                    PackageManager.FEATURE_AUTOMOTIVE, 0)) {
+                return SystemProperties.getBoolean("persist.wm.debug.shell_transit", true);
+            }
+        } catch (RemoteException re) {
+            Log.w(TAG, "Error getting system features");
+        }
+        return true;
+    }
 }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index fb8efc4..abc7b29 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -169,12 +169,17 @@
         mController = controller;
         final Bundle outSavedState = new Bundle();
         if (Flags.aeBackStackRestore()) {
-            outSavedState.setClassLoader(TaskContainer.class.getClassLoader());
+            outSavedState.setClassLoader(ParcelableTaskContainerData.class.getClassLoader());
             registerOrganizer(false /* isSystemOrganizer */, outSavedState);
         } else {
             registerOrganizer();
         }
         mBackupHelper = new BackupHelper(controller, outSavedState);
+        if (!SplitController.ENABLE_SHELL_TRANSITIONS) {
+            // TODO(b/207070762): cleanup with legacy app transition
+            // Animation will be handled by WM Shell when Shell transition is enabled.
+            overrideSplitAnimation();
+        }
     }
 
     void scheduleBackup() {
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
index 5795e8d..608a3be 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskContainer.java
@@ -16,7 +16,6 @@
 
 package androidx.window.extensions.embedding;
 
-import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
@@ -32,8 +31,6 @@
 import android.content.res.Configuration;
 import android.graphics.Rect;
 import android.os.IBinder;
-import android.os.Parcel;
-import android.os.Parcelable;
 import android.util.ArraySet;
 import android.util.DisplayMetrics;
 import android.util.Log;
@@ -57,11 +54,12 @@
 import java.util.Set;
 
 /** Represents TaskFragments and split pairs below a Task. */
-class TaskContainer implements Parcelable {
+class TaskContainer {
     private static final String TAG = TaskContainer.class.getSimpleName();
 
-    /** The unique task id. */
-    private final int mTaskId;
+    /** Parcelable data of this TaskContainer. */
+    @NonNull
+    private final ParcelableTaskContainerData mParcelableTaskContainerData;
 
     /** Active TaskFragments in this Task. */
     @NonNull
@@ -130,16 +128,15 @@
      * @param splitController The {@link SplitController}.
      */
     TaskContainer(int taskId, @NonNull Activity activityInTask,
-            @Nullable SplitController splitController) {
-        if (taskId == INVALID_TASK_ID) {
-            throw new IllegalArgumentException("Invalid Task id");
-        }
-        mTaskId = taskId;
+            @NonNull SplitController splitController) {
+        mParcelableTaskContainerData = new ParcelableTaskContainerData(taskId, this);
+
         final TaskProperties taskProperties = TaskProperties
                 .getTaskPropertiesFromActivity(activityInTask);
         mInfo = new TaskFragmentParentInfo(
                 taskProperties.getConfiguration(),
                 taskProperties.getDisplayId(),
+                taskId,
                 // Note that it is always called when there's a new Activity is started, which
                 // implies the host task is visible and has an activity in the task.
                 true /* visible */,
@@ -148,8 +145,44 @@
         mSplitController = splitController;
     }
 
+    /** This is only used when restoring it from a {@link ParcelableTaskContainerData}. */
+    TaskContainer(@NonNull ParcelableTaskContainerData data,
+            @NonNull SplitController splitController) {
+        mParcelableTaskContainerData = new ParcelableTaskContainerData(data, this);
+        mSplitController = splitController;
+        for (ParcelableTaskFragmentContainerData tfData :
+                data.getParcelableTaskFragmentContainerDataList()) {
+            final TaskFragmentContainer container =
+                    new TaskFragmentContainer(tfData, splitController, this);
+            mContainers.add(container);
+        }
+    }
+
+    @NonNull
+    ParcelableTaskContainerData getParcelableData() {
+        return mParcelableTaskContainerData;
+    }
+
+    @NonNull
+    List<ParcelableTaskFragmentContainerData> getParcelableTaskFragmentContainerDataList() {
+        final List<ParcelableTaskFragmentContainerData> data = new ArrayList<>(mContainers.size());
+        for (TaskFragmentContainer container : mContainers) {
+            data.add(container.getParcelableData());
+        }
+        return data;
+    }
+
+    @NonNull
+    List<ParcelableSplitContainerData> getParcelableSplitContainerDataList() {
+        final List<ParcelableSplitContainerData> data = new ArrayList<>(mSplitContainers.size());
+        for (SplitContainer splitContainer : mSplitContainers) {
+            data.add(splitContainer.getParcelableData());
+        }
+        return data;
+    }
+
     int getTaskId() {
-        return mTaskId;
+        return mParcelableTaskContainerData.mTaskId;
     }
 
     int getDisplayId() {
@@ -162,7 +195,8 @@
 
     void setInvisible() {
         mInfo = new TaskFragmentParentInfo(mInfo.getConfiguration(), mInfo.getDisplayId(),
-                false /* visible */, mInfo.hasDirectActivity(), mInfo.getDecorSurface());
+                mInfo.getTaskId(), false /* visible */, mInfo.hasDirectActivity(),
+                mInfo.getDecorSurface());
     }
 
     boolean hasDirectActivity() {
@@ -680,34 +714,6 @@
         return activityStacks;
     }
 
-    @Override
-    public int describeContents() {
-        return 0;
-    }
-
-    @Override
-    public void writeToParcel(Parcel dest, int flags) {
-        dest.writeInt(mTaskId);
-        // TODO(b/289875940)
-    }
-
-    protected TaskContainer(Parcel in) {
-        mTaskId = in.readInt();
-        // TODO(b/289875940)
-    }
-
-    public static final Creator<TaskContainer> CREATOR = new Creator<>() {
-        @Override
-        public TaskContainer createFromParcel(Parcel in) {
-            return new TaskContainer(in);
-        }
-
-        @Override
-        public TaskContainer[] newArray(int size) {
-            return new TaskContainer[size];
-        }
-    };
-
     /** A wrapper class which contains the information of {@link TaskContainer} */
     static final class TaskProperties {
         private final int mDisplayId;
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationAdapter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationAdapter.java
new file mode 100644
index 0000000..33220c4
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationAdapter.java
@@ -0,0 +1,224 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static android.graphics.Matrix.MTRANS_X;
+import static android.graphics.Matrix.MTRANS_Y;
+import static android.view.RemoteAnimationTarget.MODE_CLOSING;
+
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.view.Choreographer;
+import android.view.RemoteAnimationTarget;
+import android.view.SurfaceControl;
+import android.view.animation.Animation;
+import android.view.animation.Transformation;
+
+import androidx.annotation.NonNull;
+
+/**
+ * Wrapper to handle the TaskFragment animation update in one {@link SurfaceControl.Transaction}.
+ *
+ * The base adapter can be used for {@link RemoteAnimationTarget} that is simple open/close.
+ */
+class TaskFragmentAnimationAdapter {
+
+    /**
+     * If {@link #mOverrideLayer} is set to this value, we don't want to override the surface layer.
+     */
+    private static final int LAYER_NO_OVERRIDE = -1;
+
+    @NonNull
+    final Animation mAnimation;
+    @NonNull
+    final RemoteAnimationTarget mTarget;
+    @NonNull
+    final SurfaceControl mLeash;
+    /** Area in absolute coordinate that the animation surface shouldn't go beyond. */
+    @NonNull
+    private final Rect mWholeAnimationBounds = new Rect();
+    /**
+     * Area in absolute coordinate that should represent all the content to show for this window.
+     * This should be the end bounds for opening window, and start bounds for closing window in case
+     * the window is resizing during the open/close transition.
+     */
+    @NonNull
+    private final Rect mContentBounds = new Rect();
+    /** Offset relative to the window parent surface for {@link #mContentBounds}. */
+    @NonNull
+    private final Point mContentRelOffset = new Point();
+
+    @NonNull
+    final Transformation mTransformation = new Transformation();
+    @NonNull
+    final float[] mMatrix = new float[9];
+    @NonNull
+    final float[] mVecs = new float[4];
+    @NonNull
+    final Rect mRect = new Rect();
+    private boolean mIsFirstFrame = true;
+    private int mOverrideLayer = LAYER_NO_OVERRIDE;
+
+    TaskFragmentAnimationAdapter(@NonNull Animation animation,
+            @NonNull RemoteAnimationTarget target) {
+        this(animation, target, target.leash, target.screenSpaceBounds);
+    }
+
+    /**
+     * @param leash the surface to animate.
+     * @param wholeAnimationBounds  area in absolute coordinate that the animation surface shouldn't
+     *                              go beyond.
+     */
+    TaskFragmentAnimationAdapter(@NonNull Animation animation,
+            @NonNull RemoteAnimationTarget target, @NonNull SurfaceControl leash,
+            @NonNull Rect wholeAnimationBounds) {
+        mAnimation = animation;
+        mTarget = target;
+        mLeash = leash;
+        mWholeAnimationBounds.set(wholeAnimationBounds);
+        if (target.mode == MODE_CLOSING) {
+            // When it is closing, we want to show the content at the start position in case the
+            // window is resizing as well. For example, when the activities is changing from split
+            // to stack, the bottom TaskFragment will be resized to fullscreen when hiding.
+            final Rect startBounds = target.startBounds;
+            final Rect endBounds = target.screenSpaceBounds;
+            mContentBounds.set(startBounds);
+            mContentRelOffset.set(target.localBounds.left, target.localBounds.top);
+            mContentRelOffset.offset(
+                    startBounds.left - endBounds.left,
+                    startBounds.top - endBounds.top);
+        } else {
+            mContentBounds.set(target.screenSpaceBounds);
+            mContentRelOffset.set(target.localBounds.left, target.localBounds.top);
+        }
+    }
+
+    /**
+     * Surface layer to be set at the first frame of the animation. We will not set the layer if it
+     * is set to {@link #LAYER_NO_OVERRIDE}.
+     */
+    final void overrideLayer(int layer) {
+        mOverrideLayer = layer;
+    }
+
+    /** Called on frame update. */
+    final void onAnimationUpdate(@NonNull SurfaceControl.Transaction t, long currentPlayTime) {
+        if (mIsFirstFrame) {
+            t.show(mLeash);
+            if (mOverrideLayer != LAYER_NO_OVERRIDE) {
+                t.setLayer(mLeash, mOverrideLayer);
+            }
+            mIsFirstFrame = false;
+        }
+
+        // Extract the transformation to the current time.
+        mAnimation.getTransformation(Math.min(currentPlayTime, mAnimation.getDuration()),
+                mTransformation);
+        t.setFrameTimelineVsync(Choreographer.getInstance().getVsyncId());
+        onAnimationUpdateInner(t);
+    }
+
+    /** To be overridden by subclasses to adjust the animation surface change. */
+    void onAnimationUpdateInner(@NonNull SurfaceControl.Transaction t) {
+        // Update the surface position and alpha.
+        mTransformation.getMatrix().postTranslate(mContentRelOffset.x, mContentRelOffset.y);
+        t.setMatrix(mLeash, mTransformation.getMatrix(), mMatrix);
+        t.setAlpha(mLeash, mTransformation.getAlpha());
+
+        // Get current surface bounds in absolute coordinate.
+        // positionX/Y are in local coordinate, so minus the local offset to get the slide amount.
+        final int positionX = Math.round(mMatrix[MTRANS_X]);
+        final int positionY = Math.round(mMatrix[MTRANS_Y]);
+        final Rect cropRect = new Rect(mContentBounds);
+        cropRect.offset(positionX - mContentRelOffset.x, positionY - mContentRelOffset.y);
+
+        // Store the current offset of the surface top left from (0,0) in absolute coordinate.
+        final int offsetX = cropRect.left;
+        final int offsetY = cropRect.top;
+
+        // Intersect to make sure the animation happens within the whole animation bounds.
+        if (!cropRect.intersect(mWholeAnimationBounds)) {
+            // Hide the surface when it is outside of the animation area.
+            t.setAlpha(mLeash, 0);
+        }
+
+        // cropRect is in absolute coordinate, so we need to translate it to surface top left.
+        cropRect.offset(-offsetX, -offsetY);
+        t.setCrop(mLeash, cropRect);
+    }
+
+    /** Called after animation finished. */
+    final void onAnimationEnd(@NonNull SurfaceControl.Transaction t) {
+        onAnimationUpdate(t, mAnimation.getDuration());
+    }
+
+    final long getDurationHint() {
+        return mAnimation.computeDurationHint();
+    }
+
+    /**
+     * Should be used for the animation of the snapshot of a {@link RemoteAnimationTarget} that has
+     * size change.
+     */
+    static class SnapshotAdapter extends TaskFragmentAnimationAdapter {
+
+        SnapshotAdapter(@NonNull Animation animation, @NonNull RemoteAnimationTarget target) {
+            // Start leash is the snapshot of the starting surface.
+            super(animation, target, target.startLeash, target.screenSpaceBounds);
+        }
+
+        @Override
+        void onAnimationUpdateInner(@NonNull SurfaceControl.Transaction t) {
+            // Snapshot should always be placed at the top left of the animation leash.
+            mTransformation.getMatrix().postTranslate(0, 0);
+            t.setMatrix(mLeash, mTransformation.getMatrix(), mMatrix);
+            t.setAlpha(mLeash, mTransformation.getAlpha());
+        }
+    }
+
+    /**
+     * Should be used for the animation of the {@link RemoteAnimationTarget} that has size change.
+     */
+    static class BoundsChangeAdapter extends TaskFragmentAnimationAdapter {
+
+        BoundsChangeAdapter(@NonNull Animation animation, @NonNull RemoteAnimationTarget target) {
+            super(animation, target);
+        }
+
+        @Override
+        void onAnimationUpdateInner(@NonNull SurfaceControl.Transaction t) {
+            mTransformation.getMatrix().postTranslate(
+                    mTarget.localBounds.left, mTarget.localBounds.top);
+            t.setMatrix(mLeash, mTransformation.getMatrix(), mMatrix);
+            t.setAlpha(mLeash, mTransformation.getAlpha());
+
+            // The following applies an inverse scale to the clip-rect so that it crops "after" the
+            // scale instead of before.
+            mVecs[1] = mVecs[2] = 0;
+            mVecs[0] = mVecs[3] = 1;
+            mTransformation.getMatrix().mapVectors(mVecs);
+            mVecs[0] = 1.f / mVecs[0];
+            mVecs[3] = 1.f / mVecs[3];
+            final Rect clipRect = mTransformation.getClipRect();
+            mRect.left = (int) (clipRect.left * mVecs[0] + 0.5f);
+            mRect.right = (int) (clipRect.right * mVecs[0] + 0.5f);
+            mRect.top = (int) (clipRect.top * mVecs[3] + 0.5f);
+            mRect.bottom = (int) (clipRect.bottom * mVecs[3] + 0.5f);
+            t.setWindowCrop(mLeash, mRect);
+        }
+    }
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationController.java
new file mode 100644
index 0000000..d7eb9a0
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationController.java
@@ -0,0 +1,79 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_OPEN;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CHANGE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_OPEN;
+
+import android.util.Log;
+import android.view.RemoteAnimationAdapter;
+import android.view.RemoteAnimationDefinition;
+import android.window.TaskFragmentOrganizer;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+/** Controls the TaskFragment remote animations. */
+class TaskFragmentAnimationController {
+
+    private static final String TAG = "TaskFragAnimationCtrl";
+    static final boolean DEBUG = false;
+
+    private final TaskFragmentOrganizer mOrganizer;
+    private final TaskFragmentAnimationRunner mRemoteRunner = new TaskFragmentAnimationRunner();
+    @VisibleForTesting
+    final RemoteAnimationDefinition mDefinition;
+    private boolean mIsRegistered;
+
+    TaskFragmentAnimationController(@NonNull TaskFragmentOrganizer organizer) {
+        mOrganizer = organizer;
+        mDefinition = new RemoteAnimationDefinition();
+        final RemoteAnimationAdapter animationAdapter =
+                new RemoteAnimationAdapter(mRemoteRunner, 0, 0, true /* changeNeedsSnapshot */);
+        mDefinition.addRemoteAnimation(TRANSIT_OLD_ACTIVITY_OPEN, animationAdapter);
+        mDefinition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_OPEN, animationAdapter);
+        mDefinition.addRemoteAnimation(TRANSIT_OLD_ACTIVITY_CLOSE, animationAdapter);
+        mDefinition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CLOSE, animationAdapter);
+        mDefinition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CHANGE, animationAdapter);
+    }
+
+    void registerRemoteAnimations() {
+        if (DEBUG) {
+            Log.v(TAG, "registerRemoteAnimations");
+        }
+        if (mIsRegistered) {
+            return;
+        }
+        mOrganizer.registerRemoteAnimations(mDefinition);
+        mIsRegistered = true;
+    }
+
+    void unregisterRemoteAnimations() {
+        if (DEBUG) {
+            Log.v(TAG, "unregisterRemoteAnimations");
+        }
+        if (!mIsRegistered) {
+            return;
+        }
+        mOrganizer.unregisterRemoteAnimations();
+        mIsRegistered = false;
+    }
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationRunner.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationRunner.java
new file mode 100644
index 0000000..d9b73a8
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationRunner.java
@@ -0,0 +1,309 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static android.os.Process.THREAD_PRIORITY_DISPLAY;
+import static android.view.RemoteAnimationTarget.MODE_CHANGING;
+import static android.view.RemoteAnimationTarget.MODE_CLOSING;
+import static android.view.RemoteAnimationTarget.MODE_OPENING;
+import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_OPEN;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CHANGE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_OPEN;
+import static android.view.WindowManagerPolicyConstants.TYPE_LAYER_OFFSET;
+
+import android.animation.Animator;
+import android.animation.ValueAnimator;
+import android.graphics.Rect;
+import android.os.Handler;
+import android.os.HandlerThread;
+import android.os.RemoteException;
+import android.util.Log;
+import android.view.IRemoteAnimationFinishedCallback;
+import android.view.IRemoteAnimationRunner;
+import android.view.RemoteAnimationTarget;
+import android.view.SurfaceControl;
+import android.view.WindowManager;
+import android.view.animation.Animation;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.BiFunction;
+
+/** To run the TaskFragment animations. */
+class TaskFragmentAnimationRunner extends IRemoteAnimationRunner.Stub {
+
+    private static final String TAG = "TaskFragAnimationRunner";
+    private final Handler mHandler;
+    private final TaskFragmentAnimationSpec mAnimationSpec;
+
+    TaskFragmentAnimationRunner() {
+        HandlerThread animationThread = new HandlerThread(
+                "androidx.window.extensions.embedding", THREAD_PRIORITY_DISPLAY);
+        animationThread.start();
+        mHandler = animationThread.getThreadHandler();
+        mAnimationSpec = new TaskFragmentAnimationSpec(mHandler);
+    }
+
+    @Nullable
+    private Animator mAnimator;
+
+    @Override
+    public void onAnimationStart(@WindowManager.TransitionOldType int transit,
+            @NonNull RemoteAnimationTarget[] apps,
+            @NonNull RemoteAnimationTarget[] wallpapers,
+            @NonNull RemoteAnimationTarget[] nonApps,
+            @NonNull IRemoteAnimationFinishedCallback finishedCallback) {
+        if (wallpapers.length != 0 || nonApps.length != 0) {
+            throw new IllegalArgumentException("TaskFragment shouldn't handle animation with"
+                    + "wallpaper or non-app windows.");
+        }
+        if (TaskFragmentAnimationController.DEBUG) {
+            Log.v(TAG, "onAnimationStart transit=" + transit);
+        }
+        mHandler.post(() -> startAnimation(transit, apps, finishedCallback));
+    }
+
+    @Override
+    public void onAnimationCancelled() {
+        mHandler.post(this::cancelAnimation);
+    }
+
+    /** Creates and starts animation. */
+    private void startAnimation(@WindowManager.TransitionOldType int transit,
+            @NonNull RemoteAnimationTarget[] targets,
+            @NonNull IRemoteAnimationFinishedCallback finishedCallback) {
+        if (mAnimator != null) {
+            Log.w(TAG, "start new animation when the previous one is not finished yet.");
+            mAnimator.cancel();
+        }
+        mAnimator = createAnimator(transit, targets, finishedCallback);
+        mAnimator.start();
+    }
+
+    /** Cancels animation. */
+    private void cancelAnimation() {
+        if (mAnimator == null) {
+            return;
+        }
+        mAnimator.cancel();
+        mAnimator = null;
+    }
+
+    /** Creates the animator given the transition type and windows. */
+    @NonNull
+    private Animator createAnimator(@WindowManager.TransitionOldType int transit,
+            @NonNull RemoteAnimationTarget[] targets,
+            @NonNull IRemoteAnimationFinishedCallback finishedCallback) {
+        final List<TaskFragmentAnimationAdapter> adapters =
+                createAnimationAdapters(transit, targets);
+        long duration = 0;
+        for (TaskFragmentAnimationAdapter adapter : adapters) {
+            duration = Math.max(duration, adapter.getDurationHint());
+        }
+        final ValueAnimator animator = ValueAnimator.ofFloat(0, 1);
+        animator.setDuration(duration);
+        animator.addUpdateListener((anim) -> {
+            // Update all adapters in the same transaction.
+            final SurfaceControl.Transaction t = new SurfaceControl.Transaction();
+            for (TaskFragmentAnimationAdapter adapter : adapters) {
+                adapter.onAnimationUpdate(t, animator.getCurrentPlayTime());
+            }
+            t.apply();
+        });
+        animator.addListener(new Animator.AnimatorListener() {
+            @Override
+            public void onAnimationStart(Animator animation) {}
+
+            @Override
+            public void onAnimationEnd(Animator animation) {
+                final SurfaceControl.Transaction t = new SurfaceControl.Transaction();
+                for (TaskFragmentAnimationAdapter adapter : adapters) {
+                    adapter.onAnimationEnd(t);
+                }
+                t.apply();
+
+                try {
+                    finishedCallback.onAnimationFinished();
+                } catch (RemoteException e) {
+                    e.rethrowFromSystemServer();
+                }
+                mAnimator = null;
+            }
+
+            @Override
+            public void onAnimationCancel(Animator animation) {}
+
+            @Override
+            public void onAnimationRepeat(Animator animation) {}
+        });
+        return animator;
+    }
+
+    /** List of {@link TaskFragmentAnimationAdapter} to handle animations on all window targets. */
+    @NonNull
+    private List<TaskFragmentAnimationAdapter> createAnimationAdapters(
+            @WindowManager.TransitionOldType int transit,
+            @NonNull RemoteAnimationTarget[] targets) {
+        switch (transit) {
+            case TRANSIT_OLD_ACTIVITY_OPEN:
+            case TRANSIT_OLD_TASK_FRAGMENT_OPEN:
+                return createOpenAnimationAdapters(targets);
+            case TRANSIT_OLD_ACTIVITY_CLOSE:
+            case TRANSIT_OLD_TASK_FRAGMENT_CLOSE:
+                return createCloseAnimationAdapters(targets);
+            case TRANSIT_OLD_TASK_FRAGMENT_CHANGE:
+                return createChangeAnimationAdapters(targets);
+            default:
+                throw new IllegalArgumentException("Unhandled transit type=" + transit);
+        }
+    }
+
+    @NonNull
+    private List<TaskFragmentAnimationAdapter> createOpenAnimationAdapters(
+            @NonNull RemoteAnimationTarget[] targets) {
+        return createOpenCloseAnimationAdapters(targets, true /* isOpening */,
+                mAnimationSpec::loadOpenAnimation);
+    }
+
+    @NonNull
+    private List<TaskFragmentAnimationAdapter> createCloseAnimationAdapters(
+            @NonNull RemoteAnimationTarget[] targets) {
+        return createOpenCloseAnimationAdapters(targets, false /* isOpening */,
+                mAnimationSpec::loadCloseAnimation);
+    }
+
+    /**
+     * Creates {@link TaskFragmentAnimationAdapter} for OPEN and CLOSE types of transition.
+     * @param isOpening {@code true} for OPEN type, {@code false} for CLOSE type.
+     */
+    @NonNull
+    private List<TaskFragmentAnimationAdapter> createOpenCloseAnimationAdapters(
+            @NonNull RemoteAnimationTarget[] targets, boolean isOpening,
+            @NonNull BiFunction<RemoteAnimationTarget, Rect, Animation> animationProvider) {
+        // We need to know if the target window is only a partial of the whole animation screen.
+        // If so, we will need to adjust it to make the whole animation screen looks like one.
+        final List<RemoteAnimationTarget> openingTargets = new ArrayList<>();
+        final List<RemoteAnimationTarget> closingTargets = new ArrayList<>();
+        final Rect openingWholeScreenBounds = new Rect();
+        final Rect closingWholeScreenBounds = new Rect();
+        for (RemoteAnimationTarget target : targets) {
+            if (target.mode != MODE_CLOSING) {
+                openingTargets.add(target);
+                openingWholeScreenBounds.union(target.screenSpaceBounds);
+            } else {
+                closingTargets.add(target);
+                closingWholeScreenBounds.union(target.screenSpaceBounds);
+                // Union the start bounds since this may be the ClosingChanging animation.
+                closingWholeScreenBounds.union(target.startBounds);
+            }
+        }
+
+        // For OPEN transition, open windows should be above close windows.
+        // For CLOSE transition, open windows should be below close windows.
+        int offsetLayer = TYPE_LAYER_OFFSET;
+        final List<TaskFragmentAnimationAdapter> adapters = new ArrayList<>();
+        for (RemoteAnimationTarget target : openingTargets) {
+            final TaskFragmentAnimationAdapter adapter = createOpenCloseAnimationAdapter(target,
+                    animationProvider, openingWholeScreenBounds);
+            if (isOpening) {
+                adapter.overrideLayer(offsetLayer++);
+            }
+            adapters.add(adapter);
+        }
+        for (RemoteAnimationTarget target : closingTargets) {
+            final TaskFragmentAnimationAdapter adapter = createOpenCloseAnimationAdapter(target,
+                    animationProvider, closingWholeScreenBounds);
+            if (!isOpening) {
+                adapter.overrideLayer(offsetLayer++);
+            }
+            adapters.add(adapter);
+        }
+        return adapters;
+    }
+
+    @NonNull
+    private TaskFragmentAnimationAdapter createOpenCloseAnimationAdapter(
+            @NonNull RemoteAnimationTarget target,
+            @NonNull BiFunction<RemoteAnimationTarget, Rect, Animation> animationProvider,
+            @NonNull Rect wholeAnimationBounds) {
+        final Animation animation = animationProvider.apply(target, wholeAnimationBounds);
+        return new TaskFragmentAnimationAdapter(animation, target, target.leash,
+                wholeAnimationBounds);
+    }
+
+    @NonNull
+    private List<TaskFragmentAnimationAdapter> createChangeAnimationAdapters(
+            @NonNull RemoteAnimationTarget[] targets) {
+        if (shouldUseJumpCutForChangeAnimation(targets)) {
+            return new ArrayList<>();
+        }
+
+        final List<TaskFragmentAnimationAdapter> adapters = new ArrayList<>();
+        for (RemoteAnimationTarget target : targets) {
+            if (target.mode == MODE_CHANGING) {
+                // This is the target with bounds change.
+                final Animation[] animations =
+                        mAnimationSpec.createChangeBoundsChangeAnimations(target);
+                // Adapter for the starting snapshot leash.
+                adapters.add(new TaskFragmentAnimationAdapter.SnapshotAdapter(
+                        animations[0], target));
+                // Adapter for the ending bounds changed leash.
+                adapters.add(new TaskFragmentAnimationAdapter.BoundsChangeAdapter(
+                        animations[1], target));
+                continue;
+            }
+
+            // These are the other targets that don't have bounds change in the same transition.
+            final Animation animation;
+            if (target.hasAnimatingParent) {
+                // No-op if it will be covered by the changing parent window.
+                animation = TaskFragmentAnimationSpec.createNoopAnimation(target);
+            } else if (target.mode == MODE_CLOSING) {
+                animation = mAnimationSpec.createChangeBoundsCloseAnimation(target);
+            } else {
+                animation = mAnimationSpec.createChangeBoundsOpenAnimation(target);
+            }
+            adapters.add(new TaskFragmentAnimationAdapter(animation, target));
+        }
+        return adapters;
+    }
+
+    /**
+     * Whether we should use jump cut for the change transition.
+     * This normally happens when opening a new secondary with the existing primary using a
+     * different split layout. This can be complicated, like from horizontal to vertical split with
+     * new split pairs.
+     * Uses a jump cut animation to simplify.
+     */
+    private boolean shouldUseJumpCutForChangeAnimation(@NonNull RemoteAnimationTarget[] targets) {
+        boolean hasOpeningWindow = false;
+        boolean hasClosingWindow = false;
+        for (RemoteAnimationTarget target : targets) {
+            if (target.hasAnimatingParent) {
+                continue;
+            }
+            hasOpeningWindow |= target.mode == MODE_OPENING;
+            hasClosingWindow |= target.mode == MODE_CLOSING;
+        }
+        return hasOpeningWindow && hasClosingWindow;
+    }
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationSpec.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationSpec.java
new file mode 100644
index 0000000..1f866c3
--- /dev/null
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentAnimationSpec.java
@@ -0,0 +1,267 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static android.view.RemoteAnimationTarget.MODE_CLOSING;
+
+import android.app.ActivityThread;
+import android.content.ContentResolver;
+import android.content.Context;
+import android.database.ContentObserver;
+import android.graphics.Rect;
+import android.os.Handler;
+import android.provider.Settings;
+import android.view.RemoteAnimationTarget;
+import android.view.WindowManager;
+import android.view.animation.AlphaAnimation;
+import android.view.animation.Animation;
+import android.view.animation.AnimationSet;
+import android.view.animation.AnimationUtils;
+import android.view.animation.ClipRectAnimation;
+import android.view.animation.Interpolator;
+import android.view.animation.LinearInterpolator;
+import android.view.animation.ScaleAnimation;
+import android.view.animation.TranslateAnimation;
+
+import androidx.annotation.NonNull;
+
+import com.android.internal.R;
+import com.android.internal.policy.AttributeCache;
+import com.android.internal.policy.TransitionAnimation;
+
+/** Animation spec for TaskFragment transition. */
+// TODO(b/206557124): provide an easier way to customize animation
+class TaskFragmentAnimationSpec {
+
+    private static final String TAG = "TaskFragAnimationSpec";
+    private static final int CHANGE_ANIMATION_DURATION = 517;
+    private static final int CHANGE_ANIMATION_FADE_DURATION = 80;
+    private static final int CHANGE_ANIMATION_FADE_OFFSET = 30;
+
+    private final Context mContext;
+    private final TransitionAnimation mTransitionAnimation;
+    private final Interpolator mFastOutExtraSlowInInterpolator;
+    private final LinearInterpolator mLinearInterpolator;
+    private float mTransitionAnimationScaleSetting;
+
+    TaskFragmentAnimationSpec(@NonNull Handler handler) {
+        mContext = ActivityThread.currentActivityThread().getApplication();
+        mTransitionAnimation = new TransitionAnimation(mContext, false /* debug */, TAG);
+        // Initialize the AttributeCache for the TransitionAnimation.
+        AttributeCache.init(mContext);
+        mFastOutExtraSlowInInterpolator = AnimationUtils.loadInterpolator(
+                mContext, android.R.interpolator.fast_out_extra_slow_in);
+        mLinearInterpolator = new LinearInterpolator();
+
+        // The transition animation should be adjusted based on the developer option.
+        final ContentResolver resolver = mContext.getContentResolver();
+        mTransitionAnimationScaleSetting = getTransitionAnimationScaleSetting();
+        resolver.registerContentObserver(
+                Settings.Global.getUriFor(Settings.Global.TRANSITION_ANIMATION_SCALE), false,
+                new SettingsObserver(handler));
+    }
+
+    /** For target that doesn't need to be animated. */
+    @NonNull
+    static Animation createNoopAnimation(@NonNull RemoteAnimationTarget target) {
+        // Noop but just keep the target showing/hiding.
+        final float alpha = target.mode == MODE_CLOSING ? 0f : 1f;
+        return new AlphaAnimation(alpha, alpha);
+    }
+
+    /** Animation for target that is opening in a change transition. */
+    @NonNull
+    Animation createChangeBoundsOpenAnimation(@NonNull RemoteAnimationTarget target) {
+        final Rect parentBounds = target.taskInfo.configuration.windowConfiguration.getBounds();
+        final Rect bounds = target.screenSpaceBounds;
+        final int startLeft;
+        final int startTop;
+        if (parentBounds.top == bounds.top && parentBounds.bottom == bounds.bottom) {
+            // The window will be animated in from left or right depending on its position.
+            startTop = 0;
+            startLeft = parentBounds.left == bounds.left ? -bounds.width() : bounds.width();
+        } else {
+            // The window will be animated in from top or bottom depending on its position.
+            startTop = parentBounds.top == bounds.top ? -bounds.height() : bounds.height();
+            startLeft = 0;
+        }
+
+        // The position should be 0-based as we will post translate in
+        // TaskFragmentAnimationAdapter#onAnimationUpdate
+        final Animation animation = new TranslateAnimation(startLeft, 0, startTop, 0);
+        animation.setInterpolator(mFastOutExtraSlowInInterpolator);
+        animation.setDuration(CHANGE_ANIMATION_DURATION);
+        animation.initialize(bounds.width(), bounds.height(), bounds.width(), bounds.height());
+        animation.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+        return animation;
+    }
+
+    /** Animation for target that is closing in a change transition. */
+    @NonNull
+    Animation createChangeBoundsCloseAnimation(@NonNull RemoteAnimationTarget target) {
+        final Rect parentBounds = target.taskInfo.configuration.windowConfiguration.getBounds();
+        // Use startBounds if the window is closing in case it may also resize.
+        final Rect bounds = target.startBounds;
+        final int endTop;
+        final int endLeft;
+        if (parentBounds.top == bounds.top && parentBounds.bottom == bounds.bottom) {
+            // The window will be animated out to left or right depending on its position.
+            endTop = 0;
+            endLeft = parentBounds.left == bounds.left ? -bounds.width() : bounds.width();
+        } else {
+            // The window will be animated out to top or bottom depending on its position.
+            endTop = parentBounds.top == bounds.top ? -bounds.height() : bounds.height();
+            endLeft = 0;
+        }
+
+        // The position should be 0-based as we will post translate in
+        // TaskFragmentAnimationAdapter#onAnimationUpdate
+        final Animation animation = new TranslateAnimation(0, endLeft, 0, endTop);
+        animation.setInterpolator(mFastOutExtraSlowInInterpolator);
+        animation.setDuration(CHANGE_ANIMATION_DURATION);
+        animation.initialize(bounds.width(), bounds.height(), bounds.width(), bounds.height());
+        animation.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+        return animation;
+    }
+
+    /**
+     * Animation for target that is changing (bounds change) in a change transition.
+     * @return the return array always has two elements. The first one is for the start leash, and
+     *         the second one is for the end leash.
+     */
+    @NonNull
+    Animation[] createChangeBoundsChangeAnimations(@NonNull RemoteAnimationTarget target) {
+        // Both start bounds and end bounds are in screen coordinates. We will post translate
+        // to the local coordinates in TaskFragmentAnimationAdapter#onAnimationUpdate
+        final Rect startBounds = target.startBounds;
+        final Rect parentBounds = target.taskInfo.configuration.windowConfiguration.getBounds();
+        final Rect endBounds = target.screenSpaceBounds;
+        float scaleX = ((float) startBounds.width()) / endBounds.width();
+        float scaleY = ((float) startBounds.height()) / endBounds.height();
+        // Start leash is a child of the end leash. Reverse the scale so that the start leash won't
+        // be scaled up with its parent.
+        float startScaleX = 1.f / scaleX;
+        float startScaleY = 1.f / scaleY;
+
+        // The start leash will be fade out.
+        final AnimationSet startSet = new AnimationSet(false /* shareInterpolator */);
+        final Animation startAlpha = new AlphaAnimation(1f, 0f);
+        startAlpha.setInterpolator(mLinearInterpolator);
+        startAlpha.setDuration(CHANGE_ANIMATION_FADE_DURATION);
+        startAlpha.setStartOffset(CHANGE_ANIMATION_FADE_OFFSET);
+        startSet.addAnimation(startAlpha);
+        final Animation startScale = new ScaleAnimation(startScaleX, startScaleX, startScaleY,
+                startScaleY);
+        startScale.setInterpolator(mFastOutExtraSlowInInterpolator);
+        startScale.setDuration(CHANGE_ANIMATION_DURATION);
+        startSet.addAnimation(startScale);
+        startSet.initialize(startBounds.width(), startBounds.height(), endBounds.width(),
+                endBounds.height());
+        startSet.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+
+        // The end leash will be moved into the end position while scaling.
+        final AnimationSet endSet = new AnimationSet(true /* shareInterpolator */);
+        endSet.setInterpolator(mFastOutExtraSlowInInterpolator);
+        final Animation endScale = new ScaleAnimation(scaleX, 1, scaleY, 1);
+        endScale.setDuration(CHANGE_ANIMATION_DURATION);
+        endSet.addAnimation(endScale);
+        // The position should be 0-based as we will post translate in
+        // TaskFragmentAnimationAdapter#onAnimationUpdate
+        final Animation endTranslate = new TranslateAnimation(startBounds.left - endBounds.left, 0,
+                startBounds.top - endBounds.top, 0);
+        endTranslate.setDuration(CHANGE_ANIMATION_DURATION);
+        endSet.addAnimation(endTranslate);
+        // The end leash is resizing, we should update the window crop based on the clip rect.
+        final Rect startClip = new Rect(startBounds);
+        final Rect endClip = new Rect(endBounds);
+        startClip.offsetTo(0, 0);
+        endClip.offsetTo(0, 0);
+        final Animation clipAnim = new ClipRectAnimation(startClip, endClip);
+        clipAnim.setDuration(CHANGE_ANIMATION_DURATION);
+        endSet.addAnimation(clipAnim);
+        endSet.initialize(startBounds.width(), startBounds.height(), parentBounds.width(),
+                parentBounds.height());
+        endSet.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+
+        return new Animation[]{startSet, endSet};
+    }
+
+    @NonNull
+    Animation loadOpenAnimation(@NonNull RemoteAnimationTarget target,
+            @NonNull Rect wholeAnimationBounds) {
+        final boolean isEnter = target.mode != MODE_CLOSING;
+        final Animation animation;
+        // Background color on TaskDisplayArea has already been set earlier in
+        // WindowContainer#getAnimationAdapter.
+        if (target.showBackdrop) {
+            animation = mTransitionAnimation.loadDefaultAnimationRes(isEnter
+                    ? com.android.internal.R.anim.task_fragment_clear_top_open_enter
+                    : com.android.internal.R.anim.task_fragment_clear_top_open_exit);
+        } else {
+            animation = mTransitionAnimation.loadDefaultAnimationRes(isEnter
+                    ? com.android.internal.R.anim.task_fragment_open_enter
+                    : com.android.internal.R.anim.task_fragment_open_exit);
+        }
+        // Use the whole animation bounds instead of the change bounds, so that when multiple change
+        // targets are opening at the same time, the animation applied to each will be the same.
+        // Otherwise, we may see gap between the activities that are launching together.
+        animation.initialize(wholeAnimationBounds.width(), wholeAnimationBounds.height(),
+                wholeAnimationBounds.width(), wholeAnimationBounds.height());
+        animation.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+        return animation;
+    }
+
+    @NonNull
+    Animation loadCloseAnimation(@NonNull RemoteAnimationTarget target,
+            @NonNull Rect wholeAnimationBounds) {
+        final boolean isEnter = target.mode != MODE_CLOSING;
+        final Animation animation;
+        if (target.showBackdrop) {
+            animation = mTransitionAnimation.loadDefaultAnimationRes(isEnter
+                    ? com.android.internal.R.anim.task_fragment_clear_top_close_enter
+                    : com.android.internal.R.anim.task_fragment_clear_top_close_exit);
+        } else {
+            animation = mTransitionAnimation.loadDefaultAnimationRes(isEnter
+                    ? com.android.internal.R.anim.task_fragment_close_enter
+                    : com.android.internal.R.anim.task_fragment_close_exit);
+        }
+        // Use the whole animation bounds instead of the change bounds, so that when multiple change
+        // targets are closing at the same time, the animation applied to each will be the same.
+        // Otherwise, we may see gap between the activities that are finishing together.
+        animation.initialize(wholeAnimationBounds.width(), wholeAnimationBounds.height(),
+                wholeAnimationBounds.width(), wholeAnimationBounds.height());
+        animation.scaleCurrentDuration(mTransitionAnimationScaleSetting);
+        return animation;
+    }
+
+    private float getTransitionAnimationScaleSetting() {
+        return WindowManager.fixScale(Settings.Global.getFloat(mContext.getContentResolver(),
+                Settings.Global.TRANSITION_ANIMATION_SCALE, mContext.getResources().getFloat(
+                                R.dimen.config_appTransitionAnimationDurationScaleDefault)));
+    }
+
+    private class SettingsObserver extends ContentObserver {
+        SettingsObserver(@NonNull Handler handler) {
+            super(handler);
+        }
+
+        @Override
+        public void onChange(boolean selfChange) {
+            mTransitionAnimationScaleSetting = getTransitionAnimationScaleSetting();
+        }
+    }
+}
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
index dc6506b..dc1d983 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/TaskFragmentContainer.java
@@ -53,15 +53,13 @@
 class TaskFragmentContainer {
     private static final int APPEAR_EMPTY_TIMEOUT_MS = 3000;
 
+    /** Parcelable data of this TaskFragmentContainer. */
+    @NonNull
+    private final ParcelableTaskFragmentContainerData mParcelableData;
+
     @NonNull
     private final SplitController mController;
 
-    /**
-     * Client-created token that uniquely identifies the task fragment container instance.
-     */
-    @NonNull
-    private final IBinder mToken;
-
     /** Parent leaf Task. */
     @NonNull
     private final TaskContainer mTaskContainer;
@@ -103,9 +101,6 @@
      */
     private final List<IBinder> mActivitiesToFinishOnExit = new ArrayList<>();
 
-    @Nullable
-    private final String mOverlayTag;
-
     /**
      * The launch options that was used to create this container. Must not {@link Bundle#isEmpty()}
      * for {@link #isOverlay()} container.
@@ -113,29 +108,13 @@
     @NonNull
     private final Bundle mLaunchOptions = new Bundle();
 
-    /**
-     * The associated {@link Activity#getActivityToken()} of the overlay container.
-     * Must be {@code null} for non-overlay container.
-     * <p>
-     * If an overlay container is associated with an activity, this overlay container will be
-     * dismissed when the associated activity is destroyed. If the overlay container is visible,
-     * activity will be launched on top of the overlay container and expanded to fill the parent
-     * container.
-     */
-    @Nullable
-    private final IBinder mAssociatedActivityToken;
-
     /** Indicates whether the container was cleaned up after the last activity was removed. */
     private boolean mIsFinished;
 
     /**
-     * Bounds that were requested last via {@link android.window.WindowContainerTransaction}.
-     */
-    private final Rect mLastRequestedBounds = new Rect();
-
-    /**
      * Windowing mode that was requested last via {@link android.window.WindowContainerTransaction}.
      */
+    // TODO(b/289875940): review this and other field that might need to be moved in the base class.
     @WindowingMode
     private int mLastRequestedWindowingMode = WINDOWING_MODE_UNDEFINED;
 
@@ -208,17 +187,17 @@
             @NonNull SplitController controller,
             @Nullable TaskFragmentContainer pairedPrimaryContainer, @Nullable String overlayTag,
             @Nullable Bundle launchOptions, @Nullable Activity associatedActivity) {
+        mParcelableData = new ParcelableTaskFragmentContainerData(
+                new Binder("TaskFragmentContainer"), overlayTag,
+                associatedActivity != null ? associatedActivity.getActivityToken() : null);
+
         if ((pendingAppearedActivity == null && pendingAppearedIntent == null)
                 || (pendingAppearedActivity != null && pendingAppearedIntent != null)) {
             throw new IllegalArgumentException(
                     "One and only one of pending activity and intent must be non-null");
         }
         mController = controller;
-        mToken = new Binder("TaskFragmentContainer");
         mTaskContainer = taskContainer;
-        mOverlayTag = overlayTag;
-        mAssociatedActivityToken = associatedActivity != null
-                ? associatedActivity.getActivityToken() : null;
 
         if (launchOptions != null) {
             mLaunchOptions.putAll(launchOptions);
@@ -259,18 +238,26 @@
         if (overlayTag != null && pendingAppearedIntent != null
                 && associatedActivity != null && !associatedActivity.isFinishing()) {
             final IBinder associatedActivityToken = associatedActivity.getActivityToken();
-            final OverlayContainerRestoreParams params = new OverlayContainerRestoreParams(mToken,
-                    launchOptions, pendingAppearedIntent);
+            final OverlayContainerRestoreParams params = new OverlayContainerRestoreParams(
+                    mParcelableData.mToken, launchOptions, pendingAppearedIntent);
             mController.mOverlayRestoreParams.put(associatedActivityToken, params);
         }
     }
 
+    /** This is only used when restoring it from a {@link ParcelableTaskFragmentContainerData}. */
+    TaskFragmentContainer(@NonNull ParcelableTaskFragmentContainerData data,
+            @NonNull SplitController splitController, @NonNull TaskContainer taskContainer) {
+        mParcelableData = data;
+        mController = splitController;
+        mTaskContainer = taskContainer;
+    }
+
     /**
      * Returns the client-created token that uniquely identifies this container.
      */
     @NonNull
     IBinder getTaskFragmentToken() {
-        return mToken;
+        return mParcelableData.mToken;
     }
 
     /** List of non-finishing activities that belong to this container and live in this process. */
@@ -389,7 +376,8 @@
             return null;
         }
         return new ActivityStack(activities, isEmpty(),
-                ActivityStack.Token.createFromBinder(mToken), mOverlayTag);
+                ActivityStack.Token.createFromBinder(mParcelableData.mToken),
+                mParcelableData.mOverlayTag);
     }
 
     /** Adds the activity that will be reparented to this container. */
@@ -413,7 +401,7 @@
         final ActivityThread.ActivityClientRecord record = ActivityThread
                 .currentActivityThread().getActivityClient(activityToken);
         if (record != null) {
-            record.mTaskFragmentToken = mToken;
+            record.mTaskFragmentToken = mParcelableData.mToken;
         }
     }
 
@@ -469,7 +457,7 @@
         if (!isOverlayWithActivityAssociation()) {
             return;
         }
-        if (mAssociatedActivityToken == activityToken) {
+        if (mParcelableData.mAssociatedActivityToken == activityToken) {
             // If the associated activity is destroyed, also finish this overlay container.
             mController.mPresenter.cleanupContainer(wct, this, false /* shouldFinishDependent */);
         }
@@ -776,8 +764,8 @@
      * @see WindowContainerTransaction#setRelativeBounds
      */
     boolean areLastRequestedBoundsEqual(@Nullable Rect relBounds) {
-        return (relBounds == null && mLastRequestedBounds.isEmpty())
-                || mLastRequestedBounds.equals(relBounds);
+        return (relBounds == null && mParcelableData.mLastRequestedBounds.isEmpty())
+                || mParcelableData.mLastRequestedBounds.equals(relBounds);
     }
 
     /**
@@ -787,14 +775,14 @@
      */
     void setLastRequestedBounds(@Nullable Rect relBounds) {
         if (relBounds == null) {
-            mLastRequestedBounds.setEmpty();
+            mParcelableData.mLastRequestedBounds.setEmpty();
         } else {
-            mLastRequestedBounds.set(relBounds);
+            mParcelableData.mLastRequestedBounds.set(relBounds);
         }
     }
 
     @NonNull Rect getLastRequestedBounds() {
-        return mLastRequestedBounds;
+        return mParcelableData.mLastRequestedBounds;
     }
 
     /**
@@ -965,6 +953,16 @@
         return mTaskContainer.getTaskId();
     }
 
+    @NonNull
+    IBinder getToken() {
+        return mParcelableData.mToken;
+    }
+
+    @NonNull
+    ParcelableTaskFragmentContainerData getParcelableData() {
+        return mParcelableData;
+    }
+
     /** Gets the parent Task. */
     @NonNull
     TaskContainer getTaskContainer() {
@@ -1011,7 +1009,7 @@
 
     /** Returns whether this taskFragment container is an overlay container. */
     boolean isOverlay() {
-        return mOverlayTag != null;
+        return mParcelableData.mOverlayTag != null;
     }
 
     /**
@@ -1020,7 +1018,7 @@
      */
     @Nullable
     String getOverlayTag() {
-        return mOverlayTag;
+        return mParcelableData.mOverlayTag;
     }
 
     /**
@@ -1045,7 +1043,7 @@
      */
     @Nullable
     IBinder getAssociatedActivityToken() {
-        return mAssociatedActivityToken;
+        return mParcelableData.mAssociatedActivityToken;
     }
 
     /**
@@ -1053,11 +1051,11 @@
      * a non-fill-parent overlay without activity association.
      */
     boolean isAlwaysOnTopOverlay() {
-        return isOverlay() && mAssociatedActivityToken == null;
+        return isOverlay() && mParcelableData.mAssociatedActivityToken == null;
     }
 
     boolean isOverlayWithActivityAssociation() {
-        return isOverlay() && mAssociatedActivityToken != null;
+        return isOverlay() && mParcelableData.mAssociatedActivityToken != null;
     }
 
     @Override
@@ -1074,13 +1072,13 @@
     private String toString(boolean includeContainersToFinishOnExit) {
         return "TaskFragmentContainer{"
                 + " parentTaskId=" + getTaskId()
-                + " token=" + mToken
+                + " token=" + mParcelableData.mToken
                 + " topNonFinishingActivity=" + getTopNonFinishingActivity()
                 + " runningActivityCount=" + getRunningActivityCount()
                 + " isFinished=" + mIsFinished
-                + " overlayTag=" + mOverlayTag
-                + " associatedActivityToken=" + mAssociatedActivityToken
-                + " lastRequestedBounds=" + mLastRequestedBounds
+                + " overlayTag=" + mParcelableData.mOverlayTag
+                + " associatedActivityToken=" + mParcelableData.mAssociatedActivityToken
+                + " lastRequestedBounds=" + mParcelableData.mLastRequestedBounds
                 + " pendingAppearedActivities=" + mPendingAppearedActivities
                 + (includeContainersToFinishOnExit ? " containersToFinishOnExit="
                 + containersToFinishOnExitToString() : "")
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
index 8911d18..ac004c3 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/JetpackTaskFragmentOrganizerTest.java
@@ -23,6 +23,8 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
@@ -84,6 +86,24 @@
     }
 
     @Test
+    public void testUnregisterOrganizer() {
+        mOrganizer.overrideSplitAnimation();
+        mOrganizer.unregisterOrganizer();
+
+        verify(mOrganizer).unregisterRemoteAnimations();
+    }
+
+    @Test
+    public void testOverrideSplitAnimation() {
+        assertNull(mOrganizer.mAnimationController);
+
+        mOrganizer.overrideSplitAnimation();
+
+        assertNotNull(mOrganizer.mAnimationController);
+        verify(mOrganizer).registerRemoteAnimations(mOrganizer.mAnimationController.mDefinition);
+    }
+
+    @Test
     public void testExpandTaskFragment() {
         final TaskContainer taskContainer = createTestTaskContainer();
         doReturn(taskContainer).when(mSplitController).getTaskContainer(anyInt());
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
index 90eeb58..5b97e7e 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/OverlayPresentationTest.java
@@ -549,7 +549,7 @@
         assertThat(taskContainer.getTaskFragmentContainers()).containsExactly(overlayContainer);
 
         taskContainer.updateTaskFragmentParentInfo(new TaskFragmentParentInfo(Configuration.EMPTY,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */));
 
         mSplitController.updateOverlayContainer(mTransaction, overlayContainer);
@@ -618,7 +618,8 @@
         final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
                 new Configuration(taskProperties.getConfiguration()), taskProperties.getDisplayId(),
-                true /* visible */, false /* hasDirectActivity */, null /* decorSurface */);
+                TASK_ID, true /* visible */, false /* hasDirectActivity */,
+                null /* decorSurface */);
         parentInfo.getConfiguration().windowConfiguration.getBounds().offset(10, 10);
 
         mSplitController.onTaskFragmentParentInfoChanged(mTransaction, TASK_ID, parentInfo);
@@ -642,7 +643,8 @@
         final TaskContainer.TaskProperties taskProperties = taskContainer.getTaskProperties();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(
                 new Configuration(taskProperties.getConfiguration()), taskProperties.getDisplayId(),
-                true /* visible */, false /* hasDirectActivity */, null /* decorSurface */);
+                TASK_ID, true /* visible */, false /* hasDirectActivity */,
+                null /* decorSurface */);
 
         mSplitController.onTaskFragmentParentInfoChanged(mTransaction, TASK_ID, parentInfo);
 
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
index d852204..0512412 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
@@ -1164,7 +1164,7 @@
     public void testOnTransactionReady_taskFragmentParentInfoChanged() {
         final TaskFragmentTransaction transaction = new TaskFragmentTransaction();
         final TaskFragmentParentInfo parentInfo = new TaskFragmentParentInfo(Configuration.EMPTY,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */);
         transaction.addChange(new TaskFragmentTransaction.Change(
                 TYPE_TASK_FRAGMENT_PARENT_INFO_CHANGED)
@@ -1625,7 +1625,7 @@
         final TaskContainer taskContainer = mSplitController.getTaskContainer(TASK_ID);
         final Configuration configuration = new Configuration();
         final TaskFragmentParentInfo originalInfo = new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */);
         mSplitController.onTaskFragmentParentInfoChanged(mock(WindowContainerTransaction.class),
                 TASK_ID, originalInfo);
@@ -1634,7 +1634,7 @@
         // Making a public configuration change while the Task is invisible.
         configuration.densityDpi += 100;
         final TaskFragmentParentInfo invisibleInfo = new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, false /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, false /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */);
         mSplitController.onTaskFragmentParentInfoChanged(mock(WindowContainerTransaction.class),
                 TASK_ID, invisibleInfo);
@@ -1646,7 +1646,7 @@
 
         // Updates when Task to become visible
         final TaskFragmentParentInfo visibleInfo = new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */);
         mSplitController.onTaskFragmentParentInfoChanged(mock(WindowContainerTransaction.class),
                 TASK_ID, visibleInfo);
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
index 2847232..97f4d07 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskContainerTest.java
@@ -23,6 +23,7 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.TASK_ID;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTestTaskContainer;
 
 import static org.junit.Assert.assertEquals;
@@ -82,7 +83,7 @@
 
         configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
         taskContainer.updateTaskFragmentParentInfo(new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */));
 
         assertEquals(WINDOWING_MODE_MULTI_WINDOW,
@@ -90,7 +91,7 @@
 
         configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FREEFORM);
         taskContainer.updateTaskFragmentParentInfo(new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */));
 
         assertEquals(WINDOWING_MODE_FREEFORM,
@@ -111,14 +112,14 @@
 
         configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_FULLSCREEN);
         taskContainer.updateTaskFragmentParentInfo(new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */));
 
         assertFalse(taskContainer.isInPictureInPicture());
 
         configuration.windowConfiguration.setWindowingMode(WINDOWING_MODE_PINNED);
         taskContainer.updateTaskFragmentParentInfo(new TaskFragmentParentInfo(configuration,
-                DEFAULT_DISPLAY, true /* visible */, false /* hasDirectActivity */,
+                DEFAULT_DISPLAY, TASK_ID, true /* visible */, false /* hasDirectActivity */,
                 null /* decorSurface */));
 
         assertTrue(taskContainer.isInPictureInPicture());
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentAnimationControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentAnimationControllerTest.java
new file mode 100644
index 0000000..a1e9f08
--- /dev/null
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/TaskFragmentAnimationControllerTest.java
@@ -0,0 +1,88 @@
+/*
+ * 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 androidx.window.extensions.embedding;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+
+import static org.mockito.Mockito.never;
+
+import android.platform.test.annotations.Presubmit;
+import android.window.TaskFragmentOrganizer;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+/**
+ * Test class for {@link TaskFragmentAnimationController}.
+ *
+ * Build/Install/Run:
+ *  atest WMJetpackUnitTests:TaskFragmentAnimationControllerTest
+ */
+@Presubmit
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class TaskFragmentAnimationControllerTest {
+    @Rule
+    public MockitoRule rule = MockitoJUnit.rule();
+
+    @Mock
+    private TaskFragmentOrganizer mOrganizer;
+    private TaskFragmentAnimationController mAnimationController;
+
+    @Before
+    public void setup() {
+        mAnimationController = new TaskFragmentAnimationController(mOrganizer);
+    }
+
+    @Test
+    public void testRegisterRemoteAnimations() {
+        mAnimationController.registerRemoteAnimations();
+
+        verify(mOrganizer).registerRemoteAnimations(mAnimationController.mDefinition);
+
+        mAnimationController.registerRemoteAnimations();
+
+        // No extra call if it has been registered.
+        verify(mOrganizer).registerRemoteAnimations(mAnimationController.mDefinition);
+    }
+
+    @Test
+    public void testUnregisterRemoteAnimations() {
+        mAnimationController.unregisterRemoteAnimations();
+
+        // No call if it is not registered.
+        verify(mOrganizer, never()).unregisterRemoteAnimations();
+
+        mAnimationController.registerRemoteAnimations();
+        mAnimationController.unregisterRemoteAnimations();
+
+        verify(mOrganizer).unregisterRemoteAnimations();
+
+        mAnimationController.unregisterRemoteAnimations();
+
+        // No extra call if it has been unregistered.
+        verify(mOrganizer).unregisterRemoteAnimations();
+    }
+}
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index b338a2a..a79bc97 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -39,17 +39,6 @@
     path: "src",
 }
 
-// Sources that have no dependencies that can be used directly downstream of this library
-// TODO(b/322791067): move these sources to WindowManager-Shell-shared
-filegroup {
-    name: "wm_shell_util-sources",
-    srcs: [
-        "src/com/android/wm/shell/common/bubbles/*.kt",
-        "src/com/android/wm/shell/common/bubbles/*.java",
-    ],
-    path: "src",
-}
-
 // Aidls which can be used directly downstream of this library
 filegroup {
     name: "wm_shell-aidls",
@@ -184,9 +173,11 @@
         ":wm_shell-shared-aidls",
     ],
     static_libs: [
+        "androidx.core_core-animation",
         "androidx.dynamicanimation_dynamicanimation",
         "jsr330",
     ],
+    kotlincflags: ["-Xjvm-default=all"],
 }
 
 java_library {
@@ -212,7 +203,6 @@
     ],
     static_libs: [
         "androidx.appcompat_appcompat",
-        "androidx.core_core-animation",
         "androidx.core_core-ktx",
         "androidx.arch.core_core-runtime",
         "androidx.datastore_datastore",
diff --git a/libs/WindowManager/Shell/aconfig/multitasking.aconfig b/libs/WindowManager/Shell/aconfig/multitasking.aconfig
index 9de10c0..df1a98c 100644
--- a/libs/WindowManager/Shell/aconfig/multitasking.aconfig
+++ b/libs/WindowManager/Shell/aconfig/multitasking.aconfig
@@ -138,3 +138,17 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+   name: "enable_bubble_to_fullscreen"
+   namespace: "multitasking"
+   description: "Enable an option to move bubbles to fullscreen"
+   bug: "363326492"
+}
+
+flag {
+    name: "enable_flexible_split"
+    namespace: "multitasking"
+    description: "Enables flexibile split feature for split screen"
+    bug: "349828130"
+}
diff --git a/libs/WindowManager/Shell/multivalentScreenshotTests/src/com/android/wm/shell/bubbles/BubbleEducationViewScreenshotTest.kt b/libs/WindowManager/Shell/multivalentScreenshotTests/src/com/android/wm/shell/bubbles/BubbleEducationViewScreenshotTest.kt
index d35f493..f09969d 100644
--- a/libs/WindowManager/Shell/multivalentScreenshotTests/src/com/android/wm/shell/bubbles/BubbleEducationViewScreenshotTest.kt
+++ b/libs/WindowManager/Shell/multivalentScreenshotTests/src/com/android/wm/shell/bubbles/BubbleEducationViewScreenshotTest.kt
@@ -16,7 +16,7 @@
 package com.android.wm.shell.bubbles
 
 import android.view.LayoutInflater
-import com.android.wm.shell.common.bubbles.BubblePopupView
+import com.android.wm.shell.shared.bubbles.BubblePopupView
 import com.android.wm.shell.testing.goldenpathmanager.WMShellGoldenPathManager
 import com.android.wm.shell.R
 import org.junit.Rule
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 4b97451..b38d00da 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
@@ -30,7 +30,7 @@
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.R
 import com.android.wm.shell.bubbles.BubblePositioner.MAX_HEIGHT
-import com.android.wm.shell.common.bubbles.BubbleBarLocation
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation
 import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors.directExecutor
 import org.junit.Before
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleStackViewTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleStackViewTest.kt
index faadf1d..96ffa03 100644
--- a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleStackViewTest.kt
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleStackViewTest.kt
@@ -53,7 +53,7 @@
 import org.mockito.kotlin.mock
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
-import com.android.wm.shell.common.bubbles.BubbleBarLocation
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation
 import java.util.concurrent.Semaphore
 import java.util.concurrent.TimeUnit
 import java.util.function.Consumer
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinControllerTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinControllerTest.kt
index 935d129..ecb2b25 100644
--- a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinControllerTest.kt
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinControllerTest.kt
@@ -31,12 +31,12 @@
 import com.android.wm.shell.R
 import com.android.wm.shell.bubbles.BubblePositioner
 import com.android.wm.shell.bubbles.DeviceConfig
-import com.android.wm.shell.common.bubbles.BaseBubblePinController
-import com.android.wm.shell.common.bubbles.BaseBubblePinController.Companion.DROP_TARGET_ALPHA_IN_DURATION
-import com.android.wm.shell.common.bubbles.BaseBubblePinController.Companion.DROP_TARGET_ALPHA_OUT_DURATION
-import com.android.wm.shell.common.bubbles.BubbleBarLocation
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.LEFT
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.RIGHT
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController.Companion.DROP_TARGET_ALPHA_IN_DURATION
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController.Companion.DROP_TARGET_ALPHA_OUT_DURATION
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.LEFT
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.RIGHT
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
diff --git a/libs/WindowManager/Shell/res/layout/bubble_bar_manage_education.xml b/libs/WindowManager/Shell/res/layout/bubble_bar_manage_education.xml
index a0a06f1..806d026 100644
--- a/libs/WindowManager/Shell/res/layout/bubble_bar_manage_education.xml
+++ b/libs/WindowManager/Shell/res/layout/bubble_bar_manage_education.xml
@@ -14,7 +14,7 @@
   ~ See the License for the specific language governing permissions and
   ~ limitations under the License
   -->
-<com.android.wm.shell.common.bubbles.BubblePopupView
+<com.android.wm.shell.shared.bubbles.BubblePopupView
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
@@ -53,4 +53,4 @@
         android:textAlignment="center"
         android:text="@string/bubble_bar_education_manage_text"/>
 
-</com.android.wm.shell.common.bubbles.BubblePopupView>
\ No newline at end of file
+</com.android.wm.shell.shared.bubbles.BubblePopupView>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/layout/bubble_bar_stack_education.xml b/libs/WindowManager/Shell/res/layout/bubble_bar_stack_education.xml
index b489a5c..7fa586c 100644
--- a/libs/WindowManager/Shell/res/layout/bubble_bar_stack_education.xml
+++ b/libs/WindowManager/Shell/res/layout/bubble_bar_stack_education.xml
@@ -14,7 +14,7 @@
   ~ See the License for the specific language governing permissions and
   ~ limitations under the License
   -->
-<com.android.wm.shell.common.bubbles.BubblePopupView
+<com.android.wm.shell.shared.bubbles.BubblePopupView
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
@@ -53,4 +53,4 @@
         android:textAlignment="center"
         android:text="@string/bubble_bar_education_stack_text"/>
 
-</com.android.wm.shell.common.bubbles.BubblePopupView>
\ No newline at end of file
+</com.android.wm.shell.shared.bubbles.BubblePopupView>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/res/values/strings.xml b/libs/WindowManager/Shell/res/values/strings.xml
index 36d0a3c..a353db7 100644
--- a/libs/WindowManager/Shell/res/values/strings.xml
+++ b/libs/WindowManager/Shell/res/values/strings.xml
@@ -155,6 +155,8 @@
     <string name="bubbles_app_settings"><xliff:g id="notification_title" example="Android Messages">%1$s</xliff:g> settings</string>
     <!-- Text used for the bubble dismiss area. Bubbles dragged to, or flung towards, this area will go away. [CHAR LIMIT=30] -->
     <string name="bubble_dismiss_text">Dismiss bubble</string>
+    <!-- Text used to move the bubble to fullscreen. [CHAR LIMIT=30] -->
+    <string name="bubble_fullscreen_text">Move to fullscreen</string>
     <!-- Button text to stop a conversation from bubbling [CHAR LIMIT=60]-->
     <string name="bubbles_dont_bubble_conversation">Don\u2019t bubble conversation</string>
     <!-- Title text for the bubbles feature education cling shown when a bubble is on screen for the first time. [CHAR LIMIT=60]-->
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BaseBubblePinController.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BaseBubblePinController.kt
similarity index 96%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BaseBubblePinController.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BaseBubblePinController.kt
index eec2468..7086691 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BaseBubblePinController.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BaseBubblePinController.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.graphics.Point
 import android.graphics.RectF
@@ -23,9 +23,9 @@
 import androidx.core.animation.Animator
 import androidx.core.animation.AnimatorListenerAdapter
 import androidx.core.animation.ObjectAnimator
-import com.android.wm.shell.common.bubbles.BaseBubblePinController.LocationChangeListener
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.LEFT
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.RIGHT
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController.LocationChangeListener
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.LEFT
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.RIGHT
 
 /**
  * Base class for common logic shared between different bubble views to support pinning bubble bar
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.aidl
similarity index 93%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.aidl
index 3c5beeb..4fe7611 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.aidl
@@ -14,6 +14,6 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 parcelable BubbleBarLocation;
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.kt
similarity index 97%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.kt
index f0bdfde..191875d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarLocation.kt
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.os.Parcel
 import android.os.Parcelable
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarUpdate.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarUpdate.java
similarity index 98%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarUpdate.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarUpdate.java
index ec3c601..5bde1e8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarUpdate.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleBarUpdate.java
@@ -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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleConstants.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleConstants.java
similarity index 89%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleConstants.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleConstants.java
index 0329b8d..3396bc4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleConstants.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleConstants.java
@@ -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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 /**
  * Constants shared between bubbles in shell & things we have to do for bubbles in launcher.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleInfo.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleInfo.java
similarity index 92%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleInfo.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleInfo.java
index 829af08..5876682 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleInfo.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubbleInfo.java
@@ -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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -48,10 +48,11 @@
     @Nullable
     private String mAppName;
     private boolean mIsImportantConversation;
+    private boolean mShowAppBadge;
 
     public BubbleInfo(String key, int flags, @Nullable String shortcutId, @Nullable Icon icon,
             int userId, String packageName, @Nullable String title, @Nullable String appName,
-            boolean isImportantConversation) {
+            boolean isImportantConversation, boolean showAppBadge) {
         mKey = key;
         mFlags = flags;
         mShortcutId = shortcutId;
@@ -61,6 +62,7 @@
         mTitle = title;
         mAppName = appName;
         mIsImportantConversation = isImportantConversation;
+        mShowAppBadge = showAppBadge;
     }
 
     private BubbleInfo(Parcel source) {
@@ -73,6 +75,7 @@
         mTitle = source.readString();
         mAppName = source.readString();
         mIsImportantConversation = source.readBoolean();
+        mShowAppBadge = source.readBoolean();
     }
 
     public String getKey() {
@@ -115,6 +118,10 @@
         return mIsImportantConversation;
     }
 
+    public boolean showAppBadge() {
+        return mShowAppBadge;
+    }
+
     /**
      * Whether this bubble is currently being hidden from the stack.
      */
@@ -172,6 +179,7 @@
         parcel.writeString(mTitle);
         parcel.writeString(mAppName);
         parcel.writeBoolean(mIsImportantConversation);
+        parcel.writeBoolean(mShowAppBadge);
     }
 
     @NonNull
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupDrawable.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupDrawable.kt
similarity index 98%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupDrawable.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupDrawable.kt
index 887af17..8681acf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupDrawable.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupDrawable.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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.annotation.ColorInt
 import android.graphics.Canvas
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupView.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupView.kt
similarity index 95%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupView.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupView.kt
index 444fbf7..802d7d1 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubblePopupView.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/BubblePopupView.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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.content.Context
 import android.graphics.Rect
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissCircleView.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissCircleView.java
similarity index 96%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissCircleView.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissCircleView.java
index 7c5bb21..0c05156 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissCircleView.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissCircleView.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 import android.content.Context;
 import android.content.res.Configuration;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissView.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissView.kt
similarity index 98%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissView.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissView.kt
index e06de9e..2bb66b0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/DismissView.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/DismissView.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2020 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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.animation.ObjectAnimator
 import android.content.Context
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/OWNERS b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/OWNERS
similarity index 100%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/OWNERS
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/OWNERS
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RelativeTouchListener.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RelativeTouchListener.kt
similarity index 98%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RelativeTouchListener.kt
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RelativeTouchListener.kt
index 4e55ba2..b1f4e33 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RelativeTouchListener.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RelativeTouchListener.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.graphics.PointF
 import android.view.MotionEvent
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RemovedBubble.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RemovedBubble.java
similarity index 94%
rename from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RemovedBubble.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RemovedBubble.java
index f90591b..c83696c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/RemovedBubble.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/bubbles/RemovedBubble.java
@@ -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.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.wm.shell.shared.bubbles;
 
 import android.annotation.NonNull;
 import android.os.Parcel;
diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
index 47d5274..424d4bf 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/desktopmode/DesktopModeFlags.kt
@@ -48,7 +48,7 @@
     TASK_STACK_OBSERVER_IN_SHELL(Flags::enableTaskStackObserverInShell, true),
     SIZE_CONSTRAINTS(Flags::enableDesktopWindowingSizeConstraints, true),
     DISABLE_SNAP_RESIZE(Flags::disableNonResizableAppSnapResizing, true),
-    DYNAMIC_INITIAL_BOUNDS(Flags::enableWindowingDynamicInitialBounds, true),
+    DYNAMIC_INITIAL_BOUNDS(Flags::enableWindowingDynamicInitialBounds, false),
     ENABLE_DESKTOP_WINDOWING_TASK_LIMIT(Flags::enableDesktopWindowingTaskLimit, true),
     BACK_NAVIGATION(Flags::enableDesktopWindowingBackNavigation, true),
     EDGE_DRAG_RESIZE(Flags::enableWindowingEdgeDragResize, true),
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/navigationbar/RegionSamplingHelper.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/handles/RegionSamplingHelper.java
similarity index 98%
rename from packages/SystemUI/shared/src/com/android/systemui/shared/navigationbar/RegionSamplingHelper.java
rename to libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/handles/RegionSamplingHelper.java
index 9999f08..b92b8ef 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/navigationbar/RegionSamplingHelper.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/handles/RegionSamplingHelper.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.shared.navigationbar;
+package com.android.wm.shell.shared.handles;
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
@@ -194,7 +194,7 @@
             ViewRootImpl viewRootImpl = mSampledView.getViewRootImpl();
             SurfaceControl stopLayerControl = null;
             if (viewRootImpl != null) {
-                 stopLayerControl = viewRootImpl.getSurfaceControl();
+                stopLayerControl = viewRootImpl.getSurfaceControl();
             }
             if (stopLayerControl == null || !stopLayerControl.isValid()) {
                 if (!mWaitingOnDraw) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
index 33949f5..ef679da 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationController.java
@@ -74,6 +74,7 @@
 import android.window.IOnBackInvokedCallback;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
+import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
 import com.android.internal.annotations.VisibleForTesting;
@@ -1272,19 +1273,24 @@
             ComponentName openComponent = null;
             int tmpSize;
             int openTaskId = INVALID_TASK_ID;
+            WindowContainerToken openToken = null;
             for (int j = init.getChanges().size() - 1; j >= 0; --j) {
                 final TransitionInfo.Change change = init.getChanges().get(j);
                 if (change.hasFlags(FLAG_BACK_GESTURE_ANIMATED)) {
                     openComponent = findComponentName(change);
                     openTaskId = findTaskId(change);
+                    openToken = findToken(change);
                     if (change.hasFlags(FLAG_SHOW_WALLPAPER)) {
                         openShowWallpaper = true;
                     }
                     break;
                 }
             }
-            if (openComponent == null && openTaskId == INVALID_TASK_ID) {
-                // shouldn't happen.
+            if (openComponent == null && openTaskId == INVALID_TASK_ID && openToken == null) {
+                // This shouldn't happen, but if that happen, consume the initial transition anyway.
+                Log.e(TAG, "Unable to merge following transition, cannot find the gesture "
+                        + "animated target from the open transition=" + mOpenTransitionInfo);
+                mOpenTransitionInfo = null;
                 return;
             }
             // find first non-prepare open target
@@ -1315,7 +1321,7 @@
                 boolean moveToTop = false;
                 for (int j = info.getChanges().size() - 1; j >= 0; --j) {
                     final TransitionInfo.Change change = info.getChanges().get(j);
-                    if (isSameChangeTarget(openComponent, openTaskId, change)) {
+                    if (isSameChangeTarget(openComponent, openTaskId, openToken, change)) {
                         moveToTop = change.hasFlags(FLAG_MOVED_TO_TOP);
                         info.getChanges().remove(j);
                     } else if ((openShowWallpaper && change.hasFlags(FLAG_IS_WALLPAPER))
@@ -1329,7 +1335,7 @@
                     for (int i = 0; i < tmpSize; ++i) {
                         final TransitionInfo.Change change = init.getChanges().get(i);
                         if (moveToTop) {
-                            if (isSameChangeTarget(openComponent, openTaskId, change)) {
+                            if (isSameChangeTarget(openComponent, openTaskId, openToken, change)) {
                                 change.setFlags(change.getFlags() | FLAG_MOVED_TO_TOP);
                             }
                         }
@@ -1358,7 +1364,7 @@
                 if (nonBackClose && nonBackOpen) {
                     for (int j = info.getChanges().size() - 1; j >= 0; --j) {
                         final TransitionInfo.Change change = info.getChanges().get(j);
-                        if (isSameChangeTarget(openComponent, openTaskId, change)) {
+                        if (isSameChangeTarget(openComponent, openTaskId, openToken, change)) {
                             info.getChanges().remove(j);
                         } else if ((openShowWallpaper && change.hasFlags(FLAG_IS_WALLPAPER))) {
                             info.getChanges().remove(j);
@@ -1368,6 +1374,8 @@
             }
             ProtoLog.d(WM_SHELL_BACK_PREVIEW, "Back animation transition, merge pending "
                     + "transitions result=%s", info);
+            // Only handle one merge transition request.
+            mOpenTransitionInfo = null;
         }
 
         @Override
@@ -1378,7 +1386,9 @@
                 mClosePrepareTransition = null;
             }
             // try to handle unexpected transition
-            mergePendingTransitions(info);
+            if (mOpenTransitionInfo != null) {
+                mergePendingTransitions(info);
+            }
 
             if (isNotGestureBackTransition(info) || shouldCancelAnimation(info)
                     || !mCloseTransitionRequested) {
@@ -1628,6 +1638,10 @@
         return false;
     }
 
+    private static WindowContainerToken findToken(TransitionInfo.Change change) {
+        return change.getContainer();
+    }
+
     private static ComponentName findComponentName(TransitionInfo.Change change) {
         final ComponentName componentName = change.getActivityComponent();
         if (componentName != null) {
@@ -1649,11 +1663,13 @@
     }
 
     private static boolean isSameChangeTarget(ComponentName topActivity, int taskId,
-            TransitionInfo.Change change) {
+            WindowContainerToken token, TransitionInfo.Change change) {
         final ComponentName openChange = findComponentName(change);
         final int firstTaskId = findTaskId(change);
+        final WindowContainerToken openToken = findToken(change);
         return (openChange != null && openChange == topActivity)
-                || (firstTaskId != INVALID_TASK_ID && firstTaskId == taskId);
+                || (firstTaskId != INVALID_TASK_ID && firstTaskId == taskId)
+                || (openToken != null && token == openToken);
     }
 
     private static boolean canBeTransitionTarget(TransitionInfo.Change change) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
index 3e758bb..0c95934 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubble.java
@@ -53,9 +53,9 @@
 import com.android.wm.shell.Flags;
 import com.android.wm.shell.bubbles.bar.BubbleBarExpandedView;
 import com.android.wm.shell.bubbles.bar.BubbleBarLayerView;
-import com.android.wm.shell.common.bubbles.BubbleInfo;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
 import com.android.wm.shell.shared.annotations.ShellMainThread;
+import com.android.wm.shell.shared.bubbles.BubbleInfo;
 
 import java.io.PrintWriter;
 import java.util.List;
@@ -349,7 +349,8 @@
                 getPackageName(),
                 getTitle(),
                 getAppName(),
-                isImportantConversation());
+                isImportantConversation(),
+                !isAppLaunchIntent());
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index b508c1b..c545d73 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -104,14 +104,14 @@
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.TaskStackListenerCallback;
 import com.android.wm.shell.common.TaskStackListenerImpl;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
-import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.onehanded.OneHandedTransitionCallback;
 import com.android.wm.shell.pip.PinnedStackListenerForwarder;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
 import com.android.wm.shell.shared.annotations.ShellMainThread;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellController;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
index 4ad1802..709a7bd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleData.java
@@ -41,10 +41,10 @@
 import com.android.internal.util.FrameworkStatsLog;
 import com.android.wm.shell.R;
 import com.android.wm.shell.bubbles.Bubbles.DismissReason;
-import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
-import com.android.wm.shell.common.bubbles.RemovedBubble;
 import com.android.wm.shell.shared.annotations.ShellBackgroundThread;
 import com.android.wm.shell.shared.annotations.ShellMainThread;
+import com.android.wm.shell.shared.bubbles.BubbleBarUpdate;
+import com.android.wm.shell.shared.bubbles.RemovedBubble;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedViewManager.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedViewManager.kt
index 4e80e90..ec4854b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedViewManager.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedViewManager.kt
@@ -16,7 +16,7 @@
 
 package com.android.wm.shell.bubbles
 
-import com.android.wm.shell.common.bubbles.BubbleBarLocation
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation
 
 /** Manager interface for bubble expanded views. */
 interface BubbleExpandedViewManager {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePopupViewExt.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePopupViewExt.kt
index bdb09e1..fd110a276 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePopupViewExt.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubblePopupViewExt.kt
@@ -17,8 +17,8 @@
 
 import android.graphics.Color
 import com.android.wm.shell.R
-import com.android.wm.shell.common.bubbles.BubblePopupDrawable
-import com.android.wm.shell.common.bubbles.BubblePopupView
+import com.android.wm.shell.shared.bubbles.BubblePopupDrawable
+import com.android.wm.shell.shared.bubbles.BubblePopupView
 
 /**
  * A convenience method to setup the [BubblePopupView] with the correct config using local resources
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 0cf187b..c386c93 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
@@ -32,7 +32,7 @@
 import com.android.internal.protolog.ProtoLog;
 import com.android.launcher3.icons.IconNormalizer;
 import com.android.wm.shell.R;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 
 /**
  * Keeps track of display size, configuration, and specific bubble sizes. One place for all
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
index 53bbf88..2795881 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleStackView.java
@@ -24,10 +24,10 @@
 import static com.android.wm.shell.bubbles.BubblePositioner.NUM_VISIBLE_WHEN_RESTING;
 import static com.android.wm.shell.bubbles.BubblePositioner.StackPinnedEdge.LEFT;
 import static com.android.wm.shell.bubbles.BubblePositioner.StackPinnedEdge.RIGHT;
-import static com.android.wm.shell.common.bubbles.BubbleConstants.BUBBLE_EXPANDED_SCRIM_ALPHA;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_BUBBLES;
 import static com.android.wm.shell.shared.animation.Interpolators.ALPHA_IN;
 import static com.android.wm.shell.shared.animation.Interpolators.ALPHA_OUT;
+import static com.android.wm.shell.shared.bubbles.BubbleConstants.BUBBLE_EXPANDED_SCRIM_ALPHA;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -91,8 +91,8 @@
 import com.android.wm.shell.bubbles.animation.StackAnimationController;
 import com.android.wm.shell.common.FloatingContentCoordinator;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.bubbles.DismissView;
-import com.android.wm.shell.common.bubbles.RelativeTouchListener;
+import com.android.wm.shell.shared.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.RelativeTouchListener;
 import com.android.wm.shell.shared.animation.Interpolators;
 import com.android.wm.shell.shared.animation.PhysicsAnimator;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
index 9a27fb6..62895fe 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/Bubbles.java
@@ -38,9 +38,9 @@
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
 
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
-import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.shared.annotations.ExternalThread;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarUpdate;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.Target;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/DismissViewExt.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/DismissViewExt.kt
index 48692d4..00a8172 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/DismissViewExt.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/DismissViewExt.kt
@@ -18,7 +18,7 @@
 package com.android.wm.shell.bubbles
 
 import com.android.wm.shell.R
-import com.android.wm.shell.common.bubbles.DismissView
+import com.android.wm.shell.shared.bubbles.DismissView
 
 fun DismissView.setup() {
     setup(DismissView.Config(
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubbles.aidl b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubbles.aidl
index 5779a8f..1855b93 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubbles.aidl
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubbles.aidl
@@ -20,7 +20,7 @@
 import android.graphics.Rect;
 import android.content.pm.ShortcutInfo;
 import com.android.wm.shell.bubbles.IBubblesListener;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 
 /**
  * Interface that is exposed to remote callers (launcher) to manipulate the bubbles feature when
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubblesListener.aidl b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubblesListener.aidl
index 14d29cd..eb907db 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubblesListener.aidl
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/IBubblesListener.aidl
@@ -17,7 +17,7 @@
 package com.android.wm.shell.bubbles;
 
 import android.os.Bundle;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 /**
  * Listener interface that Launcher attaches to SystemUI to get bubbles callbacks.
  */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
index 6d868d2..f90b2aa 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedView.java
@@ -46,7 +46,7 @@
 import com.android.wm.shell.bubbles.BubbleTaskView;
 import com.android.wm.shell.bubbles.BubbleTaskViewHelper;
 import com.android.wm.shell.bubbles.Bubbles;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
 import com.android.wm.shell.taskview.TaskView;
 
 import java.util.function.Supplier;
@@ -228,6 +228,13 @@
             public void onDismissBubble(Bubble bubble) {
                 mManager.dismissBubble(bubble, Bubbles.DISMISS_USER_GESTURE);
             }
+
+            @Override
+            public void onMoveToFullscreen(Bubble bubble) {
+                if (mTaskView != null) {
+                    mTaskView.moveToFullscreen();
+                }
+            }
         });
         mHandleView.setOnClickListener(view -> {
             mMenuViewController.showMenu(true /* animated */);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
index eeb5c94..07463bb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarExpandedViewDragController.kt
@@ -20,8 +20,8 @@
 import android.view.MotionEvent
 import android.view.View
 import com.android.wm.shell.bubbles.BubblePositioner
-import com.android.wm.shell.common.bubbles.DismissView
-import com.android.wm.shell.common.bubbles.RelativeTouchListener
+import com.android.wm.shell.shared.bubbles.DismissView
+import com.android.wm.shell.shared.bubbles.RelativeTouchListener
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject
 
 /** Controller for handling drag interactions with [BubbleBarExpandedView] */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
index ac42453..1c9c195 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarLayerView.java
@@ -44,9 +44,9 @@
 import com.android.wm.shell.bubbles.DeviceConfig;
 import com.android.wm.shell.bubbles.DismissViewUtils;
 import com.android.wm.shell.bubbles.bar.BubbleBarExpandedViewDragController.DragListener;
-import com.android.wm.shell.common.bubbles.BaseBubblePinController;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
-import com.android.wm.shell.common.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.DismissView;
 
 import kotlin.Unit;
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarMenuViewController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarMenuViewController.java
index 0d72998..5148107 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarMenuViewController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleBarMenuViewController.java
@@ -29,6 +29,7 @@
 import androidx.dynamicanimation.animation.DynamicAnimation;
 import androidx.dynamicanimation.animation.SpringForce;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.R;
 import com.android.wm.shell.bubbles.Bubble;
 import com.android.wm.shell.shared.animation.PhysicsAnimator;
@@ -219,6 +220,21 @@
                 }
         ));
 
+        if (Flags.enableBubbleAnything() || Flags.enableBubbleToFullscreen()) {
+            menuActions.add(new BubbleBarMenuView.MenuAction(
+                    Icon.createWithResource(resources,
+                            R.drawable.desktop_mode_ic_handle_menu_fullscreen),
+                    resources.getString(R.string.bubble_fullscreen_text),
+                    tintColor,
+                    view -> {
+                        hideMenu(true /* animated */);
+                        if (mListener != null) {
+                            mListener.onMoveToFullscreen(bubble);
+                        }
+                    }
+            ));
+        }
+
         return menuActions;
     }
 
@@ -249,5 +265,10 @@
          * Dismiss bubble and remove it from the bubble stack
          */
         void onDismissBubble(Bubble bubble);
+
+        /**
+         * Move the bubble to fullscreen.
+         */
+        void onMoveToFullscreen(Bubble bubble);
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleEducationViewController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleEducationViewController.kt
index e108f7b..9fd255d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleEducationViewController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleEducationViewController.kt
@@ -34,9 +34,9 @@
 import com.android.wm.shell.bubbles.BubbleEducationController
 import com.android.wm.shell.bubbles.BubbleViewProvider
 import com.android.wm.shell.bubbles.setup
-import com.android.wm.shell.common.bubbles.BubblePopupDrawable
-import com.android.wm.shell.common.bubbles.BubblePopupView
 import com.android.wm.shell.shared.animation.PhysicsAnimator
+import com.android.wm.shell.shared.bubbles.BubblePopupDrawable
+import com.android.wm.shell.shared.bubbles.BubblePopupView
 import kotlin.math.roundToInt
 
 /** Manages bubble education presentation and animation */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinController.kt
index 651bf02..23ba2bf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/bar/BubbleExpandedViewPinController.kt
@@ -25,8 +25,8 @@
 import androidx.core.view.updateLayoutParams
 import com.android.wm.shell.R
 import com.android.wm.shell.bubbles.BubblePositioner
-import com.android.wm.shell.common.bubbles.BaseBubblePinController
-import com.android.wm.shell.common.bubbles.BubbleBarLocation
+import com.android.wm.shell.shared.bubbles.BaseBubblePinController
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation
 
 /**
  * Controller to manage pinning bubble bar to left or right when dragging starts from the bubble bar
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
index fad3dee..1929729 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/ScreenshotUtils.java
@@ -42,6 +42,7 @@
                     .setSourceCrop(crop)
                     .setCaptureSecureLayers(true)
                     .setAllowProtected(true)
+                    .setHintForSeamlessTransition(true)
                     .build()));
     }
 
@@ -78,6 +79,9 @@
             mTransaction.setColorSpace(mScreenshot, buffer.getColorSpace());
             mTransaction.reparent(mScreenshot, mParentSurfaceControl);
             mTransaction.setLayer(mScreenshot, mLayer);
+            if (buffer.containsHdrLayers()) {
+                mTransaction.setDimmingEnabled(mScreenshot, false);
+            }
             mTransaction.show(mScreenshot);
             mTransaction.apply();
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index 42937c1..4adea23 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -148,7 +148,11 @@
  * dependencies that are device/form factor SystemUI implementation specific should go into their
  * respective modules (ie. {@link WMShellModule} for handheld, {@link TvWMShellModule} for tv, etc.)
  */
-@Module(includes = WMShellConcurrencyModule.class)
+@Module(
+        includes = {
+                WMShellConcurrencyModule.class,
+                WMShellCoroutinesModule.class
+        })
 public abstract class WMShellBaseModule {
 
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellCoroutinesModule.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellCoroutinesModule.kt
index 423fe57..cc47dbb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellCoroutinesModule.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellCoroutinesModule.kt
@@ -61,20 +61,20 @@
   @WMSingleton
   @ShellMainThread
   fun provideApplicationScope(
-      @ShellMainThread applicationDispatcher: CoroutineDispatcher,
+      @ShellMainThread applicationDispatcher: MainCoroutineDispatcher,
   ): CoroutineScope = CoroutineScope(applicationDispatcher)
 
   @Provides
   @WMSingleton
   @ShellBackgroundThread
   fun provideBackgroundCoroutineScope(
-      @ShellBackgroundThread backgroundDispatcher: CoroutineDispatcher,
+      @ShellBackgroundThread backgroundDispatcher: MainCoroutineDispatcher,
   ): CoroutineScope = CoroutineScope(backgroundDispatcher)
 
   @Provides
   @WMSingleton
   @ShellBackgroundThread
   fun provideBackgroundCoroutineContext(
-      @ShellBackgroundThread backgroundDispatcher: CoroutineDispatcher
+      @ShellBackgroundThread backgroundDispatcher: MainCoroutineDispatcher
   ): CoroutineContext = backgroundDispatcher + SupervisorJob()
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 4db8a82..02ecfd9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -59,6 +59,7 @@
 import com.android.wm.shell.dagger.back.ShellBackAnimationModule;
 import com.android.wm.shell.dagger.pip.PipModule;
 import com.android.wm.shell.desktopmode.DefaultDragToDesktopTransitionHandler;
+import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
 import com.android.wm.shell.desktopmode.DesktopModeDragAndDropTransitionHandler;
 import com.android.wm.shell.desktopmode.DesktopModeEventLogger;
 import com.android.wm.shell.desktopmode.DesktopModeLoggerTransitionObserver;
@@ -72,6 +73,7 @@
 import com.android.wm.shell.desktopmode.ReturnToDragStartAnimator;
 import com.android.wm.shell.desktopmode.SpringDragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler;
+import com.android.wm.shell.desktopmode.education.AppHandleEducationController;
 import com.android.wm.shell.desktopmode.education.AppHandleEducationFilter;
 import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository;
 import com.android.wm.shell.draganddrop.DragAndDropController;
@@ -117,6 +119,8 @@
 import dagger.Module;
 import dagger.Provides;
 
+import kotlinx.coroutines.CoroutineScope;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -237,7 +241,8 @@
             InteractionJankMonitor interactionJankMonitor,
             AppToWebGenericLinksParser genericLinksParser,
             MultiInstanceHelper multiInstanceHelper,
-            Optional<DesktopTasksLimiter> desktopTasksLimiter) {
+            Optional<DesktopTasksLimiter> desktopTasksLimiter,
+            Optional<DesktopActivityOrientationChangeHandler> desktopActivityOrientationHandler) {
         if (DesktopModeStatus.canEnterDesktopMode(context)) {
             return new DesktopModeWindowDecorViewModel(
                     context,
@@ -259,7 +264,8 @@
                     interactionJankMonitor,
                     genericLinksParser,
                     multiInstanceHelper,
-                    desktopTasksLimiter);
+                    desktopTasksLimiter,
+                    desktopActivityOrientationHandler);
         }
         return new CaptionWindowDecorViewModel(
                 context,
@@ -677,6 +683,24 @@
 
     @WMSingleton
     @Provides
+    static Optional<DesktopActivityOrientationChangeHandler> provideActivityOrientationHandler(
+            Context context,
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
+            TaskStackListenerImpl taskStackListener,
+            ToggleResizeDesktopTaskTransitionHandler toggleResizeDesktopTaskTransitionHandler,
+            @DynamicOverride DesktopModeTaskRepository desktopModeTaskRepository
+    ) {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
+            return Optional.of(new DesktopActivityOrientationChangeHandler(
+                    context, shellInit, shellTaskOrganizer, taskStackListener,
+                    toggleResizeDesktopTaskTransitionHandler, desktopModeTaskRepository));
+        }
+        return Optional.empty();
+    }
+
+    @WMSingleton
+    @Provides
     static Optional<DesktopTasksTransitionObserver> provideDesktopTasksTransitionObserver(
             Context context,
             Optional<DesktopModeTaskRepository> desktopModeTaskRepository,
@@ -722,6 +746,17 @@
         return new AppHandleEducationFilter(context, appHandleEducationDatastoreRepository);
     }
 
+    @WMSingleton
+    @Provides
+    static AppHandleEducationController provideAppHandleEducationController(
+            AppHandleEducationFilter appHandleEducationFilter,
+            ShellTaskOrganizer shellTaskOrganizer,
+            AppHandleEducationDatastoreRepository appHandleEducationDatastoreRepository,
+            @ShellMainThread CoroutineScope applicationScope) {
+        return new AppHandleEducationController(appHandleEducationFilter,
+                shellTaskOrganizer, appHandleEducationDatastoreRepository, applicationScope);
+    }
+
     //
     // Drag and drop
     //
@@ -763,7 +798,8 @@
     @Provides
     static Object provideIndependentShellComponentsToCreate(
             DragAndDropController dragAndDropController,
-            Optional<DesktopTasksTransitionObserver> desktopTasksTransitionObserverOptional
+            Optional<DesktopTasksTransitionObserver> desktopTasksTransitionObserverOptional,
+            AppHandleEducationController appHandleEducationController
     ) {
         return new Object();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt
new file mode 100644
index 0000000..59e0068
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt
@@ -0,0 +1,112 @@
+/*
+ * 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.wm.shell.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.content.Context
+import android.content.pm.ActivityInfo
+import android.content.pm.ActivityInfo.ScreenOrientation
+import android.content.res.Configuration.ORIENTATION_LANDSCAPE
+import android.content.res.Configuration.ORIENTATION_PORTRAIT
+import android.graphics.Rect
+import android.util.Size
+import android.window.WindowContainerTransaction
+import com.android.window.flags.Flags
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.common.TaskStackListenerCallback
+import com.android.wm.shell.common.TaskStackListenerImpl
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.android.wm.shell.sysui.ShellInit
+
+/** Handles task resizing to respect orientation change of non-resizeable activities in desktop. */
+class DesktopActivityOrientationChangeHandler(
+    context: Context,
+    shellInit: ShellInit,
+    private val shellTaskOrganizer: ShellTaskOrganizer,
+    private val taskStackListener: TaskStackListenerImpl,
+    private val resizeHandler: ToggleResizeDesktopTaskTransitionHandler,
+    private val taskRepository: DesktopModeTaskRepository,
+) {
+
+    init {
+        if (DesktopModeStatus.canEnterDesktopMode(context)) {
+            shellInit.addInitCallback({ onInit() }, this)
+        }
+    }
+
+    private fun onInit() {
+        taskStackListener.addListener(object : TaskStackListenerCallback {
+            override fun onActivityRequestedOrientationChanged(
+                taskId: Int,
+                @ScreenOrientation requestedOrientation: Int
+            ) {
+                // Handle requested screen orientation changes at runtime.
+                handleActivityOrientationChange(taskId, requestedOrientation)
+            }
+        })
+    }
+
+    /**
+     * Triggered with onTaskInfoChanged to handle:
+     * * New activity launching from same task with different orientation
+     * * Top activity closing in same task with different orientation to previous activity
+     */
+    fun handleActivityOrientationChange(oldTask: RunningTaskInfo, newTask: RunningTaskInfo) {
+        val newTopActivityInfo = newTask.topActivityInfo ?: return
+        val oldTopActivityInfo = oldTask.topActivityInfo ?: return
+        // Check if screen orientation is different from old task info so there is no duplicated
+        // calls to handle runtime requested orientation changes.
+        if (oldTopActivityInfo.screenOrientation != newTopActivityInfo.screenOrientation) {
+            handleActivityOrientationChange(newTask.taskId, newTopActivityInfo.screenOrientation)
+        }
+    }
+
+    private fun handleActivityOrientationChange(
+        taskId: Int,
+        @ScreenOrientation requestedOrientation: Int
+    ) {
+        if (!Flags.respectOrientationChangeForUnresizeable()) return
+        val task = shellTaskOrganizer.getRunningTaskInfo(taskId) ?: return
+        if (!isDesktopModeShowing(task.displayId) || !task.isFreeform || task.isResizeable) return
+
+        val taskBounds = task.configuration.windowConfiguration.bounds
+        val taskHeight = taskBounds.height()
+        val taskWidth = taskBounds.width()
+        if (taskWidth == taskHeight) return
+        val orientation =
+            if (taskWidth > taskHeight) ORIENTATION_LANDSCAPE else ORIENTATION_PORTRAIT
+
+        // Non-resizeable activity requested opposite orientation.
+        if (orientation == ORIENTATION_PORTRAIT
+                && ActivityInfo.isFixedOrientationLandscape(requestedOrientation)
+            || orientation == ORIENTATION_LANDSCAPE
+                && ActivityInfo.isFixedOrientationPortrait(requestedOrientation)) {
+
+            val finalSize = Size(taskHeight, taskWidth)
+            // Use the center x as the resizing anchor point.
+            val left = taskBounds.centerX() - finalSize.width / 2
+            val right = left + finalSize.width
+            val finalBounds = Rect(left, taskBounds.top, right, taskBounds.top + finalSize.height)
+
+            val wct = WindowContainerTransaction().setBounds(task.token, finalBounds)
+            resizeHandler.startTransition(wct)
+        }
+    }
+
+    private fun isDesktopModeShowing(displayId: Int): Boolean =
+        taskRepository.getVisibleTaskCount(displayId) > 0
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
index 336e5e3..0637474 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserver.kt
@@ -22,6 +22,7 @@
 import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.content.Context
 import android.os.IBinder
+import android.os.SystemProperties
 import android.os.Trace
 import android.util.SparseArray
 import android.view.SurfaceControl
@@ -52,8 +53,6 @@
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
 
-const val VISIBLE_TASKS_COUNTER_NAME = "DESKTOP_MODE_VISIBLE_TASKS"
-
 /**
  * A [Transitions.TransitionObserver] that observes transitions and the proposed changes to log
  * appropriate desktop mode session log events. This observes transitions related to desktop mode
@@ -307,6 +306,8 @@
                         VISIBLE_TASKS_COUNTER_NAME,
                         postTransitionVisibleFreeformTasks.size().toLong()
                     )
+                    SystemProperties.set(VISIBLE_TASKS_COUNTER_SYSTEM_PROPERTY,
+                        postTransitionVisibleFreeformTasks.size().toString())
                 }
                 // old tasks that were resized or repositioned
                 // TODO(b/347935387): Log changes only once they are stable.
@@ -326,6 +327,8 @@
                     VISIBLE_TASKS_COUNTER_NAME,
                     postTransitionVisibleFreeformTasks.size().toLong()
                 )
+                SystemProperties.set(VISIBLE_TASKS_COUNTER_SYSTEM_PROPERTY,
+                    postTransitionVisibleFreeformTasks.size().toString())
             }
         }
     }
@@ -431,4 +434,12 @@
         return this.type == WindowManager.TRANSIT_TO_FRONT &&
             this.flags == WindowManager.TRANSIT_FLAG_IS_RECENTS
     }
+
+    companion object {
+        @VisibleForTesting
+        const val VISIBLE_TASKS_COUNTER_NAME = "desktop_mode_visible_tasks"
+        @VisibleForTesting
+        const val VISIBLE_TASKS_COUNTER_SYSTEM_PROPERTY =
+            "debug.tracing." + VISIBLE_TASKS_COUNTER_NAME
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
index b68b436..c8ffe28 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeUtils.kt
@@ -171,6 +171,18 @@
         minOf(appBounds.height(), appBounds.width()).toFloat()
 }
 
+/** Returns true if task's width or height is maximized else returns false. */
+fun isTaskWidthOrHeightEqual(taskBounds: Rect, stableBounds: Rect): Boolean {
+    return taskBounds.width() == stableBounds.width() ||
+            taskBounds.height() == stableBounds.height()
+}
+
+/** Returns true if task bound is equal to stable bounds else returns false. */
+fun isTaskBoundsEqual(taskBounds: Rect, stableBounds: Rect): Boolean {
+    return taskBounds.width() == stableBounds.width() &&
+            taskBounds.height() == stableBounds.height()
+}
+
 /**
  * Calculates the desired initial bounds for applications in desktop windowing. This is done as a
  * scale of the screen bounds.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index 2852631..1d16980 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -169,6 +169,9 @@
             }
         }
 
+    @VisibleForTesting
+    var taskbarDesktopTaskListener: TaskbarDesktopTaskListener? = null
+
     /** Task id of the task currently being dragged from fullscreen/split. */
     val draggingTaskId
         get() = dragToDesktopTransitionHandler.draggingTaskId
@@ -610,13 +613,10 @@
         val currentTaskBounds = taskInfo.configuration.windowConfiguration.bounds
         val destinationBounds = Rect()
 
-        val isMaximized = if (taskInfo.isResizeable) {
-            currentTaskBounds == stableBounds
-        } else {
-            currentTaskBounds.width() == stableBounds.width()
-                    || currentTaskBounds.height() == stableBounds.height()
-        }
-
+        val isMaximized = isTaskMaximized(taskInfo, stableBounds)
+        // If the task is currently maximized, we will toggle it not to be and vice versa. This is
+        // helpful to eliminate the current task from logic to calculate taskbar corner rounding.
+        val willMaximize = !isMaximized
         if (isMaximized) {
             // The desktop task is at the maximized width and/or height of the stable bounds.
             // If the task's pre-maximize stable bounds were saved, toggle the task to those bounds.
@@ -651,6 +651,18 @@
             }
         }
 
+
+
+        val shouldRestoreToSnap =
+            isMaximized && isTaskSnappedToHalfScreen(taskInfo, destinationBounds)
+
+        logD("willMaximize = %s", willMaximize)
+        logD("shouldRestoreToSnap = %s", shouldRestoreToSnap)
+
+        val doesAnyTaskRequireTaskbarRounding = willMaximize || shouldRestoreToSnap ||
+                doesAnyTaskRequireTaskbarRounding(taskInfo.displayId, taskInfo.taskId)
+
+        taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding)
         val wct = WindowContainerTransaction().setBounds(taskInfo.token, destinationBounds)
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             toggleResizeDesktopTaskTransitionHandler.startTransition(wct)
@@ -659,23 +671,98 @@
         }
     }
 
+    private fun isTaskMaximized(
+        taskInfo: RunningTaskInfo,
+        stableBounds: Rect
+    ): Boolean {
+        val currentTaskBounds = taskInfo.configuration.windowConfiguration.bounds
+
+        return if (taskInfo.isResizeable) {
+            isTaskBoundsEqual(currentTaskBounds, stableBounds)
+        } else {
+            isTaskWidthOrHeightEqual(currentTaskBounds, stableBounds)
+        }
+    }
+
+    private fun isMaximizedToStableBoundsEdges(
+        taskInfo: RunningTaskInfo,
+        stableBounds: Rect
+    ): Boolean {
+        val currentTaskBounds = taskInfo.configuration.windowConfiguration.bounds
+        return isTaskBoundsEqual(currentTaskBounds, stableBounds)
+    }
+
+    /** Returns if current task bound is snapped to half screen */
+    private fun isTaskSnappedToHalfScreen(
+        taskInfo: RunningTaskInfo,
+        taskBounds: Rect = taskInfo.configuration.windowConfiguration.bounds
+    ): Boolean =
+        getSnapBounds(taskInfo, SnapPosition.LEFT) == taskBounds ||
+                getSnapBounds(taskInfo, SnapPosition.RIGHT) == taskBounds
+
+    @VisibleForTesting
+    fun doesAnyTaskRequireTaskbarRounding(
+        displayId: Int,
+        excludeTaskId: Int? = null,
+    ): Boolean {
+        val doesAnyTaskRequireTaskbarRounding =
+            taskRepository.getActiveNonMinimizedOrderedTasks(displayId)
+                // exclude current task since maximize/restore transition has not taken place yet.
+                .filterNot { taskId -> taskId == excludeTaskId }
+                .any { taskId ->
+                    val taskInfo = shellTaskOrganizer.getRunningTaskInfo(taskId)!!
+                    val displayLayout = displayController.getDisplayLayout(taskInfo.displayId)
+                    val stableBounds = Rect().apply { displayLayout?.getStableBounds(this) }
+                    logD("taskInfo = %s", taskInfo)
+                    logD(
+                        "isTaskSnappedToHalfScreen(taskInfo) = %s",
+                        isTaskSnappedToHalfScreen(taskInfo)
+                    )
+                    logD(
+                        "isMaximizedToStableBoundsEdges(taskInfo, stableBounds) = %s",
+                        isMaximizedToStableBoundsEdges(taskInfo, stableBounds)
+                    )
+                    isTaskSnappedToHalfScreen(taskInfo)
+                            || isMaximizedToStableBoundsEdges(taskInfo, stableBounds)
+                }
+
+        logD("doesAnyTaskRequireTaskbarRounding = %s", doesAnyTaskRequireTaskbarRounding)
+        return doesAnyTaskRequireTaskbarRounding
+    }
+
     /**
      * Quick-resize to the right or left half of the stable bounds.
      *
      * @param taskInfo current task that is being snap-resized via dragging or maximize menu button
+     * @param taskSurface the leash of the task being dragged
      * @param currentDragBounds current position of the task leash being dragged (or current task
      *                          bounds if being snapped resize via maximize menu button)
      * @param position the portion of the screen (RIGHT or LEFT) we want to snap the task to.
      */
     fun snapToHalfScreen(
         taskInfo: RunningTaskInfo,
+        taskSurface: SurfaceControl,
         currentDragBounds: Rect,
         position: SnapPosition
     ) {
         val destinationBounds = getSnapBounds(taskInfo, position)
+        if (destinationBounds == taskInfo.configuration.windowConfiguration.bounds) {
+            // Handle the case where we attempt to snap resize when already snap resized: the task
+            // position won't need to change but we want to animate the surface going back to the
+            // snapped position from the "dragged-to-the-edge" position.
+            if (destinationBounds != currentDragBounds) {
+                returnToDragStartAnimator.start(
+                    taskInfo.taskId,
+                    taskSurface,
+                    startBounds = currentDragBounds,
+                    endBounds = destinationBounds,
+                    isResizable = taskInfo.isResizeable
+                )
+            }
+            return
+        }
 
-        if (destinationBounds == taskInfo.configuration.windowConfiguration.bounds) return
-
+        taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(true)
         val wct = WindowContainerTransaction().setBounds(taskInfo.token, destinationBounds)
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             toggleResizeDesktopTaskTransitionHandler.startTransition(wct, currentDragBounds)
@@ -703,13 +790,14 @@
                 taskInfo.taskId,
                 taskSurface,
                 startBounds = currentDragBounds,
-                endBounds = dragStartBounds
+                endBounds = dragStartBounds,
+                isResizable = taskInfo.isResizeable,
             )
         } else {
             interactionJankMonitor.begin(
                 taskSurface, context, CUJ_DESKTOP_MODE_SNAP_RESIZE, "drag_resizable"
             )
-            snapToHalfScreen(taskInfo, currentDragBounds, position)
+            snapToHalfScreen(taskInfo, taskSurface, currentDragBounds, position)
         }
     }
 
@@ -806,6 +894,10 @@
             .mapNotNull { taskId -> shellTaskOrganizer.getRunningTaskInfo(taskId) }
             .reversed() // Start from the back so the front task is brought forward last
             .forEach { task -> wct.reorder(task.token, /* onTop= */ true) }
+
+        taskbarDesktopTaskListener?.
+            onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding(displayId))
+
         return taskToMinimize
     }
 
@@ -821,6 +913,7 @@
         val intent = Intent(context, DesktopWallpaperActivity::class.java)
         val options =
             ActivityOptions.makeBasic().apply {
+                launchWindowingMode = WINDOWING_MODE_FULLSCREEN
                 pendingIntentBackgroundActivityStartMode =
                     ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
             }
@@ -1156,6 +1249,12 @@
         ) {
             wct.removeTask(task.token)
         }
+        taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
+            doesAnyTaskRequireTaskbarRounding(
+                task.displayId,
+                task.id
+            )
+        )
         return if (wct.isEmpty) null else wct
     }
 
@@ -1450,6 +1549,8 @@
         }
         // A freeform drag-move ended, remove the indicator immediately.
         releaseVisualIndicator()
+        taskbarDesktopTaskListener
+            ?.onTaskbarCornerRoundingUpdate(doesAnyTaskRequireTaskbarRounding(taskInfo.displayId))
     }
 
     /**
@@ -1686,17 +1787,39 @@
                 }
             }
 
+        private val mTaskbarDesktopTaskListener: TaskbarDesktopTaskListener =
+                object : TaskbarDesktopTaskListener {
+                    override fun onTaskbarCornerRoundingUpdate(
+                        hasTasksRequiringTaskbarRounding: Boolean) {
+                        ProtoLog.v(
+                                WM_SHELL_DESKTOP_MODE,
+                                "IDesktopModeImpl: onTaskbarCornerRoundingUpdate " +
+                                        "doesAnyTaskRequireTaskbarRounding=%s",
+                                hasTasksRequiringTaskbarRounding
+                        )
+
+                        remoteListener.call { l ->
+                            l.onTaskbarCornerRoundingUpdate(hasTasksRequiringTaskbarRounding)
+                        }
+                    }
+                }
+
         init {
             remoteListener =
                 SingleInstanceRemoteListener<DesktopTasksController, IDesktopTaskListener>(
                     controller,
                     { c ->
-                        c.taskRepository.addVisibleTasksListener(
-                            listener,
-                            c.mainExecutor
-                        )
+                        run {
+                            c.taskRepository.addVisibleTasksListener(listener, c.mainExecutor)
+                            c.taskbarDesktopTaskListener = mTaskbarDesktopTaskListener
+                        }
                     },
-                    { c -> c.taskRepository.removeVisibleTasksListener(listener) }
+                    { c ->
+                        run {
+                            c.taskRepository.removeVisibleTasksListener(listener)
+                            c.taskbarDesktopTaskListener = null
+                        }
+                    }
                 )
         }
 
@@ -1779,6 +1902,16 @@
         private const val TAG = "DesktopTasksController"
     }
 
+    /** Defines interface for classes that can listen to changes for task resize. */
+    // TODO(b/343931111): Migrate to using TransitionObservers when ready
+    interface TaskbarDesktopTaskListener {
+        /**
+         * [hasTasksRequiringTaskbarRounding] is true when a task is either maximized or snapped
+         * left/right and rounded corners are enabled.
+         */
+        fun onTaskbarCornerRoundingUpdate(hasTasksRequiringTaskbarRounding: Boolean)
+    }
+
     /** The positions on a screen that a task can snap to. */
     enum class SnapPosition {
         RIGHT,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopTaskListener.aidl b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopTaskListener.aidl
index 8ebdfdc..c2acb87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopTaskListener.aidl
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/IDesktopTaskListener.aidl
@@ -27,4 +27,10 @@
 
     /** @deprecated this is no longer supported. */
     oneway void onStashedChanged(int displayId, boolean stashed);
+
+    /**
+     * Shows taskbar corner radius when running desktop tasks are updated if
+     * [hasTasksRequiringTaskbarRounding] is true.
+     */
+    oneway void onTaskbarCornerRoundingUpdate(boolean hasTasksRequiringTaskbarRounding);
 }
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ReturnToDragStartAnimator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ReturnToDragStartAnimator.kt
index 4c5258f..f4df42c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ReturnToDragStartAnimator.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/ReturnToDragStartAnimator.kt
@@ -48,7 +48,13 @@
     }
 
     /** Builds new animator and starts animation of task leash reposition. */
-    fun start(taskId: Int, taskSurface: SurfaceControl, startBounds: Rect, endBounds: Rect) {
+    fun start(
+        taskId: Int,
+        taskSurface: SurfaceControl,
+        startBounds: Rect,
+        endBounds: Rect,
+        isResizable: Boolean
+    ) {
         val tx = transactionSupplier.get()
 
         boundsAnimator?.cancel()
@@ -81,11 +87,13 @@
                                 .apply()
                             taskRepositionAnimationListener.onAnimationEnd(taskId)
                             boundsAnimator = null
-                            Toast.makeText(
-                                context,
-                                R.string.desktop_mode_non_resizable_snap_text,
-                                Toast.LENGTH_SHORT
-                            ).show()
+                            if (!isResizable) {
+                                Toast.makeText(
+                                    context,
+                                    R.string.desktop_mode_non_resizable_snap_text,
+                                    Toast.LENGTH_SHORT
+                                ).show()
+                            }
                             interactionJankMonitor.end(Cuj.CUJ_DESKTOP_MODE_SNAP_RESIZE)
                         }
                     )
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationController.kt
new file mode 100644
index 0000000..6013e97
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationController.kt
@@ -0,0 +1,113 @@
+/*
+ * 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.wm.shell.desktopmode.education
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.os.SystemProperties
+import com.android.window.flags.Flags
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository
+import com.android.wm.shell.shared.annotations.ShellMainThread
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.flow.debounce
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+
+/**
+ * Controls app handle education end to end.
+ *
+ * Listen to the user trigger for app handle education, calls an api to check if the education
+ * should be shown and calls an api to show education.
+ */
+@OptIn(kotlinx.coroutines.FlowPreview::class)
+@kotlinx.coroutines.ExperimentalCoroutinesApi
+class AppHandleEducationController(
+    private val appHandleEducationFilter: AppHandleEducationFilter,
+    shellTaskOrganizer: ShellTaskOrganizer,
+    private val appHandleEducationDatastoreRepository: AppHandleEducationDatastoreRepository,
+    @ShellMainThread private val applicationCoroutineScope: CoroutineScope
+) {
+  init {
+    runIfEducationFeatureEnabled {
+      // TODO: b/361038716 - Use app handle state flow instead of focus task change flow
+      val focusTaskChangeFlow = focusTaskChangeFlow(shellTaskOrganizer)
+      applicationCoroutineScope.launch {
+        // Central block handling the app's educational flow end-to-end.
+        // This flow listens to the changes to the result of
+        // [WindowingEducationProto#hasEducationViewedTimestampMillis()] in datastore proto object
+        isEducationViewedFlow()
+            .flatMapLatest { isEducationViewed ->
+              if (isEducationViewed) {
+                // If the education is viewed then return emptyFlow() that completes immediately.
+                // This will help us to not listen to focus task changes after the education has
+                // been viewed already.
+                emptyFlow()
+              } else {
+                // This flow listens for focus task changes, which trigger the app handle education.
+                focusTaskChangeFlow
+                    .filter { runningTaskInfo ->
+                      runningTaskInfo.topActivityInfo?.packageName?.let {
+                        appHandleEducationFilter.shouldShowAppHandleEducation(it)
+                      } ?: false && runningTaskInfo.windowingMode != WINDOWING_MODE_FREEFORM
+                    }
+                    .distinctUntilChanged()
+              }
+            }
+            .debounce(
+                APP_HANDLE_EDUCATION_DELAY) // Wait for few seconds, if the focus task changes.
+            // During the delay then current emission will be cancelled.
+            .flowOn(Dispatchers.IO)
+            .collectLatest {
+              // Fire and forget show education suspend function, manage entire lifecycle of
+              // tooltip in UI class.
+            }
+      }
+    }
+  }
+
+  private inline fun runIfEducationFeatureEnabled(block: () -> Unit) {
+    if (Flags.enableDesktopWindowingAppHandleEducation()) block()
+  }
+
+  private fun isEducationViewedFlow(): Flow<Boolean> =
+      appHandleEducationDatastoreRepository.dataStoreFlow
+          .map { preferences -> preferences.hasEducationViewedTimestampMillis() }
+          .distinctUntilChanged()
+
+  private fun focusTaskChangeFlow(shellTaskOrganizer: ShellTaskOrganizer): Flow<RunningTaskInfo> =
+      callbackFlow {
+        val focusTaskChange = ShellTaskOrganizer.FocusListener { taskInfo -> trySend(taskInfo) }
+        shellTaskOrganizer.addFocusListener(focusTaskChange)
+        awaitClose { shellTaskOrganizer.removeFocusListener(focusTaskChange) }
+      }
+
+  private companion object {
+    val APP_HANDLE_EDUCATION_DELAY: Long
+      get() = SystemProperties.getLong("persist.windowing_app_handle_education_delay", 3000L)
+  }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
index a7fff8a..f420c5b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
@@ -25,9 +25,12 @@
 import androidx.datastore.dataStoreFile
 import com.android.framework.protobuf.InvalidProtocolBufferException
 import com.android.internal.annotations.VisibleForTesting
+import java.io.IOException
 import java.io.InputStream
 import java.io.OutputStream
 import java.time.Duration
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.first
 
 /**
@@ -46,17 +49,26 @@
           serializer = WindowingEducationProtoSerializer,
           produceFile = { context.dataStoreFile(APP_HANDLE_EDUCATION_DATASTORE_FILEPATH) }))
 
+  /** Provides dataStore.data flow and handles exceptions thrown during collection */
+  val dataStoreFlow: Flow<WindowingEducationProto> =
+      dataStore.data.catch { exception ->
+        // dataStore.data throws an IOException when an error is encountered when reading data
+        if (exception is IOException) {
+          Log.e(
+              TAG,
+              "Error in reading app handle education related data from datastore, data is " +
+                  "stored in a file named $APP_HANDLE_EDUCATION_DATASTORE_FILEPATH",
+              exception)
+        } else {
+          throw exception
+        }
+      }
+
   /**
    * Reads and returns the [WindowingEducationProto] Proto object from the DataStore. If the
    * DataStore is empty or there's an error reading, it returns the default value of Proto.
    */
-  suspend fun windowingEducationProto(): WindowingEducationProto =
-      try {
-        dataStore.data.first()
-      } catch (e: Exception) {
-        Log.e(TAG, "Unable to read from datastore")
-        WindowingEducationProto.getDefaultInstance()
-      }
+  suspend fun windowingEducationProto(): WindowingEducationProto = dataStoreFlow.first()
 
   /**
    * Updates [AppHandleEducation.appUsageStats] and
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
index 0acc7df..faa97ac 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
@@ -98,9 +98,8 @@
 ### Exposing shared code for use in Launcher
 Launcher doesn't currently build against the Shell library, but needs to have access to some shared
 AIDL interfaces and constants.  Currently, all AIDL files, and classes under the
-`com.android.wm.shell.util` package are automatically built into the `SystemUISharedLib` that
+`com.android.wm.shell.shared` package are automatically built into the `SystemUISharedLib` that
 Launcher uses.
 
-If the new code doesn't fall into those categories, they can be added explicitly in the Shell's
-[Android.bp](/libs/WindowManager/Shell/Android.bp) file under the
-`wm_shell_util-sources` filegroup.
\ No newline at end of file
+If the new code doesn't fall into those categories, they should be moved to the Shell shared
+package (`com.android.wm.shell.shared`) under the `WindowManager-Shell-shared` library.
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index e4cd10f..ab222c9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -99,6 +99,7 @@
 import java.lang.ref.WeakReference;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.StringJoiner;
 import java.util.function.Consumer;
 import java.util.function.IntConsumer;
 
@@ -831,6 +832,7 @@
                     mPictureInPictureParams.getTitle());
             mPipParamsChangedForwarder.notifySubtitleChanged(
                     mPictureInPictureParams.getSubtitle());
+            logRemoteActions(mPictureInPictureParams);
         }
 
         mPipUiEventLoggerLogger.setTaskInfo(mTaskInfo);
@@ -1112,6 +1114,7 @@
         }
         applyNewPictureInPictureParams(newParams);
         mPictureInPictureParams = newParams;
+        logRemoteActions(mPictureInPictureParams);
     }
 
     @Override
@@ -1420,6 +1423,16 @@
         }
     }
 
+    private void logRemoteActions(@NonNull PictureInPictureParams params) {
+        StringJoiner sj = new StringJoiner("|", "[", "]");
+        if (params.hasSetActions()) {
+            params.getActions().forEach((action) -> sj.add(action.getTitle()));
+        }
+
+        ProtoLog.d(ShellProtoLogGroup.WM_SHELL_PICTURE_IN_PICTURE,
+                "%s: PIP remote actions=%s", TAG, sj.toString());
+    }
+
     /**
      * Animates resizing of the pinned stack given the duration.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipDismissTargetHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipDismissTargetHandler.java
index 0d2b8e7..06d2311 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipDismissTargetHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipDismissTargetHandler.java
@@ -35,9 +35,9 @@
 import com.android.wm.shell.R;
 import com.android.wm.shell.bubbles.DismissViewUtils;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.bubbles.DismissCircleView;
-import com.android.wm.shell.common.bubbles.DismissView;
 import com.android.wm.shell.common.pip.PipUiEventLogger;
+import com.android.wm.shell.shared.bubbles.DismissCircleView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
 
 import kotlin.Unit;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java
index 8a9302b..8ebdc96 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/animation/PipEnterExitAnimator.java
@@ -22,6 +22,7 @@
 import android.annotation.IntDef;
 import android.content.Context;
 import android.graphics.Rect;
+import android.view.Surface;
 import android.view.SurfaceControl;
 
 import androidx.annotation.NonNull;
@@ -51,8 +52,10 @@
 
     @NonNull private final SurfaceControl mLeash;
     private final SurfaceControl.Transaction mStartTransaction;
-    private final int mEnterAnimationDuration;
+    private final SurfaceControl.Transaction mFinishTransaction;
+    private final int mEnterExitAnimationDuration;
     private final @BOUNDS int mDirection;
+    private final @Surface.Rotation int mRotation;
 
     // optional callbacks for tracking animation start and end
     @Nullable private Runnable mAnimationStartCallback;
@@ -62,37 +65,59 @@
     private final Rect mStartBounds = new Rect();
     private final Rect mEndBounds = new Rect();
 
+    @Nullable private final Rect mSourceRectHint;
+    private final Rect mSourceRectHintInsets = new Rect();
+    private final Rect mZeroInsets = new Rect(0, 0, 0, 0);
+
     // Bounds updated by the evaluator as animator is running.
     private final Rect mAnimatedRect = new Rect();
 
     private final PipSurfaceTransactionHelper.SurfaceControlTransactionFactory
             mSurfaceControlTransactionFactory;
     private final RectEvaluator mRectEvaluator;
+    private final RectEvaluator mInsetEvaluator;
     private final PipSurfaceTransactionHelper mPipSurfaceTransactionHelper;
 
     public PipEnterExitAnimator(Context context,
             @NonNull SurfaceControl leash,
             SurfaceControl.Transaction startTransaction,
+            SurfaceControl.Transaction finishTransaction,
             @NonNull Rect baseBounds,
             @NonNull Rect startBounds,
             @NonNull Rect endBounds,
-            @BOUNDS int direction) {
+            @Nullable Rect sourceRectHint,
+            @BOUNDS int direction,
+            @Surface.Rotation int rotation) {
         mLeash = leash;
         mStartTransaction = startTransaction;
+        mFinishTransaction = finishTransaction;
         mBaseBounds.set(baseBounds);
         mStartBounds.set(startBounds);
         mAnimatedRect.set(startBounds);
         mEndBounds.set(endBounds);
         mRectEvaluator = new RectEvaluator(mAnimatedRect);
+        mInsetEvaluator = new RectEvaluator(new Rect());
         mPipSurfaceTransactionHelper = new PipSurfaceTransactionHelper(context);
         mDirection = direction;
+        mRotation = rotation;
+
+        mSourceRectHint = sourceRectHint != null ? new Rect(sourceRectHint) : null;
+        if (mSourceRectHint != null) {
+            mSourceRectHintInsets.set(
+                    mSourceRectHint.left - mBaseBounds.left,
+                    mSourceRectHint.top - mBaseBounds.top,
+                    mBaseBounds.right - mSourceRectHint.right,
+                    mBaseBounds.bottom - mSourceRectHint.bottom
+            );
+        }
 
         mSurfaceControlTransactionFactory =
                 new PipSurfaceTransactionHelper.VsyncSurfaceControlTransactionFactory();
-        mEnterAnimationDuration = context.getResources()
+        mEnterExitAnimationDuration = context.getResources()
                 .getInteger(R.integer.config_pipEnterAnimationDuration);
 
-        setDuration(mEnterAnimationDuration);
+        setObjectValues(startBounds, endBounds);
+        setDuration(mEnterExitAnimationDuration);
         setEvaluator(mRectEvaluator);
         addListener(this);
         addUpdateListener(this);
@@ -118,6 +143,14 @@
 
     @Override
     public void onAnimationEnd(@NonNull Animator animation) {
+        if (mFinishTransaction != null) {
+            // finishTransaction might override some state (eg. corner radii) so we want to
+            // manually set the state to the end of the animation
+            mPipSurfaceTransactionHelper.scaleAndCrop(mFinishTransaction, mLeash, mSourceRectHint,
+                            mBaseBounds, mAnimatedRect, getInsets(1f), isInPipDirection(), 1f)
+                    .round(mFinishTransaction, mLeash, isInPipDirection())
+                    .shadow(mFinishTransaction, mLeash, isInPipDirection());
+        }
         if (mAnimationEndCallback != null) {
             mAnimationEndCallback.run();
         }
@@ -127,19 +160,32 @@
     public void onAnimationUpdate(@NonNull ValueAnimator animation) {
         final SurfaceControl.Transaction tx = mSurfaceControlTransactionFactory.getTransaction();
         final float fraction = getAnimatedFraction();
+        Rect insets = getInsets(fraction);
+
         // TODO (b/350801661): implement fixed rotation
 
-        mPipSurfaceTransactionHelper.scaleAndCrop(tx, mLeash, null,
-                mBaseBounds, mAnimatedRect, null, isInPipDirection(), fraction)
+        mPipSurfaceTransactionHelper.scaleAndCrop(tx, mLeash, mSourceRectHint,
+                mBaseBounds, mAnimatedRect, insets, isInPipDirection(), fraction)
                 .round(tx, mLeash, isInPipDirection())
                 .shadow(tx, mLeash, isInPipDirection());
         tx.apply();
     }
 
+    private Rect getInsets(float fraction) {
+        Rect startInsets = isInPipDirection() ? mZeroInsets : mSourceRectHintInsets;
+        Rect endInsets = isInPipDirection() ? mSourceRectHintInsets : mZeroInsets;
+
+        return mInsetEvaluator.evaluate(fraction, startInsets, endInsets);
+    }
+
     private boolean isInPipDirection() {
         return mDirection == BOUNDS_ENTER;
     }
 
+    private boolean isOutPipDirection() {
+        return mDirection == BOUNDS_EXIT;
+    }
+
     // no-ops
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipDismissTargetHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipDismissTargetHandler.java
index e04178e..b3070f2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipDismissTargetHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipDismissTargetHandler.java
@@ -35,9 +35,9 @@
 import com.android.wm.shell.R;
 import com.android.wm.shell.bubbles.DismissViewUtils;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.bubbles.DismissCircleView;
-import com.android.wm.shell.common.bubbles.DismissView;
 import com.android.wm.shell.common.pip.PipUiEventLogger;
+import com.android.wm.shell.shared.bubbles.DismissCircleView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
 
 import kotlin.Unit;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
index 7f16880..262c14d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTaskListener.java
@@ -25,6 +25,7 @@
 import android.os.Bundle;
 import android.view.SurfaceControl;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.internal.util.Preconditions;
@@ -88,6 +89,11 @@
                 : new PictureInPictureParams.Builder().build());
     }
 
+    @NonNull
+    public PictureInPictureParams getPictureInPictureParams() {
+        return mPictureInPictureParams;
+    }
+
     @Override
     public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
         PictureInPictureParams params = taskInfo.pictureInPictureParams;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
index 44baabd..f93233e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTransition.java
@@ -36,6 +36,7 @@
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.IBinder;
+import android.view.Surface;
 import android.view.SurfaceControl;
 import android.window.TransitionInfo;
 import android.window.TransitionRequestInfo;
@@ -398,17 +399,22 @@
         SurfaceControl pipLeash = mPipTransitionState.mPinnedTaskLeash;
         Preconditions.checkNotNull(pipLeash, "Leash is null for bounds transition.");
 
+        Rect sourceRectHint = null;
+        if (pipChange.getTaskInfo() != null
+                && pipChange.getTaskInfo().pictureInPictureParams != null) {
+            sourceRectHint = pipChange.getTaskInfo().pictureInPictureParams.getSourceRectHint();
+        }
+
         PipEnterExitAnimator animator = new PipEnterExitAnimator(mContext, pipLeash,
-                startTransaction, startBounds, startBounds, endBounds,
-                PipEnterExitAnimator.BOUNDS_ENTER);
+                startTransaction, finishTransaction, startBounds, startBounds, endBounds,
+                sourceRectHint, PipEnterExitAnimator.BOUNDS_ENTER, Surface.ROTATION_0);
 
         tx.addTransactionCommittedListener(mPipScheduler.getMainExecutor(),
                 this::onClientDrawAtTransitionEnd);
         finishWct.setBoundsChangeTransaction(pipTaskToken, tx);
 
-        animator.setAnimationEndCallback(() -> {
-            finishCallback.onTransitionFinished(finishWct.isEmpty() ? null : finishWct);
-        });
+        animator.setAnimationEndCallback(() ->
+                finishCallback.onTransitionFinished(finishWct));
 
         animator.start();
         return true;
@@ -452,19 +458,53 @@
 
         TransitionInfo.Change pipChange = getChangeByToken(info, pipToken);
         if (pipChange == null) {
-            return false;
+            // pipChange is null, check to see if we've reparented the PIP activity for
+            // the multi activity case. If so we should use the activity leash instead
+            for (TransitionInfo.Change change : info.getChanges()) {
+                if (change.getTaskInfo() == null
+                        && change.getLastParent() != null
+                        && change.getLastParent().equals(pipToken)) {
+                    pipChange = change;
+                    break;
+                }
+            }
+
+            // failsafe
+            if (pipChange == null) {
+                return false;
+            }
+        }
+
+        // for multi activity, we need to manually set the leash layer
+        if (pipChange.getTaskInfo() == null) {
+            TransitionInfo.Change parent = getChangeByToken(info, pipChange.getParent());
+            if (parent != null) {
+                startTransaction.setLayer(parent.getLeash(), Integer.MAX_VALUE - 1);
+            }
         }
 
         Rect startBounds = pipChange.getStartAbsBounds();
         Rect endBounds = pipChange.getEndAbsBounds();
         SurfaceControl pipLeash = pipChange.getLeash();
+        Preconditions.checkNotNull(pipLeash, "Leash is null for exit transition.");
+
+        Rect sourceRectHint = null;
+        if (pipChange.getTaskInfo() != null
+                && pipChange.getTaskInfo().pictureInPictureParams != null) {
+            // single activity
+            sourceRectHint = pipChange.getTaskInfo().pictureInPictureParams.getSourceRectHint();
+        } else if (mPipTaskListener.getPictureInPictureParams().hasSourceBoundsHint()) {
+            // multi activity
+            sourceRectHint = mPipTaskListener.getPictureInPictureParams().getSourceRectHint();
+        }
 
         PipEnterExitAnimator animator = new PipEnterExitAnimator(mContext, pipLeash,
-                startTransaction, startBounds, startBounds, endBounds,
-                PipEnterExitAnimator.BOUNDS_EXIT);
+                startTransaction, finishTransaction, endBounds, startBounds, endBounds,
+                sourceRectHint, PipEnterExitAnimator.BOUNDS_EXIT, Surface.ROTATION_0);
+
         animator.setAnimationEndCallback(() -> {
-            finishCallback.onTransitionFinished(null);
             mPipTransitionState.setState(PipTransitionState.EXITED_PIP);
+            finishCallback.onTransitionFinished(null);
         });
 
         animator.start();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
index 27ded57..cea995d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
@@ -321,8 +321,10 @@
                 0 /* enterReason */,
                 0 /* exitReason */,
                 mLastSplitRatio,
-                0 /* mainStagePosition */, 0 /* mainStageUid */,
-                0 /* sideStagePosition */, 0 /* sideStageUid */,
+                mLastMainStagePosition,
+                mLastMainStageUid,
+                mLastSideStagePosition,
+                mLastSideStageUid,
                 0 /* dragInstanceId */,
                 mLoggerSessionId.getId());
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
index b18feefe..81f444b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/SplashscreenContentDrawer.java
@@ -352,12 +352,17 @@
     /** Extract the window background color from {@code attrs}. */
     private static int peekWindowBGColor(Context context, SplashScreenWindowAttrs attrs) {
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "peekWindowBGColor");
-        final Drawable themeBGDrawable;
+        Drawable themeBGDrawable = null;
         if (attrs.mWindowBgColor != 0) {
             themeBGDrawable = new ColorDrawable(attrs.mWindowBgColor);
         } else if (attrs.mWindowBgResId != 0) {
-            themeBGDrawable = context.getDrawable(attrs.mWindowBgResId);
-        } else {
+            try {
+                themeBGDrawable = context.getDrawable(attrs.mWindowBgResId);
+            } catch (Resources.NotFoundException e) {
+                Slog.w(TAG, "Unable get drawable from resource", e);
+            }
+        }
+        if (themeBGDrawable == null) {
             themeBGDrawable = createDefaultBackgroundDrawable();
             Slog.w(TAG, "Window background does not exist, using " + themeBGDrawable);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
index a85188a..82c0aaf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskView.java
@@ -118,6 +118,13 @@
         mTaskViewTaskController.startShortcutActivity(shortcut, options, launchBounds);
     }
 
+    /**
+     * Moves the current task in taskview out of the view and back to fullscreen.
+     */
+    public void moveToFullscreen() {
+        mTaskViewTaskController.moveToFullscreen();
+    }
+
     @Override
     public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
         if (mTaskViewTaskController.isUsingShellTransitions()) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
index 9750d3e..e74342e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
@@ -17,6 +17,7 @@
 package com.android.wm.shell.taskview;
 
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
+import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.view.WindowManager.TRANSIT_CHANGE;
 
 import android.annotation.NonNull;
@@ -256,6 +257,24 @@
         mTaskViewTransitions.startInstantTransition(TRANSIT_CHANGE, wct);
     }
 
+    /**
+     * Moves the current task in TaskView out of the view and back to fullscreen.
+     */
+    public void moveToFullscreen() {
+        if (mTaskToken == null) return;
+        mShellExecutor.execute(() -> {
+            WindowContainerTransaction wct = new WindowContainerTransaction();
+            wct.setWindowingMode(mTaskToken, WINDOWING_MODE_UNDEFINED);
+            wct.setAlwaysOnTop(mTaskToken, false);
+            mTaskOrganizer.setInterceptBackPressedOnTaskRoot(mTaskToken, false);
+            mTaskViewTransitions.moveTaskViewToFullscreen(wct, this);
+            if (mListener != null) {
+                // Task is being "removed" from the clients perspective
+                mListener.onTaskRemovalStarted(mTaskInfo.taskId);
+            }
+        });
+    }
+
     private void prepareActivityOptions(ActivityOptions options, Rect launchBounds) {
         final Binder launchCookie = new Binder();
         mShellExecutor.execute(() -> {
@@ -585,7 +604,6 @@
                 });
             }
             mTaskViewBase.onTaskVanished(taskInfo);
-            mTaskOrganizer.setInterceptBackPressedOnTaskRoot(taskInfo.token, false);
         }
     }
 
@@ -699,6 +717,9 @@
             mTaskViewBase.setResizeBgColor(startTransaction, backgroundColor);
         }
 
+        // After the embedded task has appeared, set it to non-trimmable. This is important
+        // to prevent recents from trimming and removing the embedded task.
+        wct.setTaskTrimmableFromRecents(taskInfo.token, false /* isTrimmableFromRecents */);
         mTaskViewBase.onTaskAppeared(mTaskInfo, mTaskLeash);
         if (mListener != null) {
             final int taskId = mTaskInfo.taskId;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
index 15fe7ab..39648f6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
@@ -236,6 +236,12 @@
         startNextTransition();
     }
 
+    void moveTaskViewToFullscreen(@NonNull WindowContainerTransaction wct,
+            @NonNull TaskViewTaskController taskView) {
+        mPending.add(new PendingTransition(TRANSIT_CHANGE, wct, taskView, null /* cookie */));
+        startNextTransition();
+    }
+
     /** Starts a new transition to make the given {@code taskView} visible. */
     public void setTaskViewVisible(TaskViewTaskController taskView, boolean visible) {
         setTaskViewVisible(taskView, visible, false /* reorder */);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
index 4fc6c44..9b0fb20 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
@@ -18,8 +18,8 @@
 
 import static android.app.ActivityOptions.ANIM_CLIP_REVEAL;
 import static android.app.ActivityOptions.ANIM_CUSTOM;
-import static android.app.ActivityOptions.ANIM_NONE;
 import static android.app.ActivityOptions.ANIM_FROM_STYLE;
+import static android.app.ActivityOptions.ANIM_NONE;
 import static android.app.ActivityOptions.ANIM_OPEN_CROSS_PROFILE_APPS;
 import static android.app.ActivityOptions.ANIM_SCALE_UP;
 import static android.app.ActivityOptions.ANIM_SCENE_TRANSITION;
@@ -473,7 +473,7 @@
                                 change.getLeash(),
                                 startTransaction);
                     } else if (isOnlyTranslucent && TransitionUtil.isOpeningType(info.getType())
-                                && TransitionUtil.isClosingType(mode)) {
+                            && TransitionUtil.isClosingType(mode)) {
                         // If there is a closing translucent task in an OPENING transition, we will
                         // actually select a CLOSING animation, so move the closing task into
                         // the animating part of the z-order.
@@ -767,12 +767,12 @@
             a = mTransitionAnimation.loadKeyguardExitAnimation(flags,
                     (changeFlags & FLAG_SHOW_WALLPAPER) != 0);
         } else if (type == TRANSIT_KEYGUARD_UNOCCLUDE) {
-            a = mTransitionAnimation.loadKeyguardUnoccludeAnimation();
+            a = mTransitionAnimation.loadKeyguardUnoccludeAnimation(options.getUserId());
         } else if ((changeFlags & FLAG_IS_VOICE_INTERACTION) != 0) {
             if (isOpeningType) {
-                a = mTransitionAnimation.loadVoiceActivityOpenAnimation(enter);
+                a = mTransitionAnimation.loadVoiceActivityOpenAnimation(enter, options.getUserId());
             } else {
-                a = mTransitionAnimation.loadVoiceActivityExitAnimation(enter);
+                a = mTransitionAnimation.loadVoiceActivityExitAnimation(enter, options.getUserId());
             }
         } else if (changeMode == TRANSIT_CHANGE) {
             // In the absence of a specific adapter, we just want to keep everything stationary.
@@ -783,9 +783,9 @@
         } else if (overrideType == ANIM_CUSTOM
                 && (!isTask || options.getOverrideTaskTransition())) {
             a = mTransitionAnimation.loadAnimationRes(options.getPackageName(), enter
-                    ? options.getEnterResId() : options.getExitResId());
+                    ? options.getEnterResId() : options.getExitResId(), options.getUserId());
         } else if (overrideType == ANIM_OPEN_CROSS_PROFILE_APPS && enter) {
-            a = mTransitionAnimation.loadCrossProfileAppEnterAnimation();
+            a = mTransitionAnimation.loadCrossProfileAppEnterAnimation(options.getUserId());
         } else if (overrideType == ANIM_CLIP_REVEAL) {
             a = mTransitionAnimation.createClipRevealAnimationLocked(type, wallpaperTransit, enter,
                     endBounds, endBounds, options.getTransitionBounds());
@@ -905,9 +905,9 @@
         final Rect bounds = change.getEndAbsBounds();
         // Show the right drawable depending on the user we're transitioning to.
         final Drawable thumbnailDrawable = change.hasFlags(FLAG_CROSS_PROFILE_OWNER_THUMBNAIL)
-                        ? mContext.getDrawable(R.drawable.ic_account_circle)
-                        : change.hasFlags(FLAG_CROSS_PROFILE_WORK_THUMBNAIL)
-                                ? mEnterpriseThumbnailDrawable : null;
+                ? mContext.getDrawable(R.drawable.ic_account_circle)
+                : change.hasFlags(FLAG_CROSS_PROFILE_WORK_THUMBNAIL)
+                        ? mEnterpriseThumbnailDrawable : null;
         if (thumbnailDrawable == null) {
             return;
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 1f95667..c88c1e2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -97,6 +97,7 @@
 import com.android.wm.shell.common.MultiInstanceHelper;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
+import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler;
 import com.android.wm.shell.desktopmode.DesktopModeVisualIndicator;
 import com.android.wm.shell.desktopmode.DesktopTasksController;
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition;
@@ -166,6 +167,8 @@
     private TaskOperations mTaskOperations;
     private final Supplier<SurfaceControl.Transaction> mTransactionFactory;
     private final Transitions mTransitions;
+    private final Optional<DesktopActivityOrientationChangeHandler>
+            mActivityOrientationChangeHandler;
 
     private SplitScreenController mSplitScreenController;
 
@@ -215,7 +218,8 @@
             InteractionJankMonitor interactionJankMonitor,
             AppToWebGenericLinksParser genericLinksParser,
             MultiInstanceHelper multiInstanceHelper,
-            Optional<DesktopTasksLimiter> desktopTasksLimiter
+            Optional<DesktopTasksLimiter> desktopTasksLimiter,
+            Optional<DesktopActivityOrientationChangeHandler> activityOrientationChangeHandler
     ) {
         this(
                 context,
@@ -241,7 +245,8 @@
                 rootTaskDisplayAreaOrganizer,
                 new SparseArray<>(),
                 interactionJankMonitor,
-                desktopTasksLimiter);
+                desktopTasksLimiter,
+                activityOrientationChangeHandler);
     }
 
     @VisibleForTesting
@@ -269,7 +274,8 @@
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
             SparseArray<DesktopModeWindowDecoration> windowDecorByTaskId,
             InteractionJankMonitor interactionJankMonitor,
-            Optional<DesktopTasksLimiter> desktopTasksLimiter) {
+            Optional<DesktopTasksLimiter> desktopTasksLimiter,
+            Optional<DesktopActivityOrientationChangeHandler> activityOrientationChangeHandler) {
         mContext = context;
         mMainExecutor = shellExecutor;
         mMainHandler = mainHandler;
@@ -297,6 +303,7 @@
                 com.android.internal.R.string.config_systemUi);
         mInteractionJankMonitor = interactionJankMonitor;
         mDesktopTasksLimiter = desktopTasksLimiter;
+        mActivityOrientationChangeHandler = activityOrientationChangeHandler;
         mOnDisplayChangingListener = (displayId, fromRotation, toRotation, displayAreaInfo, t) -> {
             DesktopModeWindowDecoration decoration;
             RunningTaskInfo taskInfo;
@@ -388,6 +395,8 @@
             incrementEventReceiverTasks(taskInfo.displayId);
         }
         decoration.relayout(taskInfo);
+        mActivityOrientationChangeHandler.ifPresent(handler ->
+                handler.handleActivityOrientationChange(oldTaskInfo, taskInfo));
     }
 
     @Override
@@ -482,7 +491,9 @@
         } else {
             mInteractionJankMonitor.begin(decoration.mTaskSurface, mContext,
                     Cuj.CUJ_DESKTOP_MODE_SNAP_RESIZE, "maximize_menu_resizable");
-            mDesktopTasksController.snapToHalfScreen(decoration.mTaskInfo,
+            mDesktopTasksController.snapToHalfScreen(
+                    decoration.mTaskInfo,
+                    decoration.mTaskSurface,
                     decoration.mTaskInfo.configuration.windowConfiguration.getBounds(),
                     left ? SnapPosition.LEFT : SnapPosition.RIGHT);
         }
@@ -496,16 +507,16 @@
         if (decoration == null) {
             return;
         }
-        openInBrowser(uri);
+        openInBrowser(uri, decoration.getUser());
         decoration.closeHandleMenu();
         decoration.closeMaximizeMenu();
     }
 
-    private void openInBrowser(Uri uri) {
+    private void openInBrowser(Uri uri, @NonNull UserHandle userHandle) {
         final Intent intent = Intent.makeMainSelectorActivity(ACTION_MAIN, CATEGORY_APP_BROWSER)
                 .setData(uri)
                 .addFlags(FLAG_ACTIVITY_NEW_TASK);
-        mContext.startActivity(intent);
+        mContext.startActivityAsUser(intent, userHandle);
     }
 
     private void onToDesktop(int taskId, DesktopModeTransitionSource source) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
index 2bec3fa..81251b8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecoration.java
@@ -54,6 +54,7 @@
 import android.net.Uri;
 import android.os.Handler;
 import android.os.Trace;
+import android.os.UserHandle;
 import android.util.Size;
 import android.util.Slog;
 import android.view.Choreographer;
@@ -480,6 +481,10 @@
         return mGenericLink;
     }
 
+    UserHandle getUser() {
+        return mUserContext.getUser();
+    }
+
     private void updateDragResizeListener(SurfaceControl oldDecorationSurface) {
         if (!isDragResizable(mTaskInfo)) {
             if (!mTaskInfo.positionInParent.equals(mPositionInParent)) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt
index deef378..9c73e4a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/MaximizeMenu.kt
@@ -121,8 +121,14 @@
 
     /** Closes the maximize window and releases its view. */
     fun close() {
-        maximizeMenuView?.animateCloseMenu {
-            maximizeMenu?.releaseView()
+        val view = maximizeMenuView
+        val menu = maximizeMenu
+        if (view == null) {
+            menu?.releaseView()
+        } else {
+            view.animateCloseMenu {
+                menu?.releaseView()
+            }
         }
         maximizeMenu = null
         maximizeMenuView = null
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/DesktopModeFlickerScenarios.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/DesktopModeFlickerScenarios.kt
index 3fb67cd..880e021 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/DesktopModeFlickerScenarios.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/DesktopModeFlickerScenarios.kt
@@ -17,11 +17,17 @@
 package com.android.wm.shell.flicker
 
 import android.tools.flicker.AssertionInvocationGroup
+import android.tools.flicker.assertors.assertions.AppLayerIncreasesInSize
 import android.tools.flicker.assertors.assertions.AppLayerIsInvisibleAtEnd
 import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAlways
 import android.tools.flicker.assertors.assertions.AppLayerIsVisibleAtStart
 import android.tools.flicker.assertors.assertions.AppWindowBecomesVisible
+import android.tools.flicker.assertors.assertions.AppWindowCoversLeftHalfScreenAtEnd
+import android.tools.flicker.assertors.assertions.AppWindowCoversRightHalfScreenAtEnd
 import android.tools.flicker.assertors.assertions.AppWindowHasDesktopModeInitialBoundsAtTheEnd
+import android.tools.flicker.assertors.assertions.AppWindowHasMaxBoundsInOnlyOneDimension
+import android.tools.flicker.assertors.assertions.AppWindowHasMaxDisplayHeight
+import android.tools.flicker.assertors.assertions.AppWindowHasMaxDisplayWidth
 import android.tools.flicker.assertors.assertions.AppWindowHasSizeOfAtLeast
 import android.tools.flicker.assertors.assertions.AppWindowIsInvisibleAtEnd
 import android.tools.flicker.assertors.assertions.AppWindowIsVisibleAlways
@@ -51,21 +57,21 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("END_DRAG_TO_DESKTOP"),
                 extractor =
-                    ShellTransitionScenarioExtractor(
-                        transitionMatcher =
-                            object : ITransitionMatcher {
-                                override fun findAll(
-                                    transitions: Collection<Transition>
-                                ): Collection<Transition> {
-                                    return transitions.filter {
-                                        // TODO(351168217) Use jank CUJ to extract a longer trace
-                                        it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP
-                                    }
-                                }
+                ShellTransitionScenarioExtractor(
+                    transitionMatcher =
+                    object : ITransitionMatcher {
+                        override fun findAll(
+                            transitions: Collection<Transition>
+                        ): Collection<Transition> {
+                            return transitions.filter {
+                                // TODO(351168217) Use jank CUJ to extract a longer trace
+                                it.type == TransitionType.DESKTOP_MODE_END_DRAG_TO_DESKTOP
                             }
-                    ),
+                        }
+                    }
+                ),
                 assertions =
-                    AssertionTemplates.COMMON_ASSERTIONS +
+                AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
                             AppLayerIsVisibleAlways(DESKTOP_MODE_APP),
                             AppWindowOnTopAtEnd(DESKTOP_MODE_APP),
@@ -81,24 +87,24 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CLOSE_APP"),
                 extractor =
-                    ShellTransitionScenarioExtractor(
-                        transitionMatcher =
-                            object : ITransitionMatcher {
-                                override fun findAll(
-                                    transitions: Collection<Transition>
-                                ): Collection<Transition> {
-                                    // In case there are multiple windows closing, filter out the
-                                    // last window closing. It should use the CLOSE_LAST_APP
-                                    // scenario below.
-                                    return transitions
-                                        .filter { it.type == TransitionType.CLOSE }
-                                        .sortedByDescending { it.id }
-                                        .drop(1)
-                                }
-                            }
-                    ),
+                ShellTransitionScenarioExtractor(
+                    transitionMatcher =
+                    object : ITransitionMatcher {
+                        override fun findAll(
+                            transitions: Collection<Transition>
+                        ): Collection<Transition> {
+                            // In case there are multiple windows closing, filter out the
+                            // last window closing. It should use the CLOSE_LAST_APP
+                            // scenario below.
+                            return transitions
+                                .filter { it.type == TransitionType.CLOSE }
+                                .sortedByDescending { it.id }
+                                .drop(1)
+                        }
+                    }
+                ),
                 assertions =
-                    AssertionTemplates.COMMON_ASSERTIONS +
+                AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
                             AppWindowOnTopAtStart(DESKTOP_MODE_APP),
                             AppLayerIsVisibleAtStart(DESKTOP_MODE_APP),
@@ -110,22 +116,22 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CLOSE_LAST_APP"),
                 extractor =
-                    ShellTransitionScenarioExtractor(
-                        transitionMatcher =
-                            object : ITransitionMatcher {
-                                override fun findAll(
-                                    transitions: Collection<Transition>
-                                ): Collection<Transition> {
-                                    val lastTransition =
-                                        transitions
-                                            .filter { it.type == TransitionType.CLOSE }
-                                            .maxByOrNull { it.id }!!
-                                    return listOf(lastTransition)
-                                }
-                            }
-                    ),
+                ShellTransitionScenarioExtractor(
+                    transitionMatcher =
+                    object : ITransitionMatcher {
+                        override fun findAll(
+                            transitions: Collection<Transition>
+                        ): Collection<Transition> {
+                            val lastTransition =
+                                transitions
+                                    .filter { it.type == TransitionType.CLOSE }
+                                    .maxByOrNull { it.id }!!
+                            return listOf(lastTransition)
+                        }
+                    }
+                ),
                 assertions =
-                    AssertionTemplates.COMMON_ASSERTIONS +
+                AssertionTemplates.COMMON_ASSERTIONS +
                         listOf(
                             AppWindowIsInvisibleAtEnd(DESKTOP_MODE_APP),
                             LauncherWindowReplacesAppAsTopWindow(DESKTOP_MODE_APP),
@@ -138,31 +144,107 @@
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CORNER_RESIZE"),
                 extractor =
-                    TaggedScenarioExtractorBuilder()
-                        .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
-                        .setTransitionMatcher(
-                            TaggedCujTransitionMatcher(associatedTransitionRequired = false)
-                        )
-                        .build(),
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
                 assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS
             )
 
+        val EDGE_RESIZE =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("EDGE_RESIZE"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(
+                            AppLayerIncreasesInSize(DESKTOP_MODE_APP),
+                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
         val CORNER_RESIZE_TO_MINIMUM_SIZE =
             FlickerConfigEntry(
                 scenarioId = ScenarioId("CORNER_RESIZE_TO_MINIMUM_SIZE"),
                 extractor =
-                    TaggedScenarioExtractorBuilder()
-                        .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
-                        .setTransitionMatcher(
-                            TaggedCujTransitionMatcher(associatedTransitionRequired = false)
-                        )
-                        .build(),
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_RESIZE_WINDOW)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
                 assertions =
                 AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
                         listOf(AppWindowHasSizeOfAtLeast(DESKTOP_MODE_APP, 770, 700))
                             .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
 
+        val SNAP_RESIZE_LEFT_WITH_BUTTON =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("SNAP_RESIZE_LEFT_WITH_BUTTON"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_SNAP_RESIZE)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(AppWindowCoversLeftHalfScreenAtEnd(DESKTOP_MODE_APP))
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val SNAP_RESIZE_RIGHT_WITH_BUTTON =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("SNAP_RESIZE_RIGHT_WITH_BUTTON"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_SNAP_RESIZE)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(AppWindowCoversRightHalfScreenAtEnd(DESKTOP_MODE_APP))
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val SNAP_RESIZE_LEFT_WITH_DRAG =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("SNAP_RESIZE_LEFT_WITH_DRAG"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_SNAP_RESIZE)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(AppWindowCoversLeftHalfScreenAtEnd(DESKTOP_MODE_APP))
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val SNAP_RESIZE_RIGHT_WITH_DRAG =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("SNAP_RESIZE_RIGHT_WITH_DRAG"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_SNAP_RESIZE)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(AppWindowCoversRightHalfScreenAtEnd(DESKTOP_MODE_APP))
+                            .associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
         val SNAP_RESIZE_WITH_DRAG_NON_RESIZABLE =
             FlickerConfigEntry(
                 scenarioId = ScenarioId("SNAP_RESIZE_WITH_DRAG_NON_RESIZABLE"),
@@ -181,5 +263,42 @@
                     AppWindowReturnsToStartBoundsAndPosition(NON_RESIZABLE_APP)
                 ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
             )
+
+        val MAXIMIZE_APP =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("MAXIMIZE_APP"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_MAXIMIZE_WINDOW)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions = AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(
+                            AppLayerIncreasesInSize(DESKTOP_MODE_APP),
+                            AppWindowHasMaxDisplayHeight(DESKTOP_MODE_APP),
+                            AppWindowHasMaxDisplayWidth(DESKTOP_MODE_APP)
+                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
+
+        val MAXIMIZE_APP_NON_RESIZABLE =
+            FlickerConfigEntry(
+                scenarioId = ScenarioId("MAXIMIZE_APP_NON_RESIZABLE"),
+                extractor =
+                TaggedScenarioExtractorBuilder()
+                    .setTargetTag(CujType.CUJ_DESKTOP_MODE_MAXIMIZE_WINDOW)
+                    .setTransitionMatcher(
+                        TaggedCujTransitionMatcher(associatedTransitionRequired = false)
+                    )
+                    .build(),
+                assertions =
+                AssertionTemplates.DESKTOP_MODE_APP_VISIBILITY_ASSERTIONS +
+                        listOf(
+                            AppLayerIncreasesInSize(DESKTOP_MODE_APP),
+                            AppWindowMaintainsAspectRatioAlways(DESKTOP_MODE_APP),
+                            AppWindowHasMaxBoundsInOnlyOneDimension(DESKTOP_MODE_APP)
+                        ).associateBy({ it }, { AssertionInvocationGroup.BLOCKING }),
+            )
     }
 }
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppLandscape.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppLandscape.kt
new file mode 100644
index 0000000..2179566
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppLandscape.kt
@@ -0,0 +1,50 @@
+/*
+ * 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.wm.shell.flicker
+
+import android.tools.Rotation.ROTATION_90
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.MAXIMIZE_APP
+import com.android.wm.shell.scenarios.MaximizeAppWindow
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Maximize app window by pressing the maximize button on the app header.
+ *
+ * Assert that the app window keeps the same increases in size, filling the vertical and horizontal
+ * stable display bounds.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class MaximizeAppLandscape : MaximizeAppWindow(rotation = ROTATION_90) {
+    @ExpectedScenarios(["MAXIMIZE_APP"])
+    @Test
+    override fun maximizeAppWindow() = super.maximizeAppWindow()
+
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(MAXIMIZE_APP)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizableLandscape.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizableLandscape.kt
new file mode 100644
index 0000000..b173a60
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizableLandscape.kt
@@ -0,0 +1,53 @@
+/*
+ * 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.wm.shell.flicker
+
+import android.tools.Rotation.ROTATION_90
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.MAXIMIZE_APP_NON_RESIZABLE
+import com.android.wm.shell.scenarios.MaximizeAppWindow
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Maximize non-resizable app window by pressing the maximize button on the app header.
+ *
+ * Assert that the app window keeps the same increases in size, maintaining its aspect ratio, until
+ * filling the vertical or horizontal stable display bounds.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class MaximizeAppNonResizableLandscape : MaximizeAppWindow(
+    rotation = ROTATION_90,
+    isResizable = false
+) {
+    @ExpectedScenarios(["MAXIMIZE_APP_NON_RESIZABLE"])
+    @Test
+    override fun maximizeAppWindow() = super.maximizeAppWindow()
+
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(MAXIMIZE_APP_NON_RESIZABLE)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizablePortrait.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizablePortrait.kt
new file mode 100644
index 0000000..88888ee
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppNonResizablePortrait.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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.MAXIMIZE_APP_NON_RESIZABLE
+import com.android.wm.shell.scenarios.MaximizeAppWindow
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Maximize non-resizable app window by pressing the maximize button on the app header.
+ *
+ * Assert that the app window keeps the same increases in size, maintaining its aspect ratio, until
+ * filling the vertical or horizontal stable display bounds.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class MaximizeAppNonResizablePortrait : MaximizeAppWindow(isResizable = false) {
+    @ExpectedScenarios(["MAXIMIZE_APP_NON_RESIZABLE"])
+    @Test
+    override fun maximizeAppWindow() = super.maximizeAppWindow()
+
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(MAXIMIZE_APP_NON_RESIZABLE)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppPortrait.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppPortrait.kt
new file mode 100644
index 0000000..b79fd203
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/MaximizeAppPortrait.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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.MAXIMIZE_APP
+import com.android.wm.shell.scenarios.MaximizeAppWindow
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Maximize app window by pressing the maximize button on the app header.
+ *
+ * Assert that the app window keeps the same increases in size, filling the vertical and horizontal
+ * stable display bounds.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class MaximizeAppPortrait : MaximizeAppWindow() {
+    @ExpectedScenarios(["MAXIMIZE_APP"])
+    @Test
+    override fun maximizeAppWindow() = super.maximizeAppWindow()
+
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(MAXIMIZE_APP)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeMouse.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeMouse.kt
new file mode 100644
index 0000000..c3abf23
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeMouse.kt
@@ -0,0 +1,55 @@
+/*
+ * 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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.server.wm.flicker.helpers.MotionEventHelper.InputMethod
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.EDGE_RESIZE
+import com.android.wm.shell.scenarios.ResizeAppWithEdgeResize
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class ResizeAppWithEdgeResizeMouse : ResizeAppWithEdgeResize(InputMethod.MOUSE) {
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeRight() = super.resizeAppWithEdgeResizeRight()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeLeft() = super.resizeAppWithEdgeResizeLeft()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeTop() = super.resizeAppWithEdgeResizeTop()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeBottom() = super.resizeAppWithEdgeResizeBottom()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(EDGE_RESIZE)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeStylus.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeStylus.kt
new file mode 100644
index 0000000..86b0e6f
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeStylus.kt
@@ -0,0 +1,55 @@
+/*
+ * 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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.server.wm.flicker.helpers.MotionEventHelper.InputMethod
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.EDGE_RESIZE
+import com.android.wm.shell.scenarios.ResizeAppWithEdgeResize
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class ResizeAppWithEdgeResizeStylus : ResizeAppWithEdgeResize(InputMethod.STYLUS) {
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeRight() = super.resizeAppWithEdgeResizeRight()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeLeft() = super.resizeAppWithEdgeResizeLeft()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeTop() = super.resizeAppWithEdgeResizeTop()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeBottom() = super.resizeAppWithEdgeResizeBottom()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(EDGE_RESIZE)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeTouchpad.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeTouchpad.kt
new file mode 100644
index 0000000..e6bb9ef
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/ResizeAppWithEdgeResizeTouchpad.kt
@@ -0,0 +1,55 @@
+/*
+ * 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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.server.wm.flicker.helpers.MotionEventHelper.InputMethod
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.EDGE_RESIZE
+import com.android.wm.shell.scenarios.ResizeAppWithEdgeResize
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class ResizeAppWithEdgeResizeTouchpad : ResizeAppWithEdgeResize(InputMethod.TOUCHPAD) {
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeRight() = super.resizeAppWithEdgeResizeRight()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeLeft() = super.resizeAppWithEdgeResizeLeft()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeTop() = super.resizeAppWithEdgeResizeTop()
+
+    @ExpectedScenarios(["EDGE_RESIZE"])
+    @Test
+    override fun resizeAppWithEdgeResizeBottom() = super.resizeAppWithEdgeResizeBottom()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(EDGE_RESIZE)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithButton.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithButton.kt
new file mode 100644
index 0000000..b5090086
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithButton.kt
@@ -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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.SNAP_RESIZE_LEFT_WITH_BUTTON
+import com.android.wm.shell.scenarios.SnapResizeAppWindowWithButton
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Snap resize app window using the Snap Left button from the maximize menu.
+ *
+ * Assert that the app window fills the left half the display after being snap resized.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class SnapResizeAppWindowLeftWithButton : SnapResizeAppWindowWithButton(toLeft = true) {
+    @ExpectedScenarios(["SNAP_RESIZE_LEFT_WITH_BUTTON"])
+    @Test
+    override fun snapResizeAppWindowWithButton() = super.snapResizeAppWindowWithButton()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(SNAP_RESIZE_LEFT_WITH_BUTTON)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithDrag.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithDrag.kt
new file mode 100644
index 0000000..a22e760
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowLeftWithDrag.kt
@@ -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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.SNAP_RESIZE_LEFT_WITH_DRAG
+import com.android.wm.shell.scenarios.SnapResizeAppWindowWithDrag
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Snap resize app window by dragging it to the left edge of the screen.
+ *
+ * Assert that the app window fills the left half the display after being snap resized.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class SnapResizeAppWindowLeftWithDrag : SnapResizeAppWindowWithDrag(toLeft = true) {
+    @ExpectedScenarios(["SNAP_RESIZE_LEFT_WITH_DRAG"])
+    @Test
+    override fun snapResizeAppWindowWithDrag() = super.snapResizeAppWindowWithDrag()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(SNAP_RESIZE_LEFT_WITH_DRAG)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithButton.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithButton.kt
new file mode 100644
index 0000000..375a2b8
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithButton.kt
@@ -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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.SNAP_RESIZE_RIGHT_WITH_BUTTON
+import com.android.wm.shell.scenarios.SnapResizeAppWindowWithButton
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Snap resize app window using the Snap Right button from the maximize menu.
+ *
+ * Assert that the app window fills the right half the display after being snap resized.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class SnapResizeAppWindowRightWithButton : SnapResizeAppWindowWithButton(toLeft = false) {
+    @ExpectedScenarios(["SNAP_RESIZE_RIGHT_WITH_BUTTON"])
+    @Test
+    override fun snapResizeAppWindowWithButton() = super.snapResizeAppWindowWithButton()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(SNAP_RESIZE_RIGHT_WITH_BUTTON)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithDrag.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithDrag.kt
new file mode 100644
index 0000000..4a9daf7
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/flicker-service/src/com/android/wm/shell/flicker/SnapResizeAppWindowRightWithDrag.kt
@@ -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.wm.shell.flicker
+
+import android.tools.flicker.FlickerConfig
+import android.tools.flicker.annotation.ExpectedScenarios
+import android.tools.flicker.annotation.FlickerConfigProvider
+import android.tools.flicker.config.FlickerConfig
+import android.tools.flicker.config.FlickerServiceConfig
+import android.tools.flicker.junit.FlickerServiceJUnit4ClassRunner
+import com.android.wm.shell.flicker.DesktopModeFlickerScenarios.Companion.SNAP_RESIZE_RIGHT_WITH_DRAG
+import com.android.wm.shell.scenarios.SnapResizeAppWindowWithDrag
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Snap resize app window by dragging it to the right edge of the screen.
+ *
+ * Assert that the app window fills the right half the display after being snap resized.
+ */
+@RunWith(FlickerServiceJUnit4ClassRunner::class)
+class SnapResizeAppWindowRightWithDrag : SnapResizeAppWindowWithDrag(toLeft = false) {
+    @ExpectedScenarios(["SNAP_RESIZE_RIGHT_WITH_DRAG"])
+    @Test
+    override fun snapResizeAppWindowWithDrag() = super.snapResizeAppWindowWithDrag()
+
+    companion object {
+        @JvmStatic
+        @FlickerConfigProvider
+        fun flickerConfigProvider(): FlickerConfig =
+            FlickerConfig().use(FlickerServiceConfig.DEFAULT).use(SNAP_RESIZE_RIGHT_WITH_DRAG)
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/DesktopScenarioCustomAppTestBase.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/DesktopScenarioCustomAppTestBase.kt
new file mode 100644
index 0000000..5a69b27
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/DesktopScenarioCustomAppTestBase.kt
@@ -0,0 +1,55 @@
+/*
+ * 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.wm.shell.scenarios
+
+import android.app.Instrumentation
+import android.tools.traces.parsers.WindowManagerStateHelper
+import android.tools.traces.parsers.toFlickerComponent
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.UiDevice
+import com.android.launcher3.tapl.LauncherInstrumentation
+import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
+import com.android.server.wm.flicker.helpers.LetterboxAppHelper
+import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
+import com.android.server.wm.flicker.helpers.SimpleAppHelper
+import com.android.server.wm.flicker.testapp.ActivityOptions
+import org.junit.Ignore
+
+/** Base test class for desktop CUJ with customizable test app. */
+@Ignore("Base Test Class")
+abstract class DesktopScenarioCustomAppTestBase(
+    isResizeable: Boolean = true,
+    isLandscapeApp: Boolean = true
+) {
+    val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
+    val tapl = LauncherInstrumentation()
+    val wmHelper = WindowManagerStateHelper(instrumentation)
+    val device = UiDevice.getInstance(instrumentation)
+    // TODO(b/363181411): Consolidate in LetterboxAppHelper.
+    val testApp = when {
+        isResizeable && isLandscapeApp -> SimpleAppHelper(instrumentation)
+        isResizeable && !isLandscapeApp -> SimpleAppHelper(
+                instrumentation,
+                launcherName = ActivityOptions.PortraitOnlyActivity.LABEL,
+                component = ActivityOptions.PortraitOnlyActivity.COMPONENT.toFlickerComponent()
+            )
+        // NonResizeablAppHelper has no fixed orientation.
+        !isResizeable && isLandscapeApp -> NonResizeableAppHelper(instrumentation)
+        // Opens NonResizeablePortraitActivity.
+        else -> LetterboxAppHelper(instrumentation)
+    }.let { DesktopModeAppHelper(it) }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/EnterDesktopWithDrag.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/EnterDesktopWithDrag.kt
index 0f0d2df..5f759e8 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/EnterDesktopWithDrag.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/EnterDesktopWithDrag.kt
@@ -17,15 +17,8 @@
 package com.android.wm.shell.scenarios
 
 import android.platform.test.annotations.Postsubmit
-import android.app.Instrumentation
 import android.tools.NavBar
 import android.tools.Rotation
-import android.tools.traces.parsers.WindowManagerStateHelper
-import androidx.test.platform.app.InstrumentationRegistry
-import androidx.test.uiautomator.UiDevice
-import com.android.launcher3.tapl.LauncherInstrumentation
-import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
-import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.window.flags.Flags
 import com.android.wm.shell.Utils
 import org.junit.After
@@ -40,13 +33,11 @@
 @Postsubmit
 open class EnterDesktopWithDrag
 @JvmOverloads
-constructor(val rotation: Rotation = Rotation.ROTATION_0) {
-
-    private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
-    private val tapl = LauncherInstrumentation()
-    private val wmHelper = WindowManagerStateHelper(instrumentation)
-    private val device = UiDevice.getInstance(instrumentation)
-    private val testApp = DesktopModeAppHelper(SimpleAppHelper(instrumentation))
+constructor(
+    val rotation: Rotation = Rotation.ROTATION_0,
+    isResizeable: Boolean = true,
+    isLandscapeApp: Boolean = true
+) : DesktopScenarioCustomAppTestBase(isResizeable, isLandscapeApp) {
 
     @Rule @JvmField val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, rotation)
 
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ExitDesktopWithDragToTopDragZone.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ExitDesktopWithDragToTopDragZone.kt
index 533be88..b616e53 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ExitDesktopWithDragToTopDragZone.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ExitDesktopWithDragToTopDragZone.kt
@@ -17,15 +17,8 @@
 package com.android.wm.shell.scenarios
 
 import android.platform.test.annotations.Postsubmit
-import android.app.Instrumentation
 import android.tools.NavBar
 import android.tools.Rotation
-import android.tools.traces.parsers.WindowManagerStateHelper
-import androidx.test.platform.app.InstrumentationRegistry
-import androidx.test.uiautomator.UiDevice
-import com.android.launcher3.tapl.LauncherInstrumentation
-import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
-import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.window.flags.Flags
 import com.android.wm.shell.Utils
 import org.junit.After
@@ -40,13 +33,11 @@
 @Postsubmit
 open class ExitDesktopWithDragToTopDragZone
 @JvmOverloads
-constructor(val rotation: Rotation = Rotation.ROTATION_0) {
-
-    private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
-    private val tapl = LauncherInstrumentation()
-    private val wmHelper = WindowManagerStateHelper(instrumentation)
-    private val device = UiDevice.getInstance(instrumentation)
-    private val testApp = DesktopModeAppHelper(SimpleAppHelper(instrumentation))
+constructor(
+    val rotation: Rotation = Rotation.ROTATION_0,
+    isResizeable: Boolean = true,
+    isLandscapeApp: Boolean = true
+) : DesktopScenarioCustomAppTestBase(isResizeable, isLandscapeApp) {
 
     @Rule @JvmField val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, rotation)
 
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/MaximizeAppWindow.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/MaximizeAppWindow.kt
index b812c59..426f40b 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/MaximizeAppWindow.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/MaximizeAppWindow.kt
@@ -16,10 +16,11 @@
 
 package com.android.wm.shell.scenarios
 
-import android.platform.test.annotations.Postsubmit
 import android.app.Instrumentation
+import android.platform.test.annotations.Postsubmit
 import android.tools.NavBar
 import android.tools.Rotation
+import android.tools.flicker.rules.ChangeDisplayOrientationRule
 import android.tools.traces.parsers.WindowManagerStateHelper
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
@@ -36,11 +37,12 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.BlockJUnit4ClassRunner
+
 @RunWith(BlockJUnit4ClassRunner::class)
 @Postsubmit
 open class MaximizeAppWindow
 @JvmOverloads
-constructor(rotation: Rotation = Rotation.ROTATION_0, isResizable: Boolean = true) {
+constructor(private val rotation: Rotation = Rotation.ROTATION_0, isResizable: Boolean = true) {
 
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     private val tapl = LauncherInstrumentation()
@@ -57,6 +59,9 @@
     @Before
     fun setup() {
         Assume.assumeTrue(Flags.enableDesktopWindowingMode() && tapl.isTablet)
+        tapl.setEnableRotation(true)
+        tapl.setExpectedRotation(rotation.value)
+        ChangeDisplayOrientationRule.setRotation(rotation)
         testApp.enterDesktopWithDrag(wmHelper, device)
     }
 
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithCornerResize.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithCornerResize.kt
index 03d970f..42940a9 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithCornerResize.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithCornerResize.kt
@@ -43,8 +43,8 @@
 @JvmOverloads
 constructor(
     val rotation: Rotation = Rotation.ROTATION_0,
-    val horizontalChange: Int = 50,
-    val verticalChange: Int = -50,
+    val horizontalChange: Int = 200,
+    val verticalChange: Int = -200,
     val appProperty: AppProperty = AppProperty.STANDARD
 ) {
 
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithEdgeResize.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithEdgeResize.kt
new file mode 100644
index 0000000..d094967
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/ResizeAppWithEdgeResize.kt
@@ -0,0 +1,110 @@
+/*
+ * 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.wm.shell.scenarios
+
+import android.platform.test.annotations.Postsubmit
+import android.app.Instrumentation
+import android.tools.NavBar
+import android.tools.Rotation
+import android.tools.traces.parsers.WindowManagerStateHelper
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.UiDevice
+import com.android.launcher3.tapl.LauncherInstrumentation
+import com.android.server.wm.flicker.helpers.DesktopModeAppHelper
+import com.android.server.wm.flicker.helpers.MotionEventHelper
+import com.android.server.wm.flicker.helpers.SimpleAppHelper
+import com.android.window.flags.Flags
+import com.android.wm.shell.Utils
+import org.junit.After
+import org.junit.Assume
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.BlockJUnit4ClassRunner
+
+@RunWith(BlockJUnit4ClassRunner::class)
+@Postsubmit
+open class ResizeAppWithEdgeResize
+@JvmOverloads
+constructor(
+    val inputMethod: MotionEventHelper.InputMethod,
+    val rotation: Rotation = Rotation.ROTATION_90
+) {
+    private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
+    private val tapl = LauncherInstrumentation()
+    private val wmHelper = WindowManagerStateHelper(instrumentation)
+    private val device = UiDevice.getInstance(instrumentation)
+    private val testApp = DesktopModeAppHelper(SimpleAppHelper(instrumentation))
+    private val motionEventHelper = MotionEventHelper(instrumentation, inputMethod)
+
+    @Rule
+    @JvmField
+    val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, rotation)
+
+    @Before
+    fun setup() {
+        Assume.assumeTrue(
+            Flags.enableDesktopWindowingMode()
+                    && Flags.enableWindowingEdgeDragResize() && tapl.isTablet
+        )
+        tapl.setEnableRotation(true)
+        tapl.setExpectedRotation(rotation.value)
+        testApp.enterDesktopWithDrag(wmHelper, device)
+    }
+
+    @Test
+    open fun resizeAppWithEdgeResizeRight() {
+        testApp.edgeResize(
+            wmHelper,
+            motionEventHelper,
+            DesktopModeAppHelper.Edges.RIGHT
+        )
+    }
+
+    @Test
+    open fun resizeAppWithEdgeResizeLeft() {
+        testApp.edgeResize(
+            wmHelper,
+            motionEventHelper,
+            DesktopModeAppHelper.Edges.LEFT
+        )
+    }
+
+    @Test
+    open fun resizeAppWithEdgeResizeTop() {
+        testApp.edgeResize(
+            wmHelper,
+            motionEventHelper,
+            DesktopModeAppHelper.Edges.TOP
+        )
+    }
+
+    @Test
+    open fun resizeAppWithEdgeResizeBottom() {
+        testApp.edgeResize(
+            wmHelper,
+            motionEventHelper,
+            DesktopModeAppHelper.Edges.BOTTOM
+        )
+    }
+
+    @After
+    fun teardown() {
+        testApp.exit(wmHelper)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithButton.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithButton.kt
index 685a3ba..33242db 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithButton.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithButton.kt
@@ -18,6 +18,8 @@
 
 import android.app.Instrumentation
 import android.platform.test.annotations.Postsubmit
+import android.tools.NavBar
+import android.tools.Rotation
 import android.tools.traces.parsers.WindowManagerStateHelper
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
@@ -26,9 +28,11 @@
 import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.window.flags.Flags
+import com.android.wm.shell.Utils
 import org.junit.After
 import org.junit.Assume
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.BlockJUnit4ClassRunner
@@ -37,7 +41,7 @@
 @Postsubmit
 open class SnapResizeAppWindowWithButton
 @JvmOverloads
-constructor(private val toLeft: Boolean = true, private val isResizable: Boolean = true) {
+constructor(private val toLeft: Boolean = true, isResizable: Boolean = true) {
 
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     private val tapl = LauncherInstrumentation()
@@ -49,6 +53,10 @@
         DesktopModeAppHelper(NonResizeableAppHelper(instrumentation))
     }
 
+    @Rule
+    @JvmField
+    val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, Rotation.ROTATION_0)
+
     @Before
     fun setup() {
         Assume.assumeTrue(Flags.enableDesktopWindowingMode() && tapl.isTablet)
diff --git a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithDrag.kt b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithDrag.kt
index 8a4aa63..14eb779 100644
--- a/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithDrag.kt
+++ b/libs/WindowManager/Shell/tests/e2e/desktopmode/scenarios/src/com/android/wm/shell/scenarios/SnapResizeAppWindowWithDrag.kt
@@ -18,6 +18,8 @@
 
 import android.app.Instrumentation
 import android.platform.test.annotations.Postsubmit
+import android.tools.NavBar
+import android.tools.Rotation
 import android.tools.traces.parsers.WindowManagerStateHelper
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.UiDevice
@@ -26,9 +28,11 @@
 import com.android.server.wm.flicker.helpers.NonResizeableAppHelper
 import com.android.server.wm.flicker.helpers.SimpleAppHelper
 import com.android.window.flags.Flags
+import com.android.wm.shell.Utils
 import org.junit.After
 import org.junit.Assume
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.BlockJUnit4ClassRunner
@@ -37,7 +41,7 @@
 @Postsubmit
 open class SnapResizeAppWindowWithDrag
 @JvmOverloads
-constructor(private val toLeft: Boolean = true, private val isResizable: Boolean = true) {
+constructor(private val toLeft: Boolean = true, isResizable: Boolean = true) {
 
     private val instrumentation: Instrumentation = InstrumentationRegistry.getInstrumentation()
     private val tapl = LauncherInstrumentation()
@@ -49,6 +53,10 @@
         DesktopModeAppHelper(NonResizeableAppHelper(instrumentation))
     }
 
+    @Rule
+    @JvmField
+    val testSetupRule = Utils.testSetupRule(NavBar.MODE_GESTURAL, Rotation.ROTATION_0)
+
     @Before
     fun setup() {
         Assume.assumeTrue(Flags.enableDesktopWindowingMode() && tapl.isTablet)
diff --git a/libs/WindowManager/Shell/tests/e2e/utils/src/com/android/wm/shell/Utils.kt b/libs/WindowManager/Shell/tests/e2e/utils/src/com/android/wm/shell/Utils.kt
index dee67f3..c0fafef 100644
--- a/libs/WindowManager/Shell/tests/e2e/utils/src/com/android/wm/shell/Utils.kt
+++ b/libs/WindowManager/Shell/tests/e2e/utils/src/com/android/wm/shell/Utils.kt
@@ -17,6 +17,7 @@
 package com.android.wm.shell
 
 import android.app.Instrumentation
+import android.platform.test.rule.EnsureDeviceSettingsRule
 import android.platform.test.rule.NavigationModeRule
 import android.platform.test.rule.PressHomeRule
 import android.platform.test.rule.UnlockScreenRule
@@ -49,5 +50,6 @@
                 )
             )
             .around(PressHomeRule())
+            .around(EnsureDeviceSettingsRule())
     }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/Android.bp b/libs/WindowManager/Shell/tests/flicker/Android.bp
index 4abaf5b..7305f49 100644
--- a/libs/WindowManager/Shell/tests/flicker/Android.bp
+++ b/libs/WindowManager/Shell/tests/flicker/Android.bp
@@ -68,6 +68,7 @@
         "flickerlib-helpers",
         "flickerlib-trace_processor_shell",
         "platform-test-annotations",
+        "platform-test-rules",
         "wm-flicker-common-app-helpers",
         "wm-flicker-common-assertions",
         "launcher-helper-lib",
diff --git a/libs/WindowManager/Shell/tests/flicker/appcompat/OWNERS b/libs/WindowManager/Shell/tests/flicker/appcompat/OWNERS
new file mode 100644
index 0000000..a36a4f8
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/flicker/appcompat/OWNERS
@@ -0,0 +1,2 @@
+# Window Manager > App Compat
+# Bug component: 970984
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipOnGoToHomeTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipOnGoToHomeTest.kt
index b85d793..a9ed13a 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipOnGoToHomeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/AutoEnterPipOnGoToHomeTest.kt
@@ -16,10 +16,14 @@
 
 package com.android.wm.shell.flicker.pip
 
+import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.tools.flicker.junit.FlickerParametersRunnerFactory
 import android.tools.flicker.legacy.FlickerBuilder
 import android.tools.flicker.legacy.LegacyFlickerTest
+import android.tools.flicker.subject.exceptions.ExceptionMessageBuilder
+import android.tools.flicker.subject.exceptions.IncorrectRegionException
+import android.tools.flicker.subject.layers.LayerSubject
 import androidx.test.filters.FlakyTest
 import androidx.test.filters.RequiresDevice
 import com.android.wm.shell.flicker.pip.common.EnterPipTransition
@@ -29,6 +33,7 @@
 import org.junit.runner.RunWith
 import org.junit.runners.MethodSorters
 import org.junit.runners.Parameterized
+import kotlin.math.abs
 
 /**
  * Test entering pip from an app via auto-enter property when navigating to home.
@@ -67,9 +72,24 @@
 
     override val defaultTeardown: FlickerBuilder.() -> Unit = { teardown { pipApp.exit(wmHelper) } }
 
-    @FlakyTest(bugId = 293133362)
+    private val widthNotSmallerThan: LayerSubject.(LayerSubject) -> Unit = {
+        val width = visibleRegion.region.bounds.width()
+        val otherWidth = it.visibleRegion.region.bounds.width()
+        if (width < otherWidth && abs(width - otherWidth) > EPSILON) {
+            val errorMsgBuilder =
+                ExceptionMessageBuilder()
+                    .forSubject(this)
+                    .forIncorrectRegion("width. $width smaller than $otherWidth")
+                    .setExpected(width)
+                    .setActual(otherWidth)
+            throw IncorrectRegionException(errorMsgBuilder)
+        }
+    }
+
+    @Postsubmit
     @Test
     override fun pipLayerReduces() {
+        Assume.assumeFalse(flicker.scenario.isGesturalNavigation)
         flicker.assertLayers {
             val pipLayerList = this.layers { pipApp.layerMatchesAnyOf(it) && it.isVisible }
             pipLayerList.zipWithNext { previous, current ->
@@ -78,6 +98,32 @@
         }
     }
 
+    /** Checks that [pipApp] window's width is first decreasing then increasing. */
+    @Postsubmit
+    @Test
+    fun pipLayerWidthDecreasesThenIncreases() {
+        Assume.assumeTrue(flicker.scenario.isGesturalNavigation)
+        flicker.assertLayers {
+            val pipLayerList = this.layers { pipApp.layerMatchesAnyOf(it) && it.isVisible }
+            var previousLayer = pipLayerList[0]
+            var currentLayer = previousLayer
+            var i = 0
+            invoke("layer area is decreasing") {
+                if (i < pipLayerList.size - 1) {
+                    previousLayer = currentLayer
+                    currentLayer = pipLayerList[++i]
+                    previousLayer.widthNotSmallerThan(currentLayer)
+                }
+            }.then().invoke("layer are is increasing", true /* isOptional */) {
+                if (i < pipLayerList.size - 1) {
+                    previousLayer = currentLayer
+                    currentLayer = pipLayerList[++i]
+                    currentLayer.widthNotSmallerThan(previousLayer)
+                }
+            }
+        }
+    }
+
     /** Checks that [pipApp] window is animated towards default position in right bottom corner */
     @FlakyTest(bugId = 255578530)
     @Test
@@ -108,4 +154,9 @@
     override fun visibleLayersShownMoreThanOneConsecutiveEntry() {
         super.visibleLayersShownMoreThanOneConsecutiveEntry()
     }
+
+    companion object {
+        // TODO(b/363080056): A margin of error allowed on certain layer size calculations.
+        const val EPSILON = 1
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
index 70be58f..429774f 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipAspectRatioChangeTest.kt
@@ -35,7 +35,7 @@
 @FixMethodOrder(MethodSorters.NAME_ASCENDING)
 class PipAspectRatioChangeTest(flicker: LegacyFlickerTest) : PipTransition(flicker) {
     override val thisTransition: FlickerBuilder.() -> Unit = {
-        transitions { pipApp.changeAspectRatio() }
+        transitions { pipApp.changeAspectRatio(wmHelper) }
     }
 
     @Presubmit
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
index 90b9798..cbd4a52 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/PipDragThenSnapTest.kt
@@ -181,6 +181,13 @@
         super.taskBarWindowIsAlwaysVisible()
     }
 
+    // Overridden to remove @Postsubmit annotation
+    @Test
+    @FlakyTest(bugId = 294993100)
+    override fun pipLayerHasCorrectCornersAtEnd() {
+        // No rounded corners as we go back to fullscreen in new orientation.
+    }
+
     companion object {
         /**
          * Creates the test configurations.
diff --git a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinned.kt b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinned.kt
index ed2a0a7..578a9b5 100644
--- a/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinned.kt
+++ b/libs/WindowManager/Shell/tests/flicker/pip/src/com/android/wm/shell/flicker/pip/SetRequestedOrientationWhilePinned.kt
@@ -147,6 +147,12 @@
     @Test
     override fun entireScreenCovered() = super.entireScreenCovered()
 
+    @Postsubmit
+    @Test
+    override fun pipLayerHasCorrectCornersAtEnd() {
+        flicker.assertLayersEnd { hasNoRoundedCorners(pipApp) }
+    }
+
     companion object {
         @Parameterized.Parameters(name = "{0}")
         @JvmStatic
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 8cb81b4..f57335c 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
@@ -70,6 +70,11 @@
         }
     }
 
+    @Test
+    override fun pipLayerHasCorrectCornersAtEnd() {
+        // PiP might have completely faded out by this point, so corner radii not applicable.
+    }
+
     companion object {
         /**
          * Creates the test configurations.
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 0742cf9..ce84eb6 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
@@ -16,6 +16,7 @@
 
 package com.android.wm.shell.flicker.pip.common
 
+import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
 import android.tools.flicker.legacy.LegacyFlickerTest
@@ -123,6 +124,12 @@
         }
     }
 
+    @Postsubmit
+    @Test
+    override fun pipLayerHasCorrectCornersAtEnd() {
+        flicker.assertLayersEnd { hasNoRoundedCorners(pipApp) }
+    }
+
     /** {@inheritDoc} */
     @Presubmit @Test override fun entireScreenCovered() = super.entireScreenCovered()
 
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 99c1ad2..bc2bfdb 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
@@ -18,6 +18,7 @@
 
 import android.app.Instrumentation
 import android.content.Intent
+import android.platform.test.annotations.Postsubmit
 import android.platform.test.annotations.Presubmit
 import android.tools.Rotation
 import android.tools.flicker.legacy.FlickerBuilder
@@ -105,4 +106,10 @@
                 .doesNotContain(false)
         }
     }
+
+    @Postsubmit
+    @Test
+    open fun pipLayerHasCorrectCornersAtEnd() {
+        flicker.assertLayersEnd { hasRoundedCorners(pipApp) }
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestShellExecutor.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestShellExecutor.java
index 4998702..f31722d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestShellExecutor.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/TestShellExecutor.java
@@ -19,6 +19,7 @@
 import com.android.wm.shell.common.ShellExecutor;
 
 import java.util.ArrayList;
+import java.util.List;
 
 /**
  * Really basic test executor. It just gathers all events in a blob. The only option is to
@@ -52,4 +53,9 @@
             mRunnables.remove(0).run();
         }
     }
+
+    /** Returns the list of callbacks for this executor. */
+    public List<Runnable> getCallbacks() {
+        return mRunnables;
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
index 859602e..6fa3788 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleDataTest.java
@@ -50,8 +50,8 @@
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.bubbles.BubbleData.TimeSource;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
-import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarUpdate;
 
 import com.google.common.collect.ImmutableList;
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleTest.java
index 50c4a18..dca5fc4 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/bubbles/BubbleTest.java
@@ -43,7 +43,7 @@
 import com.android.wm.shell.R;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.bubbles.BubbleInfo;
+import com.android.wm.shell.shared.bubbles.BubbleInfo;
 
 import org.junit.Before;
 import org.junit.Test;
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
new file mode 100644
index 0000000..b14f163
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandlerTest.kt
@@ -0,0 +1,263 @@
+/*
+ * 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.wm.shell.desktopmode
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.content.pm.ActivityInfo
+import android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE
+import android.content.pm.ActivityInfo.SCREEN_ORIENTATION_PORTRAIT
+import android.content.pm.ActivityInfo.SCREEN_ORIENTATION_SENSOR_PORTRAIT
+import android.graphics.Rect
+import android.os.Binder
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.testing.AndroidTestingRunner
+import android.view.Display.DEFAULT_DISPLAY
+import android.window.WindowContainerTransaction
+import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.ExtendedMockito.never
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODE
+import com.android.window.flags.Flags.FLAG_RESPECT_ORIENTATION_CHANGE_FOR_UNRESIZEABLE
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.common.TaskStackListenerImpl
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
+import com.android.wm.shell.sysui.ShellInit
+import com.android.wm.shell.transition.Transitions
+import junit.framework.Assert.assertEquals
+import junit.framework.Assert.assertTrue
+import kotlin.test.assertNotNull
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.ArgumentMatchers.isNull
+import org.mockito.Mock
+import org.mockito.Mockito.anyInt
+import org.mockito.Mockito.clearInvocations
+import org.mockito.Mockito.spy
+import org.mockito.Mockito.verify
+import org.mockito.kotlin.any
+import org.mockito.kotlin.atLeastOnce
+import org.mockito.kotlin.capture
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.whenever
+import org.mockito.quality.Strictness
+
+/**
+ * Test class for {@link DesktopActivityOrientationChangeHandler}
+ *
+ * Usage: atest WMShellUnitTests:DesktopActivityOrientationChangeHandlerTest
+ */
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_MODE, FLAG_RESPECT_ORIENTATION_CHANGE_FOR_UNRESIZEABLE)
+class DesktopActivityOrientationChangeHandlerTest : ShellTestCase() {
+    @JvmField @Rule val setFlagsRule = SetFlagsRule()
+
+    @Mock lateinit var testExecutor: ShellExecutor
+    @Mock lateinit var shellTaskOrganizer: ShellTaskOrganizer
+    @Mock lateinit var transitions: Transitions
+    @Mock lateinit var resizeTransitionHandler: ToggleResizeDesktopTaskTransitionHandler
+    @Mock lateinit var taskStackListener: TaskStackListenerImpl
+
+    private lateinit var mockitoSession: StaticMockitoSession
+    private lateinit var handler: DesktopActivityOrientationChangeHandler
+    private lateinit var shellInit: ShellInit
+    private lateinit var taskRepository: DesktopModeTaskRepository
+    // Mock running tasks are registered here so we can get the list from mock shell task organizer.
+    private val runningTasks = mutableListOf<RunningTaskInfo>()
+
+    @Before
+    fun setUp() {
+        mockitoSession =
+            mockitoSession()
+                .strictness(Strictness.LENIENT)
+                .spyStatic(DesktopModeStatus::class.java)
+                .startMocking()
+        doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
+
+        shellInit = spy(ShellInit(testExecutor))
+        taskRepository = DesktopModeTaskRepository()
+        whenever(shellTaskOrganizer.getRunningTasks(anyInt())).thenAnswer { runningTasks }
+        whenever(transitions.startTransition(anyInt(), any(), isNull())).thenAnswer { Binder() }
+
+        handler = DesktopActivityOrientationChangeHandler(context, shellInit, shellTaskOrganizer,
+            taskStackListener, resizeTransitionHandler, taskRepository)
+
+        shellInit.init()
+    }
+
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
+
+        runningTasks.clear()
+    }
+
+    @Test
+    fun instantiate_addInitCallback() {
+        verify(shellInit).addInitCallback(any(), any<DesktopActivityOrientationChangeHandler>())
+    }
+
+    @Test
+    fun instantiate_cannotEnterDesktopMode_doNotAddInitCallback() {
+        whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(false)
+        clearInvocations(shellInit)
+
+        handler = DesktopActivityOrientationChangeHandler(context, shellInit, shellTaskOrganizer,
+            taskStackListener, resizeTransitionHandler, taskRepository)
+
+        verify(shellInit, never()).addInitCallback(any(),
+            any<DesktopActivityOrientationChangeHandler>())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_resizeable_doNothing() {
+        val task = setUpFreeformTask()
+
+        taskStackListener.onActivityRequestedOrientationChanged(task.taskId,
+            SCREEN_ORIENTATION_LANDSCAPE)
+
+        verify(resizeTransitionHandler, never()).startTransition(any(), any())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_nonResizeableFullscreen_doNothing() {
+        val task = createFullscreenTask()
+        task.isResizeable = false
+        val activityInfo = ActivityInfo()
+        activityInfo.screenOrientation = SCREEN_ORIENTATION_PORTRAIT
+        task.topActivityInfo = activityInfo
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        taskRepository.addActiveTask(DEFAULT_DISPLAY, task.taskId)
+        taskRepository.updateTaskVisibility(DEFAULT_DISPLAY, task.taskId, visible = true)
+        runningTasks.add(task)
+
+        taskStackListener.onActivityRequestedOrientationChanged(task.taskId,
+            SCREEN_ORIENTATION_LANDSCAPE)
+
+        verify(resizeTransitionHandler, never()).startTransition(any(), any())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_nonResizeablePortrait_requestSameOrientation_doNothing() {
+        val task = setUpFreeformTask(isResizeable = false)
+        val newTask = setUpFreeformTask(isResizeable = false,
+            orientation = SCREEN_ORIENTATION_SENSOR_PORTRAIT)
+
+        handler.handleActivityOrientationChange(task, newTask)
+
+        verify(resizeTransitionHandler, never()).startTransition(any(), any())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_notInDesktopMode_doNothing() {
+        val task = setUpFreeformTask(isResizeable = false)
+        taskRepository.updateTaskVisibility(task.displayId, task.taskId, visible = false)
+
+        taskStackListener.onActivityRequestedOrientationChanged(task.taskId,
+            SCREEN_ORIENTATION_LANDSCAPE)
+
+        verify(resizeTransitionHandler, never()).startTransition(any(), any())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_nonResizeablePortrait_respectLandscapeRequest() {
+        val task = setUpFreeformTask(isResizeable = false)
+        val oldBounds = task.configuration.windowConfiguration.bounds
+        val newTask = setUpFreeformTask(isResizeable = false,
+            orientation = SCREEN_ORIENTATION_LANDSCAPE)
+
+        handler.handleActivityOrientationChange(task, newTask)
+
+        val wct = getLatestResizeDesktopTaskWct()
+        val finalBounds = findBoundsChange(wct, newTask)
+        assertNotNull(finalBounds)
+        val finalWidth = finalBounds.width()
+        val finalHeight = finalBounds.height()
+        // Bounds is landscape.
+        assertTrue(finalWidth > finalHeight)
+        // Aspect ratio remains the same.
+        assertEquals(oldBounds.height() / oldBounds.width(), finalWidth / finalHeight)
+        // Anchor point for resizing is at the center.
+        assertEquals(oldBounds.centerX(), finalBounds.centerX())
+    }
+
+    @Test
+    fun handleActivityOrientationChange_nonResizeableLandscape_respectPortraitRequest() {
+        val oldBounds = Rect(0, 0, 500, 200)
+        val task = setUpFreeformTask(
+            isResizeable = false, orientation = SCREEN_ORIENTATION_LANDSCAPE, bounds = oldBounds
+        )
+        val newTask = setUpFreeformTask(isResizeable = false, bounds = oldBounds)
+
+        handler.handleActivityOrientationChange(task, newTask)
+
+        val wct = getLatestResizeDesktopTaskWct()
+        val finalBounds = findBoundsChange(wct, newTask)
+        assertNotNull(finalBounds)
+        val finalWidth = finalBounds.width()
+        val finalHeight = finalBounds.height()
+        // Bounds is portrait.
+        assertTrue(finalHeight > finalWidth)
+        // Aspect ratio remains the same.
+        assertEquals(oldBounds.width() / oldBounds.height(), finalHeight / finalWidth)
+        // Anchor point for resizing is at the center.
+        assertEquals(oldBounds.centerX(), finalBounds.centerX())
+    }
+
+    private fun setUpFreeformTask(
+        displayId: Int = DEFAULT_DISPLAY,
+        isResizeable: Boolean = true,
+        orientation: Int = SCREEN_ORIENTATION_PORTRAIT,
+        bounds: Rect? = Rect(0, 0, 200, 500)
+    ): RunningTaskInfo {
+        val task = createFreeformTask(displayId, bounds)
+        val activityInfo = ActivityInfo()
+        activityInfo.screenOrientation = orientation
+        task.topActivityInfo = activityInfo
+        task.isResizeable = isResizeable
+        whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        taskRepository.addActiveTask(displayId, task.taskId)
+        taskRepository.updateTaskVisibility(displayId, task.taskId, visible = true)
+        taskRepository.addOrMoveFreeformTaskToTop(displayId, task.taskId)
+        runningTasks.add(task)
+        return task
+    }
+
+    private fun getLatestResizeDesktopTaskWct(
+        currentBounds: Rect? = null
+    ): WindowContainerTransaction {
+        val arg: ArgumentCaptor<WindowContainerTransaction> =
+            ArgumentCaptor.forClass(WindowContainerTransaction::class.java)
+        verify(resizeTransitionHandler, atLeastOnce())
+            .startTransition(capture(arg), eq(currentBounds))
+        return arg.value
+    }
+
+    private fun findBoundsChange(wct: WindowContainerTransaction, task: RunningTaskInfo): Rect? =
+        wct.changes[task.token.asBinder()]?.configuration?.windowConfiguration?.bounds
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
index e49eb36..d399b20 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeLoggerTransitionObserverTest.kt
@@ -22,6 +22,8 @@
 import android.graphics.Point
 import android.graphics.Rect
 import android.os.IBinder
+import android.os.SystemProperties
+import android.os.Trace
 import android.testing.AndroidTestingRunner
 import android.view.SurfaceControl
 import android.view.WindowManager.TRANSIT_CHANGE
@@ -38,6 +40,7 @@
 import android.window.TransitionInfo.Change
 import android.window.WindowContainerToken
 import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.modules.utils.testing.ExtendedMockitoRule
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.ShellExecutor
@@ -86,7 +89,11 @@
   @JvmField
   @Rule
   val extendedMockitoRule =
-      ExtendedMockitoRule.Builder(this).mockStatic(DesktopModeStatus::class.java).build()!!
+      ExtendedMockitoRule.Builder(this)
+          .mockStatic(DesktopModeStatus::class.java)
+          .mockStatic(SystemProperties::class.java)
+          .mockStatic(Trace::class.java)
+          .build()!!
 
   private val testExecutor = mock<ShellExecutor>()
   private val mockShellInit = mock<ShellInit>()
@@ -695,6 +702,17 @@
     assertNotNull(sessionId)
     verify(desktopModeEventLogger, times(1)).logSessionEnter(eq(sessionId!!), eq(enterReason))
     verify(desktopModeEventLogger, times(1)).logTaskAdded(eq(sessionId), eq(taskUpdate))
+    ExtendedMockito.verify {
+        Trace.setCounter(
+            eq(Trace.TRACE_TAG_WINDOW_MANAGER),
+            eq(DesktopModeLoggerTransitionObserver.VISIBLE_TASKS_COUNTER_NAME),
+            eq(taskUpdate.visibleTaskCount.toLong()))
+    }
+    ExtendedMockito.verify {
+        SystemProperties.set(
+            eq(DesktopModeLoggerTransitionObserver.VISIBLE_TASKS_COUNTER_SYSTEM_PROPERTY),
+            eq(taskUpdate.visibleTaskCount.toString()))
+    }
     verifyZeroInteractions(desktopModeEventLogger)
   }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeUtilsTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeUtilsTest.kt
new file mode 100644
index 0000000..8fab410
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeUtilsTest.kt
@@ -0,0 +1,66 @@
+/*
+ * 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.wm.shell.desktopmode
+
+import android.graphics.Rect
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class DesktopModeUtilsTest {
+    @Test
+    fun isTaskBoundsEqual_stableBoundsAreEqual_returnTrue() {
+        assertThat(isTaskBoundsEqual(task2Bounds, stableBounds)).isTrue()
+    }
+
+    @Test
+    fun isTaskBoundsEqual_stableBoundsAreNotEqual_returnFalse() {
+        assertThat(isTaskBoundsEqual(task4Bounds, stableBounds)).isFalse()
+    }
+
+    @Test
+    fun isTaskWidthOrHeightEqual_stableBoundsAreEqual_returnTrue() {
+        assertThat(isTaskWidthOrHeightEqual(task2Bounds, stableBounds)).isTrue()
+    }
+
+    @Test
+    fun isTaskWidthOrHeightEqual_stableBoundWidthIsEquals_returnTrue() {
+        assertThat(isTaskWidthOrHeightEqual(task3Bounds, stableBounds)).isTrue()
+    }
+
+    @Test
+    fun isTaskWidthOrHeightEqual_stableBoundHeightIsEquals_returnTrue() {
+        assertThat(isTaskWidthOrHeightEqual(task3Bounds, stableBounds)).isTrue()
+    }
+
+    @Test
+    fun isTaskWidthOrHeightEqual_stableBoundsWidthOrHeightAreNotEquals_returnFalse() {
+        assertThat(isTaskWidthOrHeightEqual(task1Bounds, stableBounds)).isTrue()
+    }
+
+    private companion object {
+        val task1Bounds = Rect(0, 0, 0, 0)
+        val task2Bounds = Rect(1, 1, 1, 1)
+        val task3Bounds = Rect(0, 1, 0, 1)
+        val task4Bounds = Rect(1, 2, 2, 1)
+        val stableBounds = Rect(1, 1, 1, 1)
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index d248720..10557dd 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -84,6 +84,7 @@
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.common.SyncTransactionQueue
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition
+import com.android.wm.shell.desktopmode.DesktopTasksController.TaskbarDesktopTaskListener
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFullscreenTask
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createHomeTask
@@ -122,12 +123,12 @@
 import org.mockito.ArgumentMatchers.isNull
 import org.mockito.Mock
 import org.mockito.Mockito
-import org.mockito.Mockito.any
 import org.mockito.Mockito.anyInt
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.spy
 import org.mockito.Mockito.verify
+import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.atLeastOnce
 import org.mockito.kotlin.eq
@@ -175,6 +176,7 @@
   @Mock
   private lateinit var mockInteractionJankMonitor: InteractionJankMonitor
   @Mock private lateinit var mockSurface: SurfaceControl
+  @Mock private lateinit var taskbarDesktopTaskListener: TaskbarDesktopTaskListener
 
   private lateinit var mockitoSession: StaticMockitoSession
   private lateinit var controller: DesktopTasksController
@@ -237,6 +239,8 @@
     val captor = ArgumentCaptor.forClass(RecentsTransitionStateListener::class.java)
     verify(recentsTransitionHandler).addTransitionStateListener(captor.capture())
     recentsTransitionStateListener = captor.value
+
+    controller.taskbarDesktopTaskListener = taskbarDesktopTaskListener
   }
 
   private fun createController(): DesktopTasksController {
@@ -282,6 +286,52 @@
   }
 
   @Test
+  fun doesAnyTaskRequireTaskbarRounding_onlyFreeFormTaskIsRunning_returnFalse() {
+    setUpFreeformTask()
+
+    assertThat(controller.doesAnyTaskRequireTaskbarRounding(DEFAULT_DISPLAY)).isFalse()
+  }
+
+  @Test
+  fun doesAnyTaskRequireTaskbarRounding_toggleResizeOfFreeFormTask_returnTrue() {
+    val task1 = setUpFreeformTask()
+
+    val argumentCaptor = ArgumentCaptor.forClass(Boolean::class.java)
+    controller.toggleDesktopTaskSize(task1)
+    verify(taskbarDesktopTaskListener).onTaskbarCornerRoundingUpdate(argumentCaptor.capture())
+
+    assertThat(argumentCaptor.value).isTrue()
+  }
+
+  @Test
+  fun doesAnyTaskRequireTaskbarRounding_fullScreenTaskIsRunning_returnTrue() {
+    val stableBounds = Rect().apply { displayLayout.getStableBounds(this) }
+    setUpFreeformTask(bounds = stableBounds, active = true)
+    assertThat(controller.doesAnyTaskRequireTaskbarRounding(DEFAULT_DISPLAY)).isTrue()
+  }
+
+  @Test
+  fun doesAnyTaskRequireTaskbarRounding_toggleResizeOfFullScreenTask_returnFalse() {
+    val stableBounds = Rect().apply { displayLayout.getStableBounds(this) }
+    val task1 = setUpFreeformTask(bounds = stableBounds, active = true)
+
+    val argumentCaptor = ArgumentCaptor.forClass(Boolean::class.java)
+    controller.toggleDesktopTaskSize(task1)
+    verify(taskbarDesktopTaskListener).onTaskbarCornerRoundingUpdate(argumentCaptor.capture())
+
+    assertThat(argumentCaptor.value).isFalse()
+  }
+
+  @Test
+  fun doesAnyTaskRequireTaskbarRounding_splitScreenTaskIsRunning_returnTrue() {
+    val stableBounds = Rect().apply { displayLayout.getStableBounds(this) }
+    setUpFreeformTask(bounds = Rect(stableBounds.left, stableBounds.top, 500, stableBounds.bottom))
+
+    assertThat(controller.doesAnyTaskRequireTaskbarRounding(DEFAULT_DISPLAY)).isTrue()
+  }
+
+
+  @Test
   fun instantiate_cannotEnterDesktopMode_doNotAddInitCallback() {
     whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(false)
     clearInvocations(shellInit)
@@ -2809,7 +2859,7 @@
   }
 
   @Test
-  fun getSnapBounds_calculatesBoundsForResizable() {
+  fun snapToHalfScreen_getSnapBounds_calculatesBoundsForResizable() {
     val bounds = Rect(100, 100, 300, 300)
     val task = setUpFreeformTask(DEFAULT_DISPLAY, bounds).apply {
       topActivityInfo = ActivityInfo().apply {
@@ -2824,13 +2874,45 @@
       STABLE_BOUNDS.left, STABLE_BOUNDS.top, STABLE_BOUNDS.right / 2, STABLE_BOUNDS.bottom
     )
 
-    controller.snapToHalfScreen(task, currentDragBounds, SnapPosition.LEFT)
+    controller.snapToHalfScreen(task, mockSurface, currentDragBounds, SnapPosition.LEFT)
     // Assert bounds set to stable bounds
     val wct = getLatestToggleResizeDesktopTaskWct(currentDragBounds)
     assertThat(findBoundsChange(wct, task)).isEqualTo(expectedBounds)
   }
 
   @Test
+  fun snapToHalfScreen_snapBoundsWhenAlreadySnapped_animatesSurfaceWithoutWCT() {
+    assumeTrue(ENABLE_SHELL_TRANSITIONS)
+    // Set up task to already be in snapped-left bounds
+    val bounds = Rect(
+      STABLE_BOUNDS.left, STABLE_BOUNDS.top, STABLE_BOUNDS.right / 2, STABLE_BOUNDS.bottom
+    )
+    val task = setUpFreeformTask(DEFAULT_DISPLAY, bounds).apply {
+      topActivityInfo = ActivityInfo().apply {
+        screenOrientation = SCREEN_ORIENTATION_LANDSCAPE
+        configuration.windowConfiguration.appBounds = bounds
+      }
+      isResizeable = true
+    }
+
+    // Attempt to snap left again
+    val currentDragBounds = Rect(bounds).apply { offset(-100, 0) }
+    controller.snapToHalfScreen(task, mockSurface, currentDragBounds, SnapPosition.LEFT)
+
+    // Assert that task is NOT updated via WCT
+    verify(toggleResizeDesktopTaskTransitionHandler, never()).startTransition(any(), any())
+
+    // Assert that task leash is updated via Surface Animations
+    verify(mReturnToDragStartAnimator).start(
+      eq(task.taskId),
+      eq(mockSurface),
+      eq(currentDragBounds),
+      eq(bounds),
+      eq(true)
+    )
+  }
+
+  @Test
   @DisableFlags(Flags.FLAG_DISABLE_NON_RESIZABLE_APP_SNAP_RESIZING)
   fun handleSnapResizingTask_nonResizable_snapsToHalfScreen() {
     val task = setUpFreeformTask(DEFAULT_DISPLAY, Rect(0, 0, 200, 100)).apply {
@@ -2861,7 +2943,8 @@
       eq(task.taskId),
       eq(mockSurface),
       eq(currentDragBounds),
-      eq(preDragBounds)
+      eq(preDragBounds),
+      eq(false)
     )
   }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleBarLocationTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleBarLocationTest.kt
similarity index 86%
rename from libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleBarLocationTest.kt
rename to libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleBarLocationTest.kt
index 27e0b19..b9bf95b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleBarLocationTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleBarLocationTest.kt
@@ -13,14 +13,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.testing.AndroidTestingRunner
 import androidx.test.filters.SmallTest
 import com.android.wm.shell.ShellTestCase
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.DEFAULT
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.LEFT
-import com.android.wm.shell.common.bubbles.BubbleBarLocation.RIGHT
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.DEFAULT
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.LEFT
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation.RIGHT
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
 import org.junit.runner.RunWith
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleInfoTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleInfoTest.kt
similarity index 96%
rename from libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleInfoTest.kt
rename to libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleInfoTest.kt
index 5b22edd..641063c 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/bubbles/BubbleInfoTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/bubbles/BubbleInfoTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles
+package com.android.wm.shell.shared.bubbles
 
 import android.os.Parcel
 import android.os.Parcelable.PARCELABLE_WRITE_RETURN_VALUE
@@ -41,6 +41,7 @@
                 "com.some.package",
                 "title",
                 "Some app",
+                true,
                 true
             )
         val parcel = Parcel.obtain()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shared/navigationbar/RegionSamplingHelperTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/handles/RegionSamplingHelperTest.kt
similarity index 75%
rename from packages/SystemUI/tests/src/com/android/systemui/shared/navigationbar/RegionSamplingHelperTest.kt
rename to libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/handles/RegionSamplingHelperTest.kt
index 2a6754c..d3e291f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shared/navigationbar/RegionSamplingHelperTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/shared/handles/RegionSamplingHelperTest.kt
@@ -14,7 +14,8 @@
  * limitations under the License.
  */
 
-package com.android.systemui.shared.navigationbar
+package com.android.wm.shell.shared.handles
+
 
 import android.graphics.Rect
 import android.testing.TestableLooper.RunWithLooper
@@ -24,16 +25,15 @@
 import androidx.concurrent.futures.DirectExecutor
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-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.time.FakeSystemClock
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.TestShellExecutor
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.any
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.eq
 import org.mockito.Mock
@@ -42,11 +42,12 @@
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when` as whenever
 import org.mockito.junit.MockitoJUnit
+import org.mockito.kotlin.argumentCaptor
 
 @RunWith(AndroidJUnit4::class)
 @SmallTest
 @RunWithLooper
-class RegionSamplingHelperTest : SysuiTestCase() {
+class RegionSamplingHelperTest : ShellTestCase() {
 
     @Mock
     lateinit var sampledView: View
@@ -72,12 +73,16 @@
         whenever(surfaceControl.isValid).thenReturn(true)
         whenever(wrappedSurfaceControl.isValid).thenReturn(true)
         whenever(samplingCallback.isSamplingEnabled).thenReturn(true)
-        regionSamplingHelper = object : RegionSamplingHelper(sampledView, samplingCallback,
-                DirectExecutor.INSTANCE, DirectExecutor.INSTANCE, compositionListener) {
-            override fun wrap(stopLayerControl: SurfaceControl?): SurfaceControl {
-                return wrappedSurfaceControl
+        getInstrumentation().runOnMainSync(Runnable {
+            regionSamplingHelper = object : RegionSamplingHelper(
+                sampledView, samplingCallback,
+                DirectExecutor.INSTANCE, DirectExecutor.INSTANCE, compositionListener
+            ) {
+                override fun wrap(stopLayerControl: SurfaceControl?): SurfaceControl {
+                    return wrappedSurfaceControl
+                }
             }
-        }
+        })
         regionSamplingHelper.setWindowVisible(true)
     }
 
@@ -99,7 +104,7 @@
         regionSamplingHelper.setWindowHasBlurs(true)
         regionSamplingHelper.start(Rect(0, 0, 100, 100))
         verify(compositionListener, never())
-                .register(any(), anyInt(), eq(wrappedSurfaceControl), any())
+            .register(any(), anyInt(), eq(wrappedSurfaceControl), any())
     }
 
     @Test
@@ -112,35 +117,38 @@
     @Test
     fun testCompositionSamplingListener_has_nonEmptyRect() {
         // simulate race condition
-        val fakeExecutor = FakeExecutor(FakeSystemClock()) // pass in as backgroundExecutor
+        val fakeExecutor = TestShellExecutor() // pass in as backgroundExecutor
         val fakeSamplingCallback = mock(RegionSamplingHelper.SamplingCallback::class.java)
 
         whenever(fakeSamplingCallback.isSamplingEnabled).thenReturn(true)
         whenever(wrappedSurfaceControl.isValid).thenReturn(true)
-
-        regionSamplingHelper = object : RegionSamplingHelper(sampledView, fakeSamplingCallback,
-                DirectExecutor.INSTANCE, fakeExecutor, compositionListener) {
-            override fun wrap(stopLayerControl: SurfaceControl?): SurfaceControl {
-                return wrappedSurfaceControl
+        getInstrumentation().runOnMainSync(Runnable {
+            regionSamplingHelper = object : RegionSamplingHelper(
+                sampledView, fakeSamplingCallback,
+                DirectExecutor.INSTANCE, fakeExecutor, compositionListener
+            ) {
+                override fun wrap(stopLayerControl: SurfaceControl?): SurfaceControl {
+                    return wrappedSurfaceControl
+                }
             }
-        }
+        })
         regionSamplingHelper.setWindowVisible(true)
         regionSamplingHelper.start(Rect(0, 0, 100, 100))
 
         // make sure background task is enqueued
-        assertThat(fakeExecutor.numPending()).isEqualTo(1)
+        assertThat(fakeExecutor.getCallbacks().size).isEqualTo(1)
 
         // make sure regionSamplingHelper will have empty Rect
         whenever(fakeSamplingCallback.getSampledRegion(any())).thenReturn(Rect(0, 0, 0, 0))
         regionSamplingHelper.onLayoutChange(sampledView, 0, 0, 0, 0, 0, 0, 0, 0)
 
         // resume running of background thread
-        fakeExecutor.runAllReady()
+        fakeExecutor.flushAll()
 
         // grab Rect passed into compositionSamplingListener and make sure it's not empty
         val argumentGrabber = argumentCaptor<Rect>()
         verify(compositionListener).register(any(), anyInt(), eq(wrappedSurfaceControl),
-                argumentGrabber.capture())
-        assertThat(argumentGrabber.value.isEmpty).isFalse()
+            argumentGrabber.capture())
+        assertThat(argumentGrabber.firstValue.isEmpty).isFalse()
     }
-}
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
index 0434742..1984885 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
@@ -294,16 +294,6 @@
     }
 
     @Test
-    public void testUnsetOnBackPressedOnTaskRoot_legacyTransitions() {
-        assumeFalse(Transitions.ENABLE_SHELL_TRANSITIONS);
-        mTaskViewTaskController.onTaskAppeared(mTaskInfo, mLeash);
-        verify(mOrganizer).setInterceptBackPressedOnTaskRoot(eq(mTaskInfo.token), eq(true));
-
-        mTaskViewTaskController.onTaskVanished(mTaskInfo);
-        verify(mOrganizer).setInterceptBackPressedOnTaskRoot(eq(mTaskInfo.token), eq(false));
-    }
-
-    @Test
     public void testOnNewTask_noSurface() {
         assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
         WindowContainerTransaction wct = new WindowContainerTransaction();
@@ -443,19 +433,6 @@
     }
 
     @Test
-    public void testUnsetOnBackPressedOnTaskRoot() {
-        assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
-        WindowContainerTransaction wct = new WindowContainerTransaction();
-        mTaskViewTaskController.prepareOpenAnimation(true /* newTask */,
-                new SurfaceControl.Transaction(), new SurfaceControl.Transaction(), mTaskInfo,
-                mLeash, wct);
-        verify(mOrganizer).setInterceptBackPressedOnTaskRoot(eq(mTaskInfo.token), eq(true));
-
-        mTaskViewTaskController.prepareCloseAnimation();
-        verify(mOrganizer).setInterceptBackPressedOnTaskRoot(eq(mTaskInfo.token), eq(false));
-    }
-
-    @Test
     public void testSetObscuredTouchRect() {
         mTaskView.setObscuredTouchRect(
                 new Rect(/* left= */ 0, /* top= */ 10, /* right= */ 100, /* bottom= */ 120));
@@ -713,4 +690,26 @@
         verify(mViewHandler).post(any());
         verify(mTaskView).setResizeBackgroundColor(eq(Color.BLUE));
     }
+
+    @Test
+    public void testOnAppeared_setsTrimmableTask() {
+        WindowContainerTransaction wct = new WindowContainerTransaction();
+        mTaskViewTaskController.prepareOpenAnimation(true /* newTask */,
+                new SurfaceControl.Transaction(), new SurfaceControl.Transaction(), mTaskInfo,
+                mLeash, wct);
+
+        assertThat(wct.getHierarchyOps().get(0).isTrimmableFromRecents()).isFalse();
+    }
+
+    @Test
+    public void testMoveToFullscreen_callsTaskRemovalStarted() {
+        WindowContainerTransaction wct = new WindowContainerTransaction();
+        mTaskViewTaskController.prepareOpenAnimation(true /* newTask */,
+                new SurfaceControl.Transaction(), new SurfaceControl.Transaction(), mTaskInfo,
+                mLeash, wct);
+        mTaskView.surfaceCreated(mock(SurfaceHolder.class));
+        mTaskViewTaskController.moveToFullscreen();
+
+        verify(mViewListener).onTaskRemovalStarted(eq(mTaskInfo.taskId));
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
index da0aca7b..be0549b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
@@ -34,6 +34,7 @@
 import android.hardware.input.InputManager
 import android.net.Uri
 import android.os.Handler
+import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.CheckFlagsRule
@@ -79,6 +80,7 @@
 import com.android.wm.shell.common.MultiInstanceHelper
 import com.android.wm.shell.common.ShellExecutor
 import com.android.wm.shell.common.SyncTransactionQueue
+import com.android.wm.shell.desktopmode.DesktopActivityOrientationChangeHandler
 import com.android.wm.shell.desktopmode.DesktopTasksController
 import com.android.wm.shell.desktopmode.DesktopTasksController.SnapPosition
 import com.android.wm.shell.desktopmode.DesktopTasksLimiter
@@ -111,7 +113,7 @@
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
-import org.mockito.Mockito.verify
+import org.mockito.kotlin.verify
 import org.mockito.kotlin.any
 import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.argThat
@@ -162,11 +164,14 @@
     @Mock private lateinit var mockWindowManager: IWindowManager
     @Mock private lateinit var mockInteractionJankMonitor: InteractionJankMonitor
     @Mock private lateinit var mockGenericLinksParser: AppToWebGenericLinksParser
+    @Mock private lateinit var mockUserHandle: UserHandle
     @Mock private lateinit var mockToast: Toast
     private val bgExecutor = TestShellExecutor()
     @Mock private lateinit var mockMultiInstanceHelper: MultiInstanceHelper
     @Mock private lateinit var mockTasksLimiter: DesktopTasksLimiter
     @Mock private lateinit var mockFreeformTaskTransitionStarter: FreeformTaskTransitionStarter
+    @Mock private lateinit var mockActivityOrientationChangeHandler:
+            DesktopActivityOrientationChangeHandler
     private lateinit var spyContext: TestableContext
 
     private val transactionFactory = Supplier<SurfaceControl.Transaction> {
@@ -220,7 +225,8 @@
                 mockRootTaskDisplayAreaOrganizer,
                 windowDecorByTaskIdSpy,
                 mockInteractionJankMonitor,
-                Optional.of(mockTasksLimiter)
+                Optional.of(mockTasksLimiter),
+                Optional.of(mockActivityOrientationChangeHandler)
         )
         desktopModeWindowDecorViewModel.setSplitScreenController(mockSplitScreenController)
         whenever(mockDisplayController.getDisplayLayout(any())).thenReturn(mockDisplayLayout)
@@ -594,6 +600,7 @@
 
     @Test
     fun testOnDecorSnappedLeft_snapResizes() {
+        val taskSurfaceCaptor = argumentCaptor<SurfaceControl>()
         val onLeftSnapClickListenerCaptor = forClass(Function0::class.java)
                 as ArgumentCaptor<Function0<Unit>>
         val decor = createOpenTaskDecoration(
@@ -604,8 +611,13 @@
         val currentBounds = decor.mTaskInfo.configuration.windowConfiguration.bounds
         onLeftSnapClickListenerCaptor.value.invoke()
 
-        verify(mockDesktopTasksController)
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.LEFT)
+        verify(mockDesktopTasksController).snapToHalfScreen(
+            eq(decor.mTaskInfo),
+            taskSurfaceCaptor.capture(),
+            eq(currentBounds),
+            eq(SnapPosition.LEFT)
+        )
+        assertEquals(taskSurfaceCaptor.firstValue, decor.mTaskSurface)
     }
 
     @Test
@@ -626,6 +638,7 @@
     @Test
     @DisableFlags(Flags.FLAG_DISABLE_NON_RESIZABLE_APP_SNAP_RESIZING)
     fun testOnSnapResizeLeft_nonResizable_decorSnappedLeft() {
+        val taskSurfaceCaptor = argumentCaptor<SurfaceControl>()
         val onLeftSnapClickListenerCaptor = forClass(Function0::class.java)
                 as ArgumentCaptor<Function0<Unit>>
         val decor = createOpenTaskDecoration(
@@ -636,8 +649,13 @@
         val currentBounds = decor.mTaskInfo.configuration.windowConfiguration.bounds
         onLeftSnapClickListenerCaptor.value.invoke()
 
-        verify(mockDesktopTasksController)
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.LEFT)
+        verify(mockDesktopTasksController).snapToHalfScreen(
+            eq(decor.mTaskInfo),
+            taskSurfaceCaptor.capture(),
+            eq(currentBounds),
+            eq(SnapPosition.LEFT)
+        )
+        assertEquals(decor.mTaskSurface, taskSurfaceCaptor.firstValue)
     }
 
     @Test
@@ -654,12 +672,13 @@
         onLeftSnapClickListenerCaptor.value.invoke()
 
         verify(mockDesktopTasksController, never())
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.LEFT)
+            .snapToHalfScreen(eq(decor.mTaskInfo), any(), eq(currentBounds), eq(SnapPosition.LEFT))
         verify(mockToast).show()
     }
 
     @Test
     fun testOnDecorSnappedRight_snapResizes() {
+        val taskSurfaceCaptor = argumentCaptor<SurfaceControl>()
         val onRightSnapClickListenerCaptor = forClass(Function0::class.java)
                 as ArgumentCaptor<Function0<Unit>>
         val decor = createOpenTaskDecoration(
@@ -670,8 +689,13 @@
         val currentBounds = decor.mTaskInfo.configuration.windowConfiguration.bounds
         onRightSnapClickListenerCaptor.value.invoke()
 
-        verify(mockDesktopTasksController)
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.RIGHT)
+        verify(mockDesktopTasksController).snapToHalfScreen(
+            eq(decor.mTaskInfo),
+            taskSurfaceCaptor.capture(),
+            eq(currentBounds),
+            eq(SnapPosition.RIGHT)
+        )
+        assertEquals(decor.mTaskSurface, taskSurfaceCaptor.firstValue)
     }
 
     @Test
@@ -692,6 +716,7 @@
     @Test
     @DisableFlags(Flags.FLAG_DISABLE_NON_RESIZABLE_APP_SNAP_RESIZING)
     fun testOnSnapResizeRight_nonResizable_decorSnappedRight() {
+        val taskSurfaceCaptor = argumentCaptor<SurfaceControl>()
         val onRightSnapClickListenerCaptor = forClass(Function0::class.java)
                 as ArgumentCaptor<Function0<Unit>>
         val decor = createOpenTaskDecoration(
@@ -702,8 +727,13 @@
         val currentBounds = decor.mTaskInfo.configuration.windowConfiguration.bounds
         onRightSnapClickListenerCaptor.value.invoke()
 
-        verify(mockDesktopTasksController)
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.RIGHT)
+        verify(mockDesktopTasksController).snapToHalfScreen(
+            eq(decor.mTaskInfo),
+            taskSurfaceCaptor.capture(),
+            eq(currentBounds),
+            eq(SnapPosition.RIGHT)
+        )
+        assertEquals(decor.mTaskSurface, taskSurfaceCaptor.firstValue)
     }
 
     @Test
@@ -720,7 +750,7 @@
         onRightSnapClickListenerCaptor.value.invoke()
 
         verify(mockDesktopTasksController, never())
-            .snapToHalfScreen(decor.mTaskInfo, currentBounds, SnapPosition.RIGHT)
+            .snapToHalfScreen(eq(decor.mTaskInfo), any(), eq(currentBounds), eq(SnapPosition.RIGHT))
         verify(mockToast).show()
     }
 
@@ -888,12 +918,12 @@
 
         openInBrowserListenerCaptor.value.accept(uri)
 
-        verify(spyContext).startActivity(argThat { intent ->
+        verify(spyContext).startActivityAsUser(argThat { intent ->
             intent.data == uri
                     && ((intent.flags and Intent.FLAG_ACTIVITY_NEW_TASK) != 0)
                     && intent.categories.contains(Intent.CATEGORY_LAUNCHER)
                     && intent.action == Intent.ACTION_MAIN
-        })
+        }, eq(mockUserHandle))
     }
 
     @Test
@@ -1027,6 +1057,7 @@
 
     private fun createOpenTaskDecoration(
         @WindowingMode windowingMode: Int,
+        taskSurface: SurfaceControl = SurfaceControl(),
         onMaxOrRestoreListenerCaptor: ArgumentCaptor<Function0<Unit>> =
             forClass(Function0::class.java) as ArgumentCaptor<Function0<Unit>>,
         onLeftSnapClickListenerCaptor: ArgumentCaptor<Function0<Unit>> =
@@ -1045,7 +1076,7 @@
             forClass(View.OnClickListener::class.java) as ArgumentCaptor<View.OnClickListener>
     ): DesktopModeWindowDecoration {
         val decor = setUpMockDecorationForTask(createTask(windowingMode = windowingMode))
-        onTaskOpening(decor.mTaskInfo)
+        onTaskOpening(decor.mTaskInfo, taskSurface)
         verify(decor).setOnMaximizeOrRestoreClickListener(onMaxOrRestoreListenerCaptor.capture())
         verify(decor).setOnLeftSnapClickListener(onLeftSnapClickListenerCaptor.capture())
         verify(decor).setOnRightSnapClickListener(onRightSnapClickListenerCaptor.capture())
@@ -1104,6 +1135,7 @@
         ).thenReturn(decoration)
         decoration.mTaskInfo = task
         whenever(decoration.isFocused).thenReturn(task.isFocused)
+        whenever(decoration.user).thenReturn(mockUserHandle)
         if (task.windowingMode == WINDOWING_MODE_MULTI_WINDOW) {
             whenever(mockSplitScreenController.isTaskInSplitScreen(task.taskId))
                 .thenReturn(true)
diff --git a/libs/androidfw/Asset.cpp b/libs/androidfw/Asset.cpp
index 43a70c1..5a4cff0 100644
--- a/libs/androidfw/Asset.cpp
+++ b/libs/androidfw/Asset.cpp
@@ -309,9 +309,8 @@
         return NULL;
     }
 
-    // We succeeded, so relinquish control of dataMap
     pAsset->mAccessMode = mode;
-    return std::move(pAsset);
+    return pAsset;
 }
 
 /*
@@ -328,9 +327,8 @@
       return NULL;
   }
 
-  // We succeeded, so relinquish control of dataMap
   pAsset->mAccessMode = mode;
-  return std::move(pAsset);
+  return pAsset;
 }
 
 /*
diff --git a/libs/input/MouseCursorController.cpp b/libs/input/MouseCursorController.cpp
index eecc741..1afef75 100644
--- a/libs/input/MouseCursorController.cpp
+++ b/libs/input/MouseCursorController.cpp
@@ -25,6 +25,9 @@
 #include <input/Input.h>
 #include <log/log.h>
 
+#define INDENT "  "
+#define INDENT2 "    "
+
 namespace {
 // Time to spend fading out the pointer completely.
 const nsecs_t POINTER_FADE_DURATION = 500 * 1000000LL; // 500 ms
@@ -449,6 +452,24 @@
     return mLocked.resourcesLoaded;
 }
 
+std::string MouseCursorController::dump() const {
+    std::string dump = INDENT "MouseCursorController:\n";
+    std::scoped_lock lock(mLock);
+    dump += StringPrintf(INDENT2 "viewport: %s\n", mLocked.viewport.toString().c_str());
+    dump += StringPrintf(INDENT2 "stylusHoverMode: %s\n",
+                         mLocked.stylusHoverMode ? "true" : "false");
+    dump += StringPrintf(INDENT2 "pointerFadeDirection: %d\n", mLocked.pointerFadeDirection);
+    dump += StringPrintf(INDENT2 "updatePointerIcon: %s\n",
+                         mLocked.updatePointerIcon ? "true" : "false");
+    dump += StringPrintf(INDENT2 "resourcesLoaded: %s\n",
+                         mLocked.resourcesLoaded ? "true" : "false");
+    dump += StringPrintf(INDENT2 "requestedPointerType: %d\n", mLocked.requestedPointerType);
+    dump += StringPrintf(INDENT2 "resolvedPointerType: %d\n", mLocked.resolvedPointerType);
+    dump += StringPrintf(INDENT2 "skipScreenshot: %s\n", mLocked.skipScreenshot ? "true" : "false");
+    dump += StringPrintf(INDENT2 "animating: %s\n", mLocked.animating ? "true" : "false");
+    return dump;
+}
+
 bool MouseCursorController::doAnimations(nsecs_t timestamp) {
     std::scoped_lock lock(mLock);
     bool keepFading = doFadingAnimationLocked(timestamp);
diff --git a/libs/input/MouseCursorController.h b/libs/input/MouseCursorController.h
index 78f6413..8600341 100644
--- a/libs/input/MouseCursorController.h
+++ b/libs/input/MouseCursorController.h
@@ -67,6 +67,8 @@
 
     bool resourcesLoaded();
 
+    std::string dump() const;
+
 private:
     mutable std::mutex mLock;
 
diff --git a/libs/input/PointerController.cpp b/libs/input/PointerController.cpp
index 11b27a2..5ae967b 100644
--- a/libs/input/PointerController.cpp
+++ b/libs/input/PointerController.cpp
@@ -25,6 +25,7 @@
 #include <android-base/stringprintf.h>
 #include <android-base/thread_annotations.h>
 #include <ftl/enum.h>
+#include <input/PrintTools.h>
 
 #include <mutex>
 
@@ -353,6 +354,8 @@
     for (const auto& [_, spotController] : mLocked.spotControllers) {
         spotController.dump(dump, INDENT3);
     }
+    dump += INDENT2 "Cursor Controller:\n";
+    dump += addLinePrefix(mCursorController.dump(), INDENT3);
     return dump;
 }
 
diff --git a/location/api/system-current.txt b/location/api/system-current.txt
index f6e76a2..cf3f740 100644
--- a/location/api/system-current.txt
+++ b/location/api/system-current.txt
@@ -437,7 +437,7 @@
   }
 
   public class LocationManager {
-    method @RequiresPermission(allOf={android.Manifest.permission.LOCATION_HARDWARE, android.Manifest.permission.INTERACT_ACROSS_USERS}) public void addProviderRequestChangedListener(@NonNull java.util.concurrent.Executor, @NonNull android.location.provider.ProviderRequest.ChangedListener);
+    method @Deprecated @FlaggedApi("android.location.flags.deprecate_provider_request_apis") @RequiresPermission(allOf={android.Manifest.permission.LOCATION_HARDWARE, android.Manifest.permission.INTERACT_ACROSS_USERS}) public void addProviderRequestChangedListener(@NonNull java.util.concurrent.Executor, @NonNull android.location.provider.ProviderRequest.ChangedListener);
     method @Deprecated @RequiresPermission(android.Manifest.permission.LOCATION_HARDWARE) public void flushGnssBatch();
     method @Deprecated @RequiresPermission(anyOf={android.Manifest.permission.ACCESS_COARSE_LOCATION, android.Manifest.permission.ACCESS_FINE_LOCATION}) public void getCurrentLocation(@NonNull android.location.LocationRequest, @Nullable android.os.CancellationSignal, @NonNull java.util.concurrent.Executor, @NonNull java.util.function.Consumer<android.location.Location>);
     method @Nullable public String getExtraLocationControllerPackage();
@@ -452,7 +452,7 @@
     method @RequiresPermission(android.Manifest.permission.READ_DEVICE_CONFIG) public boolean isProviderPackage(@Nullable String, @NonNull String, @Nullable String);
     method @Deprecated @RequiresPermission(allOf={android.Manifest.permission.LOCATION_HARDWARE, android.Manifest.permission.UPDATE_APP_OPS_STATS}) public boolean registerGnssBatchedLocationCallback(long, boolean, @NonNull android.location.BatchedLocationCallback, @Nullable android.os.Handler);
     method @Deprecated @RequiresPermission(android.Manifest.permission.ACCESS_FINE_LOCATION) public boolean registerGnssMeasurementsCallback(@NonNull android.location.GnssRequest, @NonNull java.util.concurrent.Executor, @NonNull android.location.GnssMeasurementsEvent.Callback);
-    method @RequiresPermission(android.Manifest.permission.LOCATION_HARDWARE) public void removeProviderRequestChangedListener(@NonNull android.location.provider.ProviderRequest.ChangedListener);
+    method @Deprecated @FlaggedApi("android.location.flags.deprecate_provider_request_apis") @RequiresPermission(android.Manifest.permission.LOCATION_HARDWARE) public void removeProviderRequestChangedListener(@NonNull android.location.provider.ProviderRequest.ChangedListener);
     method @Deprecated @RequiresPermission(anyOf={android.Manifest.permission.ACCESS_COARSE_LOCATION, android.Manifest.permission.ACCESS_FINE_LOCATION}) public void requestLocationUpdates(@Nullable android.location.LocationRequest, @NonNull android.location.LocationListener, @Nullable android.os.Looper);
     method @Deprecated @RequiresPermission(anyOf={android.Manifest.permission.ACCESS_COARSE_LOCATION, android.Manifest.permission.ACCESS_FINE_LOCATION}) public void requestLocationUpdates(@Nullable android.location.LocationRequest, @NonNull java.util.concurrent.Executor, @NonNull android.location.LocationListener);
     method @Deprecated @RequiresPermission(anyOf={android.Manifest.permission.ACCESS_COARSE_LOCATION, android.Manifest.permission.ACCESS_FINE_LOCATION}) public void requestLocationUpdates(@Nullable android.location.LocationRequest, @NonNull android.app.PendingIntent);
diff --git a/location/java/android/location/LocationManager.java b/location/java/android/location/LocationManager.java
index d921730..6342489 100644
--- a/location/java/android/location/LocationManager.java
+++ b/location/java/android/location/LocationManager.java
@@ -28,6 +28,7 @@
 
 import android.Manifest;
 import android.annotation.CallbackExecutor;
+import android.annotation.FlaggedApi;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresFeature;
@@ -47,6 +48,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
+import android.location.flags.Flags;
 import android.location.provider.IProviderRequestListener;
 import android.location.provider.ProviderProperties;
 import android.location.provider.ProviderRequest;
@@ -2957,8 +2959,13 @@
      *
      * @param executor the executor that the callback runs on
      * @param listener the listener to register
+     *
+     * @deprecated Do not use - this API was intended for testing only, and may not return any
+     * results in the future.
      * @hide
      */
+    @FlaggedApi(Flags.FLAG_DEPRECATE_PROVIDER_REQUEST_APIS)
+    @Deprecated
     @SystemApi
     @RequiresPermission(allOf = {Manifest.permission.LOCATION_HARDWARE,
             Manifest.permission.INTERACT_ACROSS_USERS})
@@ -2973,8 +2980,13 @@
      * Removes a {@link ProviderRequest.ChangedListener} that has been added.
      *
      * @param listener the listener to remove.
+     *
+     * @deprecated Do not use - this API was intended for testing only, and may not return any
+     * results in the future.
      * @hide
      */
+    @FlaggedApi(Flags.FLAG_DEPRECATE_PROVIDER_REQUEST_APIS)
+    @Deprecated
     @SystemApi
     @RequiresPermission(Manifest.permission.LOCATION_HARDWARE)
     public void removeProviderRequestChangedListener(
diff --git a/location/java/android/location/flags/location.aconfig b/location/java/android/location/flags/location.aconfig
index 0edaaef..3c387d3 100644
--- a/location/java/android/location/flags/location.aconfig
+++ b/location/java/android/location/flags/location.aconfig
@@ -2,6 +2,13 @@
 container: "system"
 
 flag {
+    name: "deprecate_provider_request_apis"
+    namespace: "location"
+    description: "Deprecates LocationManager ProviderChanged APIs"
+    bug: "361811782"
+}
+
+flag {
     name: "keep_gnss_stationary_throttling"
     namespace: "location"
     description: "Keeps stationary throttling for the GNSS provider even if the disable_stationary_throttling flag is true."
diff --git a/media/java/android/media/AudioSystem.java b/media/java/android/media/AudioSystem.java
index a255f73..ebdfd3e 100644
--- a/media/java/android/media/AudioSystem.java
+++ b/media/java/android/media/AudioSystem.java
@@ -2655,7 +2655,16 @@
     /**
      * Register a native listener for system property sysprop
      * @param callback the listener which fires when the property changes
+     * @return a native handle for use in subsequent methods
      * @hide
      */
-    public static native void listenForSystemPropertyChange(String sysprop, Runnable callback);
+    public static native long listenForSystemPropertyChange(String sysprop, Runnable callback);
+
+    /**
+     * Trigger a sysprop listener update, if the property has been updated: synchronously validating
+     * there are no pending sysprop changes.
+     * @param handle the handle returned by {@link listenForSystemPropertyChange}
+     * @hide
+     */
+    public static native void triggerSystemPropertyUpdate(long handle);
 }
diff --git a/media/java/android/media/tv/tuner/Tuner.java b/media/java/android/media/tv/tuner/Tuner.java
index d14275f..92f6eaf 100644
--- a/media/java/android/media/tv/tuner/Tuner.java
+++ b/media/java/android/media/tv/tuner/Tuner.java
@@ -2533,16 +2533,19 @@
     /**
      * Request a frontend by frontend type.
      *
-     * <p> This API is used if the applications want to select a frontend with desired type when
-     * there are multiple frontends of the same type is there before {@link tune}. The applied
-     * frontend will be one of the not in-use frontend. If all frontends are in-use, this API will
-     * reclaim and apply the frontend owned by the lowest priority client if current client has
-     * higher priority. Otherwise, this API will not apply any frontend and return
-     * {@link #RESULT_UNAVAILABLE}.
+     * <p> This API is used (before {@link #tune(FrontendSettings)}) if the applications want to
+     * select a frontend of a particular type for {@link #tune(FrontendSettings)} when there are
+     * multiple frontends of the same type present, allowing the system to select which one is
+     * applied. The applied frontend will be one of the not-in-use frontends. If all frontends are
+     * in-use, this API will reclaim and apply the frontend owned by the lowest priority client if
+     * current client has higher priority. Otherwise, this API will not apply any frontend and
+     * return {@link #RESULT_UNAVAILABLE}.
      *
      * @param desiredFrontendType the Type of the desired fronted. Should be one of
      *                            {@link android.media.tv.tuner.frontend.FrontendSettings.Type}
      * @return result status of open operation.
+     * @see #applyFrontend(FrontendInfo)
+     * @see #tune(FrontendSettings)
      */
     @Result
     @FlaggedApi(Flags.FLAG_TUNER_W_APIS)
diff --git a/media/native/midi/Android.bp b/media/native/midi/Android.bp
index a991a71f..7acb8c7 100644
--- a/media/native/midi/Android.bp
+++ b/media/native/midi/Android.bp
@@ -74,8 +74,4 @@
     symbol_file: "libamidi.map.txt",
 
     first_version: "29",
-    export_header_libs: [
-        "amidi",
-    ],
-
 }
diff --git a/native/android/Android.bp b/native/android/Android.bp
index c4c4102..3eb99c3 100644
--- a/native/android/Android.bp
+++ b/native/android/Android.bp
@@ -27,9 +27,6 @@
     symbol_file: "libandroid.map.txt",
     first_version: "9",
     unversioned_until: "current",
-    export_header_libs: [
-        "libandroid_headers",
-    ],
 }
 
 cc_defaults {
diff --git a/nfc/api/current.txt b/nfc/api/current.txt
index 447e980..5b6b6c0 100644
--- a/nfc/api/current.txt
+++ b/nfc/api/current.txt
@@ -220,14 +220,15 @@
     field @Deprecated public static final String ACTION_CHANGE_DEFAULT = "android.nfc.cardemulation.action.ACTION_CHANGE_DEFAULT";
     field public static final String CATEGORY_OTHER = "other";
     field public static final String CATEGORY_PAYMENT = "payment";
-    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final String DH = "DH";
-    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final String ESE = "ESE";
     field public static final String EXTRA_CATEGORY = "category";
     field public static final String EXTRA_SERVICE_COMPONENT = "component";
+    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_DH = 0; // 0x0
+    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE = 1; // 0x1
+    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_UICC = 2; // 0x2
+    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET = -1; // 0xffffffff
     field public static final int SELECTION_MODE_ALWAYS_ASK = 1; // 0x1
     field public static final int SELECTION_MODE_ASK_IF_CONFLICT = 2; // 0x2
     field public static final int SELECTION_MODE_PREFER_DEFAULT = 0; // 0x0
-    field @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public static final String UICC = "UICC";
   }
 
   public abstract class HostApduService extends android.app.Service {
diff --git a/nfc/api/system-current.txt b/nfc/api/system-current.txt
index 25a01b9..0f97b2c 100644
--- a/nfc/api/system-current.txt
+++ b/nfc/api/system-current.txt
@@ -73,6 +73,7 @@
     method public void onApplyRouting(@NonNull java.util.function.Consumer<java.lang.Boolean>);
     method public void onBootFinished(int);
     method public void onBootStarted();
+    method public void onCardEmulationActivated(boolean);
     method public void onDisable(@NonNull java.util.function.Consumer<java.lang.Boolean>);
     method public void onDisableFinished(int);
     method public void onDisableStarted();
@@ -81,6 +82,8 @@
     method public void onEnableStarted();
     method public void onHceEventReceived(int);
     method public void onNdefRead(@NonNull java.util.function.Consumer<java.lang.Boolean>);
+    method public void onRfDiscoveryStarted(boolean);
+    method public void onRfFieldActivated(boolean);
     method public void onRoutingChanged();
     method public void onStateUpdated(int);
     method public void onTagConnected(boolean, @NonNull android.nfc.Tag);
@@ -94,7 +97,7 @@
   public final class CardEmulation {
     method @FlaggedApi("android.permission.flags.wallet_role_enabled") @Nullable @RequiresPermission(android.Manifest.permission.NFC_PREFERRED_PAYMENT_INFO) public static android.content.ComponentName getPreferredPaymentService(@NonNull android.content.Context);
     method @FlaggedApi("android.nfc.enable_nfc_mainline") @NonNull public java.util.List<android.nfc.cardemulation.ApduServiceInfo> getServices(@NonNull String, int);
-    method @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public void overrideRoutingTable(@NonNull android.app.Activity, @Nullable String, @Nullable String);
+    method @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public void overrideRoutingTable(@NonNull android.app.Activity, int, int);
     method @FlaggedApi("android.nfc.nfc_override_recover_routing_table") public void recoverRoutingTable(@NonNull android.app.Activity);
   }
 
diff --git a/nfc/java/android/nfc/INfcAdapter.aidl b/nfc/java/android/nfc/INfcAdapter.aidl
index 6c0f933..e2ec952 100644
--- a/nfc/java/android/nfc/INfcAdapter.aidl
+++ b/nfc/java/android/nfc/INfcAdapter.aidl
@@ -62,7 +62,7 @@
 
     void dispatch(in Tag tag);
 
-    void setReaderMode (IBinder b, IAppCallback callback, int flags, in Bundle extras);
+    void setReaderMode (IBinder b, IAppCallback callback, int flags, in Bundle extras, String pkg);
 
     void addNfcUnlockHandler(INfcUnlockHandler unlockHandler, in int[] techList);
     void removeNfcUnlockHandler(INfcUnlockHandler unlockHandler);
@@ -100,7 +100,7 @@
     void unregisterWlcStateListener(in INfcWlcStateListener listener);
     WlcListenerDeviceInfo getWlcListenerDeviceInfo();
 
-    void updateDiscoveryTechnology(IBinder b, int pollFlags, int listenFlags);
+    void updateDiscoveryTechnology(IBinder b, int pollFlags, int listenFlags, String pkg);
 
     void notifyPollingLoop(in PollingFrame frame);
     void notifyHceDeactivated();
diff --git a/nfc/java/android/nfc/INfcOemExtensionCallback.aidl b/nfc/java/android/nfc/INfcOemExtensionCallback.aidl
index c19a44b..b65c837 100644
--- a/nfc/java/android/nfc/INfcOemExtensionCallback.aidl
+++ b/nfc/java/android/nfc/INfcOemExtensionCallback.aidl
@@ -37,4 +37,7 @@
    void onTagDispatch(in ResultReceiver isSkipped);
    void onRoutingChanged();
    void onHceEventReceived(int action);
+   void onCardEmulationActivated(boolean isActivated);
+   void onRfFieldActivated(boolean isActivated);
+   void onRfDiscoveryStarted(boolean isDiscoveryStarted);
 }
diff --git a/nfc/java/android/nfc/NfcActivityManager.java b/nfc/java/android/nfc/NfcActivityManager.java
index 0eb846d..909eca7 100644
--- a/nfc/java/android/nfc/NfcActivityManager.java
+++ b/nfc/java/android/nfc/NfcActivityManager.java
@@ -236,7 +236,8 @@
 
     public void setReaderMode(Binder token, int flags, Bundle extras) {
         if (DBG) Log.d(TAG, "Setting reader mode");
-        NfcAdapter.callService(() -> NfcAdapter.sService.setReaderMode(token, this, flags, extras));
+        NfcAdapter.callService(() -> NfcAdapter.sService.setReaderMode(
+                token, this, flags, extras, mAdapter.getContext().getPackageName()));
     }
 
     /**
@@ -395,7 +396,8 @@
 
     private void changeDiscoveryTech(Binder token, int pollTech, int listenTech) {
         NfcAdapter.callService(
-            () -> NfcAdapter.sService.updateDiscoveryTechnology(token, pollTech, listenTech));
+                () -> NfcAdapter.sService.updateDiscoveryTechnology(
+                        token, pollTech, listenTech, mAdapter.getContext().getPackageName()));
     }
 
 }
diff --git a/nfc/java/android/nfc/NfcAdapter.java b/nfc/java/android/nfc/NfcAdapter.java
index 525e2c5..22ae612 100644
--- a/nfc/java/android/nfc/NfcAdapter.java
+++ b/nfc/java/android/nfc/NfcAdapter.java
@@ -1731,7 +1731,8 @@
         }
         Binder token = new Binder();
         int flags = enable ? ENABLE_POLLING_FLAGS : DISABLE_POLLING_FLAGS;
-        callService(() -> sService.setReaderMode(token, null, flags, null));
+        callService(() -> sService.setReaderMode(
+                token, null, flags, null, mContext.getPackageName()));
     }
 
     /**
@@ -1804,7 +1805,8 @@
                 || (listenTechnology & FLAG_SET_DEFAULT_TECH) == FLAG_SET_DEFAULT_TECH)) {
             Binder token = new Binder();
             callService( () ->
-                sService.updateDiscoveryTechnology(token, pollTechnology, listenTechnology));
+                    sService.updateDiscoveryTechnology(
+                            token, pollTechnology, listenTechnology, mContext.getPackageName()));
         } else {
             mNfcActivityManager.setDiscoveryTech(activity, pollTechnology, listenTechnology);
         }
diff --git a/nfc/java/android/nfc/NfcOemExtension.java b/nfc/java/android/nfc/NfcOemExtension.java
index 6c02edd..632f693 100644
--- a/nfc/java/android/nfc/NfcOemExtension.java
+++ b/nfc/java/android/nfc/NfcOemExtension.java
@@ -32,7 +32,9 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
@@ -40,6 +42,7 @@
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -58,10 +61,13 @@
     private static final int OEM_EXTENSION_RESPONSE_THRESHOLD_MS = 2000;
     private final NfcAdapter mAdapter;
     private final NfcOemExtensionCallback mOemNfcExtensionCallback;
+    private boolean mIsRegistered = false;
+    private final Map<Callback, Executor> mCallbackMap = new HashMap<>();
     private final Context mContext;
-    private Executor mExecutor = null;
-    private Callback mCallback = null;
     private final Object mLock = new Object();
+    private boolean mCardEmulationActivated = false;
+    private boolean mRfFieldActivated = false;
+    private boolean mRfDiscoveryStarted = false;
 
     /**
      * Event that Host Card Emulation is activated.
@@ -215,6 +221,32 @@
          * @param action Flag indicating actions to activate, start and stop cpu boost.
          */
         void onHceEventReceived(@HostCardEmulationAction int action);
+
+        /**
+        * Notifies NFC is activated in listen mode.
+        * NFC Forum NCI-2.3 ch.5.2.6 specification
+        *
+        * <p>NFCC is ready to communicate with a Card reader
+        *
+        * @param isActivated true, if card emulation activated, else de-activated.
+        */
+        void onCardEmulationActivated(boolean isActivated);
+
+        /**
+        * Notifies the Remote NFC Endpoint RF Field is activated.
+        * NFC Forum NCI-2.3 ch.5.3 specification
+        *
+        * @param isActivated true, if RF Field is ON, else RF Field is OFF.
+        */
+        void onRfFieldActivated(boolean isActivated);
+
+        /**
+        * Notifies the NFC RF discovery is started or in the IDLE state.
+        * NFC Forum NCI-2.3 ch.5.2 specification
+        *
+        * @param isDiscoveryStarted true, if RF discovery started, else RF state is Idle.
+        */
+        void onRfDiscoveryStarted(boolean isDiscoveryStarted);
     }
 
 
@@ -229,7 +261,12 @@
 
     /**
      * Register an {@link Callback} to listen for NFC oem extension callbacks
+     * Multiple clients can register and callbacks will be invoked asynchronously.
+     *
      * <p>The provided callback will be invoked by the given {@link Executor}.
+     * As part of {@link #registerCallback(Executor, Callback)} the
+     * {@link Callback} will be invoked with current NFC state
+     * before the {@link #registerCallback(Executor, Callback)} function completes.
      *
      * @param executor an {@link Executor} to execute given callback
      * @param callback oem implementation of {@link Callback}
@@ -239,15 +276,35 @@
     public void registerCallback(@NonNull @CallbackExecutor Executor executor,
             @NonNull Callback callback) {
         synchronized (mLock) {
-            if (mCallback != null) {
+            if (executor == null || callback == null) {
+                Log.e(TAG, "Executor and Callback must not be null!");
+                throw new IllegalArgumentException();
+            }
+
+            if (mCallbackMap.containsKey(callback)) {
                 Log.e(TAG, "Callback already registered. Unregister existing callback before"
                         + "registering");
                 throw new IllegalArgumentException();
             }
-            NfcAdapter.callService(() -> {
-                NfcAdapter.sService.registerOemExtensionCallback(mOemNfcExtensionCallback);
-                mCallback = callback;
-                mExecutor = executor;
+            mCallbackMap.put(callback, executor);
+            if (!mIsRegistered) {
+                NfcAdapter.callService(() -> {
+                    NfcAdapter.sService.registerOemExtensionCallback(mOemNfcExtensionCallback);
+                    mIsRegistered = true;
+                });
+            } else {
+                updateNfCState(callback, executor);
+            }
+        }
+    }
+
+    private void updateNfCState(Callback callback, Executor executor) {
+        if (callback != null) {
+            Log.i(TAG, "updateNfCState");
+            executor.execute(() -> {
+                callback.onCardEmulationActivated(mCardEmulationActivated);
+                callback.onRfFieldActivated(mRfFieldActivated);
+                callback.onRfDiscoveryStarted(mRfDiscoveryStarted);
             });
         }
     }
@@ -266,15 +323,19 @@
     @RequiresPermission(android.Manifest.permission.WRITE_SECURE_SETTINGS)
     public void unregisterCallback(@NonNull Callback callback) {
         synchronized (mLock) {
-            if (mCallback == null || mCallback != callback) {
+            if (!mCallbackMap.containsKey(callback) || !mIsRegistered) {
                 Log.e(TAG, "Callback not registered");
                 throw new IllegalArgumentException();
             }
-            NfcAdapter.callService(() -> {
-                NfcAdapter.sService.unregisterOemExtensionCallback(mOemNfcExtensionCallback);
-                mCallback = null;
-                mExecutor = null;
-            });
+            if (mCallbackMap.size() == 1) {
+                NfcAdapter.callService(() -> {
+                    NfcAdapter.sService.unregisterOemExtensionCallback(mOemNfcExtensionCallback);
+                    mIsRegistered = false;
+                    mCallbackMap.remove(callback);
+                });
+            } else {
+                mCallbackMap.remove(callback);
+            }
         }
     }
 
@@ -322,90 +383,133 @@
     }
 
     private final class NfcOemExtensionCallback extends INfcOemExtensionCallback.Stub {
+
         @Override
         public void onTagConnected(boolean connected, Tag tag) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoid2ArgCallback(connected, tag, cb::onTagConnected, ex));
+        }
+
+        @Override
+        public void onCardEmulationActivated(boolean isActivated) throws RemoteException {
+            mCardEmulationActivated = isActivated;
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(isActivated, cb::onCardEmulationActivated, ex));
+        }
+
+        @Override
+        public void onRfFieldActivated(boolean isActivated) throws RemoteException {
+            mRfFieldActivated = isActivated;
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(isActivated, cb::onRfFieldActivated, ex));
+        }
+
+        @Override
+        public void onRfDiscoveryStarted(boolean isDiscoveryStarted) throws RemoteException {
+            mRfDiscoveryStarted = isDiscoveryStarted;
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(isDiscoveryStarted, cb::onRfDiscoveryStarted, ex));
+        }
+
+        @Override
+        public void onStateUpdated(int state) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(state, cb::onStateUpdated, ex));
+        }
+
+        @Override
+        public void onApplyRouting(ResultReceiver isSkipped) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(
+                        new ReceiverWrapper(isSkipped), cb::onApplyRouting, ex));
+        }
+        @Override
+        public void onNdefRead(ResultReceiver isSkipped) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(
+                        new ReceiverWrapper(isSkipped), cb::onNdefRead, ex));
+        }
+        @Override
+        public void onEnable(ResultReceiver isAllowed) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(
+                        new ReceiverWrapper(isAllowed), cb::onEnable, ex));
+        }
+        @Override
+        public void onDisable(ResultReceiver isAllowed) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(
+                        new ReceiverWrapper(isAllowed), cb::onDisable, ex));
+        }
+        @Override
+        public void onBootStarted() throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(null, (Object input) -> cb.onBootStarted(), ex));
+        }
+        @Override
+        public void onEnableStarted() throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(null, (Object input) -> cb.onEnableStarted(), ex));
+        }
+        @Override
+        public void onDisableStarted() throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(null, (Object input) -> cb.onDisableStarted(), ex));
+        }
+        @Override
+        public void onBootFinished(int status) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(status, cb::onBootFinished, ex));
+        }
+        @Override
+        public void onEnableFinished(int status) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(status, cb::onEnableFinished, ex));
+        }
+        @Override
+        public void onDisableFinished(int status) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(status, cb::onDisableFinished, ex));
+        }
+        @Override
+        public void onTagDispatch(ResultReceiver isSkipped) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(
+                        new ReceiverWrapper(isSkipped), cb::onTagDispatch, ex));
+        }
+        @Override
+        public void onRoutingChanged() throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(null, (Object input) -> cb.onRoutingChanged(), ex));
+        }
+        @Override
+        public void onHceEventReceived(int action) throws RemoteException {
+            mCallbackMap.forEach((cb, ex) ->
+                    handleVoidCallback(action, cb::onHceEventReceived, ex));
+        }
+
+        private <T> void handleVoidCallback(
+                T input, Consumer<T> callbackMethod, Executor executor) {
             synchronized (mLock) {
-                if (mCallback == null || mExecutor == null) {
-                    return;
-                }
                 final long identity = Binder.clearCallingIdentity();
                 try {
-                    mExecutor.execute(() -> mCallback.onTagConnected(connected, tag));
+                    executor.execute(() -> callbackMethod.accept(input));
+                } catch (RuntimeException ex) {
+                    throw ex;
                 } finally {
                     Binder.restoreCallingIdentity(identity);
                 }
             }
         }
-        @Override
-        public void onStateUpdated(int state) throws RemoteException {
-            handleVoidCallback(state, mCallback::onStateUpdated);
-        }
-        @Override
-        public void onApplyRouting(ResultReceiver isSkipped) throws RemoteException {
-            handleVoidCallback(
-                    new ReceiverWrapper(isSkipped), mCallback::onApplyRouting);
-        }
-        @Override
-        public void onNdefRead(ResultReceiver isSkipped) throws RemoteException {
-            handleVoidCallback(
-                    new ReceiverWrapper(isSkipped), mCallback::onNdefRead);
-        }
-        @Override
-        public void onEnable(ResultReceiver isAllowed) throws RemoteException {
-            handleVoidCallback(
-                    new ReceiverWrapper(isAllowed), mCallback::onEnable);
-        }
-        @Override
-        public void onDisable(ResultReceiver isAllowed) throws RemoteException {
-            handleVoidCallback(
-                    new ReceiverWrapper(isAllowed), mCallback::onDisable);
-        }
-        @Override
-        public void onBootStarted() throws RemoteException {
-            handleVoidCallback(null, (Object input) -> mCallback.onBootStarted());
-        }
-        @Override
-        public void onEnableStarted() throws RemoteException {
-            handleVoidCallback(null, (Object input) -> mCallback.onEnableStarted());
-        }
-        @Override
-        public void onDisableStarted() throws RemoteException {
-            handleVoidCallback(null, (Object input) -> mCallback.onDisableStarted());
-        }
-        @Override
-        public void onBootFinished(int status) throws RemoteException {
-            handleVoidCallback(status, mCallback::onBootFinished);
-        }
-        @Override
-        public void onEnableFinished(int status) throws RemoteException {
-            handleVoidCallback(status, mCallback::onEnableFinished);
-        }
-        @Override
-        public void onDisableFinished(int status) throws RemoteException {
-            handleVoidCallback(status, mCallback::onDisableFinished);
-        }
-        @Override
-        public void onTagDispatch(ResultReceiver isSkipped) throws RemoteException {
-            handleVoidCallback(
-                    new ReceiverWrapper(isSkipped), mCallback::onTagDispatch);
-        }
-        @Override
-        public void onRoutingChanged() throws RemoteException {
-            handleVoidCallback(null, (Object input) -> mCallback.onRoutingChanged());
-        }
-        @Override
-        public void onHceEventReceived(int action) throws RemoteException {
-            handleVoidCallback(action, mCallback::onHceEventReceived);
-        }
 
-        private <T> void handleVoidCallback(T input, Consumer<T> callbackMethod) {
+        private <T1, T2> void handleVoid2ArgCallback(
+                T1 input1, T2 input2, BiConsumer<T1, T2> callbackMethod, Executor executor) {
             synchronized (mLock) {
-                if (mCallback == null || mExecutor == null) {
-                    return;
-                }
                 final long identity = Binder.clearCallingIdentity();
                 try {
-                    mExecutor.execute(() -> callbackMethod.accept(input));
+                    executor.execute(() -> callbackMethod.accept(input1, input2));
+                } catch (RuntimeException ex) {
+                    throw ex;
                 } finally {
                     Binder.restoreCallingIdentity(identity);
                 }
@@ -415,17 +519,12 @@
         private <S, T> S handleNonVoidCallbackWithInput(
                 S defaultValue, T input, Function<T, S> callbackMethod) throws RemoteException {
             synchronized (mLock) {
-                if (mCallback == null) {
-                    return defaultValue;
-                }
                 final long identity = Binder.clearCallingIdentity();
                 S result = defaultValue;
                 try {
                     ExecutorService executor = Executors.newSingleThreadExecutor();
-                    FutureTask<S> futureTask = new FutureTask<>(
-                            () -> callbackMethod.apply(input)
-                    );
-                    executor.submit(futureTask);
+                    FutureTask<S> futureTask = new FutureTask<>(() -> callbackMethod.apply(input));
+                    var unused = executor.submit(futureTask);
                     try {
                         result = futureTask.get(
                                 OEM_EXTENSION_RESPONSE_THRESHOLD_MS, TimeUnit.MILLISECONDS);
@@ -447,17 +546,12 @@
         private <T> T handleNonVoidCallbackWithoutInput(T defaultValue, Supplier<T> callbackMethod)
                 throws RemoteException {
             synchronized (mLock) {
-                if (mCallback == null) {
-                    return defaultValue;
-                }
                 final long identity = Binder.clearCallingIdentity();
                 T result = defaultValue;
                 try {
                     ExecutorService executor = Executors.newSingleThreadExecutor();
-                    FutureTask<T> futureTask = new FutureTask<>(
-                            callbackMethod::get
-                    );
-                    executor.submit(futureTask);
+                    FutureTask<T> futureTask = new FutureTask<>(callbackMethod::get);
+                    var unused = executor.submit(futureTask);
                     try {
                         result = futureTask.get(
                                 OEM_EXTENSION_RESPONSE_THRESHOLD_MS, TimeUnit.MILLISECONDS);
diff --git a/nfc/java/android/nfc/cardemulation/CardEmulation.java b/nfc/java/android/nfc/cardemulation/CardEmulation.java
index 0605dbe..497309c 100644
--- a/nfc/java/android/nfc/cardemulation/CardEmulation.java
+++ b/nfc/java/android/nfc/cardemulation/CardEmulation.java
@@ -18,12 +18,12 @@
 
 import android.Manifest;
 import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SdkConstant;
 import android.annotation.SdkConstant.SdkConstantType;
-import android.annotation.StringDef;
 import android.annotation.SystemApi;
 import android.annotation.UserHandleAware;
 import android.annotation.UserIdInt;
@@ -155,17 +155,23 @@
      * Route to Device Host (DH).
      */
     @FlaggedApi(Flags.FLAG_NFC_OVERRIDE_RECOVER_ROUTING_TABLE)
-    public static final String DH = "DH";
+    public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_DH = 0;
     /**
      * Route to eSE.
      */
     @FlaggedApi(Flags.FLAG_NFC_OVERRIDE_RECOVER_ROUTING_TABLE)
-    public static final String ESE = "ESE";
+    public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE = 1;
     /**
      * Route to UICC.
      */
     @FlaggedApi(Flags.FLAG_NFC_OVERRIDE_RECOVER_ROUTING_TABLE)
-    public static final String UICC = "UICC";
+    public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_UICC = 2;
+
+    /**
+     * Route unset.
+     */
+    @FlaggedApi(Flags.FLAG_NFC_OVERRIDE_RECOVER_ROUTING_TABLE)
+    public static final int PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET = -1;
 
     static boolean sIsInitialized = false;
     static HashMap<Context, CardEmulation> sCardEmus = new HashMap<Context, CardEmulation>();
@@ -734,7 +740,7 @@
      *
      * @return the preferred payment service description
      */
-    @RequiresPermission(android.Manifest.permission.NFC_PREFERRED_PAYMENT_INFO)
+    @RequiresPermission(Manifest.permission.NFC_PREFERRED_PAYMENT_INFO)
     @Nullable
     public CharSequence getDescriptionForPreferredPaymentService() {
         ApduServiceInfo serviceInfo = callServiceReturn(() ->
@@ -884,10 +890,12 @@
     }
 
     /** @hide */
-    @StringDef({
-        DH,
-        ESE,
-        UICC
+    @IntDef(prefix = "PROTOCOL_AND_TECHNOLOGY_ROUTE_",
+            value = {
+                    PROTOCOL_AND_TECHNOLOGY_ROUTE_DH,
+                    PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE,
+                    PROTOCOL_AND_TECHNOLOGY_ROUTE_UICC,
+                    PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface ProtocolAndTechnologyRoute {}
@@ -896,29 +904,32 @@
       * Setting NFC controller routing table, which includes Protocol Route and Technology Route,
       * while this Activity is in the foreground.
       *
-      * The parameter set to null can be used to keep current values for that entry. Either
+      * The parameter set to {@link #PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET}
+      * can be used to keep current values for that entry. Either
       * Protocol Route or Technology Route should be override when calling this API, otherwise
       * throw {@link IllegalArgumentException}.
       * <p>
       * Example usage in an Activity that requires to set proto route to "ESE" and keep tech route:
       * <pre>
       * protected void onResume() {
-      *     mNfcAdapter.overrideRoutingTable(this , "ESE" , null);
+      *     mNfcAdapter.overrideRoutingTable(
+      *         this, {@link #PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE}, null);
       * }</pre>
       * </p>
       * Also activities must call {@link #recoverRoutingTable(Activity)}
       * when it goes to the background. Only the package of the
       * currently preferred service (the service set as preferred by the current foreground
       * application via {@link CardEmulation#setPreferredService(Activity, ComponentName)} or the
-      * current Default Wallet Role Holder {@link android.app.role.RoleManager#ROLE_WALLET}),
+      * current Default Wallet Role Holder {@link RoleManager#ROLE_WALLET}),
       * otherwise a call to this method will fail and throw {@link SecurityException}.
       * @param activity The Activity that requests NFC controller routing table to be changed.
-      * @param protocol ISO-DEP route destination, which can be "DH" or "UICC" or "ESE".
-      * @param technology Tech-A, Tech-B and Tech-F route destination, which can be "DH" or "UICC"
-      * or "ESE".
+      * @param protocol ISO-DEP route destination, where the possible inputs are defined
+      *                 in {@link ProtocolAndTechnologyRoute}.
+      * @param technology Tech-A, Tech-B and Tech-F route destination, where the possible inputs
+      *                   are defined in {@link ProtocolAndTechnologyRoute}
       * @throws SecurityException if the caller is not the preferred NFC service
       * @throws IllegalArgumentException if the activity is not resumed or the caller is not in the
-      * foreground, or both protocol route and technology route are null.
+      * foreground.
       * <p>
       * This is a high risk API and only included to support mainline effort
       * @hide
@@ -926,25 +937,36 @@
     @SystemApi
     @FlaggedApi(Flags.FLAG_NFC_OVERRIDE_RECOVER_ROUTING_TABLE)
     public void overrideRoutingTable(
-            @NonNull Activity activity, @ProtocolAndTechnologyRoute @Nullable String protocol,
-            @ProtocolAndTechnologyRoute @Nullable String technology) {
+            @NonNull Activity activity, @ProtocolAndTechnologyRoute int protocol,
+            @ProtocolAndTechnologyRoute int technology) {
         if (!activity.isResumed()) {
             throw new IllegalArgumentException("Activity must be resumed.");
         }
-        if (protocol == null && technology == null) {
-            throw new IllegalArgumentException(("Both Protocol and Technology are null."));
-        }
+        String protocolRoute = switch (protocol) {
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_DH -> "DH";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE -> "ESE";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_UICC -> "UICC";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET -> null;
+            default -> throw new IllegalStateException("Unexpected value: " + protocol);
+        };
+        String technologyRoute = switch (technology) {
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_DH -> "DH";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_ESE -> "ESE";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_UICC -> "UICC";
+            case PROTOCOL_AND_TECHNOLOGY_ROUTE_UNSET -> null;
+            default -> throw new IllegalStateException("Unexpected value: " + protocol);
+        };
         callService(() ->
                 sService.overrideRoutingTable(
                     mContext.getUser().getIdentifier(),
-                    protocol,
-                    technology,
+                    protocolRoute,
+                    technologyRoute,
                     mContext.getPackageName()));
     }
 
     /**
      * Restore the NFC controller routing table,
-     * which was changed by {@link #overrideRoutingTable(Activity, String, String)}
+     * which was changed by {@link #overrideRoutingTable(Activity, int, int)}
      *
      * @param activity The Activity that requested NFC controller routing table to be changed.
      * @throws IllegalArgumentException if the caller is not in the foreground.
diff --git a/nfc/java/android/nfc/flags.aconfig b/nfc/java/android/nfc/flags.aconfig
index 0fda91d..0ade4db 100644
--- a/nfc/java/android/nfc/flags.aconfig
+++ b/nfc/java/android/nfc/flags.aconfig
@@ -141,3 +141,11 @@
     description: "Enable override and recover routing table"
     bug: "329043523"
 }
+
+flag {
+    name: "nfc_watchdog"
+    is_exported: true
+    namespace: "nfc"
+    description: "Enable watchdog for the NFC system process"
+    bug: "362937338"
+}
\ No newline at end of file
diff --git a/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_left_bk.xml b/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_left_bk.xml
index 8a25726..0d4ef3c 100644
--- a/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_left_bk.xml
+++ b/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_left_bk.xml
@@ -20,7 +20,7 @@
        xmlns:tools="http://schemas.android.com/tools"
        tools:targetApi="28"
        android:shape="rectangle">
-    <solid android:color="?androidprv:attr/colorSurface" />
+    <solid android:color="?androidprv:attr/materialColorSurfaceContainerHigh" />
     <corners
         android:topLeftRadius="?android:attr/dialogCornerRadius"
         android:topRightRadius="0dp"
diff --git a/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_right_bk.xml b/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_right_bk.xml
index 7e626e5..3072772 100644
--- a/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_right_bk.xml
+++ b/packages/SettingsLib/ActionButtonsPreference/res/drawable/half_rounded_right_bk.xml
@@ -20,7 +20,7 @@
        xmlns:tools="http://schemas.android.com/tools"
        tools:targetApi="28"
        android:shape="rectangle">
-    <solid android:color="?androidprv:attr/colorSurface" />
+    <solid android:color="?androidprv:attr/materialColorSurfaceContainerHigh" />
     <corners
         android:topLeftRadius="0dp"
         android:topRightRadius="?android:attr/dialogCornerRadius"
diff --git a/packages/SettingsLib/ActionButtonsPreference/res/drawable/rounded_bk.xml b/packages/SettingsLib/ActionButtonsPreference/res/drawable/rounded_bk.xml
index 9f4980b..f1790f9 100644
--- a/packages/SettingsLib/ActionButtonsPreference/res/drawable/rounded_bk.xml
+++ b/packages/SettingsLib/ActionButtonsPreference/res/drawable/rounded_bk.xml
@@ -20,7 +20,7 @@
        xmlns:tools="http://schemas.android.com/tools"
        tools:targetApi="28"
        android:shape="rectangle">
-    <solid android:color="?androidprv:attr/colorSurface" />
+    <solid android:color="?androidprv:attr/materialColorSurfaceContainerHigh" />
     <corners
         android:radius="?android:attr/dialogCornerRadius"
     />
diff --git a/packages/SettingsLib/ActionButtonsPreference/res/drawable/square_bk.xml b/packages/SettingsLib/ActionButtonsPreference/res/drawable/square_bk.xml
index 67b5107..f0da7b4 100644
--- a/packages/SettingsLib/ActionButtonsPreference/res/drawable/square_bk.xml
+++ b/packages/SettingsLib/ActionButtonsPreference/res/drawable/square_bk.xml
@@ -18,7 +18,7 @@
 <shape xmlns:android="http://schemas.android.com/apk/res/android"
        xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
        android:shape="rectangle">
-    <solid android:color="?androidprv:attr/colorSurface" />
+    <solid android:color="?androidprv:attr/materialColorSurfaceContainerHigh" />
     <corners
         android:radius="0dp"
     />
diff --git a/packages/SettingsLib/Spa/build.gradle.kts b/packages/SettingsLib/Spa/build.gradle.kts
index a543450..3011ce0 100644
--- a/packages/SettingsLib/Spa/build.gradle.kts
+++ b/packages/SettingsLib/Spa/build.gradle.kts
@@ -29,7 +29,7 @@
 
 allprojects {
     extra["androidTop"] = androidTop
-    extra["jetpackComposeVersion"] = "1.7.0-beta07"
+    extra["jetpackComposeVersion"] = "1.7.0-rc01"
 }
 
 subprojects {
diff --git a/packages/SettingsLib/Spa/gradle/libs.versions.toml b/packages/SettingsLib/Spa/gradle/libs.versions.toml
index 3507605..d01c0b9 100644
--- a/packages/SettingsLib/Spa/gradle/libs.versions.toml
+++ b/packages/SettingsLib/Spa/gradle/libs.versions.toml
@@ -15,7 +15,7 @@
 #
 
 [versions]
-agp = "8.5.2"
+agp = "8.6.0"
 compose-compiler = "1.5.11"
 dexmaker-mockito = "2.28.3"
 jvm = "17"
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.10-bin.zip b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.10-bin.zip
new file mode 100644
index 0000000..50432f3
--- /dev/null
+++ b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.10-bin.zip
Binary files differ
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.9-bin.zip b/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.9-bin.zip
deleted file mode 100644
index 9a97e46..0000000
--- a/packages/SettingsLib/Spa/gradle/wrapper/gradle-8.9-bin.zip
+++ /dev/null
Binary files differ
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.jar b/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.jar
index 2c35211..a4b76b9 100644
--- a/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.jar
+++ b/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.jar
Binary files differ
diff --git a/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties b/packages/SettingsLib/Spa/gradle/wrapper/gradle-wrapper.properties
index 9f29c77..9a7f4b6 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.9-bin.zip
+distributionUrl=gradle-8.10-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 e9153e3..f0c2ea6 100644
--- a/packages/SettingsLib/Spa/spa/build.gradle.kts
+++ b/packages/SettingsLib/Spa/spa/build.gradle.kts
@@ -54,13 +54,13 @@
 dependencies {
     api(project(":SettingsLibColor"))
     api("androidx.appcompat:appcompat:1.7.0")
-    api("androidx.compose.material3:material3:1.3.0-beta05")
+    api("androidx.compose.material3:material3:1.3.0-rc01")
     api("androidx.compose.material:material-icons-extended:$jetpackComposeVersion")
     api("androidx.compose.runtime:runtime-livedata:$jetpackComposeVersion")
     api("androidx.compose.ui:ui-tooling-preview:$jetpackComposeVersion")
     api("androidx.lifecycle:lifecycle-livedata-ktx")
     api("androidx.lifecycle:lifecycle-runtime-compose")
-    api("androidx.navigation:navigation-compose:2.8.0-beta07")
+    api("androidx.navigation:navigation-compose:2.8.0-rc01")
     api("com.github.PhilJay:MPAndroidChart:v3.1.0-alpha")
     api("com.google.android.material:material:1.11.0")
     debugApi("androidx.compose.ui:ui-tooling:$jetpackComposeVersion")
diff --git a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/preference/SwitchPreference.kt b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/preference/SwitchPreference.kt
index aceb545..62af08e 100644
--- a/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/preference/SwitchPreference.kt
+++ b/packages/SettingsLib/Spa/spa/src/com/android/settingslib/spa/widget/preference/SwitchPreference.kt
@@ -117,7 +117,7 @@
     val indication = LocalIndication.current
     val onChangeWithLog = wrapOnSwitchWithLog(onCheckedChange)
     val interactionSource = remember { MutableInteractionSource() }
-    val modifier = remember(checked, changeable) {
+    val modifier =
         if (checked != null && onChangeWithLog != null) {
             Modifier.toggleable(
                 value = checked,
@@ -128,7 +128,6 @@
                 onValueChange = onChangeWithLog,
             )
         } else Modifier
-    }
     BasePreference(
         title = title,
         summary = summary,
diff --git a/packages/SettingsLib/aconfig/settingslib.aconfig b/packages/SettingsLib/aconfig/settingslib.aconfig
index 2b3862f..34b597b 100644
--- a/packages/SettingsLib/aconfig/settingslib.aconfig
+++ b/packages/SettingsLib/aconfig/settingslib.aconfig
@@ -129,3 +129,13 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "member_device_lea_active_state_sync_fix"
+    namespace: "cross_device_experiences"
+    description: "Gates whether to enable fix for member device active state sync on lea profile"
+    bug: "364201289"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg.xml b/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg.xml
new file mode 100644
index 0000000..35517ea
--- /dev/null
+++ b/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg.xml
@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 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.
+  -->
+
+<shape
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
+    android:shape="rectangle">
+    <solid android:color="?androidprv:attr/colorAccentPrimary" />
+    <corners android:radius="12dp" />
+</shape>
\ No newline at end of file
diff --git a/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg_ripple.xml b/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg_ripple.xml
new file mode 100644
index 0000000..18696c6
--- /dev/null
+++ b/packages/SettingsLib/res/drawable/audio_sharing_rounded_bg_ripple.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2023 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.
+  -->
+
+<ripple xmlns:android="http://schemas.android.com/apk/res/android"
+    android:color="?android:attr/colorControlHighlight">
+    <item android:drawable="@drawable/audio_sharing_rounded_bg"/>
+</ripple>
\ No newline at end of file
diff --git a/packages/SettingsLib/res/values/config.xml b/packages/SettingsLib/res/values/config.xml
index 68b81db..3c3de04 100644
--- a/packages/SettingsLib/res/values/config.xml
+++ b/packages/SettingsLib/res/values/config.xml
@@ -31,4 +31,14 @@
     <!-- Control whether status bar should distinguish HSPA data icon form UMTS
     data icon on devices -->
     <bool name="config_hspa_data_distinguishable">false</bool>
+
+    <!-- Edit User avatar explicit package name -->
+    <string name="config_avatar_picker_package" translatable="false">
+        com.android.avatarpicker
+    </string>
+
+    <!-- Edit User avatar explicit activity class -->
+    <string name="config_avatar_picker_class" translatable="false">
+        com.android.avatarpicker.ui.AvatarPickerActivity
+    </string>
 </resources>
\ No newline at end of file
diff --git a/packages/SettingsLib/res/values/strings.xml b/packages/SettingsLib/res/values/strings.xml
index 687c728..4d771c0 100644
--- a/packages/SettingsLib/res/values/strings.xml
+++ b/packages/SettingsLib/res/values/strings.xml
@@ -1364,6 +1364,9 @@
     <!-- Keywords for setting screen for controlling apps that can schedule alarms [CHAR LIMIT=100] -->
     <string name="keywords_alarms_and_reminders">schedule, alarm, reminder, clock</string>
 
+    <!-- Priority Modes: Name of the "manual" Do Not Disturb mode. [CHAR LIMIT=50] -->
+    <string name="zen_mode_do_not_disturb_name">Do Not Disturb</string>
+
     <!-- Sound: Title for the Do not Disturb option and associated settings page. [CHAR LIMIT=50]-->
     <string name="zen_mode_settings_title">Do Not Disturb</string>
 
diff --git a/packages/SettingsLib/src/com/android/settingslib/Utils.java b/packages/SettingsLib/src/com/android/settingslib/Utils.java
index b02b0c4..1e58335 100644
--- a/packages/SettingsLib/src/com/android/settingslib/Utils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/Utils.java
@@ -495,7 +495,7 @@
                 || packageName.equals(sServicesSystemSharedLibPackageName)
                 || packageName.equals(sSharedSystemSharedLibPackageName)
                 || packageName.equals(PrintManager.PRINT_SPOOLER_PACKAGE_NAME)
-                || (updateServiceV2() && packageName.equals(getDefaultWebViewPackageName()))
+                || (updateServiceV2() && packageName.equals(getDefaultWebViewPackageName(pm)))
                 || isDeviceProvisioningPackage(resources, packageName);
     }
 
@@ -511,7 +511,7 @@
 
     /** Fetch the package name of the default WebView provider. */
     @Nullable
-    private static String getDefaultWebViewPackageName() {
+    private static String getDefaultWebViewPackageName(PackageManager pm) {
         if (sDefaultWebViewPackageName != null) {
             return sDefaultWebViewPackageName;
         }
@@ -519,9 +519,10 @@
         WebViewProviderInfo provider = null;
 
         if (android.webkit.Flags.updateServiceIpcWrapper()) {
-            WebViewUpdateManager manager = WebViewUpdateManager.getInstance();
-            if (manager != null) {
-                provider = manager.getDefaultWebViewPackage();
+            if (pm.hasSystemFeature(PackageManager.FEATURE_WEBVIEW)) {
+                // WebViewUpdateManager.getInstance() will not return null on devices with
+                // FEATURE_WEBVIEW.
+                provider = WebViewUpdateManager.getInstance().getDefaultWebViewPackage();
             }
         } else {
             try {
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/BluetoothUtils.java b/packages/SettingsLib/src/com/android/settingslib/bluetooth/BluetoothUtils.java
index 055afed..0ffb763 100644
--- a/packages/SettingsLib/src/com/android/settingslib/bluetooth/BluetoothUtils.java
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/BluetoothUtils.java
@@ -1,5 +1,6 @@
 package com.android.settingslib.bluetooth;
 
+import static com.android.settingslib.bluetooth.LocalBluetoothLeBroadcast.UNKNOWN_VALUE_PLACEHOLDER;
 import static com.android.settingslib.widget.AdaptiveOutlineDrawable.ICON_TYPE_ADVANCED;
 
 import android.annotation.SuppressLint;
@@ -704,12 +705,50 @@
         return !sourceList.isEmpty() && sourceList.stream().anyMatch(BluetoothUtils::isConnected);
     }
 
+    /**
+     * Check if {@link BluetoothDevice} has a active local broadcast source.
+     *
+     * @param device The bluetooth device to check.
+     * @param localBtManager The BT manager to provide BT functions.
+     * @return Whether the device has a active local broadcast source.
+     */
+    @WorkerThread
+    public static boolean hasActiveLocalBroadcastSourceForBtDevice(
+            @Nullable BluetoothDevice device, @Nullable LocalBluetoothManager localBtManager) {
+        LocalBluetoothLeBroadcastAssistant assistant =
+                localBtManager == null
+                        ? null
+                        : localBtManager.getProfileManager().getLeAudioBroadcastAssistantProfile();
+        LocalBluetoothLeBroadcast broadcast =
+                localBtManager == null
+                        ? null
+                        : localBtManager.getProfileManager().getLeAudioBroadcastProfile();
+        if (device == null || assistant == null || broadcast == null) {
+            Log.d(TAG, "Skip check hasActiveLocalBroadcastSourceForBtDevice due to arg is null");
+            return false;
+        }
+        List<BluetoothLeBroadcastReceiveState> sourceList = assistant.getAllSources(device);
+        int broadcastId = broadcast.getLatestBroadcastId();
+        return !sourceList.isEmpty()
+                && broadcastId != UNKNOWN_VALUE_PLACEHOLDER
+                && sourceList.stream()
+                        .anyMatch(
+                                source -> isSourceMatched(source, broadcastId));
+    }
+
     /** Checks the connectivity status based on the provided broadcast receive state. */
     @WorkerThread
     public static boolean isConnected(BluetoothLeBroadcastReceiveState state) {
         return state.getBisSyncState().stream().anyMatch(bitmap -> bitmap != 0);
     }
 
+    /** Checks if the broadcast id is matched based on the provided broadcast receive state. */
+    @WorkerThread
+    public static boolean isSourceMatched(
+            @Nullable BluetoothLeBroadcastReceiveState state, int broadcastId) {
+        return state != null && state.getBroadcastId() == broadcastId;
+    }
+
     /**
      * Checks if the Bluetooth device is an available hearing device, which means: 1) currently
      * connected 2) is Hearing Aid 3) connected profile match hearing aid related profiles (e.g.
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/LocalBluetoothLeBroadcast.java b/packages/SettingsLib/src/com/android/settingslib/bluetooth/LocalBluetoothLeBroadcast.java
index 26905b1..6f2567b 100644
--- a/packages/SettingsLib/src/com/android/settingslib/bluetooth/LocalBluetoothLeBroadcast.java
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/LocalBluetoothLeBroadcast.java
@@ -101,7 +101,7 @@
     private static final int DEFAULT_CODE_MIN = 1000;
     // Order of this profile in device profiles list
     private static final int ORDINAL = 1;
-    private static final int UNKNOWN_VALUE_PLACEHOLDER = -1;
+    static final int UNKNOWN_VALUE_PLACEHOLDER = -1;
     private static final Uri[] SETTINGS_URIS =
             new Uri[] {
                 Settings.Secure.getUriFor(Settings.Secure.BLUETOOTH_LE_BROADCAST_NAME),
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.aidl
similarity index 84%
copy from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
copy to packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.aidl
index 3c5beeb..1726036 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.aidl
@@ -14,6 +14,6 @@
  * limitations under the License.
  */
 
-package com.android.wm.shell.common.bubbles;
+package com.android.settingslib.bluetooth.devicesettings;
 
-parcelable BubbleBarLocation;
\ No newline at end of file
+parcelable DeviceSettingsProviderServiceStatus;
\ No newline at end of file
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.kt b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.kt
new file mode 100644
index 0000000..977849e
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatus.kt
@@ -0,0 +1,60 @@
+/*
+ * 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.settingslib.bluetooth.devicesettings
+
+import android.os.Bundle
+import android.os.Parcel
+import android.os.Parcelable
+
+/**
+ * A data class representing a device settings item in bluetooth device details config.
+ *
+ * @property enabled Whether the service is enabled.
+ * @property extras Extra bundle
+ */
+data class DeviceSettingsProviderServiceStatus(
+    val enabled: Boolean,
+    val extras: Bundle = Bundle.EMPTY,
+) : Parcelable {
+
+    override fun describeContents(): Int = 0
+
+    override fun writeToParcel(parcel: Parcel, flags: Int) {
+        parcel.run {
+            writeBoolean(enabled)
+            writeBundle(extras)
+        }
+    }
+
+    companion object {
+        @JvmField
+        val CREATOR: Parcelable.Creator<DeviceSettingsProviderServiceStatus> =
+            object : Parcelable.Creator<DeviceSettingsProviderServiceStatus> {
+                override fun createFromParcel(parcel: Parcel) =
+                    parcel.run {
+                        DeviceSettingsProviderServiceStatus(
+                            enabled = readBoolean(),
+                            extras = readBundle((Bundle::class.java.classLoader)) ?: Bundle.EMPTY,
+                        )
+                    }
+
+                override fun newArray(size: Int): Array<DeviceSettingsProviderServiceStatus?> {
+                    return arrayOfNulls(size)
+                }
+            }
+    }
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/IDeviceSettingsProviderService.aidl b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/IDeviceSettingsProviderService.aidl
index d5efac9..1c0a1fd 100644
--- a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/IDeviceSettingsProviderService.aidl
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/IDeviceSettingsProviderService.aidl
@@ -18,10 +18,12 @@
 
 import com.android.settingslib.bluetooth.devicesettings.DeviceInfo;
 import com.android.settingslib.bluetooth.devicesettings.DeviceSettingState;
+import com.android.settingslib.bluetooth.devicesettings.DeviceSettingsProviderServiceStatus;
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsListener;
 
-oneway interface IDeviceSettingsProviderService {
-   void registerDeviceSettingsListener(in DeviceInfo device, in IDeviceSettingsListener callback);
-   void unregisterDeviceSettingsListener(in DeviceInfo device, in IDeviceSettingsListener callback);
-   void updateDeviceSettings(in DeviceInfo device, in DeviceSettingState params);
+interface IDeviceSettingsProviderService {
+   DeviceSettingsProviderServiceStatus getServiceStatus();
+   oneway void registerDeviceSettingsListener(in DeviceInfo device, in IDeviceSettingsListener callback);
+   oneway void unregisterDeviceSettingsListener(in DeviceInfo device, in IDeviceSettingsListener callback);
+   oneway void updateDeviceSettings(in DeviceInfo device, in DeviceSettingState params);
 }
\ No newline at end of file
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/model/ServiceConnectionStatus.kt b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/model/ServiceConnectionStatus.kt
new file mode 100644
index 0000000..25080bc
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/model/ServiceConnectionStatus.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.settingslib.bluetooth.devicesettings.data.model
+
+import android.os.IInterface
+
+/** Present a service connection status. */
+sealed interface ServiceConnectionStatus<out T : IInterface> {
+    /** Service is connecting. */
+    data object Connecting : ServiceConnectionStatus<Nothing>
+
+    /** Service is connected. */
+    data class Connected<T : IInterface>(val service: T) : ServiceConnectionStatus<T>
+
+    /** Service connection failed. */
+    data object Failed : ServiceConnectionStatus<Nothing>
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingServiceConnection.kt b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingServiceConnection.kt
index d6b2862..33beb06 100644
--- a/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingServiceConnection.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingServiceConnection.kt
@@ -22,7 +22,8 @@
 import android.content.Intent
 import android.content.ServiceConnection
 import android.os.IBinder
-import com.android.internal.util.ConcurrentUtils
+import android.os.IInterface
+import android.util.Log
 import com.android.settingslib.bluetooth.BluetoothUtils
 import com.android.settingslib.bluetooth.CachedBluetoothDevice
 import com.android.settingslib.bluetooth.devicesettings.DeviceInfo
@@ -34,27 +35,28 @@
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsConfigProviderService
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsListener
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsProviderService
+import com.android.settingslib.bluetooth.devicesettings.data.model.ServiceConnectionStatus
 import java.util.concurrent.ConcurrentHashMap
-import java.util.concurrent.atomic.AtomicReference
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.async
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.emitAll
-import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.flatMapConcat
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.mapNotNull
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
@@ -84,64 +86,132 @@
             }
     }
 
-    private var config = AtomicReference<DeviceSettingsConfig?>(null)
-    private var idToSetting = AtomicReference<Flow<Map<Int, DeviceSetting>>?>(null)
-
-    /** Gets [DeviceSettingsConfig] for the device, return null when failed. */
-    suspend fun getDeviceSettingsConfig(): DeviceSettingsConfig? =
-        config.computeIfAbsent {
-            getConfigServiceBindingIntent(cachedDevice)
-                .flatMapLatest { getService(it) }
-                .map { it?.let { IDeviceSettingsConfigProviderService.Stub.asInterface(it) } }
-                .map {
-                    it?.getDeviceSettingsConfig(
-                        deviceInfo { setBluetoothAddress(cachedDevice.address) }
-                    )
+    private var isServiceEnabled =
+        coroutineScope.async(backgroundCoroutineContext, start = CoroutineStart.LAZY) {
+            val states = getSettingsProviderServices()?.values ?: return@async false
+            combine(states) { it.toList() }
+                .mapNotNull { allStatus ->
+                    if (allStatus.any { it is ServiceConnectionStatus.Failed }) {
+                        false
+                    } else if (allStatus.all { it is ServiceConnectionStatus.Connected }) {
+                        allStatus
+                            .filterIsInstance<
+                                ServiceConnectionStatus.Connected<IDeviceSettingsProviderService>
+                            >()
+                            .all { it.service.serviceStatus?.enabled == true }
+                    } else {
+                        null
+                    }
                 }
                 .first()
         }
 
+    private var config =
+        coroutineScope.async(backgroundCoroutineContext, start = CoroutineStart.LAZY) {
+            val intent =
+                tryGetEndpointFromMetadata(cachedDevice)?.toIntent()
+                    ?: run {
+                        Log.i(TAG, "Unable to read device setting metadata from $cachedDevice")
+                        return@async null
+                    }
+            getService(intent, IDeviceSettingsConfigProviderService.Stub::asInterface)
+                .flatMapConcat {
+                    when (it) {
+                        is ServiceConnectionStatus.Connected ->
+                            flowOf(
+                                it.service.getDeviceSettingsConfig(
+                                    deviceInfo { setBluetoothAddress(cachedDevice.address) }
+                                )
+                            )
+                        ServiceConnectionStatus.Connecting -> flowOf()
+                        ServiceConnectionStatus.Failed -> flowOf(null)
+                    }
+                }
+                .first()
+        }
+
+    private val settingIdToItemMapping =
+        flow {
+                if (!isServiceEnabled.await()) {
+                    Log.w(TAG, "Service is disabled")
+                    return@flow
+                }
+                getSettingsProviderServices()
+                    ?.values
+                    ?.map {
+                        it.flatMapLatest { status ->
+                            when (status) {
+                                is ServiceConnectionStatus.Connected ->
+                                    getDeviceSettingsFromService(cachedDevice, status.service)
+                                else -> flowOf(emptyList())
+                            }
+                        }
+                    }
+                    ?.let { items -> combine(items) { it.toList().flatten() } }
+                    ?.map { items -> items.associateBy { it.settingId } }
+                    ?.let { emitAll(it) }
+            }
+            .shareIn(scope = coroutineScope, started = SharingStarted.WhileSubscribed(), replay = 1)
+
+    /** Gets [DeviceSettingsConfig] for the device, return null when failed. */
+    suspend fun getDeviceSettingsConfig(): DeviceSettingsConfig? {
+        if (!isServiceEnabled.await()) {
+            Log.w(TAG, "Service is disabled")
+            return null
+        }
+        return readConfig()
+    }
+
     /** Gets all device settings for the device. */
     fun getDeviceSettingList(): Flow<List<DeviceSetting>> =
-        getSettingIdToItemMapping().map { it.values.toList() }
+        settingIdToItemMapping.map { it.values.toList() }
 
     /** Gets the device settings with the ID for the device. */
     fun getDeviceSetting(@DeviceSettingId deviceSettingId: Int): Flow<DeviceSetting?> =
-        getSettingIdToItemMapping().map { it[deviceSettingId] }
+        settingIdToItemMapping.map { it[deviceSettingId] }
 
     /** Updates the device setting state for the device. */
     suspend fun updateDeviceSettings(
         @DeviceSettingId deviceSettingId: Int,
         deviceSettingPreferenceState: DeviceSettingPreferenceState,
     ) {
-        getDeviceSettingsConfig()?.let { config ->
+        if (!isServiceEnabled.await()) {
+            Log.w(TAG, "Service is disabled")
+            return
+        }
+        readConfig()?.let { config ->
             (config.mainContentItems + config.moreSettingsItems)
                 .find { it.settingId == deviceSettingId }
                 ?.let {
                     getSettingsProviderServices()
                         ?.get(EndPoint(it.packageName, it.className, it.intentAction))
-                        ?.filterNotNull()
+                        ?.filterIsInstance<
+                            ServiceConnectionStatus.Connected<IDeviceSettingsProviderService>
+                        >()
                         ?.first()
                 }
+                ?.service
                 ?.updateDeviceSettings(
                     deviceInfo { setBluetoothAddress(cachedDevice.address) },
                     DeviceSettingState.Builder()
                         .setSettingId(deviceSettingId)
                         .setPreferenceState(deviceSettingPreferenceState)
-                        .build()
+                        .build(),
                 )
         }
     }
 
+    private suspend fun readConfig(): DeviceSettingsConfig? = config.await()
+
     private suspend fun getSettingsProviderServices():
-        Map<EndPoint, StateFlow<IDeviceSettingsProviderService?>>? =
-        getDeviceSettingsConfig()
+        Map<EndPoint, StateFlow<ServiceConnectionStatus<IDeviceSettingsProviderService>>>? =
+        readConfig()
             ?.let { config ->
                 (config.mainContentItems + config.moreSettingsItems).map {
                     EndPoint(
                         packageName = it.packageName,
                         className = it.className,
-                        intentAction = it.intentAction
+                        intentAction = it.intentAction,
                     )
                 }
             }
@@ -150,43 +220,22 @@
                 { it },
                 { endpoint ->
                     services.computeIfAbsent(endpoint) {
-                        getService(endpoint.toIntent())
-                            .map { service ->
-                                IDeviceSettingsProviderService.Stub.asInterface(service)
-                            }
-                            .stateIn(coroutineScope, SharingStarted.WhileSubscribed(), null)
+                        getService(
+                                endpoint.toIntent(),
+                                IDeviceSettingsProviderService.Stub::asInterface,
+                            )
+                            .stateIn(
+                                coroutineScope,
+                                SharingStarted.WhileSubscribed(),
+                                ServiceConnectionStatus.Connecting,
+                            )
                     }
-                }
+                },
             )
 
-    private fun getSettingIdToItemMapping(): Flow<Map<Int, DeviceSetting>> =
-        idToSetting.computeIfAbsent {
-            flow {
-                    getSettingsProviderServices()
-                        ?.values
-                        ?.map {
-                            it.flatMapLatest { service ->
-                                if (service != null) {
-                                    getDeviceSettingsFromService(cachedDevice, service)
-                                } else {
-                                    flowOf(emptyList())
-                                }
-                            }
-                        }
-                        ?.let { items -> combine(items) { it.toList().flatten() } }
-                        ?.map { items -> items.associateBy { it.settingId } }
-                        ?.let { emitAll(it) }
-                }
-                .shareIn(
-                    scope = coroutineScope,
-                    started = SharingStarted.WhileSubscribed(),
-                    replay = 1
-                )
-        }!!
-
     private fun getDeviceSettingsFromService(
         cachedDevice: CachedBluetoothDevice,
-        service: IDeviceSettingsProviderService
+        service: IDeviceSettingsProviderService,
     ): Flow<List<DeviceSetting>> {
         return callbackFlow {
                 val listener =
@@ -202,51 +251,28 @@
             .stateIn(coroutineScope, SharingStarted.WhileSubscribed(), emptyList())
     }
 
-    private fun getService(intent: Intent): Flow<IBinder?> {
+    private fun <T : IInterface> getService(
+        intent: Intent,
+        transform: ((IBinder) -> T),
+    ): Flow<ServiceConnectionStatus<T>> {
         return callbackFlow {
             val serviceConnection =
                 object : ServiceConnection {
                     override fun onServiceConnected(name: ComponentName, service: IBinder) {
-                        launch { send(service) }
+                        launch { send(ServiceConnectionStatus.Connected(transform(service))) }
                     }
 
                     override fun onServiceDisconnected(name: ComponentName?) {
-                        launch { send(null) }
+                        launch { send(ServiceConnectionStatus.Connecting) }
                     }
                 }
             if (!context.bindService(intent, serviceConnection, Context.BIND_AUTO_CREATE)) {
-                launch { send(null) }
+                launch { send(ServiceConnectionStatus.Failed) }
             }
             awaitClose { context.unbindService(serviceConnection) }
         }
     }
 
-    private fun getConfigServiceBindingIntent(cachedDevice: CachedBluetoothDevice): Flow<Intent> {
-        return callbackFlow {
-                val listener =
-                    BluetoothAdapter.OnMetadataChangedListener { device, key, _ ->
-                        if (
-                            key == METADATA_FAST_PAIR_CUSTOMIZED_FIELDS &&
-                                cachedDevice.device == device
-                        ) {
-                            launch { tryGetEndpointFromMetadata(cachedDevice)?.let { send(it) } }
-                        }
-                    }
-                bluetoothAdaptor.addOnMetadataChangedListener(
-                    cachedDevice.device,
-                    ConcurrentUtils.DIRECT_EXECUTOR,
-                    listener,
-                )
-                awaitClose {
-                    bluetoothAdaptor.removeOnMetadataChangedListener(cachedDevice.device, listener)
-                }
-            }
-            .onStart { tryGetEndpointFromMetadata(cachedDevice)?.let { emit(it) } }
-            .distinctUntilChanged()
-            .map { it.toIntent() }
-            .flowOn(backgroundCoroutineContext)
-    }
-
     private suspend fun tryGetEndpointFromMetadata(cachedDevice: CachedBluetoothDevice): EndPoint? =
         withContext(backgroundCoroutineContext) {
             val packageName =
@@ -257,29 +283,31 @@
             val className =
                 BluetoothUtils.getFastPairCustomizedField(
                     cachedDevice.device,
-                    CONFIG_SERVICE_CLASS_NAME
+                    CONFIG_SERVICE_CLASS_NAME,
                 ) ?: return@withContext null
             val intentAction =
                 BluetoothUtils.getFastPairCustomizedField(
                     cachedDevice.device,
-                    CONFIG_SERVICE_INTENT_ACTION
+                    CONFIG_SERVICE_INTENT_ACTION,
                 ) ?: return@withContext null
             EndPoint(packageName, className, intentAction)
         }
 
-    private inline fun <T> AtomicReference<T?>.computeIfAbsent(producer: () -> T): T? =
-        get() ?: producer().let { compareAndExchange(null, it) ?: it }
-
     private inline fun deviceInfo(block: DeviceInfo.Builder.() -> Unit): DeviceInfo {
         return DeviceInfo.Builder().apply { block() }.build()
     }
 
     companion object {
+        const val TAG = "DeviceSettingSrvConn"
         const val METADATA_FAST_PAIR_CUSTOMIZED_FIELDS: Int = 25
         const val CONFIG_SERVICE_PACKAGE_NAME = "DEVICE_SETTINGS_CONFIG_PACKAGE_NAME"
         const val CONFIG_SERVICE_CLASS_NAME = "DEVICE_SETTINGS_CONFIG_CLASS"
         const val CONFIG_SERVICE_INTENT_ACTION = "DEVICE_SETTINGS_CONFIG_ACTION"
 
-        val services = ConcurrentHashMap<EndPoint, StateFlow<IDeviceSettingsProviderService?>>()
+        val services =
+            ConcurrentHashMap<
+                EndPoint,
+                StateFlow<ServiceConnectionStatus<IDeviceSettingsProviderService>>,
+            >()
     }
 }
diff --git a/packages/SettingsLib/src/com/android/settingslib/fuelgauge/PowerAllowlistBackend.java b/packages/SettingsLib/src/com/android/settingslib/fuelgauge/PowerAllowlistBackend.java
index 4f2329b..47a08eb 100644
--- a/packages/SettingsLib/src/com/android/settingslib/fuelgauge/PowerAllowlistBackend.java
+++ b/packages/SettingsLib/src/com/android/settingslib/fuelgauge/PowerAllowlistBackend.java
@@ -136,12 +136,10 @@
             return true;
         }
 
-        if (android.app.admin.flags.Flags.disallowUserControlBgUsageFix()) {
-            // App is subject to DevicePolicyManager.setUserControlDisabledPackages() policy.
-            final int userId = UserHandle.getUserId(uid);
-            if (mAppContext.getPackageManager().isPackageStateProtected(pkg, userId)) {
-                return true;
-            }
+        // App is subject to DevicePolicyManager.setUserControlDisabledPackages() policy.
+        final int userId = UserHandle.getUserId(uid);
+        if (mAppContext.getPackageManager().isPackageStateProtected(pkg, userId)) {
+            return true;
         }
 
         return false;
diff --git a/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt b/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt
deleted file mode 100644
index 02d684d..0000000
--- a/packages/SettingsLib/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExt.kt
+++ /dev/null
@@ -1,100 +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.settingslib.media.data.repository
-
-import android.media.AudioManager
-import android.media.IVolumeController
-import kotlinx.coroutines.channels.awaitClose
-import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.buffer
-import kotlinx.coroutines.flow.callbackFlow
-import kotlinx.coroutines.launch
-
-/** Returns [AudioManager.setVolumeController] events as a [Flow] */
-fun AudioManager.volumeControllerEvents(): Flow<VolumeControllerEvent> =
-    callbackFlow {
-            volumeController =
-                object : IVolumeController.Stub() {
-                    override fun displaySafeVolumeWarning(flags: Int) {
-                        launch { send(VolumeControllerEvent.DisplaySafeVolumeWarning(flags)) }
-                    }
-
-                    override fun volumeChanged(streamType: Int, flags: Int) {
-                        launch { send(VolumeControllerEvent.VolumeChanged(streamType, flags)) }
-                    }
-
-                    override fun masterMuteChanged(flags: Int) {
-                        launch { send(VolumeControllerEvent.MasterMuteChanged(flags)) }
-                    }
-
-                    override fun setLayoutDirection(layoutDirection: Int) {
-                        launch { send(VolumeControllerEvent.SetLayoutDirection(layoutDirection)) }
-                    }
-
-                    override fun dismiss() {
-                        launch { send(VolumeControllerEvent.Dismiss) }
-                    }
-
-                    override fun setA11yMode(mode: Int) {
-                        launch { send(VolumeControllerEvent.SetA11yMode(mode)) }
-                    }
-
-                    override fun displayCsdWarning(
-                        csdWarning: Int,
-                        displayDurationMs: Int,
-                    ) {
-                        launch {
-                            send(
-                                VolumeControllerEvent.DisplayCsdWarning(
-                                    csdWarning,
-                                    displayDurationMs,
-                                )
-                            )
-                        }
-                    }
-                }
-            awaitClose { volumeController = null }
-        }
-        .buffer()
-
-/** Models events received via [IVolumeController] */
-sealed interface VolumeControllerEvent {
-
-    /** @see [IVolumeController.displaySafeVolumeWarning] */
-    data class DisplaySafeVolumeWarning(val flags: Int) : VolumeControllerEvent
-
-    /** @see [IVolumeController.volumeChanged] */
-    data class VolumeChanged(val streamType: Int, val flags: Int) : VolumeControllerEvent
-
-    /** @see [IVolumeController.masterMuteChanged] */
-    data class MasterMuteChanged(val flags: Int) : VolumeControllerEvent
-
-    /** @see [IVolumeController.setLayoutDirection] */
-    data class SetLayoutDirection(val layoutDirection: Int) : VolumeControllerEvent
-
-    /** @see [IVolumeController.setA11yMode] */
-    data class SetA11yMode(val mode: Int) : VolumeControllerEvent
-
-    /** @see [IVolumeController.displayCsdWarning] */
-    data class DisplayCsdWarning(
-        val csdWarning: Int,
-        val displayDurationMs: Int,
-    ) : VolumeControllerEvent
-
-    /** @see [IVolumeController.dismiss] */
-    data object Dismiss : VolumeControllerEvent
-}
\ No newline at end of file
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
index 4371f05..043219a 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/FakeZenModeRepository.kt
@@ -41,6 +41,8 @@
     override val modes: Flow<List<ZenMode>>
         get() = mutableModesFlow.asStateFlow()
 
+    override fun getModes(): List<ZenMode> = mutableModesFlow.value
+
     private val activeModesDurations = mutableMapOf<String, Duration?>()
 
     init {
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/ZenModeRepository.kt b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/ZenModeRepository.kt
index 0ff7f84..7fdbcda 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/ZenModeRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/data/repository/ZenModeRepository.kt
@@ -59,6 +59,8 @@
     /** A list of all existing priority modes. */
     val modes: Flow<List<ZenMode>>
 
+    fun getModes(): List<ZenMode>
+
     fun activateMode(zenMode: ZenMode, duration: Duration? = null)
 
     fun deactivateMode(zenMode: ZenMode)
@@ -184,6 +186,15 @@
         }
     }
 
+    /**
+     * Gets the current list of [ZenMode] instances according to the backend.
+     *
+     * This is necessary, and cannot be supplanted by making [modes] a StateFlow, because it will be
+     * called whenever we know or suspect that [modes] may not have caught up to the latest data
+     * (such as right after a user switch).
+     */
+    override fun getModes(): List<ZenMode> = backend.modes
+
     override fun activateMode(zenMode: ZenMode, duration: Duration?) {
         backend.activateMode(zenMode, duration)
     }
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
index f7492cf..8e0cdf8 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/TestModeBuilder.java
@@ -16,6 +16,7 @@
 
 package com.android.settingslib.notification.modes;
 
+import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.service.notification.ZenModeConfig.ORIGIN_UNKNOWN;
 import static android.service.notification.ZenModeConfig.ORIGIN_USER_IN_SYSTEMUI;
 
@@ -24,11 +25,13 @@
 import android.content.ComponentName;
 import android.net.Uri;
 import android.service.notification.Condition;
+import android.service.notification.SystemZenRules;
 import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenModeConfig;
 import android.service.notification.ZenPolicy;
 
 import androidx.annotation.DrawableRes;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import java.util.Random;
@@ -40,13 +43,27 @@
     private ZenModeConfig.ZenRule mConfigZenRule;
 
     public static final ZenMode EXAMPLE = new TestModeBuilder().build();
-    public static final ZenMode MANUAL_DND_ACTIVE = manualDnd(Uri.EMPTY, true);
-    public static final ZenMode MANUAL_DND_INACTIVE = manualDnd(Uri.EMPTY, false);
 
-    public static ZenMode manualDnd(Uri conditionId, boolean isActive) {
+    public static final ZenMode MANUAL_DND_ACTIVE = manualDnd(Uri.EMPTY,
+            INTERRUPTION_FILTER_PRIORITY, true);
+
+    public static final ZenMode MANUAL_DND_INACTIVE = manualDnd(Uri.EMPTY,
+            INTERRUPTION_FILTER_PRIORITY, false);
+
+    @NonNull
+    public static ZenMode manualDnd(@NotificationManager.InterruptionFilter int filter,
+            boolean isActive) {
+        return manualDnd(Uri.EMPTY, filter, isActive);
+    }
+
+    private static ZenMode manualDnd(Uri conditionId,
+            @NotificationManager.InterruptionFilter int filter, boolean isActive) {
         return ZenMode.manualDndMode(
                 new AutomaticZenRule.Builder("Do Not Disturb", conditionId)
-                        .setInterruptionFilter(NotificationManager.INTERRUPTION_FILTER_PRIORITY)
+                        .setInterruptionFilter(filter)
+                        .setType(AutomaticZenRule.TYPE_OTHER)
+                        .setManualInvocationAllowed(true)
+                        .setPackage(SystemZenRules.PACKAGE_ANDROID)
                         .setZenPolicy(new ZenPolicy.Builder().disallowAllSounds().build())
                         .build(),
                 isActive);
@@ -58,7 +75,7 @@
         mId = "rule_" + id;
         mRule = new AutomaticZenRule.Builder("Test Rule #" + id, Uri.parse("rule://" + id))
                 .setPackage("some_package")
-                .setInterruptionFilter(NotificationManager.INTERRUPTION_FILTER_PRIORITY)
+                .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
                 .setZenPolicy(new ZenPolicy.Builder().disallowAllSounds().build())
                 .build();
         mConfigZenRule = new ZenModeConfig.ZenRule();
@@ -134,7 +151,7 @@
             @NotificationManager.InterruptionFilter int interruptionFilter) {
         mRule.setInterruptionFilter(interruptionFilter);
         mConfigZenRule.zenMode = NotificationManager.zenModeFromInterruptionFilter(
-                interruptionFilter, NotificationManager.INTERRUPTION_FILTER_PRIORITY);
+                interruptionFilter, INTERRUPTION_FILTER_PRIORITY);
         return this;
     }
 
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java
index 0a0b65b..79dabf0 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconKeys.java
@@ -37,7 +37,7 @@
      * chosen one via Settings).
      */
     static final ZenIcon.Key IMPLICIT_MODE_DEFAULT = ZenIcon.Key.forSystemResource(
-            R.drawable.ic_zen_mode_type_unknown);
+            R.drawable.ic_zen_mode_type_special_dnd);
 
     private static final ImmutableMap<Integer, ZenIcon.Key> TYPE_DEFAULTS = ImmutableMap.of(
             AutomaticZenRule.TYPE_UNKNOWN,
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
index fe0f98a..43c6c50 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenIconLoader.java
@@ -31,7 +31,6 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
-import androidx.annotation.VisibleForTesting;
 
 import com.google.common.util.concurrent.FluentFuture;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -53,25 +52,22 @@
     private final LruCache<ZenIcon.Key, Drawable> mCache;
     private final ListeningExecutorService mBackgroundExecutor;
 
+    /** Obtains the singleton {@link ZenIconLoader}. */
     public static ZenIconLoader getInstance() {
         if (sInstance == null) {
-            sInstance = new ZenIconLoader();
+            sInstance = new ZenIconLoader(Executors.newFixedThreadPool(4));
         }
         return sInstance;
     }
 
-    /** Replaces the singleton instance of {@link ZenIconLoader} by the provided one. */
-    @VisibleForTesting(otherwise = VisibleForTesting.NONE)
-    public static void setInstance(@Nullable ZenIconLoader instance) {
-        sInstance = instance;
-    }
-
-    private ZenIconLoader() {
-        this(Executors.newFixedThreadPool(4));
-    }
-
-    @VisibleForTesting
-    public ZenIconLoader(ExecutorService backgroundExecutor) {
+    /**
+     * Constructs a ZenIconLoader with the specified {@code backgroundExecutor}.
+     *
+     * <p>ZenIconLoader <em>should be a singleton</em>, so this should only be used to instantiate
+     * and provide the singleton instance in a module. If the app doesn't support dependency
+     * injection, use {@link #getInstance} instead.
+     */
+    public ZenIconLoader(@NonNull ExecutorService backgroundExecutor) {
         mCache = new LruCache<>(50);
         mBackgroundExecutor =
                 MoreExecutors.listeningDecorator(backgroundExecutor);
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
index 36975c7..7b2a284 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenMode.java
@@ -18,9 +18,10 @@
 
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_TIME;
+import static android.app.NotificationManager.INTERRUPTION_FILTER_ALARMS;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALL;
+import static android.app.NotificationManager.INTERRUPTION_FILTER_NONE;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
-import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
 import static android.service.notification.SystemZenRules.getTriggerDescriptionForScheduleEvent;
 import static android.service.notification.SystemZenRules.getTriggerDescriptionForScheduleTime;
 import static android.service.notification.ZenModeConfig.tryParseCountdownConditionId;
@@ -68,23 +69,6 @@
     static final String MANUAL_DND_MODE_ID = ZenModeConfig.MANUAL_RULE_ID;
     static final String TEMP_NEW_MODE_ID = "temp_new_mode";
 
-    // Must match com.android.server.notification.ZenModeHelper#applyCustomPolicy.
-    private static final ZenPolicy POLICY_INTERRUPTION_FILTER_ALARMS =
-            new ZenPolicy.Builder()
-                    .disallowAllSounds()
-                    .allowAlarms(true)
-                    .allowMedia(true)
-                    .allowPriorityChannels(false)
-                    .build();
-
-    // Must match com.android.server.notification.ZenModeHelper#applyCustomPolicy.
-    private static final ZenPolicy POLICY_INTERRUPTION_FILTER_NONE =
-            new ZenPolicy.Builder()
-                    .disallowAllSounds()
-                    .hideAllVisualEffects()
-                    .allowPriorityChannels(false)
-                    .build();
-
     private static final Comparator<Integer> PRIORITIZED_TYPE_COMPARATOR = new Comparator<>() {
 
         private static final ImmutableList</* @AutomaticZenRule.Type */ Integer>
@@ -171,13 +155,9 @@
     }
 
     static ZenMode manualDndMode(AutomaticZenRule manualRule, boolean isActive) {
-        // Manual rule is owned by the system, so we set it here
-        AutomaticZenRule manualRuleWithPkg = new AutomaticZenRule.Builder(manualRule)
-                .setPackage(PACKAGE_ANDROID)
-                .build();
         return new ZenMode(
                 MANUAL_DND_MODE_ID,
-                manualRuleWithPkg,
+                manualRule,
                 Kind.MANUAL_DND,
                 isActive ? Status.ENABLED_AND_ACTIVE : Status.ENABLED);
     }
@@ -298,6 +278,23 @@
         }
     }
 
+    /** Returns the interruption filter of the mode. */
+    @NotificationManager.InterruptionFilter
+    public int getInterruptionFilter() {
+        return mRule.getInterruptionFilter();
+    }
+
+    /**
+     * Sets the interruption filter of the mode. This is valid for {@link AutomaticZenRule}-backed
+     * modes (and not manual DND).
+     */
+    public void setInterruptionFilter(@NotificationManager.InterruptionFilter int filter) {
+        if (isManualDnd() || !canEditPolicy()) {
+            throw new IllegalStateException("Cannot update interruption filter for mode " + this);
+        }
+        mRule.setInterruptionFilter(filter);
+    }
+
     @NonNull
     public ZenPolicy getPolicy() {
         switch (mRule.getInterruptionFilter()) {
@@ -306,10 +303,12 @@
                 return requireNonNull(mRule.getZenPolicy());
 
             case NotificationManager.INTERRUPTION_FILTER_ALARMS:
-                return POLICY_INTERRUPTION_FILTER_ALARMS;
+                return new ZenPolicy.Builder(ZenModeConfig.getDefaultZenPolicy()).build()
+                        .overwrittenWith(ZenPolicy.getBasePolicyInterruptionFilterAlarms());
 
             case NotificationManager.INTERRUPTION_FILTER_NONE:
-                return POLICY_INTERRUPTION_FILTER_NONE;
+                return new ZenPolicy.Builder(ZenModeConfig.getDefaultZenPolicy()).build()
+                        .overwrittenWith(ZenPolicy.getBasePolicyInterruptionFilterNone());
 
             case NotificationManager.INTERRUPTION_FILTER_UNKNOWN:
             default:
@@ -326,6 +325,10 @@
      */
     @SuppressLint("WrongConstant")
     public void setPolicy(@NonNull ZenPolicy policy) {
+        if (!canEditPolicy()) {
+            throw new IllegalStateException("Cannot update ZenPolicy for mode " + this);
+        }
+
         ZenPolicy currentPolicy = getPolicy();
         if (currentPolicy.equals(policy)) {
             return;
@@ -342,6 +345,12 @@
         mRule.setZenPolicy(policy);
     }
 
+    /**
+     * Returns the {@link ZenDeviceEffects} of the mode.
+     *
+     * <p>This is never {@code null}; if the backing AutomaticZenRule doesn't have effects set then
+     * a default (empty) effects set is returned.
+     */
     @NonNull
     public ZenDeviceEffects getDeviceEffects() {
         return mRule.getDeviceEffects() != null
@@ -349,6 +358,15 @@
                 : new ZenDeviceEffects.Builder().build();
     }
 
+    /** Sets the {@link ZenDeviceEffects} of the mode. */
+    public void setDeviceEffects(@NonNull ZenDeviceEffects effects) {
+        checkNotNull(effects);
+        if (!canEditPolicy()) {
+            throw new IllegalStateException("Cannot update device effects for mode " + this);
+        }
+        mRule.setDeviceEffects(effects);
+    }
+
     public void setCustomModeConditionId(Context context, Uri conditionId) {
         checkState(SystemZenRules.PACKAGE_ANDROID.equals(mRule.getPackageName()),
                 "Trying to change condition of non-system-owned rule %s (to %s)",
@@ -391,6 +409,18 @@
         return !isManualDnd();
     }
 
+    /**
+     * Whether the mode has an editable policy. Calling {@link #setPolicy},
+     * {@link #setDeviceEffects}, or {@link #setInterruptionFilter} is not valid for modes with a
+     * read-only policy.
+     */
+    public boolean canEditPolicy() {
+        // Cannot edit the policy of a temporarily active non-PRIORITY DND mode.
+        // Note that it's fine to edit the policy of an *AutomaticZenRule* with non-PRIORITY filter;
+        // the filter will we set to PRIORITY if you do.
+        return !isManualDndWithSpecialFilter();
+    }
+
     public boolean canBeDeleted() {
         return !isManualDnd();
     }
@@ -399,6 +429,12 @@
         return mKind == Kind.MANUAL_DND;
     }
 
+    private boolean isManualDndWithSpecialFilter() {
+        return isManualDnd()
+                && (mRule.getInterruptionFilter() == INTERRUPTION_FILTER_ALARMS
+                || mRule.getInterruptionFilter() == INTERRUPTION_FILTER_NONE);
+    }
+
     /**
      * A <em>custom manual</em> mode is a mode created by the user, and not yet assigned an
      * automatic trigger condition (neither time schedule nor a calendar).
@@ -414,6 +450,18 @@
         return mRule.isEnabled();
     }
 
+    /**
+     * Enables or disables the mode.
+     *
+     * <p>The DND mode cannot be disabled; trying to do so will fail.
+     */
+    public void setEnabled(boolean enabled) {
+        if (isManualDnd()) {
+            throw new IllegalStateException("Cannot update enabled for manual DND mode " + this);
+        }
+        mRule.setEnabled(enabled);
+    }
+
     public boolean isActive() {
         return mStatus == Status.ENABLED_AND_ACTIVE;
     }
diff --git a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenModesBackend.java b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenModesBackend.java
index c8a12f4..71e03c1 100644
--- a/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenModesBackend.java
+++ b/packages/SettingsLib/src/com/android/settingslib/notification/modes/ZenModesBackend.java
@@ -16,6 +16,11 @@
 
 package com.android.settingslib.notification.modes;
 
+import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
+import static android.app.NotificationManager.zenModeToInterruptionFilter;
+import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
+import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
+
 import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.app.AutomaticZenRule;
@@ -112,14 +117,22 @@
 
     private ZenMode getManualDndMode(ZenModeConfig config) {
         ZenModeConfig.ZenRule manualRule = config.manualRule;
+
+        // If DND is currently on with an interruption filter other than PRIORITY, construct the
+        // rule with that. DND will be *non-editable* while in this state.
+        int dndInterruptionFilter = config.isManualActive()
+                ? zenModeToInterruptionFilter(manualRule.zenMode)
+                : INTERRUPTION_FILTER_PRIORITY;
+
         AutomaticZenRule manualDndRule = new AutomaticZenRule.Builder(
-                mContext.getString(R.string.zen_mode_settings_title), manualRule.conditionId)
-                .setType(manualRule.type)
+                mContext.getString(R.string.zen_mode_do_not_disturb_name), manualRule.conditionId)
+                .setPackage(PACKAGE_ANDROID)
+                .setType(AutomaticZenRule.TYPE_OTHER)
                 .setZenPolicy(manualRule.zenPolicy)
                 .setDeviceEffects(manualRule.zenDeviceEffects)
-                .setManualInvocationAllowed(manualRule.allowManualInvocation)
+                .setManualInvocationAllowed(true)
                 .setConfigurationActivity(null) // No further settings
-                .setInterruptionFilter(NotificationManager.INTERRUPTION_FILTER_PRIORITY)
+                .setInterruptionFilter(dndInterruptionFilter)
                 .build();
 
         return ZenMode.manualDndMode(manualDndRule, config.isManualActive());
@@ -150,7 +163,7 @@
                 durationConditionId = ZenModeConfig.toTimeCondition(mContext,
                         (int) forDuration.toMinutes(), ActivityManager.getCurrentUser(), true).id;
             }
-            mNotificationManager.setZenMode(Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS,
+            mNotificationManager.setZenMode(ZEN_MODE_IMPORTANT_INTERRUPTIONS,
                     durationConditionId, TAG, /* fromUser= */ true);
 
         } else {
diff --git a/packages/SettingsLib/src/com/android/settingslib/users/EditUserPhotoController.java b/packages/SettingsLib/src/com/android/settingslib/users/EditUserPhotoController.java
index cdc3f12..f38e91a 100644
--- a/packages/SettingsLib/src/com/android/settingslib/users/EditUserPhotoController.java
+++ b/packages/SettingsLib/src/com/android/settingslib/users/EditUserPhotoController.java
@@ -32,6 +32,7 @@
 
 import com.android.internal.util.UserIcons;
 import com.android.settingslib.drawable.CircleFramedDrawable;
+import com.android.settingslib.R;
 import com.android.settingslib.utils.ThreadUtils;
 
 import com.google.common.util.concurrent.FutureCallback;
@@ -132,6 +133,13 @@
         intent.addCategory(Intent.CATEGORY_DEFAULT);
         if (Flags.avatarSync()) {
             intent.putExtra(EXTRA_IS_USER_NEW, isUserNew);
+            // Fix vulnerability b/341688848 by explicitly set the class name of avatar picker.
+            if (Flags.fixAvatarCrossUserLeak()) {
+                final String packageName =
+                        mActivity.getString(R.string.config_avatar_picker_package);
+                final String className = mActivity.getString(R.string.config_avatar_picker_class);
+                intent.setClassName(packageName, className);
+            }
         } else {
             // SettingsLib is used by multiple apps therefore we need to know out of all apps
             // using settingsLib which one is the one we return value to.
diff --git a/packages/SettingsLib/src/com/android/settingslib/volume/data/model/VolumeControllerEvent.kt b/packages/SettingsLib/src/com/android/settingslib/volume/data/model/VolumeControllerEvent.kt
new file mode 100644
index 0000000..0fe385b
--- /dev/null
+++ b/packages/SettingsLib/src/com/android/settingslib/volume/data/model/VolumeControllerEvent.kt
@@ -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.settingslib.volume.data.model
+
+import android.media.IVolumeController
+
+/** Models events received via [IVolumeController] */
+sealed interface VolumeControllerEvent {
+
+    /** @see [IVolumeController.displaySafeVolumeWarning] */
+    data class DisplaySafeVolumeWarning(val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.volumeChanged] */
+    data class VolumeChanged(val streamType: Int, val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.masterMuteChanged] */
+    data class MasterMuteChanged(val flags: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.setLayoutDirection] */
+    data class SetLayoutDirection(val layoutDirection: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.setA11yMode] */
+    data class SetA11yMode(val mode: Int) : VolumeControllerEvent
+
+    /** @see [IVolumeController.displayCsdWarning] */
+    data class DisplayCsdWarning(
+        val csdWarning: Int,
+        val displayDurationMs: Int,
+    ) : VolumeControllerEvent
+
+    /** @see [IVolumeController.dismiss] */
+    data object Dismiss : VolumeControllerEvent
+}
diff --git a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioRepository.kt b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioRepository.kt
index 0e71116..3e2d832 100644
--- a/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioRepository.kt
+++ b/packages/SettingsLib/src/com/android/settingslib/volume/data/repository/AudioRepository.kt
@@ -22,9 +22,12 @@
 import android.media.AudioManager
 import android.media.AudioManager.AudioDeviceCategory
 import android.media.AudioManager.OnCommunicationDeviceChangedListener
+import android.media.IVolumeController
 import android.provider.Settings
+import android.util.Log
 import androidx.concurrent.futures.DirectExecutor
 import com.android.internal.util.ConcurrentUtils
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
 import com.android.settingslib.volume.shared.AudioLogger
 import com.android.settingslib.volume.shared.AudioManagerEventsReceiver
 import com.android.settingslib.volume.shared.model.AudioManagerEvent
@@ -36,10 +39,13 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.conflate
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.emptyFlow
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.filterIsInstance
@@ -73,6 +79,11 @@
      */
     val communicationDevice: StateFlow<AudioDeviceInfo?>
 
+    /** Events from [AudioManager.setVolumeController] */
+    val volumeControllerEvents: Flow<VolumeControllerEvent>
+
+    fun init()
+
     /** State of the [AudioStream]. */
     fun getAudioStream(audioStream: AudioStream): Flow<AudioStreamModel>
 
@@ -90,8 +101,9 @@
     suspend fun setRingerMode(audioStream: AudioStream, mode: RingerMode)
 
     /** Gets audio device category. */
-    @AudioDeviceCategory
-    suspend fun getBluetoothAudioDeviceCategory(bluetoothAddress: String): Int
+    @AudioDeviceCategory suspend fun getBluetoothAudioDeviceCategory(bluetoothAddress: String): Int
+
+    suspend fun notifyVolumeControllerVisible(isVisible: Boolean)
 }
 
 class AudioRepositoryImpl(
@@ -101,8 +113,10 @@
     private val backgroundCoroutineContext: CoroutineContext,
     private val coroutineScope: CoroutineScope,
     private val logger: AudioLogger,
+    shouldUseVolumeController: Boolean,
 ) : AudioRepository {
 
+    private val volumeController = ProducingVolumeController()
     private val streamSettingNames: Map<AudioStream, String> =
         mapOf(
             AudioStream(AudioManager.STREAM_VOICE_CALL) to Settings.System.VOLUME_VOICE,
@@ -116,12 +130,19 @@
             AudioStream(AudioManager.STREAM_ASSISTANT) to Settings.System.VOLUME_ASSISTANT,
         )
 
+    override val volumeControllerEvents: Flow<VolumeControllerEvent> =
+        if (shouldUseVolumeController) {
+            volumeController.events
+        } else {
+            emptyFlow()
+        }
+
     override val mode: StateFlow<Int> =
         callbackFlow {
-            val listener = AudioManager.OnModeChangedListener { newMode -> trySend(newMode) }
-            audioManager.addOnModeChangedListener(ConcurrentUtils.DIRECT_EXECUTOR, listener)
-            awaitClose { audioManager.removeOnModeChangedListener(listener) }
-        }
+                val listener = AudioManager.OnModeChangedListener { newMode -> trySend(newMode) }
+                audioManager.addOnModeChangedListener(ConcurrentUtils.DIRECT_EXECUTOR, listener)
+                awaitClose { audioManager.removeOnModeChangedListener(listener) }
+            }
             .onStart { emit(audioManager.mode) }
             .flowOn(backgroundCoroutineContext)
             .stateIn(coroutineScope, SharingStarted.WhileSubscribed(), audioManager.mode)
@@ -141,14 +162,14 @@
     override val communicationDevice: StateFlow<AudioDeviceInfo?>
         get() =
             callbackFlow {
-                val listener = OnCommunicationDeviceChangedListener { trySend(Unit) }
-                audioManager.addOnCommunicationDeviceChangedListener(
-                    ConcurrentUtils.DIRECT_EXECUTOR,
-                    listener,
-                )
+                    val listener = OnCommunicationDeviceChangedListener { trySend(Unit) }
+                    audioManager.addOnCommunicationDeviceChangedListener(
+                        ConcurrentUtils.DIRECT_EXECUTOR,
+                        listener,
+                    )
 
-                awaitClose { audioManager.removeOnCommunicationDeviceChangedListener(listener) }
-            }
+                    awaitClose { audioManager.removeOnCommunicationDeviceChangedListener(listener) }
+                }
                 .filterNotNull()
                 .map { audioManager.communicationDevice }
                 .onStart { emit(audioManager.communicationDevice) }
@@ -159,20 +180,30 @@
                     audioManager.communicationDevice,
                 )
 
+    override fun init() {
+        try {
+            audioManager.volumeController = volumeController
+        } catch (error: SecurityException) {
+            Log.wtf("AudioManager", "Unable to set the volume controller", error)
+        }
+    }
+
     override fun getAudioStream(audioStream: AudioStream): Flow<AudioStreamModel> {
         return merge(
-            audioManagerEventsReceiver.events.filter {
-                if (it is StreamAudioManagerEvent) {
-                    it.audioStream == audioStream
-                } else {
-                    true
-                }
-            },
-            volumeSettingChanges(audioStream),
-        )
+                audioManagerEventsReceiver.events.filter {
+                    if (it is StreamAudioManagerEvent) {
+                        it.audioStream == audioStream
+                    } else {
+                        true
+                    }
+                },
+                volumeSettingChanges(audioStream),
+                volumeControllerEvents.filter { it is VolumeControllerEvent.VolumeChanged },
+            )
             .conflate()
             .map { getCurrentAudioStream(audioStream) }
             .onStart { emit(getCurrentAudioStream(audioStream)) }
+            .distinctUntilChanged()
             .onEach { logger.onVolumeUpdateReceived(audioStream, it) }
             .flowOn(backgroundCoroutineContext)
     }
@@ -228,6 +259,12 @@
         }
     }
 
+    override suspend fun notifyVolumeControllerVisible(isVisible: Boolean) {
+        withContext(backgroundCoroutineContext) {
+            audioManager.notifyVolumeControllerVisible(volumeController, isVisible)
+        }
+    }
+
     private fun getMinVolume(stream: AudioStream): Int =
         try {
             audioManager.getStreamMinVolume(stream.value)
@@ -253,3 +290,45 @@
         }
     }
 }
+
+private class ProducingVolumeController : IVolumeController.Stub() {
+
+    private val mutableEvents = MutableSharedFlow<VolumeControllerEvent>(extraBufferCapacity = 32)
+    val events = mutableEvents.asSharedFlow()
+
+    override fun displaySafeVolumeWarning(flags: Int) {
+        mutableEvents.tryEmit(VolumeControllerEvent.DisplaySafeVolumeWarning(flags))
+    }
+
+    override fun volumeChanged(streamType: Int, flags: Int) {
+        mutableEvents.tryEmit(VolumeControllerEvent.VolumeChanged(streamType, flags))
+    }
+
+    override fun masterMuteChanged(flags: Int) {
+        mutableEvents.tryEmit(VolumeControllerEvent.MasterMuteChanged(flags))
+    }
+
+    override fun setLayoutDirection(layoutDirection: Int) {
+        mutableEvents.tryEmit(VolumeControllerEvent.SetLayoutDirection(layoutDirection))
+    }
+
+    override fun dismiss() {
+        mutableEvents.tryEmit(VolumeControllerEvent.Dismiss)
+    }
+
+    override fun setA11yMode(mode: Int) {
+        mutableEvents.tryEmit(VolumeControllerEvent.SetA11yMode(mode))
+    }
+
+    override fun displayCsdWarning(
+        csdWarning: Int,
+        displayDurationMs: Int,
+    ) {
+        mutableEvents.tryEmit(
+            VolumeControllerEvent.DisplayCsdWarning(
+                csdWarning,
+                displayDurationMs,
+            )
+        )
+    }
+}
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryTest.kt b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryTest.kt
index 0e43acb..52e6391 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryTest.kt
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryTest.kt
@@ -44,6 +44,7 @@
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Captor
 import org.mockito.Mock
+import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
@@ -111,6 +112,7 @@
                 testScope.testScheduler,
                 testScope.backgroundScope,
                 logger,
+                true,
             )
     }
 
@@ -261,8 +263,8 @@
     @Test
     fun getBluetoothAudioDeviceCategory() {
         testScope.runTest {
-            `when`(audioManager.getBluetoothAudioDeviceCategory("12:34:56:78")).thenReturn(
-                AudioManager.AUDIO_DEVICE_CATEGORY_HEADPHONES)
+            `when`(audioManager.getBluetoothAudioDeviceCategory("12:34:56:78"))
+                .thenReturn(AudioManager.AUDIO_DEVICE_CATEGORY_HEADPHONES)
 
             val category = underTest.getBluetoothAudioDeviceCategory("12:34:56:78")
             runCurrent()
@@ -271,6 +273,27 @@
         }
     }
 
+    @Test
+    fun useVolumeControllerDisabled_setVolumeController_notCalled() {
+        testScope.runTest {
+            underTest =
+                AudioRepositoryImpl(
+                    eventsReceiver,
+                    audioManager,
+                    contentResolver,
+                    testScope.testScheduler,
+                    testScope.backgroundScope,
+                    logger,
+                    false,
+                )
+
+            underTest.volumeControllerEvents.launchIn(backgroundScope)
+            runCurrent()
+
+            verify(audioManager, never()).volumeController = any()
+        }
+    }
+
     private fun triggerConnectedDeviceChange(communicationDevice: AudioDeviceInfo?) {
         verify(audioManager)
             .addOnCommunicationDeviceChangedListener(
diff --git a/packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryVolumeControllerEventsTest.kt
similarity index 76%
rename from packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt
rename to packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryVolumeControllerEventsTest.kt
index 83b612d..f5c2f01 100644
--- a/packages/SettingsLib/tests/integ/src/com/android/settingslib/media/data/repository/AudioManagerVolumeControllerExtTest.kt
+++ b/packages/SettingsLib/tests/integ/src/com/android/settingslib/volume/data/repository/AudioRepositoryVolumeControllerEventsTest.kt
@@ -14,12 +14,15 @@
  * limitations under the License.
  */
 
-package com.android.settingslib.media.data.repository
+package com.android.settingslib.volume.data.repository
 
+import android.content.ContentResolver
 import android.media.AudioManager
 import android.media.IVolumeController
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
+import com.android.settingslib.volume.shared.FakeAudioManagerEventsReceiver
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.launchIn
@@ -39,16 +42,32 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-class AudioManagerVolumeControllerExtTest {
+class AudioRepositoryVolumeControllerEventsTest {
 
     private val testScope = TestScope()
 
     @Captor private lateinit var volumeControllerCaptor: ArgumentCaptor<IVolumeController>
     @Mock private lateinit var audioManager: AudioManager
+    @Mock private lateinit var contentResolver: ContentResolver
+
+    private val logger = FakeAudioRepositoryLogger()
+    private val eventsReceiver = FakeAudioManagerEventsReceiver()
+
+    private lateinit var underTest: AudioRepository
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
+        underTest =
+            AudioRepositoryImpl(
+                eventsReceiver,
+                audioManager,
+                contentResolver,
+                testScope.testScheduler,
+                testScope.backgroundScope,
+                logger,
+                true,
+            )
     }
 
     @Test
@@ -83,7 +102,7 @@
     ) =
         testScope.runTest {
             var event: VolumeControllerEvent? = null
-            audioManager.volumeControllerEvents().onEach { event = it }.launchIn(backgroundScope)
+            underTest.volumeControllerEvents.onEach { event = it }.launchIn(backgroundScope)
             runCurrent()
             verify(audioManager).volumeController = volumeControllerCaptor.capture()
 
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/BluetoothUtilsTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/BluetoothUtilsTest.java
index 3a7b0c7..a0e764a 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/BluetoothUtilsTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/BluetoothUtilsTest.java
@@ -16,6 +16,7 @@
 package com.android.settingslib.bluetooth;
 
 import static com.android.settingslib.bluetooth.BluetoothUtils.isAvailableAudioSharingMediaBluetoothDevice;
+import static com.android.settingslib.bluetooth.LocalBluetoothLeBroadcast.UNKNOWN_VALUE_PLACEHOLDER;
 import static com.android.settingslib.flags.Flags.FLAG_ENABLE_DETERMINING_ADVANCED_DETAILS_HEADER_WITH_METADATA;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -96,6 +97,7 @@
                     + "</HEARABLE_CONTROL_SLICE_WITH_WIDTH>";
     private static final String TEST_EXCLUSIVE_MANAGER_PACKAGE = "com.test.manager";
     private static final String TEST_EXCLUSIVE_MANAGER_COMPONENT = "com.test.manager/.component";
+    private static final int TEST_BROADCAST_ID = 25;
 
     @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
@@ -671,6 +673,34 @@
     }
 
     @Test
+    public void testHasActiveLocalBroadcastSourceForBtDevice_hasActiveLocalSource() {
+        when(mBroadcast.getLatestBroadcastId()).thenReturn(TEST_BROADCAST_ID);
+        when(mLeBroadcastReceiveState.getBroadcastId()).thenReturn(TEST_BROADCAST_ID);
+        List<BluetoothLeBroadcastReceiveState> sourceList = new ArrayList<>();
+        sourceList.add(mLeBroadcastReceiveState);
+        when(mAssistant.getAllSources(mBluetoothDevice)).thenReturn(sourceList);
+
+        assertThat(
+                BluetoothUtils.hasActiveLocalBroadcastSourceForBtDevice(
+                        mBluetoothDevice, mLocalBluetoothManager))
+                .isTrue();
+    }
+
+    @Test
+    public void testHasActiveLocalBroadcastSourceForBtDevice_noActiveLocalSource() {
+        when(mLeBroadcastReceiveState.getBroadcastId()).thenReturn(UNKNOWN_VALUE_PLACEHOLDER);
+        List<BluetoothLeBroadcastReceiveState> sourceList = new ArrayList<>();
+        sourceList.add(mLeBroadcastReceiveState);
+        when(mAssistant.getAllSources(mBluetoothDevice)).thenReturn(sourceList);
+
+        assertThat(
+                BluetoothUtils.hasActiveLocalBroadcastSourceForBtDevice(
+                        mBluetoothDevice, mLocalBluetoothManager))
+                .isFalse();
+    }
+
+
+    @Test
     public void isAvailableHearingDevice_isConnectedHearingAid_returnTure() {
         when(mCachedBluetoothDevice.isConnectedHearingAidDevice()).thenReturn(true);
         when(mCachedBluetoothDevice.getDevice()).thenReturn(mBluetoothDevice);
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatusTest.kt b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatusTest.kt
new file mode 100644
index 0000000..aa22fac
--- /dev/null
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/DeviceSettingsProviderServiceStatusTest.kt
@@ -0,0 +1,51 @@
+/*
+ * 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.settingslib.bluetooth.devicesettings
+
+import android.os.Bundle
+import android.os.Parcel
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@RunWith(RobolectricTestRunner::class)
+class DeviceSettingsProviderServiceStatusTest {
+
+    @Test
+    fun parcelOperation() {
+        val item =
+            DeviceSettingsProviderServiceStatus(
+                enabled = true,
+                extras = Bundle().apply { putString("key1", "value1") },
+            )
+
+        val fromParcel = writeAndRead(item)
+
+        assertThat(fromParcel.enabled).isEqualTo(item.enabled)
+        assertThat(fromParcel.extras.getString("key1")).isEqualTo(item.extras.getString("key1"))
+    }
+
+    private fun writeAndRead(
+        item: DeviceSettingsProviderServiceStatus
+    ): DeviceSettingsProviderServiceStatus {
+        val parcel = Parcel.obtain()
+        item.writeToParcel(parcel, 0)
+        parcel.setDataPosition(0)
+        return DeviceSettingsProviderServiceStatus.CREATOR.createFromParcel(parcel)
+    }
+}
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingRepositoryTest.kt b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingRepositoryTest.kt
index 95ee46e..ce155b5 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingRepositoryTest.kt
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/bluetooth/devicesettings/data/repository/DeviceSettingRepositoryTest.kt
@@ -33,6 +33,7 @@
 import com.android.settingslib.bluetooth.devicesettings.DeviceSettingItem
 import com.android.settingslib.bluetooth.devicesettings.DeviceSettingState
 import com.android.settingslib.bluetooth.devicesettings.DeviceSettingsConfig
+import com.android.settingslib.bluetooth.devicesettings.DeviceSettingsProviderServiceStatus
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsConfigProviderService
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsListener
 import com.android.settingslib.bluetooth.devicesettings.IDeviceSettingsProviderService
@@ -47,10 +48,8 @@
 import com.android.settingslib.bluetooth.devicesettings.shared.model.ToggleModel
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
@@ -59,12 +58,9 @@
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.ArgumentCaptor
 import org.mockito.ArgumentMatchers.any
 import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.ArgumentMatchers.anyString
-import org.mockito.ArgumentMatchers.eq
-import org.mockito.Captor
 import org.mockito.Mock
 import org.mockito.Mockito.doReturn
 import org.mockito.Mockito.verify
@@ -85,9 +81,6 @@
     @Mock private lateinit var configService: IDeviceSettingsConfigProviderService.Stub
     @Mock private lateinit var settingProviderService1: IDeviceSettingsProviderService.Stub
     @Mock private lateinit var settingProviderService2: IDeviceSettingsProviderService.Stub
-    @Captor
-    private lateinit var metadataChangeCaptor:
-        ArgumentCaptor<BluetoothAdapter.OnMetadataChangedListener>
 
     private lateinit var underTest: DeviceSettingRepository
     private val testScope = TestScope()
@@ -153,6 +146,12 @@
     fun getDeviceSettingsConfig_withMetadata_success() {
         testScope.runTest {
             `when`(configService.getDeviceSettingsConfig(any())).thenReturn(DEVICE_SETTING_CONFIG)
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
 
             val config = underTest.getDeviceSettingsConfig(cachedDevice)
 
@@ -161,32 +160,40 @@
     }
 
     @Test
-    fun getDeviceSettingsConfig_waitMetadataChange_success() {
+    fun getDeviceSettingsConfig_noMetadata_returnNull() {
+        testScope.runTest {
+            `when`(
+                bluetoothDevice.getMetadata(
+                    DeviceSettingServiceConnection.METADATA_FAST_PAIR_CUSTOMIZED_FIELDS))
+                .thenReturn("".toByteArray())
+            `when`(configService.getDeviceSettingsConfig(any())).thenReturn(DEVICE_SETTING_CONFIG)
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+
+            val config = underTest.getDeviceSettingsConfig(cachedDevice)
+
+            assertThat(config).isNull()
+        }
+    }
+
+    @Test
+    fun getDeviceSettingsConfig_providerServiceNotEnabled_returnNull() {
         testScope.runTest {
             `when`(configService.getDeviceSettingsConfig(any())).thenReturn(DEVICE_SETTING_CONFIG)
-            `when`(
-                    bluetoothDevice.getMetadata(
-                        DeviceSettingServiceConnection.METADATA_FAST_PAIR_CUSTOMIZED_FIELDS))
-                .thenReturn("".toByteArray())
-
-            var config: DeviceSettingConfigModel? = null
-            val job = launch { config = underTest.getDeviceSettingsConfig(cachedDevice) }
-            delay(1000)
-            verify(bluetoothAdapter)
-                .addOnMetadataChangedListener(
-                    eq(bluetoothDevice), any(), metadataChangeCaptor.capture())
-            metadataChangeCaptor.value.onMetadataChanged(
-                bluetoothDevice,
-                DeviceSettingServiceConnection.METADATA_FAST_PAIR_CUSTOMIZED_FIELDS,
-                BLUETOOTH_DEVICE_METADATA.toByteArray(),
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(false)
             )
-            `when`(
-                    bluetoothDevice.getMetadata(
-                        DeviceSettingServiceConnection.METADATA_FAST_PAIR_CUSTOMIZED_FIELDS))
-                .thenReturn(BLUETOOTH_DEVICE_METADATA.toByteArray())
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
 
-            job.join()
-            assertConfig(config!!, DEVICE_SETTING_CONFIG)
+            val config = underTest.getDeviceSettingsConfig(cachedDevice)
+
+            assertThat(config).isNull()
         }
     }
 
@@ -212,6 +219,12 @@
                     .getArgument<IDeviceSettingsListener>(1)
                     .onDeviceSettingsChanged(listOf(DEVICE_SETTING_1))
             }
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
             var setting: DeviceSettingModel? = null
 
             underTest
@@ -234,6 +247,12 @@
                     .getArgument<IDeviceSettingsListener>(1)
                     .onDeviceSettingsChanged(listOf(DEVICE_SETTING_2))
             }
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
             var setting: DeviceSettingModel? = null
 
             underTest
@@ -256,6 +275,12 @@
                     .getArgument<IDeviceSettingsListener>(1)
                     .onDeviceSettingsChanged(listOf(DEVICE_SETTING_HELP))
             }
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
             var setting: DeviceSettingModel? = null
 
             underTest
@@ -299,6 +324,12 @@
                     .getArgument<IDeviceSettingsListener>(1)
                     .onDeviceSettingsChanged(listOf(DEVICE_SETTING_1))
             }
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
             var setting: DeviceSettingModel? = null
 
             underTest
@@ -331,6 +362,12 @@
                     .getArgument<IDeviceSettingsListener>(1)
                     .onDeviceSettingsChanged(listOf(DEVICE_SETTING_2))
             }
+            `when`(settingProviderService1.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
+            `when`(settingProviderService2.serviceStatus).thenReturn(
+                DeviceSettingsProviderServiceStatus(true)
+            )
             var setting: DeviceSettingModel? = null
 
             underTest
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/data/repository/ZenModeRepositoryTest.kt b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/data/repository/ZenModeRepositoryTest.kt
index 67c73b1..c136644 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/data/repository/ZenModeRepositoryTest.kt
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/data/repository/ZenModeRepositoryTest.kt
@@ -199,6 +199,15 @@
         }
     }
 
+    @EnableFlags(android.app.Flags.FLAG_MODES_UI)
+    @Test
+    fun getModes_returnsModes() {
+        val modesList = listOf(TestModeBuilder().setId("One").build())
+        `when`(zenModesBackend.modes).thenReturn(modesList)
+
+        assertThat(underTest.getModes()).isEqualTo(modesList)
+    }
+
     private fun triggerIntent(action: String, extras: Map<String, Parcelable>? = null) {
         verify(context).registerReceiver(receiverCaptor.capture(), any(), any(), any())
         val intent = Intent(action)
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
index e64b0c6..d08d91d 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModeTest.java
@@ -24,6 +24,7 @@
 import static android.app.AutomaticZenRule.TYPE_THEATER;
 import static android.app.AutomaticZenRule.TYPE_UNKNOWN;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_ALARMS;
+import static android.app.NotificationManager.INTERRUPTION_FILTER_ALL;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_NONE;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
@@ -31,11 +32,14 @@
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
 
+import static org.junit.Assert.assertThrows;
+
 import android.app.AutomaticZenRule;
 import android.net.Uri;
 import android.os.Parcel;
 import android.service.notification.Condition;
 import android.service.notification.SystemZenRules;
+import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenModeConfig;
 import android.service.notification.ZenPolicy;
 
@@ -69,20 +73,26 @@
                     .build();
 
     @Test
-    public void testBasicMethods() {
+    public void testBasicMethods_mode() {
         ZenMode zenMode = new ZenMode("id", ZEN_RULE, zenConfigRuleFor(ZEN_RULE, true));
 
         assertThat(zenMode.getId()).isEqualTo("id");
         assertThat(zenMode.getRule()).isEqualTo(ZEN_RULE);
         assertThat(zenMode.isManualDnd()).isFalse();
         assertThat(zenMode.canEditNameAndIcon()).isTrue();
+        assertThat(zenMode.canEditPolicy()).isTrue();
         assertThat(zenMode.canBeDeleted()).isTrue();
         assertThat(zenMode.isActive()).isTrue();
+    }
 
-        ZenMode manualMode = ZenMode.manualDndMode(ZEN_RULE, false);
+    @Test
+    public void testBasicMethods_manualDnd() {
+        ZenMode manualMode = TestModeBuilder.MANUAL_DND_INACTIVE;
+
         assertThat(manualMode.getId()).isEqualTo(ZenMode.MANUAL_DND_MODE_ID);
         assertThat(manualMode.isManualDnd()).isTrue();
         assertThat(manualMode.canEditNameAndIcon()).isFalse();
+        assertThat(manualMode.canEditPolicy()).isTrue();
         assertThat(manualMode.canBeDeleted()).isFalse();
         assertThat(manualMode.isActive()).isFalse();
         assertThat(manualMode.getRule().getPackageName()).isEqualTo(PACKAGE_ANDROID);
@@ -203,7 +213,7 @@
         ZenMode zenMode = new ZenMode("id", azr, zenConfigRuleFor(azr, false));
 
         assertThat(zenMode.getPolicy()).isEqualTo(
-                new ZenPolicy.Builder()
+                new ZenPolicy.Builder(ZenModeConfig.getDefaultZenPolicy())
                         .disallowAllSounds()
                         .allowAlarms(true)
                         .allowMedia(true)
@@ -220,9 +230,8 @@
         ZenMode zenMode = new ZenMode("id", azr, zenConfigRuleFor(azr, false));
 
         assertThat(zenMode.getPolicy()).isEqualTo(
-                new ZenPolicy.Builder()
+                new ZenPolicy.Builder(ZenModeConfig.getDefaultZenPolicy())
                         .disallowAllSounds()
-                        .hideAllVisualEffects()
                         .allowPriorityChannels(false)
                         .build());
     }
@@ -243,6 +252,77 @@
     }
 
     @Test
+    public void getInterruptionFilter_returnsFilter() {
+        ZenMode mode = new TestModeBuilder().setInterruptionFilter(
+                INTERRUPTION_FILTER_ALARMS).build();
+
+        assertThat(mode.getInterruptionFilter()).isEqualTo(INTERRUPTION_FILTER_ALARMS);
+    }
+
+    @Test
+    public void setInterruptionFilter_setsFilter() {
+        ZenMode mode = new TestModeBuilder().setInterruptionFilter(
+                INTERRUPTION_FILTER_ALARMS).build();
+
+        mode.setInterruptionFilter(INTERRUPTION_FILTER_ALL);
+
+        assertThat(mode.getInterruptionFilter()).isEqualTo(INTERRUPTION_FILTER_ALL);
+    }
+
+    @Test
+    public void setInterruptionFilter_manualDnd_throws() {
+        ZenMode manualDnd = TestModeBuilder.MANUAL_DND_INACTIVE;
+
+        assertThrows(IllegalStateException.class,
+                () -> manualDnd.setInterruptionFilter(INTERRUPTION_FILTER_ALL));
+    }
+
+    @Test
+    public void canEditPolicy_onlyFalseForSpecialDnd() {
+        assertThat(TestModeBuilder.EXAMPLE.canEditPolicy()).isTrue();
+        assertThat(TestModeBuilder.MANUAL_DND_ACTIVE.canEditPolicy()).isTrue();
+        assertThat(TestModeBuilder.MANUAL_DND_INACTIVE.canEditPolicy()).isTrue();
+
+        ZenMode dndWithAlarms = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_ALARMS, true);
+        assertThat(dndWithAlarms.canEditPolicy()).isFalse();
+        ZenMode dndWithNone = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_NONE, true);
+        assertThat(dndWithNone.canEditPolicy()).isFalse();
+
+        // Note: Backend will never return an inactive manual mode with custom filter.
+        ZenMode badDndWithAlarms = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_ALARMS, false);
+        assertThat(badDndWithAlarms.canEditPolicy()).isFalse();
+        ZenMode badDndWithNone = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_NONE, false);
+        assertThat(badDndWithNone.canEditPolicy()).isFalse();
+    }
+
+    @Test
+    public void canEditPolicy_whenTrue_allowsSettingPolicyAndEffects() {
+        ZenMode normalDnd = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_PRIORITY, true);
+
+        assertThat(normalDnd.canEditPolicy()).isTrue();
+
+        ZenPolicy somePolicy = new ZenPolicy.Builder().showBadges(true).build();
+        normalDnd.setPolicy(somePolicy);
+        assertThat(normalDnd.getPolicy()).isEqualTo(somePolicy);
+
+        ZenDeviceEffects someEffects = new ZenDeviceEffects.Builder()
+                .setShouldUseNightMode(true).build();
+        normalDnd.setDeviceEffects(someEffects);
+        assertThat(normalDnd.getDeviceEffects()).isEqualTo(someEffects);
+    }
+
+    @Test
+    public void canEditPolicy_whenFalse_preventsSettingFilterPolicyOrEffects() {
+        ZenMode specialDnd = TestModeBuilder.manualDnd(INTERRUPTION_FILTER_ALARMS, true);
+
+        assertThat(specialDnd.canEditPolicy()).isFalse();
+        assertThrows(IllegalStateException.class,
+                () -> specialDnd.setPolicy(ZEN_POLICY));
+        assertThrows(IllegalStateException.class,
+                () -> specialDnd.setDeviceEffects(new ZenDeviceEffects.Builder().build()));
+    }
+
+    @Test
     public void comparator_prioritizes() {
         ZenMode manualDnd = TestModeBuilder.MANUAL_DND_INACTIVE;
         ZenMode driving1 = new TestModeBuilder().setName("b1").setType(TYPE_DRIVING).build();
@@ -359,7 +439,7 @@
     }
 
     @Test
-    public void getIconKey_implicitModeWithoutCustomIcon_isSpecialIcon() {
+    public void getIconKey_implicitModeWithoutCustomIcon_isDndIcon() {
         ZenMode mode = new TestModeBuilder()
                 .setId(ZenModeConfig.implicitRuleId("some.package"))
                 .setPackage("some_package")
@@ -370,7 +450,7 @@
 
         assertThat(iconKey.resPackage()).isNull();
         assertThat(iconKey.resId()).isEqualTo(
-                com.android.internal.R.drawable.ic_zen_mode_type_unknown);
+                com.android.internal.R.drawable.ic_zen_mode_type_special_dnd);
     }
 
     private static void assertUnparceledIsEqualToOriginal(String type, ZenMode original) {
diff --git a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModesBackendTest.java b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModesBackendTest.java
index 539519b..aae72b3 100644
--- a/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModesBackendTest.java
+++ b/packages/SettingsLib/tests/robotests/src/com/android/settingslib/notification/modes/ZenModesBackendTest.java
@@ -16,12 +16,14 @@
 
 package com.android.settingslib.notification.modes;
 
+import static android.app.NotificationManager.INTERRUPTION_FILTER_ALARMS;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.provider.Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS;
 import static android.provider.Settings.Global.ZEN_MODE_OFF;
 import static android.service.notification.Condition.SOURCE_UNKNOWN;
 import static android.service.notification.Condition.STATE_FALSE;
 import static android.service.notification.Condition.STATE_TRUE;
+import static android.service.notification.SystemZenRules.PACKAGE_ANDROID;
 import static android.service.notification.ZenAdapters.notificationPolicyToZenPolicy;
 import static android.service.notification.ZenPolicy.STATE_ALLOW;
 
@@ -47,8 +49,6 @@
 import android.service.notification.ZenModeConfig;
 import android.service.notification.ZenPolicy;
 
-import com.android.settingslib.R;
-
 import com.google.common.collect.ImmutableMap;
 
 import org.junit.Before;
@@ -172,9 +172,8 @@
         // all modes exist, but none of them are currently active
         assertThat(modes).containsExactly(
                         ZenMode.manualDndMode(
-                                new AutomaticZenRule.Builder(
-                                        mContext.getString(R.string.zen_mode_settings_title),
-                                        Uri.EMPTY)
+                                new AutomaticZenRule.Builder("Do Not Disturb", Uri.EMPTY)
+                                        .setPackage(PACKAGE_ANDROID)
                                         .setType(AutomaticZenRule.TYPE_OTHER)
                                         .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
                                         .setZenPolicy(notificationPolicyToZenPolicy(dndPolicy))
@@ -196,15 +195,59 @@
 
         ZenMode mode = mBackend.getMode(ZenMode.MANUAL_DND_MODE_ID);
 
+        assertThat(mode).isNotNull();
         assertThat(mode).isEqualTo(
                 ZenMode.manualDndMode(
-                        new AutomaticZenRule.Builder(
-                                mContext.getString(R.string.zen_mode_settings_title), Uri.EMPTY)
+                        new AutomaticZenRule.Builder("Do Not Disturb", Uri.EMPTY)
+                                .setPackage(PACKAGE_ANDROID)
                                 .setType(AutomaticZenRule.TYPE_OTHER)
                                 .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
                                 .setZenPolicy(notificationPolicyToZenPolicy(dndPolicy))
                                 .setManualInvocationAllowed(true)
                                 .build(), false));
+
+        assertThat(mode.isManualDnd()).isTrue();
+        assertThat(mode.isEnabled()).isTrue();
+        assertThat(mode.isActive()).isFalse();
+        assertThat(mode.canEditPolicy()).isTrue();
+        assertThat(mode.getPolicy()).isEqualTo(notificationPolicyToZenPolicy(dndPolicy));
+    }
+
+    @Test
+    public void getMode_dndWithOtherInterruptionFilter_returnsSpecialDndMode() {
+        ZenModeConfig config = configWithManualRule(new ZenModeConfig(), true);
+        config.manualRule.zenMode = Settings.Global.ZEN_MODE_ALARMS;
+        Policy dndPolicyForPriority = new Policy(Policy.PRIORITY_CATEGORY_ALARMS,
+                Policy.PRIORITY_SENDERS_CONTACTS, Policy.PRIORITY_SENDERS_CONTACTS);
+        config.applyNotificationPolicy(dndPolicyForPriority);
+        when(mNm.getZenModeConfig()).thenReturn(config);
+
+        ZenMode mode = mBackend.getMode(ZenMode.MANUAL_DND_MODE_ID);
+
+        assertThat(mode).isNotNull();
+        assertThat(mode).isEqualTo(
+                ZenMode.manualDndMode(
+                        new AutomaticZenRule.Builder("Do Not Disturb", Uri.EMPTY)
+                                .setPackage(PACKAGE_ANDROID)
+                                .setType(AutomaticZenRule.TYPE_OTHER)
+                                .setInterruptionFilter(INTERRUPTION_FILTER_ALARMS)
+                                .setZenPolicy(notificationPolicyToZenPolicy(dndPolicyForPriority))
+                                .setManualInvocationAllowed(true)
+                                .build(), true));
+
+        assertThat(mode.isManualDnd()).isTrue();
+        assertThat(mode.isEnabled()).isTrue();
+        assertThat(mode.isActive()).isTrue();
+
+        // Mode itself has a special fixed policy, different to the rule.
+        assertThat(mode.canEditPolicy()).isFalse();
+        assertThat(mode.getPolicy()).isEqualTo(
+                new ZenPolicy.Builder(ZenModeConfig.getDefaultZenPolicy())
+                        .disallowAllSounds()
+                        .allowAlarms(true)
+                        .allowMedia(true)
+                        .allowPriorityChannels(false)
+                        .build());
     }
 
     @Test
diff --git a/packages/SettingsProvider/src/android/provider/settings/validators/GlobalSettingsValidators.java b/packages/SettingsProvider/src/android/provider/settings/validators/GlobalSettingsValidators.java
index 65937ea..f6e1057 100644
--- a/packages/SettingsProvider/src/android/provider/settings/validators/GlobalSettingsValidators.java
+++ b/packages/SettingsProvider/src/android/provider/settings/validators/GlobalSettingsValidators.java
@@ -340,7 +340,8 @@
                         new String[] {
                             String.valueOf(Global.Wearable.PAIRED_DEVICE_OS_TYPE_UNKNOWN),
                             String.valueOf(Global.Wearable.PAIRED_DEVICE_OS_TYPE_ANDROID),
-                            String.valueOf(Global.Wearable.PAIRED_DEVICE_OS_TYPE_IOS)
+                            String.valueOf(Global.Wearable.PAIRED_DEVICE_OS_TYPE_IOS),
+                            String.valueOf(Global.Wearable.PAIRED_DEVICE_OS_TYPE_NONE)
                         }));
         VALIDATORS.put(
                 Global.Wearable.COMPANION_BLE_ROLE,
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index c1bb55c..d26a906 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -949,36 +949,37 @@
     strict_mode: false,
 }
 
-// Disable for now. TODO(b/356666754) Re-enable it
-// android_ravenwood_test {
-//     name: "SystemUiRavenTests",
-//     srcs: [
-//         ":SystemUI-tests-utils",
-//         ":SystemUI-tests-multivalent",
-//         // TODO(b/294256649): pivot to using {.aapt.jar} and re-enable
-//         // use_resource_processor: true when better supported by soong
-//         ":SystemUIRobo-stub{.aapt.srcjar}",
-//     ],
-//     static_libs: [
-//         "SystemUI-core",
-//         "SystemUI-res",
-//         "SystemUI-tests-base",
-//         "androidx.test.uiautomator_uiautomator",
-//         "androidx.core_core-animation-testing",
-//         "androidx.test.ext.junit",
-//         "kosmos",
-//         "mockito-kotlin-nodeps",
-//     ],
-//     libs: [
-//         "android.test.runner",
-//         "android.test.base",
-//         "android.test.mock",
-//     ],
-//     auto_gen_config: true,
-//     plugins: [
-//         "dagger2-compiler",
-//     ],
-// }
+android_ravenwood_test {
+    name: "SystemUiRavenTests",
+    srcs: [
+        ":SystemUI-tests-utils",
+        ":SystemUI-tests-multivalent",
+        // TODO(b/294256649): pivot to using {.aapt.jar} and re-enable
+        // use_resource_processor: true when better supported by soong
+        ":SystemUIRobo-stub{.aapt.srcjar}",
+    ],
+    static_libs: [
+        "SystemUI-core",
+        "SystemUI-res",
+        "SystemUI-tests-base",
+        "androidx.test.uiautomator_uiautomator",
+        "androidx.core_core-animation-testing",
+        "androidx.test.ext.junit",
+        "kosmos",
+        "kotlin-test",
+        "mockito-kotlin-nodeps",
+        "androidx.compose.runtime_runtime",
+    ],
+    libs: [
+        "android.test.runner",
+        "android.test.base",
+        "android.test.mock",
+    ],
+    auto_gen_config: true,
+    plugins: [
+        "dagger2-compiler",
+    ],
+}
 
 // Opt-out config for optimizing the SystemUI target using R8.
 // Disabled via `export SYSTEMUI_OPTIMIZE_JAVA=false`, or explicitly in Make via
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index f98b29a..157af7d 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -92,6 +92,7 @@
     <uses-permission android:name="android.permission.MOUNT_UNMOUNT_FILESYSTEMS" />
     <uses-permission android:name="android.permission.MASTER_CLEAR" />
     <uses-permission android:name="android.permission.VIBRATE" />
+    <uses-permission android:name="android.permission.VIBRATE_SYSTEM_CONSTANTS" />
     <uses-permission android:name="android.permission.MANAGE_SENSOR_PRIVACY" />
     <uses-permission android:name="android.permission.OBSERVE_SENSOR_PRIVACY" />
     <uses-permission android:name="android.permission.ACCESS_AMBIENT_CONTEXT_EVENT" />
@@ -370,6 +371,9 @@
 
     <uses-permission android:name="android.permission.MONITOR_STICKY_MODIFIER_STATE" />
 
+    <!-- Listen to keyboard shortcut events from input manager -->
+    <uses-permission android:name="android.permission.MANAGE_KEY_GESTURES" />
+
     <!-- To follow the grammatical gender preference -->
     <uses-permission android:name="android.permission.READ_SYSTEM_GRAMMATICAL_GENDER" />
 
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/Android.bp b/packages/SystemUI/accessibility/accessibilitymenu/Android.bp
index c60eb61..fb1f715 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/Android.bp
+++ b/packages/SystemUI/accessibility/accessibilitymenu/Android.bp
@@ -41,6 +41,7 @@
         "SettingsLibDisplayUtils",
         "SettingsLibSettingsTheme",
         "com_android_a11y_menu_flags_lib",
+        "//frameworks/libs/systemui:view_capture",
     ],
 
     optimize: {
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuAdapter.java b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuAdapter.java
index c2cf6e1..c333a7a 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuAdapter.java
+++ b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuAdapter.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.accessibility.accessibilitymenu.view;
 
-import android.content.Context;
 import android.graphics.Rect;
 import android.view.LayoutInflater;
 import android.view.TouchDelegate;
@@ -43,16 +42,14 @@
     private final int mLargeTextSize;
 
     private final AccessibilityMenuService mService;
-    private final LayoutInflater mInflater;
     private final List<A11yMenuShortcut> mShortcutDataList;
     private final ShortcutDrawableUtils mShortcutDrawableUtils;
 
     public A11yMenuAdapter(
             AccessibilityMenuService service,
-            Context displayContext, List<A11yMenuShortcut> shortcutDataList) {
+            List<A11yMenuShortcut> shortcutDataList) {
         this.mService = service;
         this.mShortcutDataList = shortcutDataList;
-        mInflater = LayoutInflater.from(displayContext);
 
         mShortcutDrawableUtils = new ShortcutDrawableUtils(service);
 
@@ -78,7 +75,8 @@
     @Override
     public View getView(int position, View convertView, ViewGroup parent) {
         if (convertView == null) {
-            convertView = mInflater.inflate(R.layout.grid_item, parent, false);
+            convertView = LayoutInflater.from(parent.getContext())
+                    .inflate(R.layout.grid_item, parent, false);
 
             configureShortcutSize(convertView,
                     A11yMenuPreferenceFragment.isLargeButtonsEnabled(mService));
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuOverlayLayout.java b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuOverlayLayout.java
index de3c472..3db61a5 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuOverlayLayout.java
+++ b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuOverlayLayout.java
@@ -22,6 +22,8 @@
 import static android.view.View.ACCESSIBILITY_LIVE_REGION_POLITE;
 import static android.view.WindowManager.LayoutParams.TYPE_ACCESSIBILITY_OVERLAY;
 
+import static com.android.app.viewcapture.ViewCaptureFactory.getViewCaptureAwareWindowManagerInstance;
+
 import static java.lang.Math.max;
 
 import android.animation.Animator;
@@ -51,7 +53,9 @@
 import android.widget.TextView;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.UiContext;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.systemui.accessibility.accessibilitymenu.AccessibilityMenuService;
 import com.android.systemui.accessibility.accessibilitymenu.Flags;
 import com.android.systemui.accessibility.accessibilitymenu.R;
@@ -101,7 +105,6 @@
     };
 
     private final AccessibilityMenuService mService;
-    private final WindowManager mWindowManager;
     private final DisplayManager mDisplayManager;
     private ViewGroup mLayout;
     private WindowManager.LayoutParams mLayoutParameter;
@@ -111,7 +114,6 @@
 
     public A11yMenuOverlayLayout(AccessibilityMenuService service) {
         mService = service;
-        mWindowManager = mService.getSystemService(WindowManager.class);
         mDisplayManager = mService.getSystemService(DisplayManager.class);
         configureLayout();
         mHandler = new Handler(Looper.getMainLooper());
@@ -134,8 +136,7 @@
         int lastVisibilityState = View.GONE;
         if (mLayout != null) {
             lastVisibilityState = mLayout.getVisibility();
-            mWindowManager.removeView(mLayout);
-            mLayout = null;
+            clearLayout();
         }
 
         if (mLayoutParameter == null) {
@@ -143,14 +144,17 @@
         }
 
         final Display display = mDisplayManager.getDisplay(DEFAULT_DISPLAY);
-        final Context context = mService.createDisplayContext(display).createWindowContext(
-                TYPE_ACCESSIBILITY_OVERLAY, null);
-        mLayout = new A11yMenuFrameLayout(context);
-        updateLayoutPosition();
-        inflateLayoutAndSetOnTouchListener(mLayout, context);
-        mA11yMenuViewPager = new A11yMenuViewPager(mService, context);
+        final Context uiContext = mService.createWindowContext(
+                display, TYPE_ACCESSIBILITY_OVERLAY, /* options= */null);
+        final ViewCaptureAwareWindowManager windowManager =
+                getViewCaptureAwareWindowManagerInstance(uiContext,
+                        com.android.systemui.Flags.enableViewCaptureTracing());
+        mLayout = new A11yMenuFrameLayout(uiContext);
+        updateLayoutPosition(uiContext);
+        inflateLayoutAndSetOnTouchListener(mLayout, uiContext);
+        mA11yMenuViewPager = new A11yMenuViewPager(mService);
         mA11yMenuViewPager.configureViewPagerAndFooter(mLayout, createShortcutList(), pageIndex);
-        mWindowManager.addView(mLayout, mLayoutParameter);
+        windowManager.addView(mLayout, mLayoutParameter);
         mLayout.setVisibility(lastVisibilityState);
         mA11yMenuViewPager.updateFooterState();
 
@@ -159,7 +163,11 @@
 
     public void clearLayout() {
         if (mLayout != null) {
-            mWindowManager.removeView(mLayout);
+            ViewCaptureAwareWindowManager windowManager = getViewCaptureAwareWindowManagerInstance(
+                    mLayout.getContext(), com.android.systemui.Flags.enableViewCaptureTracing());
+            if (windowManager != null) {
+                windowManager.removeView(mLayout);
+            }
             mLayout.setOnTouchListener(null);
             mLayout = null;
         }
@@ -170,8 +178,11 @@
         if (mLayout == null || mLayoutParameter == null) {
             return;
         }
-        updateLayoutPosition();
-        mWindowManager.updateViewLayout(mLayout, mLayoutParameter);
+        updateLayoutPosition(mLayout.getContext());
+        WindowManager windowManager = mLayout.getContext().getSystemService(WindowManager.class);
+        if (windowManager != null) {
+            windowManager.updateViewLayout(mLayout, mLayoutParameter);
+        }
     }
 
     private void initLayoutParams() {
@@ -183,8 +194,8 @@
         mLayoutParameter.setTitle(mService.getString(R.string.accessibility_menu_service_name));
     }
 
-    private void inflateLayoutAndSetOnTouchListener(ViewGroup view, Context displayContext) {
-        LayoutInflater inflater = LayoutInflater.from(displayContext);
+    private void inflateLayoutAndSetOnTouchListener(ViewGroup view, @UiContext Context uiContext) {
+        LayoutInflater inflater = LayoutInflater.from(uiContext);
         inflater.inflate(R.layout.paged_menu, view);
         view.setOnTouchListener(mService);
     }
@@ -238,7 +249,11 @@
     }
 
     /** Updates a11y menu layout position by configuring layout params. */
-    private void updateLayoutPosition() {
+    private void updateLayoutPosition(@UiContext @NonNull Context uiContext) {
+        WindowManager windowManager = uiContext.getSystemService(WindowManager.class);
+        if (windowManager == null) {
+            return;
+        }
         final Display display = mDisplayManager.getDisplay(Display.DEFAULT_DISPLAY);
         final Configuration configuration = mService.getResources().getConfiguration();
         final int orientation = configuration.orientation;
@@ -276,14 +291,13 @@
             mLayoutParameter.height = WindowManager.LayoutParams.WRAP_CONTENT;
             mLayout.setBackgroundResource(R.drawable.shadow_0deg);
         }
-
         // Adjusts the y position of a11y menu layout to make the layout not to overlap bottom
         // navigation bar window.
-        updateLayoutByWindowInsetsIfNeeded();
+        updateLayoutByWindowInsetsIfNeeded(windowManager);
         mLayout.setOnApplyWindowInsetsListener(
                 (view, insets) -> {
-                    if (updateLayoutByWindowInsetsIfNeeded()) {
-                        mWindowManager.updateViewLayout(mLayout, mLayoutParameter);
+                    if (updateLayoutByWindowInsetsIfNeeded(windowManager)) {
+                        windowManager.updateViewLayout(mLayout, mLayoutParameter);
                     }
                     return view.onApplyWindowInsets(insets);
                 });
@@ -295,9 +309,9 @@
      * This method adjusts the layout position and size to
      * make a11y menu not to overlap navigation bar window.
      */
-    private boolean updateLayoutByWindowInsetsIfNeeded() {
+    private boolean updateLayoutByWindowInsetsIfNeeded(@NonNull WindowManager windowManager) {
         boolean shouldUpdateLayout = false;
-        WindowMetrics windowMetrics = mWindowManager.getCurrentWindowMetrics();
+        WindowMetrics windowMetrics = windowManager.getCurrentWindowMetrics();
         Insets windowInsets = windowMetrics.getWindowInsets().getInsetsIgnoringVisibility(
                 WindowInsets.Type.systemBars() | WindowInsets.Type.displayCutout());
         int xOffset = max(windowInsets.left, windowInsets.right);
@@ -396,7 +410,7 @@
     }
 
     private class A11yMenuFrameLayout extends FrameLayout {
-        A11yMenuFrameLayout(@NonNull Context context) {
+        A11yMenuFrameLayout(@UiContext @NonNull Context context) {
             super(context);
         }
 
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuViewPager.java b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuViewPager.java
index 08bbf19..a29ce82 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuViewPager.java
+++ b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/A11yMenuViewPager.java
@@ -146,12 +146,8 @@
     /** The container layout for a11y menu. */
     private ViewGroup mA11yMenuLayout;
 
-    /** Display context for inflating views. */
-    private Context mDisplayContext;
-
-    public A11yMenuViewPager(AccessibilityMenuService service, Context displayContext) {
+    public A11yMenuViewPager(AccessibilityMenuService service) {
         this.mService = service;
-        this.mDisplayContext = displayContext;
     }
 
     /**
@@ -289,7 +285,8 @@
             footerLayout.getLayoutParams().height =
                     (int) (footerLayout.getLayoutParams().height / densityScale);
             // Adjust the view pager height for system bar and display cutout insets.
-            WindowManager windowManager = mService.getSystemService(WindowManager.class);
+            WindowManager windowManager = mA11yMenuLayout.getContext()
+                    .getSystemService(WindowManager.class);
             WindowMetrics windowMetric = windowManager.getCurrentWindowMetrics();
             Insets windowInsets = windowMetric.getWindowInsets().getInsetsIgnoringVisibility(
                     WindowInsets.Type.systemBars() | WindowInsets.Type.displayCutout());
diff --git a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/ViewPagerAdapter.java b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/ViewPagerAdapter.java
index 43ec956..152ff68 100644
--- a/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/ViewPagerAdapter.java
+++ b/packages/SystemUI/accessibility/accessibilitymenu/src/com/android/systemui/accessibility/accessibilitymenu/view/ViewPagerAdapter.java
@@ -57,7 +57,7 @@
     @Override
     public void onBindViewHolder(RecyclerView.ViewHolder holder, int position) {
         A11yMenuAdapter adapter = new A11yMenuAdapter(
-                mService, holder.itemView.getContext(), mShortcutList.get(position));
+                mService, mShortcutList.get(position));
         GridView gridView = (GridView) holder.itemView;
         gridView.setNumColumns(A11yMenuViewPager.GridViewParams.getGridColumnCount(mService));
         gridView.setAdapter(adapter);
diff --git a/packages/SystemUI/aconfig/biometrics_framework.aconfig b/packages/SystemUI/aconfig/biometrics_framework.aconfig
index 95e4b59..10d7352 100644
--- a/packages/SystemUI/aconfig/biometrics_framework.aconfig
+++ b/packages/SystemUI/aconfig/biometrics_framework.aconfig
@@ -3,3 +3,12 @@
 
 # NOTE: Keep alphabetized to help limit merge conflicts from multiple simultaneous editors.
 
+flag {
+  name: "bp_icon_a11y"
+  namespace: "biometrics_framework"
+  description: "Fixes biometric prompt icon not working as button with a11y"
+  bug: "359423579"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 8a1d81b..cf13621 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -139,13 +139,6 @@
 }
 
 flag {
-    name: "notifications_heads_up_refactor"
-    namespace: "systemui"
-    description: "Use HeadsUpInteractor to feed HUN updates to the NSSL."
-    bug: "325936094"
-}
-
-flag {
    name: "notification_transparent_header_fix"
    namespace: "systemui"
    description: "fix the transparent group header issue for async header inflation."
@@ -370,6 +363,13 @@
 }
 
 flag {
+   name: "status_bar_signal_policy_refactor"
+   namespace: "systemui"
+   description: "Use a settings observer for airplane mode and make StatusBarSignalPolicy startable"
+   bug: "264539100"
+}
+
+flag {
     name: "status_bar_swipe_over_chip"
     namespace: "systemui"
     description: "Allow users to swipe over the status bar chip to open the shade"
@@ -380,6 +380,17 @@
 }
 
 flag {
+    name: "status_bar_always_check_underlying_networks"
+    namespace: "systemui"
+    description: "For status bar connectivity UI, always check underlying networks for wifi and "
+        "carrier merged information, regardless of the sepcified transport type"
+    bug: "352162710"
+    metadata {
+      purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "status_bar_stop_updating_window_height"
     namespace: "systemui"
     description: "Don't have PhoneStatusBarView manually trigger an update of the height in "
@@ -391,6 +402,13 @@
 }
 
 flag {
+    name: "status_bar_ron_chips"
+    namespace: "systemui"
+    description: "Show rich ongoing notifications as chips in the status bar"
+    bug: "361346412"
+}
+
+flag {
     name: "compose_bouncer"
     namespace: "systemui"
     description: "Use the new compose bouncer in SystemUI"
@@ -588,16 +606,6 @@
 }
 
 flag {
-    name: "screenshot_private_profile_behavior_fix"
-    namespace: "systemui"
-    description: "Private profile support for screenshots"
-    bug: "327613051"
-    metadata {
-        purpose: PURPOSE_BUGFIX
-    }
-}
-
-flag {
     name: "screenshot_save_image_exporter"
     namespace: "systemui"
     description: "Save all screenshots using ImageExporter"
@@ -1348,13 +1356,6 @@
 }
 
 flag {
-    name: "new_picker_ui"
-    namespace: "systemui"
-    description: "Enables the BC25 design of the customization picker UI."
-    bug: "339081035"
-}
-
-flag {
   namespace: "systemui"
   name: "settings_ext_register_content_observer_on_bg_thread"
   description: "Register content observer in callback flow APIs on background thread in SettingsProxyExt."
@@ -1365,6 +1366,16 @@
 }
 
 flag {
+    name: "notify_password_text_view_user_activity_in_background"
+    namespace: "systemui"
+    description: "Decide whether to notify the user activity in password text view, to power manager in the background thread."
+    bug: "346882515"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "face_message_defer_update"
     namespace: "systemui"
     description: "Only analyze the last n frames when determining whether to defer a face auth help message like low light"
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/BouncerSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/BouncerSceneModule.kt
index 2b1268e..5b368df 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/BouncerSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/BouncerSceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.bouncer.ui.composable.BouncerScene
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/CommunalSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/CommunalSceneModule.kt
index 94b5db2..74ce4bb 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/CommunalSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/CommunalSceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.communal.ui.compose.CommunalScene
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/GoneSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/GoneSceneModule.kt
index bc3fef1..871ade9 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/GoneSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/GoneSceneModule.kt
@@ -16,8 +16,8 @@
 
 package com.android.systemui.scene
 
-import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.ui.composable.GoneScene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/LockscreenSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/LockscreenSceneModule.kt
index 72965fb..bfeaf92 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/LockscreenSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/LockscreenSceneModule.kt
@@ -27,7 +27,7 @@
 import com.android.systemui.keyguard.ui.composable.LockscreenSceneBlueprintModule
 import com.android.systemui.keyguard.ui.composable.blueprint.ComposableLockscreenSceneBlueprint
 import com.android.systemui.keyguard.ui.viewmodel.LockscreenContentViewModel
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.Provides
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeOverlayModule.kt
similarity index 61%
copy from packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt
copy to packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeOverlayModule.kt
index 7e09a10..e55520a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeOverlayModule.kt
@@ -14,16 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyguard
+package com.android.systemui.scene
 
-import com.android.systemui.Flags
+import com.android.systemui.notifications.ui.composable.NotificationsShadeOverlay
+import com.android.systemui.scene.ui.composable.Overlay
+import dagger.Binds
+import dagger.Module
+import dagger.multibindings.IntoSet
 
-/** Helper for reading or using the new picker UI flag. */
-@Suppress("NOTHING_TO_INLINE")
-object NewPickerUiKeyguardPreview {
+@Module
+interface NotificationsShadeOverlayModule {
 
-    /** Is the new picker UI enabled */
-    @JvmStatic
-    inline val isEnabled
-        get() = Flags.newPickerUi()
+    @Binds @IntoSet fun notificationsShade(overlay: NotificationsShadeOverlay): Overlay
 }
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt
index 9b736b8..c58df35 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/NotificationsShadeSceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.notifications.ui.composable.NotificationsShadeScene
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsSceneModule.kt
index ee1f525..d55210d 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsSceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.qs.ui.composable.QuickSettingsScene
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeOverlayModule.kt
similarity index 61%
rename from packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt
rename to packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeOverlayModule.kt
index 7e09a10..bc4adf9 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/NewPickerUiKeyguardPreview.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeOverlayModule.kt
@@ -14,16 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.systemui.keyguard
+package com.android.systemui.scene
 
-import com.android.systemui.Flags
+import com.android.systemui.qs.ui.composable.QuickSettingsShadeOverlay
+import com.android.systemui.scene.ui.composable.Overlay
+import dagger.Binds
+import dagger.Module
+import dagger.multibindings.IntoSet
 
-/** Helper for reading or using the new picker UI flag. */
-@Suppress("NOTHING_TO_INLINE")
-object NewPickerUiKeyguardPreview {
+@Module
+interface QuickSettingsShadeOverlayModule {
 
-    /** Is the new picker UI enabled */
-    @JvmStatic
-    inline val isEnabled
-        get() = Flags.newPickerUi()
+    @Binds @IntoSet fun quickSettingsShade(overlay: QuickSettingsShadeOverlay): Overlay
 }
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt
index 3d7401d..5bb6ae4 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/QuickSettingsShadeSceneModule.kt
@@ -17,7 +17,7 @@
 package com.android.systemui.scene
 
 import com.android.systemui.qs.ui.composable.QuickSettingsShadeScene
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Binds
 import dagger.Module
 import dagger.multibindings.IntoSet
diff --git a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ShadeSceneModule.kt b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ShadeSceneModule.kt
index c655d6b..186914f 100644
--- a/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ShadeSceneModule.kt
+++ b/packages/SystemUI/compose/facade/enabled/src/com/android/systemui/scene/ShadeSceneModule.kt
@@ -16,7 +16,7 @@
 
 package com.android.systemui.scene
 
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.shade.ui.composable.ShadeScene
 import dagger.Binds
 import dagger.Module
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
index d164eab..9f78d69 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/bouncer/ui/composable/BouncerContent.kt
@@ -528,7 +528,7 @@
     // Update state whenever currentSceneKey has changed.
     LaunchedEffect(state, currentSceneKey) {
         if (currentSceneKey != state.transitionState.currentScene) {
-            state.setTargetScene(currentSceneKey, coroutineScope = this)
+            state.setTargetScene(currentSceneKey, animationScope = this)
         }
     }
 
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 5a4020d..7fb88e8 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
@@ -27,14 +27,14 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.bouncer.ui.BouncerDialogFactory
-import com.android.systemui.bouncer.ui.viewmodel.BouncerSceneActionsViewModel
 import com.android.systemui.bouncer.ui.viewmodel.BouncerSceneContentViewModel
+import com.android.systemui.bouncer.ui.viewmodel.BouncerUserActionsViewModel
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.rememberViewModel
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 
@@ -54,18 +54,17 @@
 class BouncerScene
 @Inject
 constructor(
-    private val actionsViewModelFactory: BouncerSceneActionsViewModel.Factory,
+    private val actionsViewModelFactory: BouncerUserActionsViewModel.Factory,
     private val contentViewModelFactory: BouncerSceneContentViewModel.Factory,
     private val dialogFactory: BouncerDialogFactory,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
     override val key = Scenes.Bouncer
 
-    private val actionsViewModel: BouncerSceneActionsViewModel by lazy {
+    private val actionsViewModel: BouncerUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -76,7 +75,7 @@
         modifier: Modifier,
     ) =
         BouncerScene(
-            viewModel = rememberViewModel { contentViewModelFactory.create() },
+            viewModel = rememberViewModel("BouncerScene") { contentViewModelFactory.create() },
             dialogFactory = dialogFactory,
             modifier = modifier,
         )
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
index b0590e0..be6a0f9 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalHub.kt
@@ -567,16 +567,8 @@
         // Do nothing if there is no new live content
         val indexOfFirstUpdatedContent =
             newLiveContentKeys.indexOfFirst { !prevLiveContentKeys.contains(it) }
-        if (indexOfFirstUpdatedContent < 0) {
-            return@LaunchedEffect
-        }
-
-        // Scroll if the live content is not visible
-        val lastVisibleItemIndex = gridState.layoutInfo.visibleItemsInfo.lastOrNull()?.index
-        if (lastVisibleItemIndex != null && indexOfFirstUpdatedContent > lastVisibleItemIndex) {
-            // Launching with a scope to prevent the job from being canceled in the case of a
-            // recomposition during scrolling
-            coroutineScope.launch { gridState.animateScrollToItem(indexOfFirstUpdatedContent) }
+        if (indexOfFirstUpdatedContent in 0 until gridState.firstVisibleItemIndex) {
+            gridState.scrollToItem(indexOfFirstUpdatedContent)
         }
     }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
index 54ffcf4..8b6de6a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/communal/ui/compose/CommunalScene.kt
@@ -29,7 +29,7 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.statusbar.phone.SystemUIDialogFactory
 import javax.inject.Inject
 import kotlinx.coroutines.awaitCancellation
@@ -46,13 +46,13 @@
     private val dialogFactory: SystemUIDialogFactory,
     private val interactionHandler: WidgetInteractionHandler,
     private val widgetSection: CommunalAppWidgetSection,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
     override val key = Scenes.Communal
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        MutableStateFlow<Map<UserAction, UserActionResult>>(
+    override val userActions: Flow<Map<UserAction, UserActionResult>> =
+        MutableStateFlow(
                 mapOf(
-                    Swipe(SwipeDirection.End) to UserActionResult(Scenes.Lockscreen),
+                    Swipe(SwipeDirection.End) to Scenes.Lockscreen,
                 )
             )
             .asStateFlow()
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/AlternateBouncer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/AlternateBouncer.kt
index 04bcc36..c60e11e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/AlternateBouncer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/AlternateBouncer.kt
@@ -73,8 +73,7 @@
             initialValue = null
         )
 
-    // TODO (b/353955910): back handling doesn't work
-    BackHandler { alternateBouncerDependencies.viewModel.onBackRequested() }
+    BackHandler(enabled = isVisible) { alternateBouncerDependencies.viewModel.onBackRequested() }
 
     AnimatedVisibility(
         visible = isVisible,
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
index 672b8a7..dbe7538 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenContent.kt
@@ -50,7 +50,7 @@
     fun SceneScope.Content(
         modifier: Modifier = Modifier,
     ) {
-        val viewModel = rememberViewModel { viewModelFactory.create() }
+        val viewModel = rememberViewModel("LockscreenContent") { viewModelFactory.create() }
         val isContentVisible: Boolean by viewModel.isContentVisible.collectAsStateWithLifecycle()
         if (!isContentVisible) {
             // If the content isn't supposed to be visible, show a large empty box as it's needed
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
index 2029e9e..c7c29f9 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/LockscreenScene.kt
@@ -22,13 +22,13 @@
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
-import com.android.compose.animation.scene.animateSceneFloatAsState
+import com.android.compose.animation.scene.animateContentFloatAsState
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.keyguard.ui.viewmodel.LockscreenSceneActionsViewModel
+import com.android.systemui.keyguard.ui.viewmodel.LockscreenUserActionsViewModel
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.qs.ui.composable.QuickSettings
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Lazy
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
@@ -38,17 +38,16 @@
 class LockscreenScene
 @Inject
 constructor(
-    actionsViewModelFactory: LockscreenSceneActionsViewModel.Factory,
+    actionsViewModelFactory: LockscreenUserActionsViewModel.Factory,
     private val lockscreenContent: Lazy<LockscreenContent>,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
     override val key = Scenes.Lockscreen
 
-    private val actionsViewModel: LockscreenSceneActionsViewModel by lazy {
+    private val actionsViewModel: LockscreenUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -70,7 +69,7 @@
     lockscreenContent: Lazy<LockscreenContent>,
     modifier: Modifier = Modifier,
 ) {
-    animateSceneFloatAsState(
+    animateContentFloatAsState(
         value = QuickSettings.SharedValues.SquishinessValues.LockscreenSceneStarting,
         key = QuickSettings.SharedValues.TilesSquishiness,
     )
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
index 5f4dc6e..18e1092 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/NotificationSection.kt
@@ -83,7 +83,7 @@
     fun SceneScope.HeadsUpNotifications() {
         SnoozeableHeadsUpNotificationSpace(
             stackScrollView = stackScrollView.get(),
-            viewModel = rememberViewModel { viewModelFactory.create() },
+            viewModel = rememberViewModel("HeadsUpNotifications") { viewModelFactory.create() },
         )
     }
 
@@ -107,7 +107,7 @@
 
         ConstrainedNotificationStack(
             stackScrollView = stackScrollView.get(),
-            viewModel = rememberViewModel { viewModelFactory.create() },
+            viewModel = rememberViewModel("Notifications") { viewModelFactory.create() },
             modifier =
                 modifier
                     .fillMaxWidth()
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
index 0eeb79b..97d89a2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/TopAreaSection.kt
@@ -93,7 +93,7 @@
         // Update state whenever currentSceneKey has changed.
         LaunchedEffect(state, currentScene) {
             if (currentScene != state.transitionState.currentScene) {
-                state.setTargetScene(currentScene, coroutineScope = this)
+                state.setTargetScene(currentScene, animationScope = this)
             }
         }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
index 808e666..5f7b1ad 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/media/controls/ui/composable/MediaCarousel.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.media.controls.ui.view.MediaHost
 import com.android.systemui.res.R
 import com.android.systemui.util.animation.MeasurementInput
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 object MediaCarousel {
     object Elements {
@@ -46,6 +47,7 @@
     }
 }
 
+@ExperimentalCoroutinesApi
 @Composable
 fun SceneScope.MediaCarousel(
     isVisible: Boolean,
@@ -54,7 +56,7 @@
     carouselController: MediaCarouselController,
     offsetProvider: (() -> IntOffset)? = null,
 ) {
-    if (!isVisible) {
+    if (!isVisible || carouselController.isLockedAndHidden()) {
         return
     }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
index 9e292d0..a2beba8 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/Notifications.kt
@@ -321,7 +321,9 @@
     val minScrimOffset: () -> Float = { minScrimTop - maxScrimTop() }
 
     // The height of the scrim visible on screen when it is in its resting (collapsed) state.
-    val minVisibleScrimHeight: () -> Float = { screenHeight - maxScrimTop() }
+    val minVisibleScrimHeight: () -> Float = {
+        screenHeight - maxScrimTop() - with(density) { navBarHeight.toPx() }
+    }
 
     // we are not scrolled to the top unless the scrim is at its maximum offset.
     LaunchedEffect(viewModel, scrimOffset) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeOverlay.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeOverlay.kt
new file mode 100644
index 0000000..a22becc
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeOverlay.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.notifications.ui.composable
+
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.padding
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.unit.dp
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import com.android.systemui.battery.BatteryMeterViewController
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.lifecycle.rememberViewModel
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeOverlayActionsViewModel
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeOverlayContentViewModel
+import com.android.systemui.scene.session.ui.composable.SaveableSession
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.shade.shared.model.ShadeMode
+import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
+import com.android.systemui.shade.ui.composable.OverlayShade
+import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
+import com.android.systemui.statusbar.phone.ui.StatusBarIconController
+import com.android.systemui.statusbar.phone.ui.TintedIconManager
+import dagger.Lazy
+import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
+
+@SysUISingleton
+class NotificationsShadeOverlay
+@Inject
+constructor(
+    private val actionsViewModelFactory: NotificationsShadeOverlayActionsViewModel.Factory,
+    private val contentViewModelFactory: NotificationsShadeOverlayContentViewModel.Factory,
+    private val tintedIconManagerFactory: TintedIconManager.Factory,
+    private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
+    private val statusBarIconController: StatusBarIconController,
+    private val shadeSession: SaveableSession,
+    private val stackScrollView: Lazy<NotificationScrollView>,
+) : Overlay {
+
+    override val key = Overlays.NotificationsShade
+
+    private val actionsViewModel: NotificationsShadeOverlayActionsViewModel by lazy {
+        actionsViewModelFactory.create()
+    }
+
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
+
+    override suspend fun activate(): Nothing {
+        actionsViewModel.activate()
+    }
+
+    @Composable
+    override fun ContentScope.Content(
+        modifier: Modifier,
+    ) {
+        val viewModel =
+            rememberViewModel("NotificationsShadeOverlay-viewModel") {
+                contentViewModelFactory.create()
+            }
+        val placeholderViewModel =
+            rememberViewModel("NotificationsShadeOverlay-notifPlaceholderViewModel") {
+                viewModel.notificationsPlaceholderViewModelFactory.create()
+            }
+
+        OverlayShade(
+            modifier = modifier,
+            onScrimClicked = viewModel::onScrimClicked,
+        ) {
+            Column {
+                ExpandedShadeHeader(
+                    viewModelFactory = viewModel.shadeHeaderViewModelFactory,
+                    createTintedIconManager = tintedIconManagerFactory::create,
+                    createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
+                    statusBarIconController = statusBarIconController,
+                    modifier = Modifier.padding(horizontal = 16.dp),
+                )
+
+                NotificationScrollingStack(
+                    shadeSession = shadeSession,
+                    stackScrollView = stackScrollView.get(),
+                    viewModel = placeholderViewModel,
+                    maxScrimTop = { 0f },
+                    shouldPunchHoleBehindScrim = false,
+                    shouldFillMaxSize = false,
+                    shouldReserveSpaceForNavBar = false,
+                    shadeMode = ShadeMode.Dual,
+                    modifier = Modifier.fillMaxWidth(),
+                )
+
+                // Communicates the bottom position of the drawable area within the shade to NSSL.
+                NotificationStackCutoffGuideline(
+                    stackScrollView = stackScrollView.get(),
+                    viewModel = placeholderViewModel,
+                )
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt
index 62213bd..1f4cd04 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/notifications/ui/composable/NotificationsShadeScene.kt
@@ -27,24 +27,21 @@
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.battery.BatteryMeterViewController
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.keyguard.ui.composable.LockscreenContent
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.rememberViewModel
-import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeSceneActionsViewModel
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeUserActionsViewModel
 import com.android.systemui.scene.session.ui.composable.SaveableSession
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
 import com.android.systemui.shade.ui.composable.OverlayShade
-import com.android.systemui.shade.ui.viewmodel.OverlayShadeViewModel
 import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
 import com.android.systemui.statusbar.phone.ui.StatusBarIconController
 import com.android.systemui.statusbar.phone.ui.TintedIconManager
 import dagger.Lazy
-import java.util.Optional
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 
@@ -52,8 +49,7 @@
 class NotificationsShadeScene
 @Inject
 constructor(
-    private val actionsViewModelFactory: NotificationsShadeSceneActionsViewModel.Factory,
-    private val overlayShadeViewModelFactory: OverlayShadeViewModel.Factory,
+    private val actionsViewModelFactory: NotificationsShadeUserActionsViewModel.Factory,
     private val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
     private val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
     private val tintedIconManagerFactory: TintedIconManager.Factory,
@@ -61,17 +57,15 @@
     private val statusBarIconController: StatusBarIconController,
     private val shadeSession: SaveableSession,
     private val stackScrollView: Lazy<NotificationScrollView>,
-    private val lockscreenContent: Lazy<Optional<LockscreenContent>>,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
 
     override val key = Scenes.NotificationsShade
 
-    private val actionsViewModel: NotificationsShadeSceneActionsViewModel by lazy {
+    private val actionsViewModel: NotificationsShadeUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -81,14 +75,14 @@
     override fun SceneScope.Content(
         modifier: Modifier,
     ) {
-        val notificationsPlaceholderViewModel = rememberViewModel {
-            notificationsPlaceholderViewModelFactory.create()
-        }
+        val notificationsPlaceholderViewModel =
+            rememberViewModel("NotificationsShadeScene") {
+                notificationsPlaceholderViewModelFactory.create()
+            }
 
         OverlayShade(
             modifier = modifier,
-            viewModelFactory = overlayShadeViewModelFactory,
-            lockscreenContent = lockscreenContent,
+            onScrimClicked = {},
         ) {
             Column {
                 ExpandedShadeHeader(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
index 1921624..671b012 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettings.kt
@@ -98,7 +98,7 @@
                 else -> QSSceneAdapter.State.CLOSED
             }
         }
-        is TransitionState.Transition.ChangeCurrentScene ->
+        is TransitionState.Transition.ChangeScene ->
             with(transitionState) {
                 when {
                     isSplitShade -> UnsquishingQS(squishiness)
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
index f11f8bb..d34295e 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsScene.kt
@@ -93,12 +93,12 @@
 import com.android.systemui.qs.footer.ui.compose.FooterActionsWithAnimatedVisibility
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaLandscapeTopOffset
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaOffset.InQS
-import com.android.systemui.qs.ui.viewmodel.QuickSettingsSceneActionsViewModel
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsSceneContentViewModel
+import com.android.systemui.qs.ui.viewmodel.QuickSettingsUserActionsViewModel
 import com.android.systemui.res.R
 import com.android.systemui.scene.session.ui.composable.SaveableSession
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.shade.ui.composable.CollapsedShadeHeader
 import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
@@ -113,9 +113,11 @@
 import javax.inject.Inject
 import javax.inject.Named
 import kotlin.math.roundToInt
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 
 /** The Quick Settings (AKA "QS") scene shows the quick setting tiles. */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class QuickSettingsScene
 @Inject
@@ -123,22 +125,21 @@
     private val shadeSession: SaveableSession,
     private val notificationStackScrollView: Lazy<NotificationScrollView>,
     private val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
-    private val actionsViewModelFactory: QuickSettingsSceneActionsViewModel.Factory,
+    private val actionsViewModelFactory: QuickSettingsUserActionsViewModel.Factory,
     private val contentViewModelFactory: QuickSettingsSceneContentViewModel.Factory,
     private val tintedIconManagerFactory: TintedIconManager.Factory,
     private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
     private val statusBarIconController: StatusBarIconController,
     private val mediaCarouselController: MediaCarouselController,
     @Named(MediaModule.QS_PANEL) private val mediaHost: MediaHost,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
     override val key = Scenes.QuickSettings
 
-    private val actionsViewModel: QuickSettingsSceneActionsViewModel by lazy {
+    private val actionsViewModel: QuickSettingsUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -152,7 +153,9 @@
             notificationStackScrollView = notificationStackScrollView.get(),
             viewModelFactory = contentViewModelFactory,
             notificationsPlaceholderViewModel =
-                rememberViewModel { notificationsPlaceholderViewModelFactory.create() },
+                rememberViewModel("QuickSettingsScene-notifPlaceholderViewModel") {
+                    notificationsPlaceholderViewModelFactory.create()
+                },
             createTintedIconManager = tintedIconManagerFactory::create,
             createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
             statusBarIconController = statusBarIconController,
@@ -179,10 +182,11 @@
 ) {
     val cutoutLocation = LocalDisplayCutout.current.location
 
-    val viewModel = rememberViewModel { viewModelFactory.create() }
-    val brightnessMirrorViewModel = rememberViewModel {
-        viewModel.brightnessMirrorViewModelFactory.create()
-    }
+    val viewModel = rememberViewModel("QuickSettingsScene-viewModel") { viewModelFactory.create() }
+    val brightnessMirrorViewModel =
+        rememberViewModel("QuickSettingsScene-brightnessMirrorViewModel") {
+            viewModel.brightnessMirrorViewModelFactory.create()
+        }
     val brightnessMirrorShowing by brightnessMirrorViewModel.isShowing.collectAsStateWithLifecycle()
     val contentAlpha by
         animateFloatAsState(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt
new file mode 100644
index 0000000..f8d0588
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeOverlay.kt
@@ -0,0 +1,165 @@
+/*
+ * 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.qs.ui.composable
+
+import androidx.compose.animation.AnimatedContent
+import androidx.compose.animation.core.tween
+import androidx.compose.animation.fadeIn
+import androidx.compose.animation.fadeOut
+import androidx.compose.animation.togetherWith
+import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Column
+import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.height
+import androidx.compose.foundation.layout.heightIn
+import androidx.compose.foundation.layout.padding
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.ui.Alignment
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.unit.dp
+import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import com.android.systemui.battery.BatteryMeterViewController
+import com.android.systemui.brightness.ui.compose.BrightnessSliderContainer
+import com.android.systemui.compose.modifiers.sysuiResTag
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.lifecycle.rememberViewModel
+import com.android.systemui.qs.panels.ui.compose.EditMode
+import com.android.systemui.qs.panels.ui.compose.TileGrid
+import com.android.systemui.qs.ui.viewmodel.QuickSettingsContainerViewModel
+import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeOverlayActionsViewModel
+import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeOverlayContentViewModel
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
+import com.android.systemui.shade.ui.composable.OverlayShade
+import com.android.systemui.statusbar.phone.ui.StatusBarIconController
+import com.android.systemui.statusbar.phone.ui.TintedIconManager
+import javax.inject.Inject
+import kotlinx.coroutines.flow.Flow
+
+@SysUISingleton
+class QuickSettingsShadeOverlay
+@Inject
+constructor(
+    private val actionsViewModelFactory: QuickSettingsShadeOverlayActionsViewModel.Factory,
+    private val contentViewModelFactory: QuickSettingsShadeOverlayContentViewModel.Factory,
+    private val tintedIconManagerFactory: TintedIconManager.Factory,
+    private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
+    private val statusBarIconController: StatusBarIconController,
+) : Overlay {
+
+    override val key = Overlays.QuickSettingsShade
+
+    private val actionsViewModel: QuickSettingsShadeOverlayActionsViewModel by lazy {
+        actionsViewModelFactory.create()
+    }
+
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
+
+    override suspend fun activate(): Nothing {
+        actionsViewModel.activate()
+    }
+
+    @Composable
+    override fun ContentScope.Content(
+        modifier: Modifier,
+    ) {
+        val viewModel =
+            rememberViewModel("QuickSettingsShadeOverlay") { contentViewModelFactory.create() }
+
+        OverlayShade(
+            modifier = modifier,
+            onScrimClicked = viewModel::onScrimClicked,
+        ) {
+            Column {
+                ExpandedShadeHeader(
+                    viewModelFactory = viewModel.shadeHeaderViewModelFactory,
+                    createTintedIconManager = tintedIconManagerFactory::create,
+                    createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
+                    statusBarIconController = statusBarIconController,
+                    modifier = Modifier.padding(QuickSettingsShade.Dimensions.Padding),
+                )
+
+                ShadeBody(
+                    viewModel = viewModel.quickSettingsContainerViewModel,
+                )
+            }
+        }
+    }
+}
+
+@Composable
+fun ShadeBody(
+    viewModel: QuickSettingsContainerViewModel,
+) {
+    val isEditing by viewModel.editModeViewModel.isEditing.collectAsStateWithLifecycle()
+
+    AnimatedContent(
+        targetState = isEditing,
+        transitionSpec = { fadeIn(tween(500)) togetherWith fadeOut(tween(500)) }
+    ) { editing ->
+        if (editing) {
+            EditMode(
+                viewModel = viewModel.editModeViewModel,
+                modifier = Modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding)
+            )
+        } else {
+            QuickSettingsLayout(
+                viewModel = viewModel,
+                modifier = Modifier.sysuiResTag("quick_settings_panel")
+            )
+        }
+    }
+}
+
+@Composable
+private fun QuickSettingsLayout(
+    viewModel: QuickSettingsContainerViewModel,
+    modifier: Modifier = Modifier,
+) {
+    Column(
+        verticalArrangement = Arrangement.spacedBy(QuickSettingsShade.Dimensions.Padding),
+        horizontalAlignment = Alignment.CenterHorizontally,
+        modifier = modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding),
+    ) {
+        BrightnessSliderContainer(
+            viewModel = viewModel.brightnessSliderViewModel,
+            modifier =
+                Modifier.fillMaxWidth()
+                    .height(QuickSettingsShade.Dimensions.BrightnessSliderHeight),
+        )
+        TileGrid(
+            viewModel = viewModel.tileGridViewModel,
+            modifier =
+                Modifier.fillMaxWidth().heightIn(max = QuickSettingsShade.Dimensions.GridMaxHeight),
+            viewModel.editModeViewModel::startEditing,
+        )
+    }
+}
+
+object QuickSettingsShade {
+
+    object Dimensions {
+        val Padding = 16.dp
+        val BrightnessSliderHeight = 64.dp
+        val GridMaxHeight = 800.dp
+    }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt
index b25773b..e27c7e2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/qs/ui/composable/QuickSettingsShadeScene.kt
@@ -16,51 +16,26 @@
 
 package com.android.systemui.qs.ui.composable
 
-import androidx.compose.animation.AnimatedContent
-import androidx.compose.animation.EnterTransition
-import androidx.compose.animation.ExitTransition
-import androidx.compose.animation.core.tween
-import androidx.compose.animation.fadeIn
-import androidx.compose.animation.fadeOut
-import androidx.compose.animation.togetherWith
-import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Column
-import androidx.compose.foundation.layout.fillMaxWidth
-import androidx.compose.foundation.layout.height
-import androidx.compose.foundation.layout.heightIn
 import androidx.compose.foundation.layout.padding
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.unit.dp
-import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.battery.BatteryMeterViewController
-import com.android.systemui.brightness.ui.compose.BrightnessSliderContainer
-import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.keyguard.ui.composable.LockscreenContent
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.rememberViewModel
-import com.android.systemui.qs.panels.ui.compose.EditMode
-import com.android.systemui.qs.panels.ui.compose.TileGrid
-import com.android.systemui.qs.ui.composable.QuickSettingsShade.Transitions.QuickSettingsLayoutEnter
-import com.android.systemui.qs.ui.composable.QuickSettingsShade.Transitions.QuickSettingsLayoutExit
-import com.android.systemui.qs.ui.viewmodel.QuickSettingsContainerViewModel
-import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeSceneActionsViewModel
 import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeSceneContentViewModel
+import com.android.systemui.qs.ui.viewmodel.QuickSettingsShadeUserActionsViewModel
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.shade.ui.composable.ExpandedShadeHeader
 import com.android.systemui.shade.ui.composable.OverlayShade
 import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
 import com.android.systemui.statusbar.phone.ui.StatusBarIconController
 import com.android.systemui.statusbar.phone.ui.TintedIconManager
-import dagger.Lazy
-import java.util.Optional
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
 
@@ -68,23 +43,21 @@
 class QuickSettingsShadeScene
 @Inject
 constructor(
-    private val actionsViewModelFactory: QuickSettingsShadeSceneActionsViewModel.Factory,
+    private val actionsViewModelFactory: QuickSettingsShadeUserActionsViewModel.Factory,
     private val contentViewModelFactory: QuickSettingsShadeSceneContentViewModel.Factory,
-    private val lockscreenContent: Lazy<Optional<LockscreenContent>>,
     private val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
     private val tintedIconManagerFactory: TintedIconManager.Factory,
     private val batteryMeterViewControllerFactory: BatteryMeterViewController.Factory,
     private val statusBarIconController: StatusBarIconController,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
 
     override val key = Scenes.QuickSettingsShade
 
-    private val actionsViewModel: QuickSettingsShadeSceneActionsViewModel by lazy {
+    private val actionsViewModel: QuickSettingsShadeUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -94,11 +67,12 @@
     override fun SceneScope.Content(
         modifier: Modifier,
     ) {
-        val viewModel = rememberViewModel { contentViewModelFactory.create() }
+        val viewModel =
+            rememberViewModel("QuickSettingsShadeScene") { contentViewModelFactory.create() }
+
         OverlayShade(
-            viewModelFactory = viewModel.overlayShadeViewModelFactory,
-            lockscreenContent = lockscreenContent,
             modifier = modifier,
+            onScrimClicked = {},
         ) {
             Column {
                 ExpandedShadeHeader(
@@ -116,68 +90,3 @@
         }
     }
 }
-
-@Composable
-fun ShadeBody(
-    viewModel: QuickSettingsContainerViewModel,
-) {
-    val isEditing by viewModel.editModeViewModel.isEditing.collectAsStateWithLifecycle()
-
-    AnimatedContent(
-        targetState = isEditing,
-        transitionSpec = { QuickSettingsLayoutEnter togetherWith QuickSettingsLayoutExit }
-    ) { editing ->
-        if (editing) {
-            EditMode(
-                viewModel = viewModel.editModeViewModel,
-                modifier = Modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding)
-            )
-        } else {
-            QuickSettingsLayout(
-                viewModel = viewModel,
-                modifier = Modifier.sysuiResTag("quick_settings_panel")
-            )
-        }
-    }
-}
-
-@Composable
-private fun QuickSettingsLayout(
-    viewModel: QuickSettingsContainerViewModel,
-    modifier: Modifier = Modifier,
-) {
-    Column(
-        verticalArrangement = Arrangement.spacedBy(QuickSettingsShade.Dimensions.Padding),
-        horizontalAlignment = Alignment.CenterHorizontally,
-        modifier = modifier.fillMaxWidth().padding(QuickSettingsShade.Dimensions.Padding),
-    ) {
-        BrightnessSliderContainer(
-            viewModel = viewModel.brightnessSliderViewModel,
-            modifier =
-                Modifier.fillMaxWidth()
-                    .height(QuickSettingsShade.Dimensions.BrightnessSliderHeight),
-        )
-        TileGrid(
-            viewModel = viewModel.tileGridViewModel,
-            modifier =
-                Modifier.fillMaxWidth().heightIn(max = QuickSettingsShade.Dimensions.GridMaxHeight),
-            viewModel.editModeViewModel::startEditing,
-        )
-    }
-}
-
-object QuickSettingsShade {
-
-    object Dimensions {
-        val Padding = 16.dp
-        val BrightnessSliderHeight = 64.dp
-        val GridMaxHeight = 800.dp
-    }
-
-    object Transitions {
-        val QuickSettingsLayoutEnter: EnterTransition = fadeIn(tween(500))
-        val QuickSettingsLayoutExit: ExitTransition = fadeOut(tween(500))
-        val QuickSettingsEditorEnter: EnterTransition = fadeIn(tween(500))
-        val QuickSettingsEditorExit: ExitTransition = fadeOut(tween(500))
-    }
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ActionableContent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ActionableContent.kt
new file mode 100644
index 0000000..8fe6893
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ActionableContent.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.scene.ui.composable
+
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import kotlinx.coroutines.flow.Flow
+
+/** Defines interface for content that can respond to user-actions. */
+interface ActionableContent {
+    /**
+     * The mapping between [UserAction] and destination [UserActionResult]s.
+     *
+     * When the scene framework detects a user action, if the current scene has a map entry for that
+     * user action, the framework starts a transition to the content specified in the map.
+     *
+     * Once the content is shown, the scene framework will read this property and set up a collector
+     * to watch for new mapping values. For each map entry, the scene framework will set up user
+     * input handling for its [UserAction] and, if such a user action is detected, initiate a
+     * transition to the specified [UserActionResult].
+     *
+     * Note that reading from this method does _not_ mean that any user action has occurred.
+     * Instead, the property is read before any user action/gesture is detected so that the
+     * framework can decide whether to set up gesture/input detectors/listeners in case user actions
+     * of the given types ever occur.
+     *
+     * A missing value for a specific [UserAction] means that the user action of the given type is
+     * not currently active in the top-most content (in z-index order) and should be ignored by the
+     * framework until the top-most content changes.
+     */
+    val userActions: Flow<Map<UserAction, UserActionResult>>
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
deleted file mode 100644
index 3da6a02..0000000
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt
+++ /dev/null
@@ -1,27 +0,0 @@
-/*
- * Copyright (C) 2023 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.scene.ui.composable
-
-import androidx.compose.runtime.Composable
-import androidx.compose.ui.Modifier
-import com.android.compose.animation.scene.SceneScope
-import com.android.systemui.scene.shared.model.Scene
-
-/** Compose-capable extension of [Scene]. */
-interface ComposableScene : Scene {
-    @Composable fun SceneScope.Content(modifier: Modifier)
-}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
index d489d73..ae5dd8a 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/GoneScene.kt
@@ -23,8 +23,8 @@
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
-import com.android.compose.animation.scene.animateSceneDpAsState
-import com.android.compose.animation.scene.animateSceneFloatAsState
+import com.android.compose.animation.scene.animateContentDpAsState
+import com.android.compose.animation.scene.animateContentFloatAsState
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.rememberViewModel
@@ -33,7 +33,7 @@
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaLandscapeTopOffset
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaOffset.Default
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.viewmodel.GoneSceneActionsViewModel
+import com.android.systemui.scene.ui.viewmodel.GoneUserActionsViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
 import dagger.Lazy
@@ -50,14 +50,13 @@
 constructor(
     private val notificationStackScrolLView: Lazy<NotificationScrollView>,
     private val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
-    private val viewModelFactory: GoneSceneActionsViewModel.Factory,
-) : ExclusiveActivatable(), ComposableScene {
+    private val viewModelFactory: GoneUserActionsViewModel.Factory,
+) : ExclusiveActivatable(), Scene {
     override val key = Scenes.Gone
 
-    private val actionsViewModel: GoneSceneActionsViewModel by lazy { viewModelFactory.create() }
+    private val actionsViewModel: GoneUserActionsViewModel by lazy { viewModelFactory.create() }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     override suspend fun onActivated(): Nothing {
         actionsViewModel.activate()
@@ -67,15 +66,18 @@
     override fun SceneScope.Content(
         modifier: Modifier,
     ) {
-        animateSceneFloatAsState(
+        animateContentFloatAsState(
             value = QuickSettings.SharedValues.SquishinessValues.GoneSceneStarting,
             key = QuickSettings.SharedValues.TilesSquishiness,
         )
-        animateSceneDpAsState(value = Default, key = MediaLandscapeTopOffset, canOverflow = false)
+        animateContentDpAsState(value = Default, key = MediaLandscapeTopOffset, canOverflow = false)
         Spacer(modifier.fillMaxSize())
         SnoozeableHeadsUpNotificationSpace(
             stackScrollView = notificationStackScrolLView.get(),
-            viewModel = rememberViewModel { notificationsPlaceholderViewModelFactory.create() },
+            viewModel =
+                rememberViewModel("GoneScene") {
+                    notificationsPlaceholderViewModelFactory.create()
+                },
         )
     }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt
new file mode 100644
index 0000000..609ce90
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Overlay.kt
@@ -0,0 +1,50 @@
+/*
+ * 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.scene.ui.composable
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.OverlayKey
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import com.android.systemui.lifecycle.Activatable
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOf
+
+/**
+ * Defines interface for classes that can describe an "overlay".
+ *
+ * In the scene framework, there can be multiple overlays in a single scene "container". The
+ * container takes care of rendering any current overlays and allowing overlays to be shown, hidden,
+ * or replaced based on a user action.
+ */
+interface Overlay : Activatable, ActionableContent {
+    /** Uniquely-identifying key for this overlay. The key must be unique within its container. */
+    val key: OverlayKey
+
+    /**
+     * The user actions supported by this overlay.
+     *
+     * @see [ActionableContent.userActions]
+     */
+    override val userActions: Flow<Map<UserAction, UserActionResult>>
+        get() = flowOf(mapOf(Back to UserActionResult.HideOverlay(key)))
+
+    @Composable fun ContentScope.Content(modifier: Modifier)
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Scene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Scene.kt
new file mode 100644
index 0000000..8d8ab8e
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Scene.kt
@@ -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 com.android.systemui.scene.ui.composable
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SceneScope
+import com.android.systemui.lifecycle.Activatable
+
+/**
+ * Defines interface for classes that can describe a "scene".
+ *
+ * In the scene framework, there can be multiple scenes in a single scene "container". The container
+ * takes care of rendering the current scene and allowing scenes to be switched from one to another
+ * based on either user action (for example, swiping down while on the lock screen scene may switch
+ * to the shade scene).
+ */
+interface Scene : Activatable, ActionableContent {
+
+    /** Uniquely-identifying key for this scene. The key must be unique within its container. */
+    val key: SceneKey
+
+    @Composable fun SceneScope.Content(modifier: Modifier)
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
index e17cb31..a7e41ce 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainer.kt
@@ -31,7 +31,9 @@
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.input.pointer.pointerInput
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.SceneTransitionLayout
 import com.android.compose.animation.scene.UserAction
@@ -52,17 +54,23 @@
  * containers.
  *
  * @param viewModel The UI state holder for this container.
- * @param sceneByKey Mapping of [ComposableScene] by [SceneKey], ordered by z-order such that the
- *   last scene is rendered on top of all other scenes. It's critical that this map contains exactly
- *   and only the scenes on this container. In other words: (a) there should be no scene in this map
- *   that is not in the configuration for this container and (b) all scenes in the configuration
- *   must have entries in this map.
+ * @param sceneByKey Mapping of [Scene] by [SceneKey], ordered by z-order such that the last scene
+ *   is rendered on top of all other scenes. It's critical that this map contains exactly and only
+ *   the scenes on this container. In other words: (a) there should be no scene in this map that is
+ *   not in the configuration for this container and (b) all scenes in the configuration must have
+ *   entries in this map.
+ * @param overlayByKey Mapping of [Overlay] by [OverlayKey], ordered by z-order such that the last
+ *   overlay is rendered on top of all other overlays. It's critical that this map contains exactly
+ *   and only the overlays on this container. In other words: (a) there should be no overlay in this
+ *   map that is not in the configuration for this container and (b) all overlays in the
+ *   configuration must have entries in this map.
  * @param modifier A modifier.
  */
 @Composable
 fun SceneContainer(
     viewModel: SceneContainerViewModel,
-    sceneByKey: Map<SceneKey, ComposableScene>,
+    sceneByKey: Map<SceneKey, Scene>,
+    overlayByKey: Map<OverlayKey, Overlay>,
     initialSceneKey: SceneKey,
     dataSourceDelegator: SceneDataSourceDelegator,
     modifier: Modifier = Modifier,
@@ -89,16 +97,19 @@
         onDispose { viewModel.setTransitionState(null) }
     }
 
-    val userActionsBySceneKey: MutableMap<SceneKey, Map<UserAction, UserActionResult>> = remember {
-        mutableStateMapOf()
-    }
+    val userActionsByContentKey: MutableMap<ContentKey, Map<UserAction, UserActionResult>> =
+        remember {
+            mutableStateMapOf()
+        }
+    // TODO(b/359173565): Add overlay user actions when the API is final.
     LaunchedEffect(currentSceneKey) {
         try {
-            sceneByKey[currentSceneKey]?.destinationScenes?.collectLatest { userActions ->
-                userActionsBySceneKey[currentSceneKey] = viewModel.resolveSceneFamilies(userActions)
+            sceneByKey[currentSceneKey]?.userActions?.collectLatest { userActions ->
+                userActionsByContentKey[currentSceneKey] =
+                    viewModel.resolveSceneFamilies(userActions)
             }
         } finally {
-            userActionsBySceneKey[currentSceneKey] = emptyMap()
+            userActionsByContentKey[currentSceneKey] = emptyMap()
         }
     }
 
@@ -112,22 +123,31 @@
             },
     ) {
         SceneTransitionLayout(state = state, modifier = modifier.fillMaxSize()) {
-            sceneByKey.forEach { (sceneKey, composableScene) ->
+            sceneByKey.forEach { (sceneKey, scene) ->
                 scene(
                     key = sceneKey,
-                    userActions = userActionsBySceneKey.getOrDefault(sceneKey, emptyMap())
+                    userActions = userActionsByContentKey.getOrDefault(sceneKey, emptyMap())
                 ) {
                     // Activate the scene.
-                    LaunchedEffect(composableScene) { composableScene.activate() }
+                    LaunchedEffect(scene) { scene.activate() }
 
                     // Render the scene.
-                    with(composableScene) {
+                    with(scene) {
                         this@scene.Content(
                             modifier = Modifier.element(sceneKey.rootElementKey).fillMaxSize(),
                         )
                     }
                 }
             }
+            overlayByKey.forEach { (overlayKey, composableOverlay) ->
+                overlay(
+                    key = overlayKey,
+                    userActions = userActionsByContentKey.getOrDefault(overlayKey, emptyMap())
+                ) {
+                    // Render the overlay.
+                    with(composableOverlay) { this@overlay.Content(Modifier) }
+                }
+            }
         }
 
         BottomRightCornerRibbon(
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
index 39fc7ef..a0ebca2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneContainerTransitions.kt
@@ -1,13 +1,11 @@
 package com.android.systemui.scene.ui.composable
 
 import androidx.compose.foundation.gestures.Orientation
-import com.android.compose.animation.scene.Edge
 import com.android.compose.animation.scene.ProgressConverter
 import com.android.compose.animation.scene.transitions
 import com.android.systemui.bouncer.ui.composable.Bouncer
 import com.android.systemui.notifications.ui.composable.Notifications
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.shared.model.TransitionKeys.OpenBottomShade
 import com.android.systemui.scene.shared.model.TransitionKeys.SlightlyFasterShadeCollapse
 import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
 import com.android.systemui.scene.ui.composable.transitions.bouncerToGoneTransition
@@ -48,18 +46,8 @@
     // Scene transitions
 
     from(Scenes.Bouncer, to = Scenes.Gone) { bouncerToGoneTransition() }
-    from(Scenes.Gone, to = Scenes.NotificationsShade) {
-        goneToNotificationsShadeTransition(Edge.Top)
-    }
-    from(Scenes.Gone, to = Scenes.NotificationsShade, key = OpenBottomShade) {
-        goneToNotificationsShadeTransition(Edge.Bottom)
-    }
-    from(Scenes.Gone, to = Scenes.QuickSettingsShade) {
-        goneToQuickSettingsShadeTransition(Edge.Top)
-    }
-    from(Scenes.Gone, to = Scenes.QuickSettingsShade, key = OpenBottomShade) {
-        goneToQuickSettingsShadeTransition(Edge.Bottom)
-    }
+    from(Scenes.Gone, to = Scenes.NotificationsShade) { goneToNotificationsShadeTransition() }
+    from(Scenes.Gone, to = Scenes.QuickSettingsShade) { goneToQuickSettingsShadeTransition() }
     from(Scenes.Gone, to = Scenes.Shade) { goneToShadeTransition() }
     from(Scenes.Gone, to = Scenes.Shade, key = ToSplitShade) { goneToSplitShadeTransition() }
     from(Scenes.Gone, to = Scenes.Shade, key = SlightlyFasterShadeCollapse) {
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
index 4b4b7ed..6738b97 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/SceneTransitionLayoutDataSource.kt
@@ -18,7 +18,9 @@
 
 package com.android.systemui.scene.ui.composable
 
+import androidx.compose.runtime.snapshotFlow
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.compose.animation.scene.observableTransitionState
@@ -52,6 +54,14 @@
                 initialValue = state.transitionState.currentScene,
             )
 
+    override val currentOverlays: StateFlow<Set<OverlayKey>> =
+        snapshotFlow { state.currentOverlays }
+            .stateIn(
+                scope = coroutineScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptySet(),
+            )
+
     override fun changeScene(
         toScene: SceneKey,
         transitionKey: TransitionKey?,
@@ -59,7 +69,7 @@
         state.setTargetScene(
             targetScene = toScene,
             transitionKey = transitionKey,
-            coroutineScope = coroutineScope,
+            animationScope = coroutineScope,
         )
     }
 
@@ -68,4 +78,29 @@
             scene = toScene,
         )
     }
+
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        state.showOverlay(
+            overlay = overlay,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        state.hideOverlay(
+            overlay = overlay,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        state.replaceOverlay(
+            from = from,
+            to = to,
+            animationScope = coroutineScope,
+            transitionKey = transitionKey,
+        )
+    }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
index fb41374..48ec198 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromGoneToNotificationsShadeTransition.kt
@@ -16,12 +16,10 @@
 
 package com.android.systemui.scene.ui.composable.transitions
 
-import com.android.compose.animation.scene.Edge
 import com.android.compose.animation.scene.TransitionBuilder
 
 fun TransitionBuilder.goneToNotificationsShadeTransition(
-    edge: Edge = Edge.Top,
     durationScale: Double = 1.0,
 ) {
-    toNotificationsShadeTransition(edge, durationScale)
+    toNotificationsShadeTransition(durationScale)
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToBouncerTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToBouncerTransition.kt
index 1fee874..022eb1f 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToBouncerTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/FromLockscreenToBouncerTransition.kt
@@ -5,10 +5,16 @@
 import com.android.compose.animation.scene.TransitionBuilder
 import com.android.systemui.bouncer.ui.composable.Bouncer
 
+const val FROM_LOCK_SCREEN_TO_BOUNCER_FADE_FRACTION = 0.5f
+
 fun TransitionBuilder.lockscreenToBouncerTransition() {
     spec = tween(durationMillis = 500)
 
     translate(Bouncer.Elements.Content, y = 300.dp)
-    fractionRange(end = 0.5f) { fade(Bouncer.Elements.Background) }
-    fractionRange(start = 0.5f) { fade(Bouncer.Elements.Content) }
+    fractionRange(end = FROM_LOCK_SCREEN_TO_BOUNCER_FADE_FRACTION) {
+        fade(Bouncer.Elements.Background)
+    }
+    fractionRange(start = FROM_LOCK_SCREEN_TO_BOUNCER_FADE_FRACTION) {
+        fade(Bouncer.Elements.Content)
+    }
 }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToNotificationsShadeTransition.kt b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToNotificationsShadeTransition.kt
index 05949b2..337f53a5 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToNotificationsShadeTransition.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/transitions/ToNotificationsShadeTransition.kt
@@ -19,12 +19,9 @@
 import androidx.compose.animation.core.Spring
 import androidx.compose.animation.core.spring
 import androidx.compose.animation.core.tween
-import androidx.compose.foundation.gestures.Orientation
-import androidx.compose.ui.unit.IntSize
 import com.android.compose.animation.scene.Edge
 import com.android.compose.animation.scene.TransitionBuilder
 import com.android.compose.animation.scene.UserActionDistance
-import com.android.compose.animation.scene.UserActionDistanceScope
 import com.android.systemui.notifications.ui.composable.Notifications
 import com.android.systemui.shade.ui.composable.OverlayShade
 import com.android.systemui.shade.ui.composable.Shade
@@ -32,11 +29,6 @@
 import kotlin.time.Duration.Companion.milliseconds
 
 fun TransitionBuilder.toNotificationsShadeTransition(
-    /**
-     * The edge where the shade will animate from. This is statically determined (i.e. doesn't
-     * change during runtime).
-     */
-    edge: Edge = Edge.Top,
     durationScale: Double = 1.0,
 ) {
     spec = tween(durationMillis = (DefaultDuration * durationScale).inWholeMilliseconds.toInt())
@@ -45,17 +37,11 @@
             stiffness = Spring.StiffnessMediumLow,
             visibilityThreshold = Shade.Dimensions.ScrimVisibilityThreshold,
         )
-    distance =
-        object : UserActionDistance {
-            override fun UserActionDistanceScope.absoluteDistance(
-                fromSceneSize: IntSize,
-                orientation: Orientation,
-            ): Float {
-                return fromSceneSize.height.toFloat() * 2 / 3f
-            }
-        }
+    distance = UserActionDistance { fromSceneSize, orientation ->
+        fromSceneSize.height.toFloat() * 2 / 3f
+    }
 
-    translate(OverlayShade.Elements.Panel, edge)
+    translate(OverlayShade.Elements.Panel, Edge.Top)
 
     fractionRange(end = .5f) { fade(OverlayShade.Elements.Scrim) }
 
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
index 445ffcb..8922224 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/OverlayShade.kt
@@ -40,56 +40,30 @@
 import androidx.compose.material3.windowsizeclass.WindowWidthSizeClass
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.ReadOnlyComposable
-import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.draw.clip
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.platform.LocalLayoutDirection
 import androidx.compose.ui.unit.dp
-import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.LowestZIndexContentPicker
 import com.android.compose.animation.scene.SceneScope
 import com.android.compose.windowsizeclass.LocalWindowSizeClass
-import com.android.systemui.keyguard.ui.composable.LockscreenContent
-import com.android.systemui.lifecycle.rememberViewModel
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.shared.model.ShadeAlignment
-import com.android.systemui.shade.ui.viewmodel.OverlayShadeViewModel
-import com.android.systemui.util.kotlin.getOrNull
-import dagger.Lazy
-import java.util.Optional
 
-/** The overlay shade renders a lightweight shade UI container on top of a background scene. */
+/** Renders a lightweight shade UI container, as an overlay. */
 @Composable
 fun SceneScope.OverlayShade(
-    viewModelFactory: OverlayShadeViewModel.Factory,
-    lockscreenContent: Lazy<Optional<LockscreenContent>>,
+    onScrimClicked: () -> Unit,
     modifier: Modifier = Modifier,
     content: @Composable () -> Unit,
 ) {
-    val viewModel = rememberViewModel { viewModelFactory.create() }
-    val backgroundScene by viewModel.backgroundScene.collectAsStateWithLifecycle()
-
     Box(modifier) {
-        if (backgroundScene == Scenes.Lockscreen) {
-            // Lockscreen content is optionally injected, because variants of System UI without a
-            // lockscreen cannot provide it.
-            val lockscreenContentOrNull = lockscreenContent.get().getOrNull()
-            lockscreenContentOrNull?.apply { Content(Modifier.fillMaxSize()) }
-        }
-
-        Scrim(onClicked = viewModel::onScrimClicked)
+        Scrim(onClicked = onScrimClicked)
 
         Box(
             modifier = Modifier.fillMaxSize().panelPadding(),
-            contentAlignment =
-                if (viewModel.panelAlignment == ShadeAlignment.Top) {
-                    Alignment.TopEnd
-                } else {
-                    Alignment.BottomEnd
-                },
+            contentAlignment = Alignment.TopEnd,
         ) {
             Panel(
                 modifier = Modifier.element(OverlayShade.Elements.Panel).panelSize(),
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
index 8c53740..05a0119 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeHeader.kt
@@ -129,7 +129,7 @@
     statusBarIconController: StatusBarIconController,
     modifier: Modifier = Modifier,
 ) {
-    val viewModel = rememberViewModel { viewModelFactory.create() }
+    val viewModel = rememberViewModel("CollapsedShadeHeader") { viewModelFactory.create() }
     val isDisabled by viewModel.isDisabled.collectAsStateWithLifecycle()
     if (isDisabled) {
         return
@@ -287,7 +287,7 @@
     statusBarIconController: StatusBarIconController,
     modifier: Modifier = Modifier,
 ) {
-    val viewModel = rememberViewModel { viewModelFactory.create() }
+    val viewModel = rememberViewModel("ExpandedShadeHeader") { viewModelFactory.create() }
     val isDisabled by viewModel.isDisabled.collectAsStateWithLifecycle()
     if (isDisabled) {
         return
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
index f8513a8..a03bf43 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/ShadeScene.kt
@@ -30,7 +30,7 @@
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.WindowInsets
 import androidx.compose.foundation.layout.asPaddingValues
-import androidx.compose.foundation.layout.displayCutoutPadding
+import androidx.compose.foundation.layout.displayCutout
 import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
@@ -47,8 +47,9 @@
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.mutableIntStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.graphics.CompositingStrategy
@@ -99,17 +100,16 @@
 import com.android.systemui.notifications.ui.composable.NotificationStackCutoffGuideline
 import com.android.systemui.qs.footer.ui.compose.FooterActionsWithAnimatedVisibility
 import com.android.systemui.qs.ui.composable.BrightnessMirror
-import com.android.systemui.qs.ui.composable.QSMediaMeasurePolicy
 import com.android.systemui.qs.ui.composable.QuickSettings
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaLandscapeTopOffset
 import com.android.systemui.qs.ui.composable.QuickSettings.SharedValues.MediaOffset.InQQS
 import com.android.systemui.res.R
 import com.android.systemui.scene.session.ui.composable.SaveableSession
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.shade.shared.model.ShadeMode
-import com.android.systemui.shade.ui.viewmodel.ShadeSceneActionsViewModel
 import com.android.systemui.shade.ui.viewmodel.ShadeSceneContentViewModel
+import com.android.systemui.shade.ui.viewmodel.ShadeUserActionsViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.NotificationScrollView
 import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
 import com.android.systemui.statusbar.phone.StatusBarLocation
@@ -120,6 +120,7 @@
 import javax.inject.Inject
 import javax.inject.Named
 import kotlin.math.roundToInt
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 
 object Shade {
@@ -146,13 +147,14 @@
 }
 
 /** The shade scene shows scrolling list of notifications and some of the quick setting tiles. */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class ShadeScene
 @Inject
 constructor(
     private val shadeSession: SaveableSession,
     private val notificationStackScrollView: Lazy<NotificationScrollView>,
-    private val actionsViewModelFactory: ShadeSceneActionsViewModel.Factory,
+    private val actionsViewModelFactory: ShadeUserActionsViewModel.Factory,
     private val contentViewModelFactory: ShadeSceneContentViewModel.Factory,
     private val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
     private val tintedIconManagerFactory: TintedIconManager.Factory,
@@ -161,11 +163,11 @@
     private val mediaCarouselController: MediaCarouselController,
     @Named(QUICK_QS_PANEL) private val qqsMediaHost: MediaHost,
     @Named(QS_PANEL) private val qsMediaHost: MediaHost,
-) : ExclusiveActivatable(), ComposableScene {
+) : ExclusiveActivatable(), Scene {
 
     override val key = Scenes.Shade
 
-    private val actionsViewModel: ShadeSceneActionsViewModel by lazy {
+    private val actionsViewModel: ShadeUserActionsViewModel by lazy {
         actionsViewModelFactory.create()
     }
 
@@ -173,8 +175,7 @@
         actionsViewModel.activate()
     }
 
-    override val destinationScenes: Flow<Map<UserAction, UserActionResult>> =
-        actionsViewModel.actions
+    override val userActions: Flow<Map<UserAction, UserActionResult>> = actionsViewModel.actions
 
     @Composable
     override fun SceneScope.Content(
@@ -182,9 +183,12 @@
     ) =
         ShadeScene(
             notificationStackScrollView.get(),
-            viewModel = rememberViewModel { contentViewModelFactory.create() },
+            viewModel =
+                rememberViewModel("ShadeScene-viewModel") { contentViewModelFactory.create() },
             notificationsPlaceholderViewModel =
-                rememberViewModel { notificationsPlaceholderViewModelFactory.create() },
+                rememberViewModel("ShadeScene-notifPlaceholderViewModel") {
+                    notificationsPlaceholderViewModelFactory.create()
+                },
             createTintedIconManager = tintedIconManagerFactory::create,
             createBatteryMeterViewController = batteryMeterViewControllerFactory::create,
             statusBarIconController = statusBarIconController,
@@ -266,13 +270,14 @@
     shadeSession: SaveableSession,
 ) {
     val cutoutLocation = LocalDisplayCutout.current.location
+    val cutoutInsets = WindowInsets.Companion.displayCutout
     val isLandscape = LocalWindowSizeClass.current.heightSizeClass == WindowHeightSizeClass.Compact
     val usingCollapsedLandscapeMedia =
         Utils.useCollapsedMediaInLandscape(LocalContext.current.resources)
     val isExpanded = !usingCollapsedLandscapeMedia || !isLandscape
     mediaHost.expansion = if (isExpanded) EXPANDED else COLLAPSED
 
-    val maxNotifScrimTop = remember { mutableStateOf(0f) }
+    var maxNotifScrimTop by remember { mutableIntStateOf(0) }
     val tileSquishiness by
         animateSceneFloatAsState(
             value = 1f,
@@ -298,6 +303,24 @@
             viewModel.qsSceneAdapter,
         )
     }
+    val shadeMeasurePolicy =
+        remember(mediaInRow) {
+            SingleShadeMeasurePolicy(
+                isMediaInRow = mediaInRow,
+                mediaOffset = { mediaOffset.roundToPx() },
+                onNotificationsTopChanged = { maxNotifScrimTop = it },
+                mediaZIndex = {
+                    if (MediaContentPicker.shouldElevateMedia(layoutState)) 1f else 0f
+                },
+                cutoutInsetsProvider = {
+                    if (cutoutLocation == CutoutLocation.CENTER) {
+                        null
+                    } else {
+                        cutoutInsets
+                    }
+                }
+            )
+        }
 
     Box(
         modifier =
@@ -315,101 +338,54 @@
                     .background(colorResource(R.color.shade_scrim_background_dark)),
         )
         Layout(
-            contents =
-                listOf(
-                    {
-                        Column(
-                            horizontalAlignment = Alignment.CenterHorizontally,
-                            modifier =
-                                Modifier.fillMaxWidth()
-                                    .thenIf(isEmptySpaceClickable) {
-                                        Modifier.clickable(
-                                            onClick = { viewModel.onEmptySpaceClicked() }
-                                        )
-                                    }
-                                    .thenIf(cutoutLocation != CutoutLocation.CENTER) {
-                                        Modifier.displayCutoutPadding()
-                                    },
-                        ) {
-                            CollapsedShadeHeader(
-                                viewModelFactory = viewModel.shadeHeaderViewModelFactory,
-                                createTintedIconManager = createTintedIconManager,
-                                createBatteryMeterViewController = createBatteryMeterViewController,
-                                statusBarIconController = statusBarIconController,
-                            )
-
-                            val content: @Composable () -> Unit = {
-                                Box(
-                                    Modifier.element(QuickSettings.Elements.QuickQuickSettings)
-                                        .layoutId(QSMediaMeasurePolicy.LayoutId.QS)
-                                ) {
-                                    QuickSettings(
-                                        viewModel.qsSceneAdapter,
-                                        { viewModel.qsSceneAdapter.qqsHeight },
-                                        isSplitShade = false,
-                                        squishiness = { tileSquishiness },
-                                    )
-                                }
-
-                                ShadeMediaCarousel(
-                                    isVisible = isMediaVisible,
-                                    mediaHost = mediaHost,
-                                    mediaOffsetProvider = mediaOffsetProvider,
-                                    modifier =
-                                        Modifier.layoutId(QSMediaMeasurePolicy.LayoutId.Media),
-                                    carouselController = mediaCarouselController,
-                                )
-                            }
-                            val landscapeQsMediaMeasurePolicy = remember {
-                                QSMediaMeasurePolicy(
-                                    { viewModel.qsSceneAdapter.qqsHeight },
-                                    { mediaOffset.roundToPx() },
-                                )
-                            }
-                            if (mediaInRow) {
-                                Layout(
-                                    content = content,
-                                    measurePolicy = landscapeQsMediaMeasurePolicy,
-                                )
-                            } else {
-                                content()
-                            }
-                        }
-                    },
-                    {
-                        NotificationScrollingStack(
-                            shadeSession = shadeSession,
-                            stackScrollView = notificationStackScrollView,
-                            viewModel = notificationsPlaceholderViewModel,
-                            maxScrimTop = { maxNotifScrimTop.value },
-                            shadeMode = ShadeMode.Single,
-                            shouldPunchHoleBehindScrim = shouldPunchHoleBehindScrim,
-                            onEmptySpaceClick =
-                                viewModel::onEmptySpaceClicked.takeIf { isEmptySpaceClickable },
-                        )
-                    },
+            modifier =
+                Modifier.thenIf(isEmptySpaceClickable) {
+                    Modifier.clickable { viewModel.onEmptySpaceClicked() }
+                },
+            content = {
+                CollapsedShadeHeader(
+                    viewModelFactory = viewModel.shadeHeaderViewModelFactory,
+                    createTintedIconManager = createTintedIconManager,
+                    createBatteryMeterViewController = createBatteryMeterViewController,
+                    statusBarIconController = statusBarIconController,
+                    modifier = Modifier.layoutId(SingleShadeMeasurePolicy.LayoutId.ShadeHeader),
                 )
-        ) { measurables, constraints ->
-            check(measurables.size == 2)
-            check(measurables[0].size == 1)
-            check(measurables[1].size == 1)
 
-            val quickSettingsPlaceable = measurables[0][0].measure(constraints)
-            val notificationsPlaceable = measurables[1][0].measure(constraints)
+                Box(
+                    Modifier.element(QuickSettings.Elements.QuickQuickSettings)
+                        .layoutId(SingleShadeMeasurePolicy.LayoutId.QuickSettings)
+                ) {
+                    QuickSettings(
+                        viewModel.qsSceneAdapter,
+                        { viewModel.qsSceneAdapter.qqsHeight },
+                        isSplitShade = false,
+                        squishiness = { tileSquishiness },
+                    )
+                }
 
-            maxNotifScrimTop.value = quickSettingsPlaceable.height.toFloat()
+                ShadeMediaCarousel(
+                    isVisible = isMediaVisible,
+                    isInRow = mediaInRow,
+                    mediaHost = mediaHost,
+                    mediaOffsetProvider = mediaOffsetProvider,
+                    carouselController = mediaCarouselController,
+                    modifier = Modifier.layoutId(SingleShadeMeasurePolicy.LayoutId.Media),
+                )
 
-            layout(constraints.maxWidth, constraints.maxHeight) {
-                val qsZIndex =
-                    if (MediaContentPicker.shouldElevateMedia(layoutState)) {
-                        1f
-                    } else {
-                        0f
-                    }
-                quickSettingsPlaceable.placeRelative(x = 0, y = 0, zIndex = qsZIndex)
-                notificationsPlaceable.placeRelative(x = 0, y = maxNotifScrimTop.value.roundToInt())
-            }
-        }
+                NotificationScrollingStack(
+                    shadeSession = shadeSession,
+                    stackScrollView = notificationStackScrollView,
+                    viewModel = notificationsPlaceholderViewModel,
+                    maxScrimTop = { maxNotifScrimTop.toFloat() },
+                    shadeMode = ShadeMode.Single,
+                    shouldPunchHoleBehindScrim = shouldPunchHoleBehindScrim,
+                    onEmptySpaceClick =
+                        viewModel::onEmptySpaceClicked.takeIf { isEmptySpaceClickable },
+                    modifier = Modifier.layoutId(SingleShadeMeasurePolicy.LayoutId.Notifications),
+                )
+            },
+            measurePolicy = shadeMeasurePolicy,
+        )
         Box(
             modifier =
                 Modifier.align(Alignment.BottomCenter)
@@ -493,9 +469,10 @@
         }
     }
 
-    val brightnessMirrorViewModel = rememberViewModel {
-        viewModel.brightnessMirrorViewModelFactory.create()
-    }
+    val brightnessMirrorViewModel =
+        rememberViewModel("SplitShade-brightnessMirrorViewModel") {
+            viewModel.brightnessMirrorViewModelFactory.create()
+        }
     val brightnessMirrorShowing by brightnessMirrorViewModel.isShowing.collectAsStateWithLifecycle()
     val contentAlpha by
         animateFloatAsState(
@@ -596,6 +573,7 @@
 
                             ShadeMediaCarousel(
                                 isVisible = isMediaVisible,
+                                isInRow = false,
                                 mediaHost = mediaHost,
                                 mediaOffsetProvider = mediaOffsetProvider,
                                 modifier =
@@ -653,6 +631,7 @@
 @Composable
 private fun SceneScope.ShadeMediaCarousel(
     isVisible: Boolean,
+    isInRow: Boolean,
     mediaHost: MediaHost,
     carouselController: MediaCarouselController,
     mediaOffsetProvider: ShadeMediaOffsetProvider,
@@ -664,7 +643,7 @@
         mediaHost = mediaHost,
         carouselController = carouselController,
         offsetProvider =
-            if (MediaContentPicker.shouldElevateMedia(layoutState)) {
+            if (isInRow || MediaContentPicker.shouldElevateMedia(layoutState)) {
                 null
             } else {
                 { mediaOffsetProvider.offset }
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/SingleShadeMeasurePolicy.kt b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/SingleShadeMeasurePolicy.kt
new file mode 100644
index 0000000..6275ac3
--- /dev/null
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/shade/ui/composable/SingleShadeMeasurePolicy.kt
@@ -0,0 +1,155 @@
+/*
+ * 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.shade.ui.composable
+
+import androidx.compose.foundation.layout.WindowInsets
+import androidx.compose.ui.layout.Measurable
+import androidx.compose.ui.layout.MeasurePolicy
+import androidx.compose.ui.layout.MeasureResult
+import androidx.compose.ui.layout.MeasureScope
+import androidx.compose.ui.layout.Placeable
+import androidx.compose.ui.layout.layoutId
+import androidx.compose.ui.unit.Constraints
+import androidx.compose.ui.unit.offset
+import androidx.compose.ui.util.fastFirst
+import androidx.compose.ui.util.fastFirstOrNull
+import com.android.systemui.shade.ui.composable.SingleShadeMeasurePolicy.LayoutId
+import kotlin.math.max
+
+/**
+ * Lays out elements from the [LayoutId] in the shade. This policy supports the case when the QS and
+ * UMO share the same row and when they should be one below another.
+ */
+class SingleShadeMeasurePolicy(
+    private val isMediaInRow: Boolean,
+    private val mediaOffset: MeasureScope.() -> Int,
+    private val onNotificationsTopChanged: (Int) -> Unit,
+    private val mediaZIndex: () -> Float,
+    private val cutoutInsetsProvider: () -> WindowInsets?,
+) : MeasurePolicy {
+
+    enum class LayoutId {
+        QuickSettings,
+        Media,
+        Notifications,
+        ShadeHeader,
+    }
+
+    override fun MeasureScope.measure(
+        measurables: List<Measurable>,
+        constraints: Constraints,
+    ): MeasureResult {
+        val cutoutInsets: WindowInsets? = cutoutInsetsProvider()
+        val constraintsWithCutout = applyCutout(constraints, cutoutInsets)
+        val insetsLeft = cutoutInsets?.getLeft(this, layoutDirection) ?: 0
+        val insetsTop = cutoutInsets?.getTop(this) ?: 0
+
+        val shadeHeaderPlaceable =
+            measurables
+                .fastFirst { it.layoutId == LayoutId.ShadeHeader }
+                .measure(constraintsWithCutout)
+        val mediaPlaceable =
+            measurables
+                .fastFirstOrNull { it.layoutId == LayoutId.Media }
+                ?.measure(applyMediaConstraints(constraintsWithCutout, isMediaInRow))
+        val quickSettingsPlaceable =
+            measurables
+                .fastFirst { it.layoutId == LayoutId.QuickSettings }
+                .measure(constraintsWithCutout)
+        val notificationsPlaceable =
+            measurables.fastFirst { it.layoutId == LayoutId.Notifications }.measure(constraints)
+
+        val notificationsTop =
+            calculateNotificationsTop(
+                statusBarHeaderPlaceable = shadeHeaderPlaceable,
+                quickSettingsPlaceable = quickSettingsPlaceable,
+                mediaPlaceable = mediaPlaceable,
+                insetsTop = insetsTop,
+                isMediaInRow = isMediaInRow,
+            )
+        onNotificationsTopChanged(notificationsTop)
+
+        return layout(constraints.maxWidth, constraints.maxHeight) {
+            shadeHeaderPlaceable.placeRelative(x = insetsLeft, y = insetsTop)
+            quickSettingsPlaceable.placeRelative(
+                x = insetsLeft,
+                y = insetsTop + shadeHeaderPlaceable.height,
+            )
+
+            if (isMediaInRow) {
+                mediaPlaceable?.placeRelative(
+                    x = insetsLeft + constraintsWithCutout.maxWidth / 2,
+                    y = mediaOffset() + insetsTop + shadeHeaderPlaceable.height,
+                    zIndex = mediaZIndex(),
+                )
+            } else {
+                mediaPlaceable?.placeRelative(
+                    x = insetsLeft,
+                    y = insetsTop + shadeHeaderPlaceable.height + quickSettingsPlaceable.height,
+                    zIndex = mediaZIndex(),
+                )
+            }
+
+            // Notifications don't need to accommodate for horizontal insets
+            notificationsPlaceable.placeRelative(x = 0, y = notificationsTop)
+        }
+    }
+
+    private fun calculateNotificationsTop(
+        statusBarHeaderPlaceable: Placeable,
+        quickSettingsPlaceable: Placeable,
+        mediaPlaceable: Placeable?,
+        insetsTop: Int,
+        isMediaInRow: Boolean,
+    ): Int {
+        val mediaHeight = mediaPlaceable?.height ?: 0
+        return insetsTop +
+            statusBarHeaderPlaceable.height +
+            if (isMediaInRow) {
+                max(quickSettingsPlaceable.height, mediaHeight)
+            } else {
+                quickSettingsPlaceable.height + mediaHeight
+            }
+    }
+
+    private fun applyMediaConstraints(
+        constraints: Constraints,
+        isMediaInRow: Boolean,
+    ): Constraints {
+        return if (isMediaInRow) {
+            constraints.copy(maxWidth = constraints.maxWidth / 2)
+        } else {
+            constraints
+        }
+    }
+
+    private fun MeasureScope.applyCutout(
+        constraints: Constraints,
+        cutoutInsets: WindowInsets?,
+    ): Constraints {
+        return if (cutoutInsets == null) {
+            constraints
+        } else {
+            val left = cutoutInsets.getLeft(this, layoutDirection)
+            val top = cutoutInsets.getTop(this)
+            val right = cutoutInsets.getRight(this, layoutDirection)
+            val bottom = cutoutInsets.getBottom(this)
+
+            constraints.offset(horizontal = -(left + right), vertical = -(top + bottom))
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/mediaoutput/ui/composable/MediaOutputComponent.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/mediaoutput/ui/composable/MediaOutputComponent.kt
index 1db96cf..c9b8013 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/mediaoutput/ui/composable/MediaOutputComponent.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/component/mediaoutput/ui/composable/MediaOutputComponent.kt
@@ -92,7 +92,12 @@
                         true
                     }
                 },
-            color = MaterialTheme.colorScheme.surface,
+            color =
+                if (enabled) {
+                    MaterialTheme.colorScheme.surface
+                } else {
+                    MaterialTheme.colorScheme.surfaceContainerHighest
+                },
             shape = RoundedCornerShape(28.dp),
             onClick =
                 if (enabled) {
@@ -119,7 +124,7 @@
                 modifier = Modifier.basicMarquee(),
                 text = connectedDeviceViewModel.label.toString(),
                 style = MaterialTheme.typography.labelMedium,
-                color = MaterialTheme.colorScheme.onSurfaceVariant,
+                color = connectedDeviceViewModel.labelColor.toColor(),
                 maxLines = 1,
             )
             connectedDeviceViewModel.deviceName?.let {
@@ -127,7 +132,7 @@
                     modifier = Modifier.basicMarquee(),
                     text = it.toString(),
                     style = MaterialTheme.typography.titleMedium,
-                    color = MaterialTheme.colorScheme.onSurface,
+                    color = connectedDeviceViewModel.deviceNameColor.toColor(),
                     maxLines = 1,
                 )
             }
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 072e91a..d4f3b5b 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
@@ -23,7 +23,6 @@
 import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.animation.core.animateDpAsState
 import androidx.compose.animation.core.tween
-import androidx.compose.animation.core.updateTransition
 import androidx.compose.animation.expandVertically
 import androidx.compose.animation.fadeIn
 import androidx.compose.animation.fadeOut
@@ -78,7 +77,6 @@
     modifier: Modifier = Modifier,
 ) {
     require(viewModels.isNotEmpty())
-    val transition = updateTransition(isExpanded, label = "CollapsableSliders")
     Column(modifier = modifier) {
         Box(
             modifier = Modifier.fillMaxWidth(),
@@ -106,8 +104,9 @@
                 sliderColors = sliderColors,
             )
         }
-        transition.AnimatedVisibility(
-            visible = { it || !isExpandable },
+        AnimatedVisibility(
+            visible = isExpanded || !isExpandable,
+            label = "CollapsableSliders",
             enter =
                 expandVertically(animationSpec = tween(durationMillis = EXPAND_DURATION_MILLIS)),
             exit =
@@ -120,23 +119,31 @@
                     for (index in 1..viewModels.lastIndex) {
                         val sliderViewModel: SliderViewModel = viewModels[index]
                         val sliderState by sliderViewModel.slider.collectAsStateWithLifecycle()
-                        transition.AnimatedVisibility(
-                            modifier = Modifier.padding(top = 16.dp),
-                            visible = { it || !isExpandable },
-                            enter = enterTransition(index = index, totalCount = viewModels.size),
-                            exit = exitTransition(index = index, totalCount = viewModels.size)
-                        ) {
-                            VolumeSlider(
-                                modifier = Modifier.fillMaxWidth(),
-                                state = sliderState,
-                                onValueChange = { newValue: Float ->
-                                    sliderViewModel.onValueChanged(sliderState, newValue)
-                                },
-                                onValueChangeFinished = { sliderViewModel.onValueChangeFinished() },
-                                onIconTapped = { sliderViewModel.toggleMuted(sliderState) },
-                                sliderColors = sliderColors,
-                            )
-                        }
+
+                        VolumeSlider(
+                            modifier =
+                                Modifier.padding(top = 16.dp)
+                                    .fillMaxWidth()
+                                    .animateEnterExit(
+                                        enter =
+                                            enterTransition(
+                                                index = index,
+                                                totalCount = viewModels.size,
+                                            ),
+                                        exit =
+                                            exitTransition(
+                                                index = index,
+                                                totalCount = viewModels.size,
+                                            ),
+                                    ),
+                            state = sliderState,
+                            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/ui/composable/VolumePanelRoot.kt b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
index 5ffb6f8..1cc0fb2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/volume/panel/ui/composable/VolumePanelRoot.kt
@@ -25,13 +25,13 @@
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
-import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.res.stringResource
 import androidx.compose.ui.semantics.paneTitle
 import androidx.compose.ui.semantics.semantics
 import androidx.compose.ui.unit.dp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
+import com.android.compose.theme.PlatformTheme
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.res.R
 import com.android.systemui.volume.panel.ui.layout.ComponentsLayout
@@ -43,7 +43,6 @@
 private val padding = 24.dp
 
 @Composable
-@OptIn(ExperimentalComposeUiApi::class)
 fun VolumePanelRoot(
     viewModel: VolumePanelViewModel,
     modifier: Modifier = Modifier,
@@ -54,18 +53,20 @@
 
     with(VolumePanelComposeScope(state)) {
         components?.let { componentsState ->
-            Components(
-                componentsState,
-                modifier
-                    .sysuiResTag(VolumePanelTestTag)
-                    .semantics { paneTitle = accessibilityTitle }
-                    .padding(
-                        start = padding,
-                        top = padding,
-                        end = padding,
-                        bottom = 20.dp,
-                    )
-            )
+            PlatformTheme {
+                Components(
+                    componentsState,
+                    modifier
+                        .sysuiResTag(VolumePanelTestTag)
+                        .semantics { paneTitle = accessibilityTitle }
+                        .padding(
+                            start = padding,
+                            top = padding,
+                            end = padding,
+                            bottom = 20.dp,
+                        )
+                )
+            }
         }
     }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateContent.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateContent.kt
index b166737..d876606 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateContent.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateContent.kt
@@ -21,9 +21,7 @@
 import androidx.compose.animation.core.SpringSpec
 import com.android.compose.animation.scene.content.state.TransitionState
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
 import kotlinx.coroutines.Job
-import kotlinx.coroutines.launch
 
 internal fun CoroutineScope.animateContent(
     layoutState: MutableSceneTransitionLayoutStateImpl,
@@ -31,37 +29,24 @@
     oneOffAnimation: OneOffAnimation,
     targetProgress: Float,
     chain: Boolean = true,
-) {
-    // Start the transition. This will compute the TransformationSpec associated to [transition],
-    // which we need to initialize the Animatable that will actually animate it.
-    layoutState.startTransition(transition, chain)
-
-    // The transition now contains the transformation spec that we should use to instantiate the
-    // Animatable.
-    val animationSpec = transition.transformationSpec.progressSpec
-    val visibilityThreshold =
-        (animationSpec as? SpringSpec)?.visibilityThreshold ?: ProgressVisibilityThreshold
-    val replacedTransition = transition.replacedTransition
-    val initialProgress = replacedTransition?.progress ?: 0f
-    val initialVelocity = replacedTransition?.progressVelocity ?: 0f
-    val animatable =
-        Animatable(initialProgress, visibilityThreshold = visibilityThreshold).also {
-            oneOffAnimation.animatable = it
-        }
-
-    // Animate the progress to its target value.
-    //
-    // Important: We start atomically to make sure that we start the coroutine even if it is
-    // cancelled right after it is launched, so that finishTransition() is correctly called.
-    // Otherwise, this transition will never be stopped and we will never settle to Idle.
-    oneOffAnimation.job =
-        launch(start = CoroutineStart.ATOMIC) {
-            try {
-                animatable.animateTo(targetProgress, animationSpec, initialVelocity)
-            } finally {
-                layoutState.finishTransition(transition)
+): Job {
+    oneOffAnimation.onRun = {
+        // Animate the progress to its target value.
+        val animationSpec = transition.transformationSpec.progressSpec
+        val visibilityThreshold =
+            (animationSpec as? SpringSpec)?.visibilityThreshold ?: ProgressVisibilityThreshold
+        val replacedTransition = transition.replacedTransition
+        val initialProgress = replacedTransition?.progress ?: 0f
+        val initialVelocity = replacedTransition?.progressVelocity ?: 0f
+        val animatable =
+            Animatable(initialProgress, visibilityThreshold = visibilityThreshold).also {
+                oneOffAnimation.animatable = it
             }
-        }
+
+        animatable.animateTo(targetProgress, animationSpec, initialVelocity)
+    }
+
+    return layoutState.startTransitionImmediately(animationScope = this, transition, chain)
 }
 
 internal class OneOffAnimation {
@@ -74,8 +59,8 @@
      */
     lateinit var animatable: Animatable<Float, AnimationVector1D>
 
-    /** The job that is animating [animatable]. */
-    lateinit var job: Job
+    /** The runnable to run for this animation. */
+    lateinit var onRun: suspend () -> Unit
 
     val progress: Float
         get() = animatable.value
@@ -83,7 +68,13 @@
     val progressVelocity: Float
         get() = animatable.velocity
 
-    fun finish(): Job = job
+    suspend fun run() {
+        onRun()
+    }
+
+    fun freezeAndAnimateToCurrentState() {
+        // Do nothing, the state of one-off animations never change and we directly animate to it.
+    }
 }
 
 // TODO(b/290184746): Compute a good default visibility threshold that depends on the layout size
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateOverlay.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateOverlay.kt
index e020f14..28116cb 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateOverlay.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateOverlay.kt
@@ -18,7 +18,6 @@
 
 import com.android.compose.animation.scene.content.state.TransitionState
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Job
 
 /** Trigger a one-off transition to show or hide an overlay. */
 internal fun CoroutineScope.showOrHideOverlay(
@@ -120,7 +119,13 @@
     override val isInitiatedByUserInput: Boolean = false
     override val isUserInputOngoing: Boolean = false
 
-    override fun finish(): Job = oneOffAnimation.finish()
+    override suspend fun run() {
+        oneOffAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        oneOffAnimation.freezeAndAnimateToCurrentState()
+    }
 }
 
 private class OneOffOverlayReplacingTransition(
@@ -140,5 +145,11 @@
     override val isInitiatedByUserInput: Boolean = false
     override val isUserInputOngoing: Boolean = false
 
-    override fun finish(): Job = oneOffAnimation.finish()
+    override suspend fun run() {
+        oneOffAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        oneOffAnimation.freezeAndAnimateToCurrentState()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
index 4aa50b5..b30f2b7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateSharedAsState.kt
@@ -399,26 +399,53 @@
 
         val fromValue = sharedValue[transition.fromContent]
         val toValue = sharedValue[transition.toContent]
-        return if (fromValue != null && toValue != null) {
-            if (fromValue == toValue) {
-                // Optimization: avoid reading progress if the values are the same, so we don't
-                // relayout/redraw for nothing.
-                fromValue
-            } else {
-                val overscrollSpec = transition.currentOverscrollSpec
-                val progress =
-                    when {
-                        overscrollSpec == null -> {
-                            if (canOverflow) transition.progress
-                            else transition.progress.fastCoerceIn(0f, 1f)
-                        }
-                        overscrollSpec.content == transition.toContent -> 1f
-                        else -> 0f
-                    }
+        if (fromValue == null && toValue == null) {
+            return null
+        }
 
-                sharedValue.type.lerp(fromValue, toValue, progress)
+        if (fromValue != null && toValue != null) {
+            return interpolateSharedValue(fromValue, toValue, transition, sharedValue)
+        }
+
+        if (transition is TransitionState.Transition.ReplaceOverlay) {
+            val currentSceneValue = sharedValue[transition.currentScene]
+            if (currentSceneValue != null) {
+                return interpolateSharedValue(
+                    fromValue = fromValue ?: currentSceneValue,
+                    toValue = toValue ?: currentSceneValue,
+                    transition,
+                    sharedValue,
+                )
             }
-        } else fromValue ?: toValue
+        }
+
+        return fromValue ?: toValue
+    }
+
+    private fun interpolateSharedValue(
+        fromValue: T,
+        toValue: T,
+        transition: TransitionState.Transition,
+        sharedValue: SharedValue<T, *>,
+    ): T? {
+        if (fromValue == toValue) {
+            // Optimization: avoid reading progress if the values are the same, so we don't
+            // relayout/redraw for nothing.
+            return fromValue
+        }
+
+        val overscrollSpec = transition.currentOverscrollSpec
+        val progress =
+            when {
+                overscrollSpec == null -> {
+                    if (canOverflow) transition.progress
+                    else transition.progress.fastCoerceIn(0f, 1f)
+                }
+                overscrollSpec.content == transition.toContent -> 1f
+                else -> 0f
+            }
+
+        return sharedValue.type.lerp(fromValue, toValue, progress)
     }
 
     private fun transition(sharedValue: SharedValue<T, Delta>): TransitionState.Transition? {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
index abe079a..86be4a4 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/AnimateToScene.kt
@@ -28,7 +28,7 @@
     layoutState: MutableSceneTransitionLayoutStateImpl,
     target: SceneKey,
     transitionKey: TransitionKey?,
-): TransitionState.Transition.ChangeCurrentScene? {
+): Pair<TransitionState.Transition.ChangeScene, Job>? {
     val transitionState = layoutState.transitionState
     if (transitionState.currentScene == target) {
         // This can happen in 3 different situations, for which there isn't anything else to do:
@@ -55,7 +55,7 @@
                 replacedTransition = null,
             )
         }
-        is TransitionState.Transition.ChangeCurrentScene -> {
+        is TransitionState.Transition.ChangeScene -> {
             val isInitiatedByUserInput = transitionState.isInitiatedByUserInput
 
             // A transition is currently running: first check whether `transition.toScene` or
@@ -139,7 +139,7 @@
     reversed: Boolean = false,
     fromScene: SceneKey = layoutState.transitionState.currentScene,
     chain: Boolean = true,
-): TransitionState.Transition.ChangeCurrentScene {
+): Pair<TransitionState.Transition.ChangeScene, Job> {
     val oneOffAnimation = OneOffAnimation()
     val targetProgress = if (reversed) 0f else 1f
     val transition =
@@ -165,15 +165,16 @@
             )
         }
 
-    animateContent(
-        layoutState = layoutState,
-        transition = transition,
-        oneOffAnimation = oneOffAnimation,
-        targetProgress = targetProgress,
-        chain = chain,
-    )
+    val job =
+        animateContent(
+            layoutState = layoutState,
+            transition = transition,
+            oneOffAnimation = oneOffAnimation,
+            targetProgress = targetProgress,
+            chain = chain,
+        )
 
-    return transition
+    return transition to job
 }
 
 private class OneOffSceneTransition(
@@ -184,7 +185,7 @@
     override val isInitiatedByUserInput: Boolean,
     replacedTransition: TransitionState.Transition?,
     private val oneOffAnimation: OneOffAnimation,
-) : TransitionState.Transition.ChangeCurrentScene(fromScene, toScene, replacedTransition) {
+) : TransitionState.Transition.ChangeScene(fromScene, toScene, replacedTransition) {
     override val progress: Float
         get() = oneOffAnimation.progress
 
@@ -193,5 +194,11 @@
 
     override val isUserInputOngoing: Boolean = false
 
-    override fun finish(): Job = oneOffAnimation.finish()
+    override suspend fun run() {
+        oneOffAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        oneOffAnimation.freezeAndAnimateToCurrentState()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
index 71ff8a8..24fef71 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/DraggableHandler.kt
@@ -28,7 +28,6 @@
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
 import com.android.compose.nestedscroll.PriorityNestedScrollConnection
 import kotlin.math.absoluteValue
-import kotlinx.coroutines.CoroutineScope
 
 internal interface DraggableHandler {
     /**
@@ -63,7 +62,6 @@
 internal class DraggableHandlerImpl(
     internal val layoutImpl: SceneTransitionLayoutImpl,
     internal val orientation: Orientation,
-    internal val coroutineScope: CoroutineScope,
 ) : DraggableHandler {
     internal val nestedScrollKey = Any()
     /** The [DraggableHandler] can only have one active [DragController] at a time. */
@@ -101,15 +99,10 @@
 
         val swipeAnimation = dragController.swipeAnimation
 
-        // Don't intercept a transition that is finishing.
-        if (swipeAnimation.isFinishing) {
-            return false
-        }
-
         // Only intercept the current transition if one of the 2 swipes results is also a transition
         // between the same pair of contents.
         val swipes = computeSwipes(startedPosition, pointersDown = 1)
-        val fromContent = swipeAnimation.currentContent
+        val fromContent = layoutImpl.content(swipeAnimation.currentContent)
         val (upOrLeft, downOrRight) = swipes.computeSwipesResults(fromContent)
         val currentScene = layoutImpl.state.currentScene
         val contentTransition = swipeAnimation.contentTransition
@@ -140,12 +133,13 @@
             // This [transition] was already driving the animation: simply take over it.
             // Stop animating and start from the current offset.
             val oldSwipeAnimation = oldDragController.swipeAnimation
-            oldSwipeAnimation.cancelOffsetAnimation()
 
             // We need to recompute the swipe results since this is a new gesture, and the
             // fromScene.userActions may have changed.
             val swipes = oldDragController.swipes
-            swipes.updateSwipesResults(fromContent = oldSwipeAnimation.fromContent)
+            swipes.updateSwipesResults(
+                fromContent = layoutImpl.content(oldSwipeAnimation.fromContent)
+            )
 
             // A new gesture should always create a new SwipeAnimation. This way there cannot be
             // different gestures controlling the same transition.
@@ -155,8 +149,10 @@
 
         val swipes = computeSwipes(startedPosition, pointersDown)
         val fromContent = layoutImpl.contentForUserActions()
+
+        swipes.updateSwipesResults(fromContent)
         val result =
-            swipes.findUserActionResult(fromContent, overSlop, updateSwipesResults = true)
+            swipes.findUserActionResult(overSlop)
                 // As we were unable to locate a valid target scene, the initial SwipeAnimation
                 // cannot be defined. Consequently, a simple NoOp Controller will be returned.
                 ?: return NoOpDragController
@@ -269,16 +265,14 @@
 
     fun updateTransition(newTransition: SwipeAnimation<*>, force: Boolean = false) {
         if (force || isDrivingTransition) {
-            layoutState.startTransition(newTransition.contentTransition)
+            layoutState.startTransitionImmediately(
+                animationScope = draggableHandler.layoutImpl.animationScope,
+                newTransition.contentTransition,
+                true
+            )
         }
 
-        val previous = swipeAnimation
         swipeAnimation = newTransition
-
-        // Finish the previous transition.
-        if (previous != newTransition) {
-            layoutState.finishTransition(previous.contentTransition)
-        }
     }
 
     /**
@@ -291,8 +285,8 @@
         return onDrag(delta, swipeAnimation)
     }
 
-    private fun <T : Content> onDrag(delta: Float, swipeAnimation: SwipeAnimation<T>): Float {
-        if (delta == 0f || !isDrivingTransition || swipeAnimation.isFinishing) {
+    private fun <T : ContentKey> onDrag(delta: Float, swipeAnimation: SwipeAnimation<T>): Float {
+        if (delta == 0f || !isDrivingTransition || swipeAnimation.isAnimatingOffset()) {
             return 0f
         }
 
@@ -304,12 +298,12 @@
         fun hasReachedToSceneUpOrLeft() =
             distance < 0 &&
                 desiredOffset <= distance &&
-                swipes.upOrLeftResult?.toContent(layoutState.currentScene) == toContent.key
+                swipes.upOrLeftResult?.toContent(layoutState.currentScene) == toContent
 
         fun hasReachedToSceneDownOrRight() =
             distance > 0 &&
                 desiredOffset >= distance &&
-                swipes.downOrRightResult?.toContent(layoutState.currentScene) == toContent.key
+                swipes.downOrRightResult?.toContent(layoutState.currentScene) == toContent
 
         // Considering accelerated swipe: Change fromContent in the case where the user quickly
         // swiped multiple times in the same direction to accelerate the transition from A => B then
@@ -321,7 +315,7 @@
             swipeAnimation.currentContent == toContent &&
                 (hasReachedToSceneUpOrLeft() || hasReachedToSceneDownOrRight())
 
-        val fromContent: Content
+        val fromContent: ContentKey
         val currentTransitionOffset: Float
         val newOffset: Float
         val consumedDelta: Float
@@ -357,12 +351,10 @@
 
         swipeAnimation.dragOffset = currentTransitionOffset
 
-        val result =
-            swipes.findUserActionResult(
-                fromContent = fromContent,
-                directionOffset = newOffset,
-                updateSwipesResults = hasReachedToContent
-            )
+        if (hasReachedToContent) {
+            swipes.updateSwipesResults(draggableHandler.layoutImpl.content(fromContent))
+        }
+        val result = swipes.findUserActionResult(directionOffset = newOffset)
 
         if (result == null) {
             onStop(velocity = delta, canChangeContent = true)
@@ -371,7 +363,7 @@
 
         val needNewTransition =
             hasReachedToContent ||
-                result.toContent(layoutState.currentScene) != swipeAnimation.toContent.key ||
+                result.toContent(layoutState.currentScene) != swipeAnimation.toContent ||
                 result.transitionKey != swipeAnimation.contentTransition.key
 
         if (needNewTransition) {
@@ -390,7 +382,7 @@
         return onStop(velocity, canChangeContent, swipeAnimation)
     }
 
-    private fun <T : Content> onStop(
+    private fun <T : ContentKey> onStop(
         velocity: Float,
         canChangeContent: Boolean,
 
@@ -401,13 +393,12 @@
         swipeAnimation: SwipeAnimation<T>,
     ): Float {
         // The state was changed since the drag started; don't do anything.
-        if (!isDrivingTransition || swipeAnimation.isFinishing) {
+        if (!isDrivingTransition || swipeAnimation.isAnimatingOffset()) {
             return 0f
         }
 
         fun animateTo(targetContent: T) {
             swipeAnimation.animateOffset(
-                coroutineScope = draggableHandler.coroutineScope,
                 initialVelocity = velocity,
                 targetContent = targetContent,
             )
@@ -518,6 +509,14 @@
         return upOrLeftResult to downOrRightResult
     }
 
+    /**
+     * Update the swipes results.
+     *
+     * Usually we don't want to update them while doing a drag, because this could change the target
+     * content (jump cutting) to a different content, when some system state changed the targets the
+     * background. However, an update is needed any time we calculate the targets for a new
+     * fromContent.
+     */
     fun updateSwipesResults(fromContent: Content) {
         val (upOrLeftResult, downOrRightResult) = computeSwipesResults(fromContent)
 
@@ -526,31 +525,17 @@
     }
 
     /**
-     * Returns the [UserActionResult] from [fromContent] in the direction of [directionOffset].
+     * Returns the [UserActionResult] in the direction of [directionOffset].
      *
-     * @param fromContent the content from which we look for the target
      * @param directionOffset signed float that indicates the direction. Positive is down or right
      *   negative is up or left.
-     * @param updateSwipesResults whether the swipe results should be updated to the current values
-     *   held in the user actions map. Usually we don't want to update them while doing a drag,
-     *   because this could change the target content (jump cutting) to a different content, when
-     *   some system state changed the targets the background. However, an update is needed any time
-     *   we calculate the targets for a new fromContent.
      * @return null when there are no targets in either direction. If one direction is null and you
      *   drag into the null direction this function will return the opposite direction, assuming
      *   that the users intention is to start the drag into the other direction eventually. If
      *   [directionOffset] is 0f and both direction are available, it will default to
      *   [upOrLeftResult].
      */
-    fun findUserActionResult(
-        fromContent: Content,
-        directionOffset: Float,
-        updateSwipesResults: Boolean,
-    ): UserActionResult? {
-        if (updateSwipesResults) {
-            updateSwipesResults(fromContent)
-        }
-
+    fun findUserActionResult(directionOffset: Float): UserActionResult? {
         return when {
             upOrLeftResult == null && downOrRightResult == null -> null
             (directionOffset < 0f && upOrLeftResult != null) || downOrRightResult == null ->
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index 9b1740d..4c0feb8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -313,10 +313,27 @@
         // If this element is not supposed to be laid out now, either because it is not part of any
         // ongoing transition or the other content of its transition is overscrolling, then lay out
         // the element normally and don't place it.
-        val overscrollScene = transition?.currentOverscrollSpec?.content
-        val isOtherSceneOverscrolling = overscrollScene != null && overscrollScene != content.key
-        if (isOtherSceneOverscrolling) {
-            return doNotPlace(measurable, constraints)
+        val overscrollContent = transition?.currentOverscrollSpec?.content
+        if (overscrollContent != null && overscrollContent != content.key) {
+            when (transition) {
+                is TransitionState.Transition.ChangeScene ->
+                    return doNotPlace(measurable, constraints)
+
+                // If we are overscrolling an overlay that does not contain an element that is in
+                // the current scene, place it in that scene otherwise the element won't be placed
+                // at all.
+                is TransitionState.Transition.ShowOrHideOverlay,
+                is TransitionState.Transition.ReplaceOverlay -> {
+                    if (
+                        content.key == transition.currentScene &&
+                            overscrollContent !in element.stateByContent
+                    ) {
+                        return placeNormally(measurable, constraints)
+                    } else {
+                        return doNotPlace(measurable, constraints)
+                    }
+                }
+            }
         }
 
         val placeable =
@@ -1230,17 +1247,30 @@
     // elements follow the finger direction.
     val isSharedElement = fromState != null && toState != null
     if (isSharedElement && isSharedElementEnabled(element.key, transition)) {
-        val start = contentValue(fromState!!)
-        val end = contentValue(toState!!)
+        return interpolateSharedElement(
+            transition = transition,
+            contentValue = contentValue,
+            fromState = fromState!!,
+            toState = toState!!,
+            isSpecified = isSpecified,
+            lerp = lerp,
+        )
+    }
 
-        // TODO(b/316901148): Remove checks to isSpecified() once the lookahead pass runs for all
-        // nodes before the intermediate layout pass.
-        if (!isSpecified(start)) return end
-        if (!isSpecified(end)) return start
-
-        // Make sure we don't read progress if values are the same and we don't need to interpolate,
-        // so we don't invalidate the phase where this is read.
-        return if (start == end) start else lerp(start, end, transition.progress)
+    // If we are replacing an overlay and the element is both in a single overlay and in the current
+    // scene, interpolate the state of the element using the current scene as the other scene.
+    if (!isSharedElement && transition is TransitionState.Transition.ReplaceOverlay) {
+        val currentSceneState = element.stateByContent[transition.currentScene]
+        if (currentSceneState != null) {
+            return interpolateSharedElement(
+                transition = transition,
+                contentValue = contentValue,
+                fromState = fromState ?: currentSceneState,
+                toState = toState ?: currentSceneState,
+                isSpecified = isSpecified,
+                lerp = lerp,
+            )
+        }
     }
 
     // Get the transformed value, i.e. the target value at the beginning (for entering elements) or
@@ -1383,3 +1413,24 @@
         lerp(idleValue, targetValue, rangeProgress)
     }
 }
+
+private inline fun <T> interpolateSharedElement(
+    transition: TransitionState.Transition,
+    contentValue: (Element.State) -> T,
+    fromState: Element.State,
+    toState: Element.State,
+    isSpecified: (T) -> Boolean,
+    lerp: (T, T, Float) -> T
+): T {
+    val start = contentValue(fromState)
+    val end = contentValue(toState)
+
+    // TODO(b/316901148): Remove checks to isSpecified() once the lookahead pass runs for all
+    // nodes before the intermediate layout pass.
+    if (!isSpecified(start)) return end
+    if (!isSpecified(end)) return start
+
+    // Make sure we don't read progress if values are the same and we don't need to interpolate,
+    // so we don't invalidate the phase where this is read.
+    return if (start == end) start else lerp(start, end, transition.progress)
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt
index 6181cfb..cb18c67 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/InterruptionHandler.kt
@@ -37,7 +37,7 @@
      * @see InterruptionResult
      */
     fun onInterruption(
-        interrupted: TransitionState.Transition.ChangeCurrentScene,
+        interrupted: TransitionState.Transition.ChangeScene,
         newTargetScene: SceneKey,
     ): InterruptionResult?
 }
@@ -76,7 +76,7 @@
  */
 object DefaultInterruptionHandler : InterruptionHandler {
     override fun onInterruption(
-        interrupted: TransitionState.Transition.ChangeCurrentScene,
+        interrupted: TransitionState.Transition.ChangeScene,
         newTargetScene: SceneKey,
     ): InterruptionResult {
         return InterruptionResult(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
index 3f8f5e7..ced177c 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Key.kt
@@ -153,4 +153,15 @@
     override fun toString(): String {
         return "TransitionKey(debugName=$debugName)"
     }
+
+    companion object {
+        /**
+         * A special transition key indicating that the associated transition should be used for
+         * Predictive Back gestures.
+         *
+         * Use this key when defining a transition that you want to be specifically triggered when
+         * the user performs a Predictive Back gesture.
+         */
+        val PredictiveBack = TransitionKey("PredictiveBack")
+    }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
index fd4c310..5780c08 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/MultiPointerDraggable.kt
@@ -56,7 +56,7 @@
 import com.android.compose.ui.util.SpaceVectorConverter
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.math.sign
-import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 
@@ -143,8 +143,8 @@
     CompositionLocalConsumerModifierNode,
     ObserverModifierNode,
     SpaceVectorConverter {
-    private val pointerInputHandler: suspend PointerInputScope.() -> Unit = { pointerInput() }
-    private val delegate = delegate(SuspendingPointerInputModifierNode(pointerInputHandler))
+    private val pointerTracker = delegate(SuspendingPointerInputModifierNode { pointerTracker() })
+    private val pointerInput = delegate(SuspendingPointerInputModifierNode { pointerInput() })
     private val velocityTracker = VelocityTracker()
     private var previousEnabled: Boolean = false
 
@@ -153,7 +153,7 @@
             // Reset the pointer input whenever enabled changed.
             if (value != field) {
                 field = value
-                delegate.resetPointerInputHandler()
+                pointerInput.resetPointerInputHandler()
             }
         }
 
@@ -173,7 +173,7 @@
             if (value != field) {
                 field = value
                 converter = SpaceVectorConverter(value)
-                delegate.resetPointerInputHandler()
+                pointerInput.resetPointerInputHandler()
             }
         }
 
@@ -186,19 +186,26 @@
         observeReads {
             val newEnabled = enabled()
             if (newEnabled != previousEnabled) {
-                delegate.resetPointerInputHandler()
+                pointerInput.resetPointerInputHandler()
             }
             previousEnabled = newEnabled
         }
     }
 
-    override fun onCancelPointerInput() = delegate.onCancelPointerInput()
+    override fun onCancelPointerInput() {
+        pointerTracker.onCancelPointerInput()
+        pointerInput.onCancelPointerInput()
+    }
 
     override fun onPointerEvent(
         pointerEvent: PointerEvent,
         pass: PointerEventPass,
         bounds: IntSize
-    ) = delegate.onPointerEvent(pointerEvent, pass, bounds)
+    ) {
+        // The order is important here: the tracker is always called first.
+        pointerTracker.onPointerEvent(pointerEvent, pass, bounds)
+        pointerInput.onPointerEvent(pointerEvent, pass, bounds)
+    }
 
     private var startedPosition: Offset? = null
     private var pointersDown: Int = 0
@@ -211,81 +218,77 @@
         )
     }
 
+    private suspend fun PointerInputScope.pointerTracker() {
+        val currentContext = currentCoroutineContext()
+        awaitPointerEventScope {
+            // Intercepts pointer inputs and exposes [PointersInfo], via
+            // [requireAncestorPointersInfoOwner], to our descendants.
+            while (currentContext.isActive) {
+                // During the Initial pass, we receive the event after our ancestors.
+                val pointers = awaitPointerEvent(PointerEventPass.Initial).changes
+                pointersDown = pointers.countDown()
+                if (pointersDown == 0) {
+                    // There are no more pointers down
+                    startedPosition = null
+                } else if (startedPosition == null) {
+                    startedPosition = pointers.first().position
+                    if (enabled()) {
+                        onFirstPointerDown()
+                    }
+                }
+            }
+        }
+    }
+
     private suspend fun PointerInputScope.pointerInput() {
         if (!enabled()) {
             return
         }
 
-        coroutineScope {
-            launch {
-                // Intercepts pointer inputs and exposes [PointersInfo], via
-                // [requireAncestorPointersInfoOwner], to our descendants.
-                awaitPointerEventScope {
-                    while (isActive) {
-                        // During the Initial pass, we receive the event after our ancestors.
-                        val pointers = awaitPointerEvent(PointerEventPass.Initial).changes
-
-                        pointersDown = pointers.countDown()
-                        if (pointersDown == 0) {
-                            // There are no more pointers down
-                            startedPosition = null
-                        } else if (startedPosition == null) {
-                            startedPosition = pointers.first().position
-                            onFirstPointerDown()
-                        }
-                    }
-                }
-            }
-
-            // The order is important here: we want to make sure that the previous PointerEventScope
-            // is initialized first. This ensures that the following PointerEventScope doesn't
-            // receive more events than the first one.
-            launch {
-                awaitPointerEventScope {
-                    while (isActive) {
-                        try {
-                            detectDragGestures(
-                                orientation = orientation,
-                                startDragImmediately = startDragImmediately,
-                                onDragStart = { startedPosition, overSlop, pointersDown ->
-                                    velocityTracker.resetTracking()
-                                    onDragStarted(startedPosition, overSlop, pointersDown)
-                                },
-                                onDrag = { controller, change, amount ->
-                                    velocityTracker.addPointerInputChange(change)
-                                    dispatchScrollEvents(
-                                        availableOnPreScroll = amount,
-                                        onScroll = { controller.onDrag(it) },
-                                        source = NestedScrollSource.UserInput,
-                                    )
-                                },
-                                onDragEnd = { controller ->
-                                    startFlingGesture(
-                                        initialVelocity =
-                                            currentValueOf(LocalViewConfiguration)
-                                                .maximumFlingVelocity
-                                                .let {
-                                                    val maxVelocity = Velocity(it, it)
-                                                    velocityTracker.calculateVelocity(maxVelocity)
-                                                }
-                                                .toFloat(),
-                                        onFling = { controller.onStop(it, canChangeContent = true) }
-                                    )
-                                },
-                                onDragCancel = { controller ->
-                                    startFlingGesture(
-                                        initialVelocity = 0f,
-                                        onFling = { controller.onStop(it, canChangeContent = true) }
-                                    )
-                                },
-                                swipeDetector = swipeDetector,
+        val currentContext = currentCoroutineContext()
+        awaitPointerEventScope {
+            while (currentContext.isActive) {
+                try {
+                    detectDragGestures(
+                        orientation = orientation,
+                        startDragImmediately = startDragImmediately,
+                        onDragStart = { startedPosition, overSlop, pointersDown ->
+                            velocityTracker.resetTracking()
+                            onDragStarted(startedPosition, overSlop, pointersDown)
+                        },
+                        onDrag = { controller, change, amount ->
+                            velocityTracker.addPointerInputChange(change)
+                            dispatchScrollEvents(
+                                availableOnPreScroll = amount,
+                                onScroll = { controller.onDrag(it) },
+                                source = NestedScrollSource.UserInput,
                             )
-                        } catch (exception: CancellationException) {
-                            // If the coroutine scope is active, we can just restart the drag cycle.
-                            if (!isActive) {
-                                throw exception
-                            }
-                        }
+                        },
+                        onDragEnd = { controller ->
+                            startFlingGesture(
+                                initialVelocity =
+                                    currentValueOf(LocalViewConfiguration)
+                                        .maximumFlingVelocity
+                                        .let {
+                                            val maxVelocity = Velocity(it, it)
+                                            velocityTracker.calculateVelocity(maxVelocity)
+                                        }
+                                        .toFloat(),
+                                onFling = { controller.onStop(it, canChangeContent = true) }
+                            )
+                        },
+                        onDragCancel = { controller ->
+                            startFlingGesture(
+                                initialVelocity = 0f,
+                                onFling = { controller.onStop(it, canChangeContent = true) }
+                            )
+                        },
+                        swipeDetector = swipeDetector,
+                    )
+                } catch (exception: CancellationException) {
+                    // If the coroutine scope is active, we can just restart the drag cycle.
+                    if (!currentContext.isActive) {
+                        throw exception
                     }
                 }
             }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
index 236e202..bd21a69 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
@@ -31,9 +31,6 @@
  *    layout or Compose drawing phases.
  * 2. [ObservableTransitionState] values are backed by Kotlin [Flow]s and can be collected by
  *    non-Compose code to observe value changes.
- * 3. [ObservableTransitionState.Transition.fromScene] and
- *    [ObservableTransitionState.Transition.toScene] will never be equal, while
- *    [TransitionState.Transition.fromScene] and [TransitionState.Transition.toScene] can be equal.
  */
 sealed interface ObservableTransitionState {
     /**
@@ -43,7 +40,7 @@
     fun currentScene(): Flow<SceneKey> {
         return when (this) {
             is Idle -> flowOf(currentScene)
-            is Transition.ChangeCurrentScene -> currentScene
+            is Transition.ChangeScene -> currentScene
             is Transition.ShowOrHideOverlay -> flowOf(currentScene)
             is Transition.ReplaceOverlay -> flowOf(currentScene)
         }
@@ -54,9 +51,8 @@
 
     /** There is a transition animating between two scenes. */
     sealed class Transition(
-        // TODO(b/353679003): Rename these to fromContent and toContent.
-        open val fromScene: ContentKey,
-        open val toScene: ContentKey,
+        val fromContent: ContentKey,
+        val toContent: ContentKey,
         val currentOverlays: Flow<Set<OverlayKey>>,
         val progress: Flow<Float>,
 
@@ -86,17 +82,17 @@
     ) : ObservableTransitionState {
         override fun toString(): String =
             """Transition
-                |(from=$fromScene,
-                | to=$toScene,
+                |(from=$fromContent,
+                | to=$toContent,
                 | isInitiatedByUserInput=$isInitiatedByUserInput,
                 | isUserInputOngoing=$isUserInputOngoing
                 |)"""
                 .trimMargin()
 
         /** A transition animating between [fromScene] and [toScene]. */
-        class ChangeCurrentScene(
-            override val fromScene: SceneKey,
-            override val toScene: SceneKey,
+        class ChangeScene(
+            val fromScene: SceneKey,
+            val toScene: SceneKey,
             val currentScene: Flow<SceneKey>,
             currentOverlays: Flow<Set<OverlayKey>>,
             progress: Flow<Float>,
@@ -174,8 +170,8 @@
                 previewProgress: Flow<Float> = flowOf(0f),
                 isInPreviewStage: Flow<Boolean> = flowOf(false),
                 currentOverlays: Flow<Set<OverlayKey>> = flowOf(emptySet()),
-            ): ChangeCurrentScene {
-                return ChangeCurrentScene(
+            ): ChangeScene {
+                return ChangeScene(
                     fromScene,
                     toScene,
                     currentScene,
@@ -190,14 +186,14 @@
         }
     }
 
-    fun isIdle(scene: SceneKey?): Boolean {
+    fun isIdle(scene: SceneKey? = null): Boolean {
         return this is Idle && (scene == null || this.currentScene == scene)
     }
 
     fun isTransitioning(from: ContentKey? = null, to: ContentKey? = null): Boolean {
         return this is Transition &&
-            (from == null || this.fromScene == from) &&
-            (to == null || this.toScene == to)
+            (from == null || this.fromContent == from) &&
+            (to == null || this.toContent == to)
     }
 }
 
@@ -210,8 +206,8 @@
     return snapshotFlow {
             when (val state = transitionState) {
                 is TransitionState.Idle -> ObservableTransitionState.Idle(state.currentScene)
-                is TransitionState.Transition.ChangeCurrentScene -> {
-                    ObservableTransitionState.Transition.ChangeCurrentScene(
+                is TransitionState.Transition.ChangeScene -> {
+                    ObservableTransitionState.Transition.ChangeScene(
                         fromScene = state.fromScene,
                         toScene = state.toScene,
                         currentScene = snapshotFlow { state.currentScene },
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
index e7e6b2a..8480d3a 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
@@ -18,120 +18,71 @@
 
 import androidx.activity.BackEventCompat
 import androidx.activity.compose.PredictiveBackHandler
-import androidx.compose.animation.core.Animatable
-import androidx.compose.animation.core.AnimationVector1D
+import androidx.compose.animation.core.snap
+import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.Composable
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableFloatStateOf
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
-import com.android.compose.animation.scene.content.state.TransitionState
-import kotlin.coroutines.cancellation.CancellationException
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.Job
+import com.android.compose.animation.scene.UserActionResult.ChangeScene
+import com.android.compose.animation.scene.UserActionResult.HideOverlay
+import com.android.compose.animation.scene.UserActionResult.ReplaceByOverlay
+import com.android.compose.animation.scene.UserActionResult.ShowOverlay
+import com.android.compose.animation.scene.transition.animateProgress
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.flow.map
 
 @Composable
 internal fun PredictiveBackHandler(
-    state: MutableSceneTransitionLayoutStateImpl,
-    coroutineScope: CoroutineScope,
-    targetSceneForBack: SceneKey? = null,
+    layoutImpl: SceneTransitionLayoutImpl,
+    result: UserActionResult?,
 ) {
     PredictiveBackHandler(
-        enabled = targetSceneForBack != null,
-    ) { progress: Flow<BackEventCompat> ->
-        val fromScene = state.transitionState.currentScene
-        if (targetSceneForBack == null || targetSceneForBack == fromScene) {
+        enabled = result != null,
+    ) { events: Flow<BackEventCompat> ->
+        if (result == null) {
             // Note: We have to collect progress otherwise PredictiveBackHandler will throw.
-            progress.first()
+            events.first()
             return@PredictiveBackHandler
         }
 
-        val transition =
-            PredictiveBackTransition(state, coroutineScope, fromScene, toScene = targetSceneForBack)
-        state.startTransition(transition)
-        try {
-            progress.collect { backEvent -> transition.dragProgress = backEvent.progress }
+        val animation =
+            createSwipeAnimation(
+                layoutImpl,
+                if (result.transitionKey != null) {
+                    result
+                } else {
+                    result.copy(transitionKey = TransitionKey.PredictiveBack)
+                },
+                isUpOrLeft = false,
+                // Note that the orientation does not matter here given that it's only used to
+                // compute the distance. In our case the distance is always 1f.
+                orientation = Orientation.Horizontal,
+                distance = 1f,
+            )
 
-            // Back gesture successful.
-            transition.animateTo(targetSceneForBack)
-        } catch (e: CancellationException) {
-            // Back gesture cancelled.
-            transition.animateTo(fromScene)
-        }
+        animateProgress(
+            state = layoutImpl.state,
+            animation = animation,
+            progress = events.map { it.progress },
+
+            // Use the transformationSpec.progressSpec. We will lazily access it later once the
+            // transition has been started, because at this point the transformation spec of the
+            // transition is not computed yet.
+            commitSpec = null,
+
+            // The predictive back APIs will automatically animate the progress for us in this case
+            // so there is no need to animate it.
+            cancelSpec = snap(),
+        )
     }
 }
 
-private class PredictiveBackTransition(
-    val state: MutableSceneTransitionLayoutStateImpl,
-    val coroutineScope: CoroutineScope,
-    fromScene: SceneKey,
-    toScene: SceneKey,
-) : TransitionState.Transition.ChangeCurrentScene(fromScene, toScene) {
-    override var currentScene by mutableStateOf(fromScene)
-        private set
-
-    /** The animated progress once the gesture was committed or cancelled. */
-    private var progressAnimatable by mutableStateOf<Animatable<Float, AnimationVector1D>?>(null)
-    var dragProgress: Float by mutableFloatStateOf(0f)
-
-    override val previewProgress: Float
-        get() = dragProgress
-
-    override val previewProgressVelocity: Float
-        get() = 0f // Currently, velocity is not exposed by predictive back API
-
-    override val isInPreviewStage: Boolean
-        get() = previewTransformationSpec != null && currentScene == fromScene
-
-    override val progress: Float
-        get() = progressAnimatable?.value ?: previewTransformationSpec?.let { 0f } ?: dragProgress
-
-    override val progressVelocity: Float
-        get() = progressAnimatable?.velocity ?: 0f
-
-    override val isInitiatedByUserInput: Boolean
-        get() = true
-
-    override val isUserInputOngoing: Boolean
-        get() = progressAnimatable == null
-
-    private var animationJob: Job? = null
-
-    override fun finish(): Job = animateTo(currentScene)
-
-    fun animateTo(scene: SceneKey): Job {
-        check(scene == fromScene || scene == toScene)
-        animationJob?.let {
-            return it
-        }
-
-        if (scene != currentScene && state.transitionState == this && state.canChangeScene(scene)) {
-            currentScene = scene
-        }
-
-        val targetProgress =
-            when (currentScene) {
-                fromScene -> 0f
-                toScene -> 1f
-                else -> error("scene $currentScene should be either $fromScene or $toScene")
-            }
-        val startProgress = if (previewTransformationSpec != null) 0f else dragProgress
-        val animatable = Animatable(startProgress).also { progressAnimatable = it }
-
-        // Important: We start atomically to make sure that we start the coroutine even if it is
-        // cancelled right after it is launched, so that finishTransition() is correctly called.
-        return coroutineScope
-            .launch(start = CoroutineStart.ATOMIC) {
-                try {
-                    animatable.animateTo(targetProgress)
-                } finally {
-                    state.finishTransition(this@PredictiveBackTransition)
-                }
-            }
-            .also { animationJob = it }
+private fun UserActionResult.copy(
+    transitionKey: TransitionKey? = this.transitionKey
+): UserActionResult {
+    return when (this) {
+        is ChangeScene -> copy(transitionKey = transitionKey)
+        is ShowOverlay -> copy(transitionKey = transitionKey)
+        is HideOverlay -> copy(transitionKey = transitionKey)
+        is ReplaceByOverlay -> copy(transitionKey = transitionKey)
     }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
index b3f74f7..094f20e 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayout.kt
@@ -503,19 +503,19 @@
     }
 
     /** A [UserActionResult] that shows [overlay]. */
-    class ShowOverlay(
+    data class ShowOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = overlay
     }
 
     /** A [UserActionResult] that hides [overlay]. */
-    class HideOverlay(
+    data class HideOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = currentScene
     }
@@ -526,10 +526,10 @@
      * Note: This result can only be used for user actions of overlays and an exception will be
      * thrown if it is used for a scene.
      */
-    class ReplaceByOverlay(
+    data class ReplaceByOverlay(
         val overlay: OverlayKey,
-        transitionKey: TransitionKey? = null,
-        requiresFullDistanceSwipe: Boolean = false,
+        override val transitionKey: TransitionKey? = null,
+        override val requiresFullDistanceSwipe: Boolean = false,
     ) : UserActionResult(transitionKey, requiresFullDistanceSwipe) {
         override fun toContent(currentScene: SceneKey): ContentKey = overlay
     }
@@ -597,7 +597,7 @@
 ) {
     val density = LocalDensity.current
     val layoutDirection = LocalLayoutDirection.current
-    val coroutineScope = rememberCoroutineScope()
+    val animationScope = rememberCoroutineScope()
     val layoutImpl = remember {
         SceneTransitionLayoutImpl(
                 state = state as MutableSceneTransitionLayoutStateImpl,
@@ -606,7 +606,7 @@
                 swipeSourceDetector = swipeSourceDetector,
                 transitionInterceptionThreshold = transitionInterceptionThreshold,
                 builder = builder,
-                coroutineScope = coroutineScope,
+                animationScope = animationScope,
             )
             .also { onLayoutImpl?.invoke(it) }
     }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
index 258be81..f36c0fa 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutImpl.kt
@@ -59,7 +59,13 @@
     internal var swipeSourceDetector: SwipeSourceDetector,
     internal var transitionInterceptionThreshold: Float,
     builder: SceneTransitionLayoutScope.() -> Unit,
-    internal val coroutineScope: CoroutineScope,
+
+    /**
+     * The scope that should be used by *animations started by this layout only*, i.e. animations
+     * triggered by gestures set up on this layout in [swipeToScene] or interruption decay
+     * animations.
+     */
+    internal val animationScope: CoroutineScope,
 ) {
     /**
      * The map of [Scene]s.
@@ -142,18 +148,10 @@
         // DraggableHandlerImpl must wait for the scenes to be initialized, in order to access the
         // current scene (required for SwipeTransition).
         horizontalDraggableHandler =
-            DraggableHandlerImpl(
-                layoutImpl = this,
-                orientation = Orientation.Horizontal,
-                coroutineScope = coroutineScope,
-            )
+            DraggableHandlerImpl(layoutImpl = this, orientation = Orientation.Horizontal)
 
         verticalDraggableHandler =
-            DraggableHandlerImpl(
-                layoutImpl = this,
-                orientation = Orientation.Vertical,
-                coroutineScope = coroutineScope,
-            )
+            DraggableHandlerImpl(layoutImpl = this, orientation = Orientation.Vertical)
 
         // Make sure that the state is created on the same thread (most probably the main thread)
         // than this STLImpl.
@@ -353,19 +351,8 @@
 
     @Composable
     private fun BackHandler() {
-        val targetSceneForBack =
-            when (val result = contentForUserActions().userActions[Back.Resolved]) {
-                null -> null
-                is UserActionResult.ChangeScene -> result.toScene
-                is UserActionResult.ShowOverlay,
-                is UserActionResult.HideOverlay,
-                is UserActionResult.ReplaceByOverlay -> {
-                    // TODO(b/353679003): Support overlay transitions when going back
-                    null
-                }
-            }
-
-        PredictiveBackHandler(state, coroutineScope, targetSceneForBack)
+        val result = contentForUserActions().userActions[Back.Resolved]
+        PredictiveBackHandler(layoutImpl = this, result = result)
     }
 
     @Composable
@@ -389,7 +376,7 @@
                 // Compose the new scene we are going to first.
                 transitions.fastForEachReversed { transition ->
                     when (transition) {
-                        is TransitionState.Transition.ChangeCurrentScene -> {
+                        is TransitionState.Transition.ChangeScene -> {
                             maybeAdd(transition.toScene)
                             maybeAdd(transition.fromScene)
                         }
@@ -439,7 +426,7 @@
 
                     transitions.fastForEach { transition ->
                         when (transition) {
-                            is TransitionState.Transition.ChangeCurrentScene -> {}
+                            is TransitionState.Transition.ChangeScene -> {}
                             is TransitionState.Transition.ShowOrHideOverlay ->
                                 maybeAdd(transition.overlay)
                             is TransitionState.Transition.ReplaceOverlay -> {
@@ -495,7 +482,7 @@
         val width: Int
         val height: Int
         val transition =
-            layoutImpl.state.currentTransition as? TransitionState.Transition.ChangeCurrentScene
+            layoutImpl.state.currentTransition as? TransitionState.Transition.ChangeScene
         if (transition == null) {
             width = placeable.width
             height = placeable.height
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
index 47065c7..c2d5dd05 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitionLayoutState.kt
@@ -30,6 +30,10 @@
 import com.android.compose.animation.scene.transition.link.StateLink
 import kotlin.math.absoluteValue
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.CoroutineStart
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.launch
 
 /**
  * The state of a [SceneTransitionLayout].
@@ -108,24 +112,24 @@
      * If [targetScene] is different than the [currentScene][TransitionState.currentScene] of
      * [transitionState], then this will animate to [targetScene]. The associated
      * [TransitionState.Transition] will be returned and will be set as the current
-     * [transitionState] of this [MutableSceneTransitionLayoutState].
+     * [transitionState] of this [MutableSceneTransitionLayoutState]. The [Job] in which the
+     * transition runs will be returned, allowing you to easily [join][Job.join] or
+     * [cancel][Job.cancel] the animation.
      *
      * Note that because a non-null [TransitionState.Transition] is returned does not mean that the
      * transition will finish and that we will settle to [targetScene]. The returned transition
      * might still be interrupted, for instance by another call to [setTargetScene] or by a user
      * gesture.
      *
-     * If [this] [CoroutineScope] is cancelled during the transition and that the transition was
-     * still active, then the [transitionState] of this [MutableSceneTransitionLayoutState] will be
-     * set to `TransitionState.Idle(targetScene)`.
-     *
-     * TODO(b/318794193): Add APIs to await() and cancel() any [TransitionState.Transition].
+     * If [animationScope] is cancelled during the transition and that the transition was still
+     * active, then the [transitionState] of this [MutableSceneTransitionLayoutState] will be set to
+     * `TransitionState.Idle(targetScene)`.
      */
     fun setTargetScene(
         targetScene: SceneKey,
-        coroutineScope: CoroutineScope,
+        animationScope: CoroutineScope,
         transitionKey: TransitionKey? = null,
-    ): TransitionState.Transition?
+    ): Pair<TransitionState.Transition, Job>?
 
     /** Immediately snap to the given [scene]. */
     fun snapToScene(
@@ -297,12 +301,12 @@
 
     override fun setTargetScene(
         targetScene: SceneKey,
-        coroutineScope: CoroutineScope,
+        animationScope: CoroutineScope,
         transitionKey: TransitionKey?,
-    ): TransitionState.Transition.ChangeCurrentScene? {
+    ): Pair<TransitionState.Transition.ChangeScene, Job>? {
         checkThread()
 
-        return coroutineScope.animateToScene(
+        return animationScope.animateToScene(
             layoutState = this@MutableSceneTransitionLayoutStateImpl,
             target = targetScene,
             transitionKey = transitionKey,
@@ -310,17 +314,67 @@
     }
 
     /**
+     * Instantly start a [transition], running it in [animationScope].
+     *
+     * This call returns immediately and [transition] will be the [currentTransition] of this
+     * [MutableSceneTransitionLayoutState].
+     *
+     * @see startTransition
+     */
+    internal fun startTransitionImmediately(
+        animationScope: CoroutineScope,
+        transition: TransitionState.Transition,
+        chain: Boolean = true,
+    ): Job {
+        // Note that we start with UNDISPATCHED so that startTransition() is called directly and
+        // transition becomes the current [transitionState] right after this call.
+        return animationScope.launch(
+            start = CoroutineStart.UNDISPATCHED,
+        ) {
+            startTransition(transition, chain)
+        }
+    }
+
+    /**
      * Start a new [transition].
      *
      * If [chain] is `true`, then the transitions will simply be added to [currentTransitions] and
      * will run in parallel to the current transitions. If [chain] is `false`, then the list of
      * [currentTransitions] will be cleared and [transition] will be the only running transition.
      *
-     * Important: you *must* call [finishTransition] once the transition is finished.
+     * If any transition is currently ongoing, it will be interrupted and forced to animate to its
+     * current state.
+     *
+     * This method returns when [transition] is done running, i.e. when the call to
+     * [run][TransitionState.Transition.run] returns.
      */
-    internal fun startTransition(transition: TransitionState.Transition, chain: Boolean = true) {
+    internal suspend fun startTransition(
+        transition: TransitionState.Transition,
+        chain: Boolean = true,
+    ) {
         checkThread()
 
+        try {
+            // Keep a reference to the previous transition (if any).
+            val previousTransition = currentTransition
+
+            // Start the transition.
+            startTransitionInternal(transition, chain)
+
+            // Handle transition links.
+            previousTransition?.let { cancelActiveTransitionLinks(it) }
+            if (stateLinks.isNotEmpty()) {
+                coroutineScope { setupTransitionLinks(transition) }
+            }
+
+            // Run the transition until it is finished.
+            transition.run()
+        } finally {
+            finishTransition(transition)
+        }
+    }
+
+    private fun startTransitionInternal(transition: TransitionState.Transition, chain: Boolean) {
         // Set the current scene and overlays on the transition.
         val currentState = transitionState
         transition.currentSceneWhenTransitionStarted = currentState.currentScene
@@ -349,10 +403,6 @@
             transition.updateOverscrollSpecs(fromSpec = null, toSpec = null)
         }
 
-        // Handle transition links.
-        currentTransition?.let { cancelActiveTransitionLinks(it) }
-        setupTransitionLinks(transition)
-
         if (!enableInterruptions) {
             // Set the current transition.
             check(transitionStates.size == 1)
@@ -367,9 +417,8 @@
                 transitionStates = listOf(transition)
             }
             is TransitionState.Transition -> {
-                // Force the current transition to finish to currentScene. The transition will call
-                // [finishTransition] once it's finished.
-                currentState.finish()
+                // Force the current transition to finish to currentScene.
+                currentState.freezeAndAnimateToCurrentState()
 
                 val tooManyTransitions = transitionStates.size >= MAX_CONCURRENT_TRANSITIONS
                 val clearCurrentTransitions = !chain || tooManyTransitions
@@ -423,7 +472,7 @@
         transition.activeTransitionLinks.clear()
     }
 
-    private fun setupTransitionLinks(transition: TransitionState.Transition) {
+    private fun CoroutineScope.setupTransitionLinks(transition: TransitionState.Transition) {
         stateLinks.fastForEach { stateLink ->
             val matchingLinks =
                 stateLink.transitionLinks.fastFilter { it.isMatchingLink(transition) }
@@ -443,7 +492,11 @@
                     key = matchingLink.targetTransitionKey,
                 )
 
-            stateLink.target.startTransition(linkedTransition)
+            // Start with UNDISPATCHED so that startTransition is called directly and the new linked
+            // transition is observable directly.
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                stateLink.target.startTransition(linkedTransition)
+            }
             transition.activeTransitionLinks[stateLink] = linkedTransition
         }
     }
@@ -453,7 +506,7 @@
      * [currentScene][TransitionState.currentScene]. This will do nothing if [transition] was
      * interrupted since it was started.
      */
-    internal fun finishTransition(transition: TransitionState.Transition) {
+    private fun finishTransition(transition: TransitionState.Transition) {
         checkThread()
 
         if (finishedTransitions.contains(transition)) {
@@ -461,6 +514,10 @@
             return
         }
 
+        // Make sure that this transition settles in case it was force finished, for instance by
+        // calling snapToScene().
+        transition.freezeAndAnimateToCurrentState()
+
         val transitionStates = this.transitionStates
         if (!transitionStates.contains(transition)) {
             // This transition was already removed from transitionStates.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index cefcff7..e65ed9b 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -90,10 +90,19 @@
             return relaxedSpec
         }
 
-        return transition(from, to, key) {
+        val relaxedReversed =
+            transition(from, to, key) {
                 (it.from == to && it.to == null) || (it.to == from && it.from == null)
             }
-            ?.reversed() ?: defaultTransition(from, to)
+        if (relaxedReversed != null) {
+            return relaxedReversed.reversed()
+        }
+
+        return if (key != null) {
+            findSpec(from, to, null)
+        } else {
+            defaultTransition(from, to)
+        }
     }
 
     private fun transition(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
index 8ca90f1..2a09a77 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SwipeAnimation.kt
@@ -17,6 +17,7 @@
 package com.android.compose.animation.scene
 
 import androidx.compose.animation.core.Animatable
+import androidx.compose.animation.core.AnimationSpec
 import androidx.compose.animation.core.AnimationVector1D
 import androidx.compose.foundation.gestures.Orientation
 import androidx.compose.runtime.getValue
@@ -24,41 +25,96 @@
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.unit.IntSize
-import com.android.compose.animation.scene.content.Content
-import com.android.compose.animation.scene.content.Overlay
-import com.android.compose.animation.scene.content.Scene
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.content.state.TransitionState.HasOverscrollProperties.Companion.DistanceUnspecified
 import kotlin.math.absoluteValue
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.CompletableDeferred
+
+internal fun createSwipeAnimation(
+    layoutState: MutableSceneTransitionLayoutStateImpl,
+    result: UserActionResult,
+    isUpOrLeft: Boolean,
+    orientation: Orientation,
+    distance: Float,
+): SwipeAnimation<*> {
+    return createSwipeAnimation(
+        layoutState,
+        result,
+        isUpOrLeft,
+        orientation,
+        distance = { distance },
+        contentForUserActions = {
+            error("Computing contentForUserActions requires a SceneTransitionLayoutImpl")
+        },
+    )
+}
 
 internal fun createSwipeAnimation(
     layoutImpl: SceneTransitionLayoutImpl,
     result: UserActionResult,
     isUpOrLeft: Boolean,
     orientation: Orientation,
+    distance: Float = DistanceUnspecified
 ): SwipeAnimation<*> {
-    fun <T : Content> swipeAnimation(fromContent: T, toContent: T): SwipeAnimation<T> {
+    var lastDistance = distance
+
+    fun distance(animation: SwipeAnimation<*>): Float {
+        if (lastDistance != DistanceUnspecified) {
+            return lastDistance
+        }
+
+        val absoluteDistance =
+            with(animation.contentTransition.transformationSpec.distance ?: DefaultSwipeDistance) {
+                layoutImpl.userActionDistanceScope.absoluteDistance(
+                    layoutImpl.content(animation.fromContent).targetSize,
+                    orientation,
+                )
+            }
+
+        if (absoluteDistance <= 0f) {
+            return DistanceUnspecified
+        }
+
+        val distance = if (isUpOrLeft) -absoluteDistance else absoluteDistance
+        lastDistance = distance
+        return distance
+    }
+
+    return createSwipeAnimation(
+        layoutImpl.state,
+        result,
+        isUpOrLeft,
+        orientation,
+        distance = ::distance,
+        contentForUserActions = { layoutImpl.contentForUserActions().key },
+    )
+}
+
+private fun createSwipeAnimation(
+    layoutState: MutableSceneTransitionLayoutStateImpl,
+    result: UserActionResult,
+    isUpOrLeft: Boolean,
+    orientation: Orientation,
+    distance: (SwipeAnimation<*>) -> Float,
+    contentForUserActions: () -> ContentKey,
+): SwipeAnimation<*> {
+    fun <T : ContentKey> swipeAnimation(fromContent: T, toContent: T): SwipeAnimation<T> {
         return SwipeAnimation(
-            layoutImpl = layoutImpl,
+            layoutState = layoutState,
             fromContent = fromContent,
             toContent = toContent,
-            userActionDistanceScope = layoutImpl.userActionDistanceScope,
             orientation = orientation,
             isUpOrLeft = isUpOrLeft,
             requiresFullDistanceSwipe = result.requiresFullDistanceSwipe,
+            distance = distance,
         )
     }
 
-    val layoutState = layoutImpl.state
     return when (result) {
         is UserActionResult.ChangeScene -> {
-            val fromScene = layoutImpl.scene(layoutState.currentScene)
-            val toScene = layoutImpl.scene(result.toScene)
-            ChangeCurrentSceneSwipeTransition(
+            val fromScene = layoutState.currentScene
+            val toScene = result.toScene
+            ChangeSceneSwipeTransition(
                     layoutState = layoutState,
                     swipeAnimation = swipeAnimation(fromContent = fromScene, toContent = toScene),
                     key = result.transitionKey,
@@ -67,12 +123,12 @@
                 .swipeAnimation
         }
         is UserActionResult.ShowOverlay -> {
-            val fromScene = layoutImpl.scene(layoutState.currentScene)
-            val overlay = layoutImpl.overlay(result.overlay)
+            val fromScene = layoutState.currentScene
+            val overlay = result.overlay
             ShowOrHideOverlaySwipeTransition(
                     layoutState = layoutState,
-                    _fromOrToScene = fromScene,
-                    _overlay = overlay,
+                    fromOrToScene = fromScene,
+                    overlay = overlay,
                     swipeAnimation = swipeAnimation(fromContent = fromScene, toContent = overlay),
                     key = result.transitionKey,
                     replacedTransition = null,
@@ -80,12 +136,12 @@
                 .swipeAnimation
         }
         is UserActionResult.HideOverlay -> {
-            val toScene = layoutImpl.scene(layoutState.currentScene)
-            val overlay = layoutImpl.overlay(result.overlay)
+            val toScene = layoutState.currentScene
+            val overlay = result.overlay
             ShowOrHideOverlaySwipeTransition(
                     layoutState = layoutState,
-                    _fromOrToScene = toScene,
-                    _overlay = overlay,
+                    fromOrToScene = toScene,
+                    overlay = overlay,
                     swipeAnimation = swipeAnimation(fromContent = overlay, toContent = toScene),
                     key = result.transitionKey,
                     replacedTransition = null,
@@ -93,8 +149,14 @@
                 .swipeAnimation
         }
         is UserActionResult.ReplaceByOverlay -> {
-            val fromOverlay = layoutImpl.contentForUserActions() as Overlay
-            val toOverlay = layoutImpl.overlay(result.overlay)
+            val fromOverlay =
+                when (val contentForUserActions = contentForUserActions()) {
+                    is SceneKey ->
+                        error("ReplaceByOverlay can only be called when an overlay is shown")
+                    is OverlayKey -> contentForUserActions
+                }
+
+            val toOverlay = result.overlay
             ReplaceOverlaySwipeTransition(
                     layoutState = layoutState,
                     swipeAnimation =
@@ -109,9 +171,8 @@
 
 internal fun createSwipeAnimation(old: SwipeAnimation<*>): SwipeAnimation<*> {
     return when (val transition = old.contentTransition) {
-        is TransitionState.Transition.ChangeCurrentScene -> {
-            ChangeCurrentSceneSwipeTransition(transition as ChangeCurrentSceneSwipeTransition)
-                .swipeAnimation
+        is TransitionState.Transition.ChangeScene -> {
+            ChangeSceneSwipeTransition(transition as ChangeSceneSwipeTransition).swipeAnimation
         }
         is TransitionState.Transition.ShowOrHideOverlay -> {
             ShowOrHideOverlaySwipeTransition(transition as ShowOrHideOverlaySwipeTransition)
@@ -125,29 +186,42 @@
 }
 
 /** A helper class that contains the main logic for swipe transitions. */
-internal class SwipeAnimation<T : Content>(
-    val layoutImpl: SceneTransitionLayoutImpl,
+internal class SwipeAnimation<T : ContentKey>(
+    val layoutState: MutableSceneTransitionLayoutStateImpl,
     val fromContent: T,
     val toContent: T,
-    private val userActionDistanceScope: UserActionDistanceScope,
     override val orientation: Orientation,
     override val isUpOrLeft: Boolean,
     val requiresFullDistanceSwipe: Boolean,
-    private var lastDistance: Float = DistanceUnspecified,
+    private val distance: (SwipeAnimation<T>) -> Float,
     currentContent: T = fromContent,
     dragOffset: Float = 0f,
 ) : TransitionState.HasOverscrollProperties {
     /** The [TransitionState.Transition] whose implementation delegates to this [SwipeAnimation]. */
     lateinit var contentTransition: TransitionState.Transition
 
-    var currentContent by mutableStateOf(currentContent)
+    private var _currentContent by mutableStateOf(currentContent)
+    var currentContent: T
+        get() = _currentContent
+        set(value) {
+            check(!isAnimatingOffset()) {
+                "currentContent can not be changed once we are animating the offset"
+            }
+            _currentContent = value
+        }
 
     val progress: Float
         get() {
             // Important: If we are going to return early because distance is equal to 0, we should
             // still make sure we read the offset before returning so that the calling code still
             // subscribes to the offset value.
-            val offset = offsetAnimation?.animatable?.value ?: dragOffset
+            val animatable = offsetAnimation
+            val offset =
+                when {
+                    isInPreviewStage -> 0f
+                    animatable != null -> animatable.value
+                    else -> dragOffset
+                }
 
             return computeProgress(offset)
         }
@@ -162,7 +236,7 @@
 
     val progressVelocity: Float
         get() {
-            val animatable = offsetAnimation?.animatable ?: return 0f
+            val animatable = offsetAnimation ?: return 0f
             val distance = distance()
             if (distance == DistanceUnspecified) {
                 return 0f
@@ -172,48 +246,60 @@
             return velocityInDistanceUnit / distance.absoluteValue
         }
 
+    val previewProgress: Float
+        get() {
+            val offset =
+                if (isInPreviewStage) {
+                    offsetAnimation?.value ?: dragOffset
+                } else {
+                    dragOffset
+                }
+            return computeProgress(offset)
+        }
+
+    val previewProgressVelocity: Float
+        get() = 0f
+
+    val isInPreviewStage: Boolean
+        get() = contentTransition.previewTransformationSpec != null && currentContent == fromContent
+
     override var bouncingContent: ContentKey? = null
 
     /** The current offset caused by the drag gesture. */
     var dragOffset by mutableFloatStateOf(dragOffset)
 
     /** The offset animation that animates the offset once the user lifts their finger. */
-    private var offsetAnimation: OffsetAnimation? by mutableStateOf(null)
+    private var offsetAnimation: Animatable<Float, AnimationVector1D>? by mutableStateOf(null)
+    private val offsetAnimationRunnable = CompletableDeferred<(suspend () -> Unit)?>()
 
     val isUserInputOngoing: Boolean
         get() = offsetAnimation == null
 
-    override val overscrollScope: OverscrollScope =
-        object : OverscrollScope {
-            override val density: Float
-                get() = layoutImpl.density.density
-
-            override val fontScale: Float
-                get() = layoutImpl.density.fontScale
-
-            override val absoluteDistance: Float
-                get() = distance().absoluteValue
-        }
-
-    /** Whether [finish] was called on this animation. */
-    var isFinishing = false
-        private set
+    override val absoluteDistance: Float
+        get() = distance().absoluteValue
 
     constructor(
         other: SwipeAnimation<T>
     ) : this(
-        layoutImpl = other.layoutImpl,
+        layoutState = other.layoutState,
         fromContent = other.fromContent,
         toContent = other.toContent,
-        userActionDistanceScope = other.userActionDistanceScope,
         orientation = other.orientation,
         isUpOrLeft = other.isUpOrLeft,
         requiresFullDistanceSwipe = other.requiresFullDistanceSwipe,
-        lastDistance = other.lastDistance,
+        distance = other.distance,
         currentContent = other.currentContent,
-        dragOffset = other.dragOffset,
+        dragOffset = other.offsetAnimation?.value ?: other.dragOffset,
     )
 
+    suspend fun run() {
+        // This animation will first be driven by finger, then when the user lift their finger we
+        // start an animation to the target offset (progress = 1f or progress = 0f). We await() for
+        // offsetAnimationRunnable to be completed and then run it.
+        val runAnimation = offsetAnimationRunnable.await() ?: return
+        runAnimation()
+    }
+
     /**
      * The signed distance between [fromContent] and [toContent]. It is negative if [fromContent] is
      * above or to the left of [toContent].
@@ -222,51 +308,17 @@
      * transition when the distance depends on the size or position of an element that is composed
      * in the content we are going to.
      */
-    fun distance(): Float {
-        if (lastDistance != DistanceUnspecified) {
-            return lastDistance
-        }
+    fun distance(): Float = distance(this)
 
-        val absoluteDistance =
-            with(contentTransition.transformationSpec.distance ?: DefaultSwipeDistance) {
-                userActionDistanceScope.absoluteDistance(
-                    fromContent.targetSize,
-                    orientation,
-                )
-            }
-
-        if (absoluteDistance <= 0f) {
-            return DistanceUnspecified
-        }
-
-        val distance = if (isUpOrLeft) -absoluteDistance else absoluteDistance
-        lastDistance = distance
-        return distance
-    }
-
-    /** Ends any previous [offsetAnimation] and runs the new [animation]. */
-    private fun startOffsetAnimation(animation: () -> OffsetAnimation): OffsetAnimation {
-        cancelOffsetAnimation()
-        return animation().also { offsetAnimation = it }
-    }
-
-    /** Cancel any ongoing offset animation. */
-    // TODO(b/317063114) This should be a suspended function to avoid multiple jobs running at
-    // the same time.
-    fun cancelOffsetAnimation() {
-        val animation = offsetAnimation ?: return
-        offsetAnimation = null
-
-        dragOffset = animation.animatable.value
-        animation.job.cancel()
-    }
+    fun isAnimatingOffset(): Boolean = offsetAnimation != null
 
     fun animateOffset(
-        // TODO(b/317063114) The CoroutineScope should be removed.
-        coroutineScope: CoroutineScope,
         initialVelocity: Float,
         targetContent: T,
-    ): OffsetAnimation {
+        spec: AnimationSpec<Float>? = null,
+    ) {
+        check(!isAnimatingOffset()) { "SwipeAnimation.animateOffset() can only be called once" }
+
         val initialProgress = progress
         // Skip the animation if we have already reached the target content and the overscroll does
         // not animate anything.
@@ -303,85 +355,93 @@
             currentContent = targetContent
         }
 
-        return startOffsetAnimation {
-            val animatable = Animatable(dragOffset, OffsetVisibilityThreshold)
-            val isTargetGreater = targetOffset > animatable.value
-            val startedWhenOvercrollingTargetContent =
-                if (targetContent == fromContent) initialProgress < 0f else initialProgress > 1f
-            val job =
-                coroutineScope
-                    // Important: We start atomically to make sure that we start the coroutine even
-                    // if it is cancelled right after it is launched, so that snapToContent() is
-                    // correctly called. Otherwise, this transition will never be stopped and we
-                    // will never settle to Idle.
-                    .launch(start = CoroutineStart.ATOMIC) {
-                        // TODO(b/327249191): Refactor the code so that we don't even launch a
-                        // coroutine if we don't need to animate.
-                        if (skipAnimation) {
-                            snapToContent(targetContent)
-                            dragOffset = targetOffset
-                            return@launch
-                        }
+        val startProgress =
+            if (contentTransition.previewTransformationSpec != null && targetContent == toContent) {
+                0f
+            } else {
+                dragOffset
+            }
 
-                        try {
-                            val swipeSpec =
-                                contentTransition.transformationSpec.swipeSpec
-                                    ?: layoutImpl.state.transitions.defaultSwipeSpec
-                            animatable.animateTo(
-                                targetValue = targetOffset,
-                                animationSpec = swipeSpec,
-                                initialVelocity = initialVelocity,
-                            ) {
-                                if (bouncingContent == null) {
-                                    val isBouncing =
-                                        if (isTargetGreater) {
-                                            if (startedWhenOvercrollingTargetContent) {
-                                                value >= targetOffset
-                                            } else {
-                                                value > targetOffset
-                                            }
-                                        } else {
-                                            if (startedWhenOvercrollingTargetContent) {
-                                                value <= targetOffset
-                                            } else {
-                                                value < targetOffset
-                                            }
-                                        }
+        val animatable =
+            Animatable(startProgress, OffsetVisibilityThreshold).also { offsetAnimation = it }
 
-                                    if (isBouncing) {
-                                        bouncingContent = targetContent.key
+        check(isAnimatingOffset())
 
-                                        // Immediately stop this transition if we are bouncing on a
-                                        // content that does not bounce.
-                                        if (!contentTransition.isWithinProgressRange(progress)) {
-                                            snapToContent(targetContent)
-                                        }
-                                    }
+        // Note: we still create the animatable and set it on offsetAnimation even when
+        // skipAnimation is true, just so that isUserInputOngoing and isAnimatingOffset() are
+        // unchanged even despite this small skip-optimization (which is just an implementation
+        // detail).
+        if (skipAnimation) {
+            // Unblock the job.
+            offsetAnimationRunnable.complete(null)
+            return
+        }
+
+        val isTargetGreater = targetOffset > animatable.value
+        val startedWhenOvercrollingTargetContent =
+            if (targetContent == fromContent) initialProgress < 0f else initialProgress > 1f
+
+        val swipeSpec =
+            spec
+                ?: contentTransition.transformationSpec.swipeSpec
+                ?: layoutState.transitions.defaultSwipeSpec
+
+        offsetAnimationRunnable.complete {
+            try {
+                animatable.animateTo(
+                    targetValue = targetOffset,
+                    animationSpec = swipeSpec,
+                    initialVelocity = initialVelocity,
+                ) {
+                    if (bouncingContent == null) {
+                        val isBouncing =
+                            if (isTargetGreater) {
+                                if (startedWhenOvercrollingTargetContent) {
+                                    value >= targetOffset
+                                } else {
+                                    value > targetOffset
+                                }
+                            } else {
+                                if (startedWhenOvercrollingTargetContent) {
+                                    value <= targetOffset
+                                } else {
+                                    value < targetOffset
                                 }
                             }
-                        } finally {
-                            snapToContent(targetContent)
+
+                        if (isBouncing) {
+                            bouncingContent = targetContent
+
+                            // Immediately stop this transition if we are bouncing on a content that
+                            // does not bounce.
+                            if (!contentTransition.isWithinProgressRange(progress)) {
+                                throw SnapException()
+                            }
                         }
                     }
-
-            OffsetAnimation(animatable, job)
+                }
+            } catch (_: SnapException) {
+                /* Ignore. */
+            }
         }
     }
 
-    private fun canChangeContent(targetContent: Content): Boolean {
-        val layoutState = layoutImpl.state
+    /** An exception thrown during the animation to stop it immediately. */
+    private class SnapException : Exception()
+
+    private fun canChangeContent(targetContent: ContentKey): Boolean {
         return when (val transition = contentTransition) {
-            is TransitionState.Transition.ChangeCurrentScene ->
-                layoutState.canChangeScene(targetContent.key as SceneKey)
+            is TransitionState.Transition.ChangeScene ->
+                layoutState.canChangeScene(targetContent as SceneKey)
             is TransitionState.Transition.ShowOrHideOverlay -> {
-                if (targetContent.key == transition.overlay) {
+                if (targetContent == transition.overlay) {
                     layoutState.canShowOverlay(transition.overlay)
                 } else {
                     layoutState.canHideOverlay(transition.overlay)
                 }
             }
             is TransitionState.Transition.ReplaceOverlay -> {
-                val to = targetContent.key as OverlayKey
+                val to = targetContent as OverlayKey
                 val from =
                     if (to == transition.toOverlay) transition.fromOverlay else transition.toOverlay
                 layoutState.canReplaceOverlay(from, to)
@@ -389,39 +449,11 @@
         }
     }
 
-    private fun snapToContent(content: T) {
-        cancelOffsetAnimation()
-        check(currentContent == content)
-        layoutImpl.state.finishTransition(contentTransition)
+    fun freezeAndAnimateToCurrentState() {
+        if (isAnimatingOffset()) return
+
+        animateOffset(initialVelocity = 0f, targetContent = currentContent)
     }
-
-    fun finish(): Job {
-        if (isFinishing) return requireNotNull(offsetAnimation).job
-        isFinishing = true
-
-        // If we were already animating the offset, simply return the job.
-        offsetAnimation?.let {
-            return it.job
-        }
-
-        // Animate to the current content.
-        val animation =
-            animateOffset(
-                coroutineScope = layoutImpl.coroutineScope,
-                initialVelocity = 0f,
-                targetContent = currentContent,
-            )
-        check(offsetAnimation == animation)
-        return animation.job
-    }
-
-    internal class OffsetAnimation(
-        /** The animatable used to animate the offset. */
-        val animatable: Animatable<Float, AnimationVector1D>,
-
-        /** The job in which [animatable] is animated. */
-        val job: Job,
-    )
 }
 
 private object DefaultSwipeDistance : UserActionDistance {
@@ -436,21 +468,21 @@
     }
 }
 
-private class ChangeCurrentSceneSwipeTransition(
+private class ChangeSceneSwipeTransition(
     val layoutState: MutableSceneTransitionLayoutStateImpl,
-    val swipeAnimation: SwipeAnimation<Scene>,
+    val swipeAnimation: SwipeAnimation<SceneKey>,
     override val key: TransitionKey?,
-    replacedTransition: ChangeCurrentSceneSwipeTransition?,
+    replacedTransition: ChangeSceneSwipeTransition?,
 ) :
-    TransitionState.Transition.ChangeCurrentScene(
-        swipeAnimation.fromContent.key,
-        swipeAnimation.toContent.key,
+    TransitionState.Transition.ChangeScene(
+        swipeAnimation.fromContent,
+        swipeAnimation.toContent,
         replacedTransition,
     ),
     TransitionState.HasOverscrollProperties by swipeAnimation {
 
     constructor(
-        other: ChangeCurrentSceneSwipeTransition
+        other: ChangeSceneSwipeTransition
     ) : this(
         layoutState = other.layoutState,
         swipeAnimation = SwipeAnimation(other.swipeAnimation),
@@ -463,7 +495,7 @@
     }
 
     override val currentScene: SceneKey
-        get() = swipeAnimation.currentContent.key
+        get() = swipeAnimation.currentContent
 
     override val progress: Float
         get() = swipeAnimation.progress
@@ -471,27 +503,42 @@
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
         get() = swipeAnimation.isUserInputOngoing
 
-    override fun finish(): Job = swipeAnimation.finish()
+    override suspend fun run() {
+        swipeAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        swipeAnimation.freezeAndAnimateToCurrentState()
+    }
 }
 
 private class ShowOrHideOverlaySwipeTransition(
     val layoutState: MutableSceneTransitionLayoutStateImpl,
-    val swipeAnimation: SwipeAnimation<Content>,
-    val _overlay: Overlay,
-    val _fromOrToScene: Scene,
+    val swipeAnimation: SwipeAnimation<ContentKey>,
+    overlay: OverlayKey,
+    fromOrToScene: SceneKey,
     override val key: TransitionKey?,
     replacedTransition: ShowOrHideOverlaySwipeTransition?,
 ) :
     TransitionState.Transition.ShowOrHideOverlay(
-        _overlay.key,
-        _fromOrToScene.key,
-        swipeAnimation.fromContent.key,
-        swipeAnimation.toContent.key,
+        overlay,
+        fromOrToScene,
+        swipeAnimation.fromContent,
+        swipeAnimation.toContent,
         replacedTransition,
     ),
     TransitionState.HasOverscrollProperties by swipeAnimation {
@@ -500,8 +547,8 @@
     ) : this(
         layoutState = other.layoutState,
         swipeAnimation = SwipeAnimation(other.swipeAnimation),
-        _overlay = other._overlay,
-        _fromOrToScene = other._fromOrToScene,
+        overlay = other.overlay,
+        fromOrToScene = other.fromOrToScene,
         key = other.key,
         replacedTransition = other,
     )
@@ -511,7 +558,7 @@
     }
 
     override val isEffectivelyShown: Boolean
-        get() = swipeAnimation.currentContent == _overlay
+        get() = swipeAnimation.currentContent == overlay
 
     override val progress: Float
         get() = swipeAnimation.progress
@@ -519,23 +566,38 @@
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
         get() = swipeAnimation.isUserInputOngoing
 
-    override fun finish(): Job = swipeAnimation.finish()
+    override suspend fun run() {
+        swipeAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        swipeAnimation.freezeAndAnimateToCurrentState()
+    }
 }
 
 private class ReplaceOverlaySwipeTransition(
     val layoutState: MutableSceneTransitionLayoutStateImpl,
-    val swipeAnimation: SwipeAnimation<Overlay>,
+    val swipeAnimation: SwipeAnimation<OverlayKey>,
     override val key: TransitionKey?,
     replacedTransition: ReplaceOverlaySwipeTransition?,
 ) :
     TransitionState.Transition.ReplaceOverlay(
-        swipeAnimation.fromContent.key,
-        swipeAnimation.toContent.key,
+        swipeAnimation.fromContent,
+        swipeAnimation.toContent,
         replacedTransition,
     ),
     TransitionState.HasOverscrollProperties by swipeAnimation {
@@ -553,7 +615,7 @@
     }
 
     override val effectivelyShownOverlay: OverlayKey
-        get() = swipeAnimation.currentContent.key
+        get() = swipeAnimation.currentContent
 
     override val progress: Float
         get() = swipeAnimation.progress
@@ -561,10 +623,25 @@
     override val progressVelocity: Float
         get() = swipeAnimation.progressVelocity
 
+    override val previewProgress: Float
+        get() = swipeAnimation.previewProgress
+
+    override val previewProgressVelocity: Float
+        get() = swipeAnimation.previewProgressVelocity
+
+    override val isInPreviewStage: Boolean
+        get() = swipeAnimation.isInPreviewStage
+
     override val isInitiatedByUserInput: Boolean = true
 
     override val isUserInputOngoing: Boolean
         get() = swipeAnimation.isUserInputOngoing
 
-    override fun finish(): Job = swipeAnimation.finish()
+    override suspend fun run() {
+        swipeAnimation.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        swipeAnimation.freezeAndAnimateToCurrentState()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
index 2b5953c..1f82e0b 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -104,23 +104,23 @@
     ): TransitionSpec
 
     /**
-     * Define the animation to be played when the [scene] is overscrolled in the given
+     * Define the animation to be played when the [content] is overscrolled in the given
      * [orientation].
      *
      * The overscroll animation always starts from a progress of 0f, and reaches 1f when moving the
      * [distance] down/right, -1f when moving in the opposite direction.
      */
     fun overscroll(
-        scene: SceneKey,
+        content: ContentKey,
         orientation: Orientation,
         builder: OverscrollBuilder.() -> Unit,
     ): OverscrollSpec
 
     /**
-     * Prevents overscroll the [scene] in the given [orientation], allowing ancestors to eventually
-     * consume the remaining gesture.
+     * Prevents overscroll the [content] in the given [orientation], allowing ancestors to
+     * eventually consume the remaining gesture.
      */
-    fun overscrollDisabled(scene: SceneKey, orientation: Orientation): OverscrollSpec
+    fun overscrollDisabled(content: ContentKey, orientation: Orientation): OverscrollSpec
 }
 
 interface BaseTransitionBuilder : PropertyTransformationBuilder {
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index 18e356f..da4c8d8 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -84,30 +84,30 @@
     }
 
     override fun overscroll(
-        scene: SceneKey,
+        content: ContentKey,
         orientation: Orientation,
         builder: OverscrollBuilder.() -> Unit
     ): OverscrollSpec {
         val impl = OverscrollBuilderImpl().apply(builder)
         check(impl.transformations.isNotEmpty()) {
             "This method does not allow empty transformations. " +
-                "Use overscrollDisabled($scene, $orientation) instead."
+                "Use overscrollDisabled($content, $orientation) instead."
         }
-        return overscrollSpec(scene, orientation, impl)
+        return overscrollSpec(content, orientation, impl)
     }
 
-    override fun overscrollDisabled(scene: SceneKey, orientation: Orientation): OverscrollSpec {
-        return overscrollSpec(scene, orientation, OverscrollBuilderImpl())
+    override fun overscrollDisabled(content: ContentKey, orientation: Orientation): OverscrollSpec {
+        return overscrollSpec(content, orientation, OverscrollBuilderImpl())
     }
 
     private fun overscrollSpec(
-        scene: SceneKey,
+        content: ContentKey,
         orientation: Orientation,
         impl: OverscrollBuilderImpl,
     ): OverscrollSpec {
         val spec =
             OverscrollSpecImpl(
-                content = scene,
+                content = content,
                 orientation = orientation,
                 transformationSpec =
                     TransformationSpecImpl(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
index fdb019f..a47caaa 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/content/state/TransitionState.kt
@@ -26,7 +26,6 @@
 import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
 import com.android.compose.animation.scene.OverlayKey
-import com.android.compose.animation.scene.OverscrollScope
 import com.android.compose.animation.scene.OverscrollSpecImpl
 import com.android.compose.animation.scene.ProgressVisibilityThreshold
 import com.android.compose.animation.scene.SceneKey
@@ -36,7 +35,6 @@
 import com.android.compose.animation.scene.TransitionKey
 import com.android.compose.animation.scene.transition.link.LinkedTransition
 import com.android.compose.animation.scene.transition.link.StateLink
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
 
 /** The state associated to a [SceneTransitionLayout] at some specific point in time. */
@@ -75,7 +73,7 @@
         val replacedTransition: Transition? = null,
     ) : TransitionState {
         /** A transition animating between [fromScene] and [toScene]. */
-        abstract class ChangeCurrentScene(
+        abstract class ChangeScene(
             /** The scene this transition is starting from. Can't be the same as toScene */
             val fromScene: SceneKey,
 
@@ -301,19 +299,19 @@
             return fromContent == content || toContent == content
         }
 
+        /** Run this transition and return once it is finished. */
+        internal abstract suspend fun run()
+
         /**
-         * Force this transition to finish and animate to an [Idle] state.
+         * Freeze this transition state so that neither [currentScene] nor [currentOverlays] will
+         * change in the future, and animate the progress towards that state. For instance, a
+         * [Transition.ChangeScene] should animate the progress to 0f if its [currentScene] is equal
+         * to its [fromScene][Transition.ChangeScene.fromScene] or animate it to 1f if its equal to
+         * its [toScene][Transition.ChangeScene.toScene].
          *
-         * Important: Once this is called, the effective state of the transition should remain
-         * unchanged. For instance, in the case of a [TransitionState.Transition], its
-         * [currentScene][TransitionState.Transition.currentScene] should never change once [finish]
-         * is called.
-         *
-         * @return the [Job] that animates to the idle state. It can be used to wait until the
-         *   animation is complete or cancel it to snap the animation. Calling [finish] multiple
-         *   times will return the same [Job].
+         * This is called when this transition is interrupted (replaced) by another transition.
          */
-        internal abstract fun finish(): Job
+        internal abstract fun freezeAndAnimateToCurrentState()
 
         internal fun updateOverscrollSpecs(
             fromSpec: OverscrollSpecImpl?,
@@ -351,7 +349,7 @@
 
             fun create(): Animatable<Float, AnimationVector1D> {
                 val animatable = Animatable(1f, visibilityThreshold = ProgressVisibilityThreshold)
-                layoutImpl.coroutineScope.launch {
+                layoutImpl.animationScope.launch {
                     val swipeSpec = layoutImpl.state.transitions.defaultSwipeSpec
                     val progressSpec =
                         spring(
@@ -386,10 +384,10 @@
         val orientation: Orientation
 
         /**
-         * Scope which can be used in the Overscroll DSL to define a transformation based on the
-         * distance between [Transition.fromContent] and [Transition.toContent].
+         * Return the absolute distance between fromScene and toScene, if available, otherwise
+         * [DistanceUnspecified].
          */
-        val overscrollScope: OverscrollScope
+        val absoluteDistance: Float
 
         /**
          * The content (scene or overlay) around which the transition is currently bouncing. When
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
index 59bca50..8f84586 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transformation/Translate.kt
@@ -17,6 +17,7 @@
 package com.android.compose.animation.scene.transformation
 
 import androidx.compose.ui.geometry.Offset
+import androidx.compose.ui.unit.Density
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
 import com.android.compose.animation.scene.ContentKey
@@ -53,6 +54,8 @@
     val x: OverscrollScope.() -> Float = { 0f },
     val y: OverscrollScope.() -> Float = { 0f },
 ) : PropertyTransformation<Offset> {
+    private val cachedOverscrollScope = CachedOverscrollScope()
+
     override fun transform(
         layoutImpl: SceneTransitionLayoutImpl,
         content: ContentKey,
@@ -65,10 +68,47 @@
         // OverscrollSpec only when the transition implements HasOverscrollProperties, we can assume
         // that this method was invoked after performing this check.
         val overscrollProperties = transition as TransitionState.HasOverscrollProperties
+        val overscrollScope =
+            cachedOverscrollScope.getFromCacheOrCompute(layoutImpl.density, overscrollProperties)
 
         return Offset(
-            x = value.x + overscrollProperties.overscrollScope.x(),
-            y = value.y + overscrollProperties.overscrollScope.y(),
+            x = value.x + overscrollScope.x(),
+            y = value.y + overscrollScope.y(),
         )
     }
 }
+
+/**
+ * A helper class to cache a [OverscrollScope] given a [Density] and
+ * [TransitionState.HasOverscrollProperties]. This helps avoid recreating a scope every frame
+ * whenever an overscroll transition is computed.
+ */
+private class CachedOverscrollScope() {
+    private var previousScope: OverscrollScope? = null
+    private var previousDensity: Density? = null
+    private var previousOverscrollProperties: TransitionState.HasOverscrollProperties? = null
+
+    fun getFromCacheOrCompute(
+        density: Density,
+        overscrollProperties: TransitionState.HasOverscrollProperties,
+    ): OverscrollScope {
+        if (
+            previousScope == null ||
+                density != previousDensity ||
+                previousOverscrollProperties != overscrollProperties
+        ) {
+            val scope =
+                object : OverscrollScope, Density by density {
+                    override val absoluteDistance: Float
+                        get() = overscrollProperties.absoluteDistance
+                }
+
+            previousScope = scope
+            previousDensity = density
+            previousOverscrollProperties = overscrollProperties
+            return scope
+        }
+
+        return checkNotNull(previousScope)
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt
new file mode 100644
index 0000000..715d979
--- /dev/null
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/Seek.kt
@@ -0,0 +1,187 @@
+/*
+ * 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.compose.animation.scene.transition
+
+import androidx.annotation.FloatRange
+import androidx.compose.animation.core.AnimationSpec
+import androidx.compose.foundation.gestures.Orientation
+import androidx.compose.ui.util.fastCoerceIn
+import com.android.compose.animation.scene.ContentKey
+import com.android.compose.animation.scene.MutableSceneTransitionLayoutState
+import com.android.compose.animation.scene.MutableSceneTransitionLayoutStateImpl
+import com.android.compose.animation.scene.OverlayKey
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SwipeAnimation
+import com.android.compose.animation.scene.TransitionKey
+import com.android.compose.animation.scene.UserActionResult
+import com.android.compose.animation.scene.createSwipeAnimation
+import kotlin.coroutines.cancellation.CancellationException
+import kotlinx.coroutines.coroutineScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.launch
+
+/**
+ * Seek to the given [scene] using [progress].
+ *
+ * This will start a transition from the
+ * [current scene][MutableSceneTransitionLayoutState.currentScene] to [scene], driven by the
+ * progress in [progress]. Once [progress] stops emitting, we will animate progress to 1f (using
+ * [animationSpec]) if it stopped normally or to 0f if it stopped with a
+ * [kotlin.coroutines.cancellation.CancellationException].
+ */
+suspend fun MutableSceneTransitionLayoutState.seekToScene(
+    scene: SceneKey,
+    @FloatRange(0.0, 1.0) progress: Flow<Float>,
+    transitionKey: TransitionKey? = null,
+    animationSpec: AnimationSpec<Float>? = null,
+) {
+    require(scene != currentScene) {
+        "seekToScene($scene) has to be called with a different scene than the current scene"
+    }
+
+    seek(UserActionResult.ChangeScene(scene, transitionKey), progress, animationSpec)
+}
+
+/**
+ * Seek to show the given [overlay] using [progress].
+ *
+ * This will start a transition to show [overlay] from the
+ * [current scene][MutableSceneTransitionLayoutState.currentScene], driven by the progress in
+ * [progress]. Once [progress] stops emitting, we will animate progress to 1f (using
+ * [animationSpec]) if it stopped normally or to 0f if it stopped with a
+ * [kotlin.coroutines.cancellation.CancellationException].
+ */
+suspend fun MutableSceneTransitionLayoutState.seekToShowOverlay(
+    overlay: OverlayKey,
+    @FloatRange(0.0, 1.0) progress: Flow<Float>,
+    transitionKey: TransitionKey? = null,
+    animationSpec: AnimationSpec<Float>? = null,
+) {
+    require(overlay in currentOverlays) {
+        "seekToShowOverlay($overlay) can be called only when the overlay is in currentOverlays"
+    }
+
+    seek(UserActionResult.ShowOverlay(overlay, transitionKey), progress, animationSpec)
+}
+
+/**
+ * Seek to hide the given [overlay] using [progress].
+ *
+ * This will start a transition to hide [overlay] to the
+ * [current scene][MutableSceneTransitionLayoutState.currentScene], driven by the progress in
+ * [progress]. Once [progress] stops emitting, we will animate progress to 1f (using
+ * [animationSpec]) if it stopped normally or to 0f if it stopped with a
+ * [kotlin.coroutines.cancellation.CancellationException].
+ */
+suspend fun MutableSceneTransitionLayoutState.seekToHideOverlay(
+    overlay: OverlayKey,
+    @FloatRange(0.0, 1.0) progress: Flow<Float>,
+    transitionKey: TransitionKey? = null,
+    animationSpec: AnimationSpec<Float>? = null,
+) {
+    require(overlay !in currentOverlays) {
+        "seekToHideOverlay($overlay) can be called only when the overlay is *not* in " +
+            "currentOverlays"
+    }
+
+    seek(UserActionResult.HideOverlay(overlay, transitionKey), progress, animationSpec)
+}
+
+private suspend fun MutableSceneTransitionLayoutState.seek(
+    result: UserActionResult,
+    progress: Flow<Float>,
+    animationSpec: AnimationSpec<Float>?,
+) {
+    val layoutState =
+        when (this) {
+            is MutableSceneTransitionLayoutStateImpl -> this
+        }
+
+    val swipeAnimation =
+        createSwipeAnimation(
+            layoutState = layoutState,
+            result = result,
+
+            // We are animating progress, so distance is always 1f.
+            distance = 1f,
+
+            // The orientation and isUpOrLeft don't matter here given that they are only used during
+            // overscroll, which is disabled for progress-based transitions.
+            orientation = Orientation.Horizontal,
+            isUpOrLeft = false,
+        )
+
+    animateProgress(
+        state = layoutState,
+        animation = swipeAnimation,
+        progress = progress,
+        commitSpec = animationSpec,
+        cancelSpec = animationSpec,
+    )
+}
+
+internal suspend fun <T : ContentKey> animateProgress(
+    state: MutableSceneTransitionLayoutStateImpl,
+    animation: SwipeAnimation<T>,
+    progress: Flow<Float>,
+    commitSpec: AnimationSpec<Float>?,
+    cancelSpec: AnimationSpec<Float>?,
+) {
+    fun animateOffset(targetContent: T, spec: AnimationSpec<Float>?) {
+        if (state.transitionState != animation.contentTransition || animation.isAnimatingOffset()) {
+            return
+        }
+
+        animation.animateOffset(
+            initialVelocity = 0f,
+            targetContent = targetContent,
+
+            // Important: we have to specify a spec that correctly animates *progress* (low
+            // visibility threshold) and not *offset* (higher visibility threshold).
+            spec = spec ?: animation.contentTransition.transformationSpec.progressSpec,
+        )
+    }
+
+    coroutineScope {
+        val collectionJob = launch {
+            try {
+                progress.collectLatest { progress ->
+                    // Progress based animation should never overscroll given that the
+                    // absoluteDistance exposed to overscroll builders is always 1f and will not
+                    // lead to any noticeable transformation.
+                    animation.dragOffset = progress.fastCoerceIn(0f, 1f)
+                }
+
+                // Transition committed.
+                animateOffset(animation.toContent, commitSpec)
+            } catch (e: CancellationException) {
+                // Transition cancelled.
+                animateOffset(animation.fromContent, cancelSpec)
+            }
+        }
+
+        // Start the transition.
+        state.startTransition(animation.contentTransition)
+
+        // The transition is done. Cancel the collection in case the transition was finished because
+        // it was interrupted by another transition.
+        if (collectionJob.isActive) {
+            collectionJob.cancel()
+        }
+    }
+}
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/link/LinkedTransition.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/link/LinkedTransition.kt
index 59ddb13..42ba9ba 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/link/LinkedTransition.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/transition/link/LinkedTransition.kt
@@ -19,7 +19,6 @@
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.compose.animation.scene.content.state.TransitionState
-import kotlinx.coroutines.Job
 
 /** A linked transition which is driven by a [originalTransition]. */
 internal class LinkedTransition(
@@ -27,7 +26,7 @@
     fromScene: SceneKey,
     toScene: SceneKey,
     override val key: TransitionKey? = null,
-) : TransitionState.Transition.ChangeCurrentScene(fromScene, toScene) {
+) : TransitionState.Transition.ChangeScene(fromScene, toScene) {
 
     override val currentScene: SceneKey
         get() {
@@ -50,5 +49,11 @@
     override val progressVelocity: Float
         get() = originalTransition.progressVelocity
 
-    override fun finish(): Job = originalTransition.finish()
+    override suspend fun run() {
+        originalTransition.run()
+    }
+
+    override fun freezeAndAnimateToCurrentState() {
+        originalTransition.freezeAndAnimateToCurrentState()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
index 8ebb42a..a491349 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/AnimatedSharedAsStateTest.kt
@@ -39,7 +39,10 @@
 import com.android.compose.animation.scene.TestScenes.SceneB
 import com.android.compose.animation.scene.TestScenes.SceneC
 import com.android.compose.animation.scene.TestScenes.SceneD
+import com.android.compose.test.setContentAndCreateMainScope
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertThrows
 import org.junit.Rule
@@ -406,30 +409,33 @@
             }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                // foo goes from 0f to 100f in A => B.
-                scene(SceneA) { animateFloat(0f, foo) }
-                scene(SceneB) { animateFloat(100f, foo) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    // foo goes from 0f to 100f in A => B.
+                    scene(SceneA) { animateFloat(0f, foo) }
+                    scene(SceneB) { animateFloat(100f, foo) }
 
-                // bar goes from 0f to 10f in C => D.
-                scene(SceneC) { animateFloat(0f, bar) }
-                scene(SceneD) { animateFloat(10f, bar) }
+                    // bar goes from 0f to 10f in C => D.
+                    scene(SceneC) { animateFloat(0f, bar) }
+                    scene(SceneD) { animateFloat(10f, bar) }
+                }
             }
-        }
 
-        rule.runOnUiThread {
-            // A => B is at 30%.
+        // A => B is at 30%.
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneA,
                     to = SceneB,
                     progress = { 0.3f },
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
+        }
 
-            // C => D is at 70%.
+        // C => D is at 70%.
+        scope.launch {
             state.startTransition(transition(from = SceneC, to = SceneD, progress = { 0.7f }))
         }
         rule.waitForIdle()
@@ -466,17 +472,18 @@
             }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneA) { animateFloat(0f, key) }
-                scene(SceneB) { animateFloat(100f, key) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) { animateFloat(0f, key) }
+                    scene(SceneB) { animateFloat(100f, key) }
+                }
             }
-        }
 
         // Overscroll on A at -100%: value should be interpolated given that there is no overscroll
         // defined for scene A.
         var progress by mutableStateOf(-1f)
-        rule.runOnIdle {
+        scope.launch {
             state.startTransition(transition(from = SceneA, to = SceneB, progress = { progress }))
         }
         rule.waitForIdle()
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
index 9fa4722..79f82c9 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/DraggableHandlerTest.kt
@@ -41,9 +41,9 @@
 import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.test.MonotonicClockTestScope
 import com.android.compose.test.runMonotonicClockTest
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.launch
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -132,7 +132,10 @@
                     swipeSourceDetector = DefaultEdgeDetector,
                     transitionInterceptionThreshold = transitionInterceptionThreshold,
                     builder = scenesBuilder,
-                    coroutineScope = testScope,
+
+                    // Use testScope and not backgroundScope here because backgroundScope does not
+                    // work well with advanceUntilIdle(), which is used by some tests.
+                    animationScope = testScope,
                 )
                 .apply { setContentsAndLayoutTargetSizeForTest(LAYOUT_SIZE) }
 
@@ -197,6 +200,8 @@
             fromScene: SceneKey? = null,
             toScene: SceneKey? = null,
             progress: Float? = null,
+            previewProgress: Float? = null,
+            isInPreviewStage: Boolean? = null,
             isUserInputOngoing: Boolean? = null
         ): Transition {
             val transition = assertThat(transitionState).isSceneTransition()
@@ -204,6 +209,10 @@
             fromScene?.let { assertThat(transition).hasFromScene(it) }
             toScene?.let { assertThat(transition).hasToScene(it) }
             progress?.let { assertThat(transition).hasProgress(it) }
+            previewProgress?.let { assertThat(transition).hasPreviewProgress(it) }
+            isInPreviewStage?.let {
+                assertThat(transition).run { if (it) isInPreviewStage() else isNotInPreviewStage() }
+            }
             isUserInputOngoing?.let { assertThat(transition).hasIsUserInputOngoing(it) }
             return transition
         }
@@ -301,8 +310,20 @@
         runMonotonicClockTest {
             val testGestureScope = TestGestureScope(testScope = this)
 
-            // run the test
-            testGestureScope.block()
+            try {
+                // Run the test.
+                testGestureScope.block()
+            } finally {
+                // Make sure we stop the last transition if it was not explicitly stopped, otherwise
+                // tests will time out after 10s given that the transitions are now started on the
+                // test scope. We don't use backgroundScope when starting the test transitions
+                // because coroutines started on the background scope don't work well with
+                // advanceUntilIdle(), which is used in a few tests.
+                if (testGestureScope.draggableHandler.isDrivingTransition) {
+                    (testGestureScope.layoutState.transitionState as Transition)
+                        .freezeAndAnimateToCurrentState()
+                }
+            }
         }
     }
 
@@ -338,6 +359,32 @@
     }
 
     @Test
+    fun onDragStoppedAfterDrag_velocityLowerThanThreshold_remainSameScene_previewAnimated() =
+        runGestureTest {
+            layoutState.transitions = transitions {
+                // set a preview for the transition
+                from(SceneA, to = SceneC, preview = {}) {}
+            }
+            val dragController = onDragStarted(overSlop = down(fractionOfScreen = 0.1f))
+            assertTransition(currentScene = SceneA)
+
+            dragController.onDragStopped(velocity = velocityThreshold - 0.01f)
+            runCurrent()
+
+            // verify that transition remains in preview stage and animates back to fromScene
+            assertTransition(
+                currentScene = SceneA,
+                isInPreviewStage = true,
+                previewProgress = 0.1f,
+                progress = 0f
+            )
+
+            // wait for the stop animation
+            advanceUntilIdle()
+            assertIdle(currentScene = SceneA)
+        }
+
+    @Test
     fun onDragStoppedAfterDrag_velocityAtLeastThreshold_goToNextScene() = runGestureTest {
         val dragController = onDragStarted(overSlop = down(fractionOfScreen = 0.1f))
         assertTransition(currentScene = SceneA)
@@ -940,7 +987,7 @@
     }
 
     @Test
-    fun finish() = runGestureTest {
+    fun freezeAndAnimateToCurrentState() = runGestureTest {
         // Start at scene C.
         navigateToSceneC()
 
@@ -952,35 +999,25 @@
         // The current transition can be intercepted.
         assertThat(draggableHandler.shouldImmediatelyIntercept(middle)).isTrue()
 
-        // Finish the transition.
+        // Freeze the transition.
         val transition = transitionState as Transition
-        val job = transition.finish()
+        transition.freezeAndAnimateToCurrentState()
         assertTransition(isUserInputOngoing = false)
-
-        // The current transition can not be intercepted anymore.
-        assertThat(draggableHandler.shouldImmediatelyIntercept(middle)).isFalse()
-
-        // Calling finish() multiple times returns the same Job.
-        assertThat(transition.finish()).isSameInstanceAs(job)
-        assertThat(transition.finish()).isSameInstanceAs(job)
-        assertThat(transition.finish()).isSameInstanceAs(job)
-
-        // We can join the job to wait for the animation to end.
-        assertTransition()
-        job.join()
+        advanceUntilIdle()
         assertIdle(SceneC)
     }
 
     @Test
-    fun finish_cancelled() = runGestureTest {
-        // Swipe up from the middle to transition to scene B.
-        val middle = Offset(SCREEN_SIZE / 2f, SCREEN_SIZE / 2f)
-        onDragStarted(startedPosition = middle, overSlop = up(0.1f))
-        assertTransition(fromScene = SceneA, toScene = SceneB)
+    fun interruptedTransitionCanNotBeImmediatelyIntercepted() = runGestureTest {
+        assertThat(draggableHandler.shouldImmediatelyIntercept(startedPosition = null)).isFalse()
+        onDragStarted(overSlop = up(0.1f))
+        assertThat(draggableHandler.shouldImmediatelyIntercept(startedPosition = null)).isTrue()
 
-        // Finish the transition and cancel the returned job.
-        (transitionState as Transition).finish().cancelAndJoin()
-        assertIdle(SceneA)
+        layoutState.startTransitionImmediately(
+            animationScope = testScope.backgroundScope,
+            transition(SceneA, SceneB)
+        )
+        assertThat(draggableHandler.shouldImmediatelyIntercept(startedPosition = null)).isFalse()
     }
 
     @Test
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index 770c0f8..60596de 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -71,10 +71,11 @@
 import com.android.compose.animation.scene.TestScenes.SceneC
 import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.test.assertSizeIsEqualTo
+import com.android.compose.test.setContentAndCreateMainScope
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertThrows
 import org.junit.Ignore
 import org.junit.Rule
@@ -504,7 +505,7 @@
     }
 
     @Test
-    fun elementModifierNodeIsRecycledInLazyLayouts() = runTest {
+    fun elementModifierNodeIsRecycledInLazyLayouts() {
         val nPages = 2
         val pagerState = PagerState(currentPage = 0) { nPages }
         var nullableLayoutImpl: SceneTransitionLayoutImpl? = null
@@ -630,18 +631,19 @@
                 )
             }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneA) { Box(Modifier.element(TestElements.Foo).size(20.dp)) }
-                scene(SceneB) {}
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) { Box(Modifier.element(TestElements.Foo).size(20.dp)) }
+                    scene(SceneB) {}
+                }
             }
-        }
 
         // Pause the clock to block recompositions.
         rule.mainClock.autoAdvance = false
 
         // Change the current transition.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(transition(from = SceneA, to = SceneB, progress = { 0.5f }))
         }
 
@@ -1296,7 +1298,7 @@
     }
 
     @Test
-    fun interruption() = runTest {
+    fun interruption() {
         // 4 frames of animation.
         val duration = 4 * 16
 
@@ -1336,37 +1338,41 @@
         val valueInC = 200f
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(
-                state,
-                Modifier.size(layoutSize),
-                onLayoutImpl = { layoutImpl = it },
-            ) {
-                // In scene A, Foo is aligned at the TopStart.
-                scene(SceneA) {
-                    Box(Modifier.fillMaxSize()) {
-                        Foo(sizeInA, valueInA, Modifier.align(Alignment.TopStart))
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(
+                    state,
+                    Modifier.size(layoutSize),
+                    onLayoutImpl = { layoutImpl = it },
+                ) {
+                    // In scene A, Foo is aligned at the TopStart.
+                    scene(SceneA) {
+                        Box(Modifier.fillMaxSize()) {
+                            Foo(sizeInA, valueInA, Modifier.align(Alignment.TopStart))
+                        }
                     }
-                }
 
-                // In scene C, Foo is aligned at the BottomEnd, so it moves vertically when coming
-                // from B. We put it before (below) scene B so that we can check that interruptions
-                // values and deltas are properly cleared once all transitions are done.
-                scene(SceneC) {
-                    Box(Modifier.fillMaxSize()) {
-                        Foo(sizeInC, valueInC, Modifier.align(Alignment.BottomEnd))
+                    // In scene C, Foo is aligned at the BottomEnd, so it moves vertically when
+                    // coming
+                    // from B. We put it before (below) scene B so that we can check that
+                    // interruptions
+                    // values and deltas are properly cleared once all transitions are done.
+                    scene(SceneC) {
+                        Box(Modifier.fillMaxSize()) {
+                            Foo(sizeInC, valueInC, Modifier.align(Alignment.BottomEnd))
+                        }
                     }
-                }
 
-                // In scene B, Foo is aligned at the TopEnd, so it moves horizontally when coming
-                // from A.
-                scene(SceneB) {
-                    Box(Modifier.fillMaxSize()) {
-                        Foo(sizeInB, valueInB, Modifier.align(Alignment.TopEnd))
+                    // In scene B, Foo is aligned at the TopEnd, so it moves horizontally when
+                    // coming
+                    // from A.
+                    scene(SceneB) {
+                        Box(Modifier.fillMaxSize()) {
+                            Foo(sizeInB, valueInB, Modifier.align(Alignment.TopEnd))
+                        }
                     }
                 }
             }
-        }
 
         // The offset of Foo when idle in A, B or C.
         val offsetInA = DpOffset.Zero
@@ -1390,12 +1396,12 @@
                 from = SceneA,
                 to = SceneB,
                 progress = { aToBProgress },
-                onFinish = neverFinish(),
+                onFreezeAndAnimate = { /* never finish */ },
             )
         val offsetInAToB = lerp(offsetInA, offsetInB, aToBProgress)
         val sizeInAToB = lerp(sizeInA, sizeInB, aToBProgress)
         val valueInAToB = lerp(valueInA, valueInB, aToBProgress)
-        rule.runOnUiThread { state.startTransition(aToB) }
+        scope.launch { state.startTransition(aToB) }
         rule
             .onNode(isElement(TestElements.Foo, SceneB))
             .assertSizeIsEqualTo(sizeInAToB)
@@ -1415,7 +1421,7 @@
                 progress = { bToCProgress },
                 interruptionProgress = { interruptionProgress },
             )
-        rule.runOnUiThread { state.startTransition(bToC) }
+        scope.launch { state.startTransition(bToC) }
 
         // The interruption deltas, which will be multiplied by the interruption progress then added
         // to the current transition offset and size.
@@ -1476,10 +1482,8 @@
             .assertSizeIsEqualTo(sizeInC)
 
         // Manually finish the transition.
-        rule.runOnUiThread {
-            state.finishTransition(aToB)
-            state.finishTransition(bToC)
-        }
+        aToB.finish()
+        bToC.finish()
         rule.waitForIdle()
         assertThat(state.transitionState).isIdle()
 
@@ -1498,7 +1502,7 @@
     }
 
     @Test
-    fun interruption_sharedTransitionDisabled() = runTest {
+    fun interruption_sharedTransitionDisabled() {
         // 4 frames of animation.
         val duration = 4 * 16
         val layoutSize = DpSize(200.dp, 100.dp)
@@ -1524,21 +1528,22 @@
             Box(modifier.element(TestElements.Foo).size(fooSize))
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state, Modifier.size(layoutSize)) {
-                scene(SceneA) {
-                    Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.TopStart)) }
-                }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state, Modifier.size(layoutSize)) {
+                    scene(SceneA) {
+                        Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.TopStart)) }
+                    }
 
-                scene(SceneB) {
-                    Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.TopEnd)) }
-                }
+                    scene(SceneB) {
+                        Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.TopEnd)) }
+                    }
 
-                scene(SceneC) {
-                    Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.BottomEnd)) }
+                    scene(SceneC) {
+                        Box(Modifier.fillMaxSize()) { Foo(Modifier.align(Alignment.BottomEnd)) }
+                    }
                 }
             }
-        }
 
         // The offset of Foo when idle in A, B or C.
         val offsetInA = DpOffset.Zero
@@ -1547,7 +1552,12 @@
 
         // State is a transition A => B at 50% interrupted by B => C at 30%.
         val aToB =
-            transition(from = SceneA, to = SceneB, progress = { 0.5f }, onFinish = neverFinish())
+            transition(
+                from = SceneA,
+                to = SceneB,
+                progress = { 0.5f },
+                onFreezeAndAnimate = { /* never finish */ },
+            )
         var bToCInterruptionProgress by mutableStateOf(1f)
         val bToC =
             transition(
@@ -1555,11 +1565,11 @@
                 to = SceneC,
                 progress = { 0.3f },
                 interruptionProgress = { bToCInterruptionProgress },
-                onFinish = neverFinish(),
+                onFreezeAndAnimate = { /* never finish */ },
             )
-        rule.runOnUiThread { state.startTransition(aToB) }
+        scope.launch { state.startTransition(aToB) }
         rule.waitForIdle()
-        rule.runOnUiThread { state.startTransition(bToC) }
+        scope.launch { state.startTransition(bToC) }
 
         // Foo is placed in both B and C given that the shared transition is disabled. In B, its
         // offset is impacted by the interruption but in C it is not.
@@ -1579,7 +1589,8 @@
 
         // Manually finish A => B so only B => C is remaining.
         bToCInterruptionProgress = 0f
-        rule.runOnUiThread { state.finishTransition(aToB) }
+        aToB.finish()
+
         rule
             .onNode(isElement(TestElements.Foo, SceneB))
             .assertPositionInRootIsEqualTo(offsetInB.x, offsetInB.y)
@@ -1595,7 +1606,7 @@
                 progress = { 0.7f },
                 interruptionProgress = { 1f },
             )
-        rule.runOnUiThread { state.startTransition(bToA) }
+        scope.launch { state.startTransition(bToA) }
 
         // Foo should have the position it had in B right before the interruption.
         rule
@@ -1609,32 +1620,35 @@
         val state =
             rule.runOnUiThread {
                 MutableSceneTransitionLayoutStateImpl(
-                        SceneA,
-                        transitions { overscrollDisabled(SceneA, Orientation.Horizontal) }
-                    )
-                    .apply {
-                        startTransition(
-                            transition(
-                                from = SceneA,
-                                to = SceneB,
-                                progress = { -1f },
-                                orientation = Orientation.Horizontal
-                            )
-                        )
-                    }
+                    SceneA,
+                    transitions { overscrollDisabled(SceneA, Orientation.Horizontal) }
+                )
             }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(
-                state,
-                Modifier.size(100.dp),
-                onLayoutImpl = { layoutImpl = it },
-            ) {
-                scene(SceneA) {}
-                scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(
+                    state,
+                    Modifier.size(100.dp),
+                    onLayoutImpl = { layoutImpl = it },
+                ) {
+                    scene(SceneA) {}
+                    scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
+                }
             }
+
+        scope.launch {
+            state.startTransition(
+                transition(
+                    from = SceneA,
+                    to = SceneB,
+                    progress = { -1f },
+                    orientation = Orientation.Horizontal
+                )
+            )
         }
+        rule.waitForIdle()
 
         assertThat(layoutImpl.elements).containsKey(TestElements.Foo)
         val foo = layoutImpl.elements.getValue(TestElements.Foo)
@@ -1647,7 +1661,7 @@
     }
 
     @Test
-    fun lastAlphaIsNotSetByOutdatedLayer() = runTest {
+    fun lastAlphaIsNotSetByOutdatedLayer() {
         val state =
             rule.runOnUiThread {
                 MutableSceneTransitionLayoutStateImpl(
@@ -1657,23 +1671,24 @@
             }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
-                scene(SceneA) {}
-                scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
-                scene(SceneC) { Box(Modifier.element(TestElements.Foo)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
+                    scene(SceneA) {}
+                    scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
+                    scene(SceneC) { Box(Modifier.element(TestElements.Foo)) }
+                }
             }
-        }
 
         // Start A => B at 0.5f.
         var aToBProgress by mutableStateOf(0.5f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneA,
                     to = SceneB,
                     progress = { aToBProgress },
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
         }
@@ -1692,7 +1707,7 @@
         assertThat(fooInB.lastAlpha).isEqualTo(0.7f)
 
         // Start B => C at 0.3f.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(transition(from = SceneB, to = SceneC, progress = { 0.3f }))
         }
         rule.waitForIdle()
@@ -1720,16 +1735,17 @@
             }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
-                scene(SceneA) {}
-                scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
+                    scene(SceneA) {}
+                    scene(SceneB) { Box(Modifier.element(TestElements.Foo)) }
+                }
             }
-        }
 
         // Start A => B at 60%.
         var interruptionProgress by mutableStateOf(1f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneA,
@@ -1774,19 +1790,20 @@
             Box(Modifier.element(TestElements.Foo).size(10.dp))
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneA) { Foo() }
-                scene(SceneB) { Foo() }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) { Foo() }
+                    scene(SceneB) { Foo() }
+                }
             }
-        }
 
         rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsDisplayed()
         rule.onNode(isElement(TestElements.Foo, SceneB)).assertDoesNotExist()
 
         // A => B while overscrolling at scene B.
         var progress by mutableStateOf(2f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(transition(from = SceneA, to = SceneB, progress = { progress }))
         }
         rule.waitForIdle()
@@ -1827,19 +1844,20 @@
             MovableElement(key, modifier) { content { Text(text) } }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneA) { MovableFoo(text = fooInA) }
-                scene(SceneB) { MovableFoo(text = fooInB) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) { MovableFoo(text = fooInA) }
+                    scene(SceneB) { MovableFoo(text = fooInB) }
+                }
             }
-        }
 
         rule.onNode(hasText(fooInA)).assertIsDisplayed()
         rule.onNode(hasText(fooInB)).assertDoesNotExist()
 
         // A => B while overscrolling at scene B.
         var progress by mutableStateOf(2f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(transition(from = SceneA, to = SceneB, progress = { progress }))
         }
         rule.waitForIdle()
@@ -1858,7 +1876,7 @@
     }
 
     @Test
-    fun interruptionThenOverscroll() = runTest {
+    fun interruptionThenOverscroll() {
         val state =
             rule.runOnUiThread {
                 MutableSceneTransitionLayoutStateImpl(
@@ -1879,22 +1897,23 @@
             }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state, Modifier.size(200.dp)) {
-                scene(SceneA) { SceneWithFoo(offset = DpOffset.Zero) }
-                scene(SceneB) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 0.dp)) }
-                scene(SceneC) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 40.dp)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                    scene(SceneA) { SceneWithFoo(offset = DpOffset.Zero) }
+                    scene(SceneB) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 0.dp)) }
+                    scene(SceneC) { SceneWithFoo(offset = DpOffset(x = 40.dp, y = 40.dp)) }
+                }
             }
-        }
 
         // Start A => B at 75%.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneA,
                     to = SceneB,
                     progress = { 0.75f },
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
         }
@@ -1907,7 +1926,7 @@
         // Interrupt A => B with B => C at 0%.
         var progress by mutableStateOf(0f)
         var interruptionProgress by mutableStateOf(1f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneB,
@@ -1915,7 +1934,7 @@
                     progress = { progress },
                     interruptionProgress = { interruptionProgress },
                     orientation = Orientation.Vertical,
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
         }
@@ -1963,12 +1982,13 @@
         }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
-                scene(SceneA) { NestedFooBar() }
-                scene(SceneB) { NestedFooBar() }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
+                    scene(SceneA) { NestedFooBar() }
+                    scene(SceneB) { NestedFooBar() }
+                }
             }
-        }
 
         // Idle on A: composed and placed only in B.
         rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsDisplayed()
@@ -1997,7 +2017,7 @@
         assertThat(barInA.lastScale).isNotEqualTo(Scale.Unspecified)
 
         // A => B: composed in both and placed only in B.
-        rule.runOnUiThread { state.startTransition(transition(from = SceneA, to = SceneB)) }
+        scope.launch { state.startTransition(transition(from = SceneA, to = SceneB)) }
         rule.onNode(isElement(TestElements.Foo, SceneA)).assertExists().assertIsNotDisplayed()
         rule.onNode(isElement(TestElements.Bar, SceneA)).assertExists().assertIsNotDisplayed()
         rule.onNode(isElement(TestElements.Foo, SceneB)).assertIsDisplayed()
@@ -2024,7 +2044,7 @@
     }
 
     @Test
-    fun currentTransitionSceneIsUsedToComputeElementValues() = runTest {
+    fun currentTransitionSceneIsUsedToComputeElementValues() {
         val state =
             rule.runOnIdle {
                 MutableSceneTransitionLayoutStateImpl(
@@ -2044,23 +2064,31 @@
             }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state, Modifier.size(200.dp)) {
-                scene(SceneA) { Foo() }
-                scene(SceneB) {}
-                scene(SceneC) { Foo() }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                    scene(SceneA) { Foo() }
+                    scene(SceneB) {}
+                    scene(SceneC) { Foo() }
+                }
             }
-        }
 
         // We have 2 transitions:
         //  - A => B at 100%
         //  - B => C at 0%
         // So Foo should have a size of (40dp, 60dp) in both A and C given that it is scaling its
         // size in B => C.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
-                transition(from = SceneA, to = SceneB, progress = { 1f }, onFinish = neverFinish())
+                transition(
+                    from = SceneA,
+                    to = SceneB,
+                    progress = { 1f },
+                    onFreezeAndAnimate = { /* never finish */ },
+                )
             )
+        }
+        scope.launch {
             state.startTransition(transition(from = SceneB, to = SceneC, progress = { 0f }))
         }
 
@@ -2069,7 +2097,7 @@
     }
 
     @Test
-    fun interruptionDeltasAreProperlyCleaned() = runTest {
+    fun interruptionDeltasAreProperlyCleaned() {
         val state = rule.runOnIdle { MutableSceneTransitionLayoutStateImpl(SceneA) }
 
         @Composable
@@ -2079,18 +2107,24 @@
             }
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state, Modifier.size(200.dp)) {
-                scene(SceneA) { Foo(offset = 0.dp) }
-                scene(SceneB) { Foo(offset = 20.dp) }
-                scene(SceneC) { Foo(offset = 40.dp) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                    scene(SceneA) { Foo(offset = 0.dp) }
+                    scene(SceneB) { Foo(offset = 20.dp) }
+                    scene(SceneC) { Foo(offset = 40.dp) }
+                }
             }
-        }
 
         // Start A => B at 50%.
         val aToB =
-            transition(from = SceneA, to = SceneB, progress = { 0.5f }, onFinish = neverFinish())
-        rule.runOnUiThread { state.startTransition(aToB) }
+            transition(
+                from = SceneA,
+                to = SceneB,
+                progress = { 0.5f },
+                onFreezeAndAnimate = { /* never finish */ },
+            )
+        scope.launch { state.startTransition(aToB) }
         rule.onNode(isElement(TestElements.Foo, SceneB)).assertPositionInRootIsEqualTo(10.dp, 10.dp)
 
         // Start B => C at 0%. This will compute an interruption delta of (-10dp, -10dp) so that the
@@ -2103,9 +2137,9 @@
                 current = { SceneB },
                 progress = { 0f },
                 interruptionProgress = { interruptionProgress },
-                onFinish = neverFinish(),
+                onFreezeAndAnimate = { /* never finish */ },
             )
-        rule.runOnUiThread { state.startTransition(bToC) }
+        scope.launch { state.startTransition(bToC) }
         rule.onNode(isElement(TestElements.Foo, SceneC)).assertPositionInRootIsEqualTo(10.dp, 10.dp)
 
         // Finish the interruption and leave the transition progress at 0f. We should be at the same
@@ -2116,9 +2150,9 @@
         // Finish both transitions but directly start a new one B => A with interruption progress
         // 100%. We should be at (20dp, 20dp), unless the interruption deltas have not been
         // correctly cleaned.
-        rule.runOnUiThread {
-            state.finishTransition(aToB)
-            state.finishTransition(bToC)
+        aToB.finish()
+        bToC.finish()
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneB,
@@ -2132,7 +2166,7 @@
     }
 
     @Test
-    fun lastSizeIsUnspecifiedWhenOverscrollingOtherScene() = runTest {
+    fun lastSizeIsUnspecifiedWhenOverscrollingOtherScene() {
         val state =
             rule.runOnIdle {
                 MutableSceneTransitionLayoutStateImpl(
@@ -2147,17 +2181,23 @@
         }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
-                scene(SceneA) { Foo() }
-                scene(SceneB) { Foo() }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
+                    scene(SceneA) { Foo() }
+                    scene(SceneB) { Foo() }
+                }
             }
-        }
 
         // Overscroll A => B on A.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
-                transition(from = SceneA, to = SceneB, progress = { -1f }, onFinish = neverFinish())
+                transition(
+                    from = SceneA,
+                    to = SceneB,
+                    progress = { -1f },
+                    onFreezeAndAnimate = { /* never finish */ },
+                )
             )
         }
         rule.waitForIdle()
@@ -2173,7 +2213,7 @@
     }
 
     @Test
-    fun transparentElementIsNotImpactingInterruption() = runTest {
+    fun transparentElementIsNotImpactingInterruption() {
         val state =
             rule.runOnIdle {
                 MutableSceneTransitionLayoutStateImpl(
@@ -2200,23 +2240,24 @@
             Box(modifier.element(TestElements.Foo).size(10.dp))
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneB) { Foo(Modifier.offset(40.dp, 60.dp)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneB) { Foo(Modifier.offset(40.dp, 60.dp)) }
 
-                // Define A after B so that Foo is placed in A during A <=> B.
-                scene(SceneA) { Foo() }
+                    // Define A after B so that Foo is placed in A during A <=> B.
+                    scene(SceneA) { Foo() }
+                }
             }
-        }
 
         // Start A => B at 70%.
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneA,
                     to = SceneB,
                     progress = { 0.7f },
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
         }
@@ -2227,14 +2268,14 @@
         // Start B => A at 50% with interruptionProgress = 100%. Foo is placed in A and should still
         // be at (40dp, 60dp) given that it was fully transparent in A before the interruption.
         var interruptionProgress by mutableStateOf(1f)
-        rule.runOnUiThread {
+        scope.launch {
             state.startTransition(
                 transition(
                     from = SceneB,
                     to = SceneA,
                     progress = { 0.5f },
                     interruptionProgress = { interruptionProgress },
-                    onFinish = neverFinish(),
+                    onFreezeAndAnimate = { /* never finish */ },
                 )
             )
         }
@@ -2250,7 +2291,7 @@
     }
 
     @Test
-    fun replacedTransitionDoesNotTriggerInterruption() = runTest {
+    fun replacedTransitionDoesNotTriggerInterruption() {
         val state = rule.runOnIdle { MutableSceneTransitionLayoutStateImpl(SceneA) }
 
         @Composable
@@ -2258,17 +2299,23 @@
             Box(modifier.element(TestElements.Foo).size(10.dp))
         }
 
-        rule.setContent {
-            SceneTransitionLayout(state) {
-                scene(SceneA) { Foo() }
-                scene(SceneB) { Foo(Modifier.offset(40.dp, 60.dp)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state) {
+                    scene(SceneA) { Foo() }
+                    scene(SceneB) { Foo(Modifier.offset(40.dp, 60.dp)) }
+                }
             }
-        }
 
         // Start A => B at 50%.
         val aToB1 =
-            transition(from = SceneA, to = SceneB, progress = { 0.5f }, onFinish = neverFinish())
-        rule.runOnUiThread { state.startTransition(aToB1) }
+            transition(
+                from = SceneA,
+                to = SceneB,
+                progress = { 0.5f },
+                onFreezeAndAnimate = { /* never finish */ },
+            )
+        scope.launch { state.startTransition(aToB1) }
         rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
         rule.onNode(isElement(TestElements.Foo, SceneB)).assertPositionInRootIsEqualTo(20.dp, 30.dp)
 
@@ -2282,7 +2329,7 @@
                 interruptionProgress = { 1f },
                 replacedTransition = aToB1,
             )
-        rule.runOnUiThread { state.startTransition(aToB2) }
+        scope.launch { state.startTransition(aToB2) }
         rule.onNode(isElement(TestElements.Foo, SceneA)).assertIsNotDisplayed()
         rule.onNode(isElement(TestElements.Foo, SceneB)).assertPositionInRootIsEqualTo(40.dp, 60.dp)
     }
@@ -2428,12 +2475,13 @@
         }
 
         lateinit var layoutImpl: SceneTransitionLayoutImpl
-        rule.setContent {
-            SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
-                scene(from) { Box { exitingElements.forEach { Foo(it) } } }
-                scene(to) { Box { enteringElements.forEach { Foo(it) } } }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayoutForTesting(state, onLayoutImpl = { layoutImpl = it }) {
+                    scene(from) { Box { exitingElements.forEach { Foo(it) } } }
+                    scene(to) { Box { enteringElements.forEach { Foo(it) } } }
+                }
             }
-        }
 
         val bToA =
             transition(
@@ -2443,7 +2491,7 @@
                 previewProgress = { previewProgress },
                 isInPreviewStage = { isInPreviewStage }
             )
-        rule.runOnUiThread { state.startTransition(bToA) }
+        scope.launch { state.startTransition(bToA) }
         rule.waitForIdle()
         return layoutImpl
     }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt
index f4e60a2..bc929bd 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/InterruptionHandlerTest.kt
@@ -25,9 +25,9 @@
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.test.runMonotonicClockTest
+import com.android.compose.test.transition
 import com.google.common.truth.Correspondence
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.launch
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -44,8 +44,8 @@
                 transitions { /* default interruption handler */ },
             )
 
-        state.setTargetScene(SceneB, coroutineScope = this)
-        state.setTargetScene(SceneC, coroutineScope = this)
+        state.setTargetScene(SceneB, animationScope = this)
+        state.setTargetScene(SceneC, animationScope = this)
 
         assertThat(state.currentTransitions)
             .comparingElementsUsing(FromToCurrentTriple)
@@ -69,7 +69,7 @@
                     interruptionHandler =
                         object : InterruptionHandler {
                             override fun onInterruption(
-                                interrupted: TransitionState.Transition.ChangeCurrentScene,
+                                interrupted: TransitionState.Transition.ChangeScene,
                                 newTargetScene: SceneKey
                             ): InterruptionResult {
                                 return InterruptionResult(
@@ -81,8 +81,8 @@
                 },
             )
 
-        state.setTargetScene(SceneB, coroutineScope = this)
-        state.setTargetScene(SceneC, coroutineScope = this)
+        state.setTargetScene(SceneB, animationScope = this)
+        state.setTargetScene(SceneC, animationScope = this)
 
         assertThat(state.currentTransitions)
             .comparingElementsUsing(FromToCurrentTriple)
@@ -104,7 +104,7 @@
                     interruptionHandler =
                         object : InterruptionHandler {
                             override fun onInterruption(
-                                interrupted: TransitionState.Transition.ChangeCurrentScene,
+                                interrupted: TransitionState.Transition.ChangeScene,
                                 newTargetScene: SceneKey
                             ): InterruptionResult {
                                 return InterruptionResult(
@@ -124,10 +124,10 @@
 
         // Animate to B and advance the transition a little bit so that progress > visibility
         // threshold and that reversing from B back to A won't immediately snap to A.
-        state.setTargetScene(SceneB, coroutineScope = this)
+        state.setTargetScene(SceneB, animationScope = this)
         testScheduler.advanceTimeBy(duration / 2L)
 
-        state.setTargetScene(SceneC, coroutineScope = this)
+        state.setTargetScene(SceneC, animationScope = this)
 
         assertThat(state.currentTransitions)
             .comparingElementsUsing(FromToCurrentTriple)
@@ -155,13 +155,21 @@
                 // Progress must be > visibility threshold otherwise we will directly snap to A.
                 progress = { 0.5f },
                 progressVelocity = { progressVelocity },
-                onFinish = { launch {} },
             )
-        state.startTransition(aToB)
+        state.startTransitionImmediately(animationScope = backgroundScope, aToB)
 
         // Animate back to A. The previous transition is reversed, i.e. it has the same (from, to)
         // pair, and its velocity is used when animating the progress back to 0.
-        val bToA = checkNotNull(state.setTargetScene(SceneA, coroutineScope = this))
+        val bToA =
+            checkNotNull(
+                    state.setTargetScene(
+                        SceneA,
+                        // We use testScope here and not backgroundScope because setTargetScene
+                        // needs the monotonic clock that is only available in the test scope.
+                        animationScope = this,
+                    )
+                )
+                .first
         testScheduler.runCurrent()
         assertThat(bToA).hasFromScene(SceneA)
         assertThat(bToA).hasToScene(SceneB)
@@ -181,13 +189,21 @@
                 to = SceneB,
                 current = { SceneA },
                 progressVelocity = { progressVelocity },
-                onFinish = { launch {} },
             )
-        state.startTransition(aToB)
+        state.startTransitionImmediately(animationScope = backgroundScope, aToB)
 
         // Animate to B. The previous transition is reversed, i.e. it has the same (from, to) pair,
         // and its velocity is used when animating the progress to 1.
-        val bToA = checkNotNull(state.setTargetScene(SceneB, coroutineScope = this))
+        val bToA =
+            checkNotNull(
+                    state.setTargetScene(
+                        SceneB,
+                        // We use testScope here and not backgroundScope because setTargetScene
+                        // needs the monotonic clock that is only available in the test scope.
+                        animationScope = this,
+                    )
+                )
+                .first
         testScheduler.runCurrent()
         assertThat(bToA).hasFromScene(SceneA)
         assertThat(bToA).hasToScene(SceneB)
@@ -198,7 +214,7 @@
     companion object {
         val FromToCurrentTriple =
             Correspondence.transforming(
-                { transition: TransitionState.Transition.ChangeCurrentScene? ->
+                { transition: TransitionState.Transition.ChangeScene? ->
                     Triple(transition?.fromScene, transition?.toScene, transition?.currentScene)
                 },
                 "(from, to, current) triple"
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementContentPickerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementContentPickerTest.kt
index e1d0945..c8e7e65 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementContentPickerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementContentPickerTest.kt
@@ -17,6 +17,7 @@
 package com.android.compose.animation.scene
 
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import org.junit.Assert.assertThrows
 import org.junit.Test
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
index a549d03..e4879d9 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MovableElementTest.kt
@@ -163,7 +163,7 @@
                             fromContentZIndex: Float,
                             toContentZIndex: Float
                         ): ContentKey {
-                            transition as TransitionState.Transition.ChangeCurrentScene
+                            transition as TransitionState.Transition.ChangeScene
                             assertThat(transition).hasFromScene(SceneA)
                             assertThat(transition).hasToScene(SceneB)
                             assertThat(fromContentZIndex).isEqualTo(0)
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
index 0543e7f..f3161f3 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
@@ -29,6 +29,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.animation.scene.TestScenes.SceneB
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.launch
@@ -89,16 +90,16 @@
             at(0) {
                 val state = observableState()
                 assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
-                assertThat((state as ObservableTransitionState.Transition).fromScene)
+                assertThat((state as ObservableTransitionState.Transition).fromContent)
                     .isEqualTo(SceneA)
-                assertThat(state.toScene).isEqualTo(SceneB)
+                assertThat(state.toContent).isEqualTo(SceneB)
                 assertThat(state.progress()).isEqualTo(0f)
             }
             at(TestTransitionDuration / 2) {
                 val state = observableState()
-                assertThat((state as ObservableTransitionState.Transition).fromScene)
+                assertThat((state as ObservableTransitionState.Transition).fromContent)
                     .isEqualTo(SceneA)
-                assertThat(state.toScene).isEqualTo(SceneB)
+                assertThat(state.toContent).isEqualTo(SceneB)
                 assertThat(state.progress()).isEqualTo(0.5f)
             }
             after {
@@ -139,7 +140,7 @@
         var transitionCurrentScene by mutableStateOf(SceneA)
         val transition =
             transition(from = SceneA, to = SceneB, current = { transitionCurrentScene })
-        state.startTransition(transition)
+        state.startTransitionImmediately(animationScope = backgroundScope, transition)
         assertThat(currentScene.value).isEqualTo(SceneA)
 
         // Change the transition current scene.
@@ -199,7 +200,7 @@
 
         var state = observableState()
         assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
-        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat((state as ObservableTransitionState.Transition).fromContent).isEqualTo(SceneA)
         assertThat(state.previewProgress()).isEqualTo(0.4f)
         assertThat(state.isInPreviewStage()).isEqualTo(true)
 
@@ -217,7 +218,7 @@
         }
         state = observableState()
         assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
-        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat((state as ObservableTransitionState.Transition).fromContent).isEqualTo(SceneA)
         assertThat(state.previewProgress()).isEqualTo(0.4f)
         assertThat(state.isInPreviewStage()).isEqualTo(false)
 
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
index bec2bb2..c25478b 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/OverlayTest.kt
@@ -22,10 +22,12 @@
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.platform.testTag
@@ -42,6 +44,8 @@
 import com.android.compose.animation.scene.TestOverlays.OverlayB
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.test.assertSizeIsEqualTo
+import com.android.compose.test.subjects.assertThat
+import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import org.junit.Rule
 import org.junit.Test
@@ -524,4 +528,124 @@
             }
         }
     }
+
+    @Test
+    fun replaceAnimation_elementInCurrentSceneAndOneOverlay() {
+        val sharedIntKey = ValueKey("sharedInt")
+        val sharedIntValueByContent = mutableMapOf<ContentKey, Int>()
+
+        @Composable
+        fun SceneScope.animateContentInt(targetValue: Int) {
+            val animatedValue = animateContentIntAsState(targetValue, sharedIntKey)
+            LaunchedEffect(animatedValue) {
+                try {
+                    snapshotFlow { animatedValue.value }
+                        .collect { sharedIntValueByContent[contentKey] = it }
+                } finally {
+                    sharedIntValueByContent.remove(contentKey)
+                }
+            }
+        }
+
+        rule.testReplaceOverlayTransition(
+            currentSceneContent = {
+                Box(Modifier.size(width = 180.dp, height = 120.dp)) {
+                    animateContentInt(targetValue = 1_000)
+                    Foo(width = 60.dp, height = 40.dp)
+                }
+            },
+            fromContent = {},
+            fromAlignment = Alignment.TopStart,
+            toContent = {
+                animateContentInt(targetValue = 2_000)
+                Foo(width = 100.dp, height = 80.dp)
+            },
+            transition = {
+                // 4 frames of animation
+                spec = tween(4 * 16, easing = LinearEasing)
+            },
+        ) {
+            // Foo moves from (0,0) with a size of 60x40dp to centered (in a 180x120dp Box) with a
+            // size of 100x80dp, so at (40,20).
+            //
+            // The animated Int goes from 1_000 to 2_000.
+            before {
+                rule
+                    .onNode(isElement(TestElements.Foo, content = SceneA))
+                    .assertSizeIsEqualTo(60.dp, 40.dp)
+                    .assertPositionInRootIsEqualTo(0.dp, 0.dp)
+                rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+                rule.onNode(isElement(TestElements.Foo, content = OverlayB)).assertDoesNotExist()
+
+                assertThat(sharedIntValueByContent).containsEntry(SceneA, 1_000)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayA)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayB)
+            }
+
+            at(16) {
+                rule
+                    .onNode(isElement(TestElements.Foo, content = SceneA))
+                    .assertExists()
+                    .assertIsNotDisplayed()
+                rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+                rule
+                    .onNode(isElement(TestElements.Foo, content = OverlayB))
+                    .assertSizeIsEqualTo(70.dp, 50.dp)
+                    .assertPositionInRootIsEqualTo(10.dp, 5.dp)
+
+                assertThat(sharedIntValueByContent).containsEntry(SceneA, 1_250)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayA)
+                assertThat(sharedIntValueByContent).containsEntry(OverlayB, 1_250)
+            }
+
+            at(32) {
+                rule
+                    .onNode(isElement(TestElements.Foo, content = SceneA))
+                    .assertExists()
+                    .assertIsNotDisplayed()
+                rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+                rule
+                    .onNode(isElement(TestElements.Foo, content = OverlayB))
+                    .assertSizeIsEqualTo(80.dp, 60.dp)
+                    .assertPositionInRootIsEqualTo(20.dp, 10.dp)
+
+                assertThat(sharedIntValueByContent).containsEntry(SceneA, 1_500)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayA)
+                assertThat(sharedIntValueByContent).containsEntry(OverlayB, 1_500)
+            }
+
+            at(48) {
+                rule
+                    .onNode(isElement(TestElements.Foo, content = SceneA))
+                    .assertExists()
+                    .assertIsNotDisplayed()
+                rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+                rule
+                    .onNode(isElement(TestElements.Foo, content = OverlayB))
+                    .assertSizeIsEqualTo(90.dp, 70.dp)
+                    .assertPositionInRootIsEqualTo(30.dp, 15.dp)
+
+                assertThat(sharedIntValueByContent).containsEntry(SceneA, 1_750)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayA)
+                assertThat(sharedIntValueByContent).containsEntry(OverlayB, 1_750)
+            }
+
+            after {
+                rule
+                    .onNode(isElement(TestElements.Foo, content = SceneA))
+                    .assertExists()
+                    .assertIsNotDisplayed()
+                rule.onNode(isElement(TestElements.Foo, content = OverlayA)).assertDoesNotExist()
+                rule
+                    .onNode(isElement(TestElements.Foo, content = OverlayB))
+                    .assertSizeIsEqualTo(100.dp, 80.dp)
+                    .assertPositionInRootIsEqualTo(40.dp, 20.dp)
+
+                // Outside of transitions, the value is equal to the target value in each content.
+                assertThat(sharedIntValueByContent).containsEntry(SceneA, 1_000)
+                assertThat(sharedIntValueByContent).doesNotContainKey(OverlayA)
+                assertThat(sharedIntValueByContent).containsEntry(OverlayB, 2_000)
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
index 00c7588..9284ffd 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/PredictiveBackHandlerTest.kt
@@ -18,12 +18,20 @@
 
 import androidx.activity.BackEventCompat
 import androidx.activity.ComponentActivity
+import androidx.compose.animation.core.LinearEasing
+import androidx.compose.animation.core.tween
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.size
 import androidx.compose.runtime.rememberCoroutineScope
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.assertIsDisplayed
+import androidx.compose.ui.test.hasTestTag
 import androidx.compose.ui.test.junit4.createAndroidComposeRule
+import androidx.compose.ui.unit.dp
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.compose.animation.scene.TestOverlays.OverlayA
+import com.android.compose.animation.scene.TestOverlays.OverlayB
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.animation.scene.TestScenes.SceneB
 import com.android.compose.animation.scene.TestScenes.SceneC
@@ -59,7 +67,23 @@
 
     @Test
     fun testPredictiveBack() {
-        val layoutState = rule.runOnUiThread { MutableSceneTransitionLayoutState(SceneA) }
+        val transitionFrames = 2
+        val layoutState =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    transitions =
+                        transitions {
+                            from(SceneA, to = SceneB) {
+                                spec =
+                                    tween(
+                                        durationMillis = transitionFrames * 16,
+                                        easing = LinearEasing
+                                    )
+                            }
+                        }
+                )
+            }
         rule.setContent {
             SceneTransitionLayout(layoutState) {
                 scene(SceneA, mapOf(Back to SceneB)) { Box(Modifier.fillMaxSize()) }
@@ -88,12 +112,27 @@
         assertThat(layoutState.transitionState).hasCurrentScene(SceneA)
         assertThat(layoutState.transitionState).isIdle()
 
+        rule.mainClock.autoAdvance = false
+
         // Start again and commit it.
         rule.runOnUiThread {
             dispatcher.dispatchOnBackStarted(backEvent())
             dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
             dispatcher.onBackPressed()
         }
+        rule.mainClock.advanceTimeByFrame()
+        rule.waitForIdle()
+        val transition2 = assertThat(layoutState.transitionState).isSceneTransition()
+        // verify that transition picks up progress from preview
+        assertThat(transition2).hasProgress(0.4f, tolerance = 0.0001f)
+
+        rule.mainClock.advanceTimeByFrame()
+        rule.waitForIdle()
+        // verify that transition is half way between preview-end-state (0.4f) and target-state (1f)
+        // after one frame
+        assertThat(transition2).hasProgress(0.7f, tolerance = 0.0001f)
+
+        rule.mainClock.autoAdvance = true
         rule.waitForIdle()
         assertThat(layoutState.transitionState).hasCurrentScene(SceneB)
         assertThat(layoutState.transitionState).isIdle()
@@ -198,6 +237,42 @@
         assertThat(canChangeSceneCalled).isFalse()
     }
 
+    @Test
+    fun backDismissesOverlayWithHighestZIndexByDefault() {
+        val state =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    initialOverlays = setOf(OverlayA, OverlayB)
+                )
+            }
+
+        rule.setContent {
+            SceneTransitionLayout(state, Modifier.size(200.dp)) {
+                scene(SceneA) { Box(Modifier.fillMaxSize()) }
+                overlay(OverlayA) { Box(Modifier.fillMaxSize()) }
+                overlay(OverlayB) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        // Initial state.
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertIsDisplayed()
+
+        // Press back. This should hide overlay B because it has a higher zIndex than overlay A.
+        rule.runOnUiThread { rule.activity.onBackPressedDispatcher.onBackPressed() }
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertDoesNotExist()
+
+        // Press back again. This should hide overlay A.
+        rule.runOnUiThread { rule.activity.onBackPressedDispatcher.onBackPressed() }
+        rule.onNode(hasTestTag(SceneA.testTag)).assertIsDisplayed()
+        rule.onNode(hasTestTag(OverlayA.testTag)).assertDoesNotExist()
+        rule.onNode(hasTestTag(OverlayB.testTag)).assertDoesNotExist()
+    }
+
     private fun backEvent(progress: Float = 0f): BackEventCompat {
         return BackEventCompat(
             touchX = 0f,
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
index 69f2cba..cd20a29a 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutStateTest.kt
@@ -30,14 +30,18 @@
 import com.android.compose.animation.scene.content.state.TransitionState
 import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.animation.scene.transition.link.StateLink
+import com.android.compose.animation.scene.transition.seekToScene
+import com.android.compose.test.MonotonicClockTestScope
+import com.android.compose.test.TestTransition
 import com.android.compose.test.runMonotonicClockTest
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
+import kotlin.coroutines.cancellation.CancellationException
 import kotlinx.coroutines.CoroutineStart
-import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancelAndJoin
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.flow.consumeAsFlow
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.test.runTest
 import org.junit.Rule
 import org.junit.Test
@@ -58,9 +62,12 @@
     }
 
     @Test
-    fun isTransitioningTo_transition() {
+    fun isTransitioningTo_transition() = runTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
-        state.startTransition(transition(from = SceneA, to = SceneB))
+        state.startTransitionImmediately(
+            animationScope = backgroundScope,
+            transition(from = SceneA, to = SceneB)
+        )
 
         assertThat(state.isTransitioning()).isTrue()
         assertThat(state.isTransitioning(from = SceneA)).isTrue()
@@ -73,17 +80,16 @@
     @Test
     fun setTargetScene_idleToSameScene() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutState(SceneA)
-        assertThat(state.setTargetScene(SceneA, coroutineScope = this)).isNull()
+        assertThat(state.setTargetScene(SceneA, animationScope = this)).isNull()
     }
 
     @Test
     fun setTargetScene_idleToDifferentScene() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutState(SceneA)
-        val transition = state.setTargetScene(SceneB, coroutineScope = this)
-        assertThat(transition).isNotNull()
+        val (transition, job) = checkNotNull(state.setTargetScene(SceneB, animationScope = this))
         assertThat(state.transitionState).isEqualTo(transition)
 
-        transition!!.finish().join()
+        job.join()
         assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneB))
     }
 
@@ -91,11 +97,10 @@
     fun setTargetScene_transitionToSameScene() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutState(SceneA)
 
-        val transition = state.setTargetScene(SceneB, coroutineScope = this)
-        assertThat(transition).isNotNull()
-        assertThat(state.setTargetScene(SceneB, coroutineScope = this)).isNull()
+        val (_, job) = checkNotNull(state.setTargetScene(SceneB, animationScope = this))
+        assertThat(state.setTargetScene(SceneB, animationScope = this)).isNull()
 
-        transition!!.finish().join()
+        job.join()
         assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneB))
     }
 
@@ -103,11 +108,10 @@
     fun setTargetScene_transitionToDifferentScene() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutState(SceneA)
 
-        assertThat(state.setTargetScene(SceneB, coroutineScope = this)).isNotNull()
-        val transition = state.setTargetScene(SceneC, coroutineScope = this)
-        assertThat(transition).isNotNull()
+        assertThat(state.setTargetScene(SceneB, animationScope = this)).isNotNull()
+        val (_, job) = checkNotNull(state.setTargetScene(SceneC, animationScope = this))
 
-        transition!!.finish().join()
+        job.join()
         assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneC))
     }
 
@@ -118,7 +122,7 @@
         lateinit var transition: TransitionState.Transition
         val job =
             launch(start = CoroutineStart.UNDISPATCHED) {
-                transition = state.setTargetScene(SceneB, coroutineScope = this)!!
+                transition = checkNotNull(state.setTargetScene(SceneB, animationScope = this)).first
             }
         assertThat(state.transitionState).isEqualTo(transition)
 
@@ -127,18 +131,6 @@
         assertThat(state.transitionState).isEqualTo(TransitionState.Idle(SceneB))
     }
 
-    @Test
-    fun transition_finishReturnsTheSameJobWhenCalledMultipleTimes() = runMonotonicClockTest {
-        val state = MutableSceneTransitionLayoutState(SceneA)
-        val transition = state.setTargetScene(SceneB, coroutineScope = this)
-        assertThat(transition).isNotNull()
-
-        val job = transition!!.finish()
-        assertThat(transition.finish()).isSameInstanceAs(job)
-        assertThat(transition.finish()).isSameInstanceAs(job)
-        assertThat(transition.finish()).isSameInstanceAs(job)
-    }
-
     private fun setupLinkedStates(
         parentInitialScene: SceneKey = SceneC,
         childInitialScene: SceneKey = SceneA,
@@ -163,22 +155,24 @@
     }
 
     @Test
-    fun linkedTransition_startsLinkAndFinishesLinkInToState() {
+    fun linkedTransition_startsLinkAndFinishesLinkInToState() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
         val childTransition = transition(SceneA, SceneB)
 
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneA, SceneB)).isTrue()
         assertThat(parentState.isTransitioning(SceneC, SceneD)).isTrue()
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneB))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneD))
     }
 
     @Test
-    fun linkedTransition_transitiveLink() {
+    fun linkedTransition_transitiveLink() = runTest {
         val parentParentState =
             MutableSceneTransitionLayoutState(SceneB) as MutableSceneTransitionLayoutStateImpl
         val parentLink =
@@ -204,25 +198,27 @@
 
         val childTransition = transition(SceneA, SceneB)
 
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneA, SceneB)).isTrue()
         assertThat(parentState.isTransitioning(SceneC, SceneD)).isTrue()
         assertThat(parentParentState.isTransitioning(SceneB, SceneC)).isTrue()
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneB))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneD))
         assertThat(parentParentState.transitionState).isEqualTo(TransitionState.Idle(SceneC))
     }
 
     @Test
-    fun linkedTransition_linkProgressIsEqual() {
+    fun linkedTransition_linkProgressIsEqual() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
         var progress = 0f
         val childTransition = transition(SceneA, SceneB, progress = { progress })
 
-        childState.startTransition(childTransition)
+        childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(parentState.currentTransition?.progress).isEqualTo(0f)
 
         progress = .5f
@@ -230,28 +226,32 @@
     }
 
     @Test
-    fun linkedTransition_reverseTransitionIsNotLinked() {
+    fun linkedTransition_reverseTransitionIsNotLinked() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
         val childTransition = transition(SceneB, SceneA, current = { SceneB })
 
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneB, SceneA)).isTrue()
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneC))
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneB))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneC))
     }
 
     @Test
-    fun linkedTransition_startsLinkAndFinishesLinkInFromState() {
+    fun linkedTransition_startsLinkAndFinishesLinkInFromState() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
         val childTransition = transition(SceneA, SceneB, current = { SceneA })
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneA))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneC))
     }
@@ -260,22 +260,14 @@
     fun linkedTransition_startsLinkButLinkedStateIsTakenOver() = runTest {
         val (parentState, childState) = setupLinkedStates()
 
-        val childTransition =
-            transition(
-                SceneA,
-                SceneB,
-                onFinish = { launch { /* Do nothing. */ } },
-            )
-        val parentTransition =
-            transition(
-                SceneC,
-                SceneA,
-                onFinish = { launch { /* Do nothing. */ } },
-            )
-        childState.startTransition(childTransition)
-        parentState.startTransition(parentTransition)
+        val childTransition = transition(SceneA, SceneB)
+        val parentTransition = transition(SceneC, SceneA)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
+        parentState.startTransitionImmediately(animationScope = backgroundScope, parentTransition)
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneB))
         assertThat(parentState.transitionState).isEqualTo(parentTransition)
     }
@@ -297,11 +289,11 @@
             )
 
         // Default transition from A to B.
-        assertThat(state.setTargetScene(SceneB, coroutineScope = this)).isNotNull()
+        assertThat(state.setTargetScene(SceneB, animationScope = this)).isNotNull()
         assertThat(state.currentTransition?.transformationSpec?.transformations).hasSize(1)
 
         // Go back to A.
-        state.setTargetScene(SceneA, coroutineScope = this)
+        state.setTargetScene(SceneA, animationScope = this)
         testScheduler.advanceUntilIdle()
         assertThat(state.transitionState).isIdle()
         assertThat(state.transitionState).hasCurrentScene(SceneA)
@@ -310,7 +302,7 @@
         assertThat(
                 state.setTargetScene(
                     SceneB,
-                    coroutineScope = this,
+                    animationScope = this,
                     transitionKey = transitionkey,
                 )
             )
@@ -321,7 +313,8 @@
     @Test
     fun snapToIdleIfClose_snapToStart() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
-        state.startTransition(
+        state.startTransitionImmediately(
+            animationScope = backgroundScope,
             transition(from = SceneA, to = SceneB, current = { SceneA }, progress = { 0.2f })
         )
         assertThat(state.isTransitioning()).isTrue()
@@ -339,7 +332,10 @@
     @Test
     fun snapToIdleIfClose_snapToEnd() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
-        state.startTransition(transition(from = SceneA, to = SceneB, progress = { 0.8f }))
+        state.startTransitionImmediately(
+            animationScope = backgroundScope,
+            transition(from = SceneA, to = SceneB, progress = { 0.8f })
+        )
         assertThat(state.isTransitioning()).isTrue()
 
         // Ignore the request if the progress is not close to 0 or 1, using the threshold.
@@ -356,18 +352,12 @@
     fun snapToIdleIfClose_multipleTransitions() = runMonotonicClockTest {
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
 
-        val aToB =
-            transition(
-                from = SceneA,
-                to = SceneB,
-                progress = { 0.5f },
-                onFinish = { launch { /* do nothing */ } },
-            )
-        state.startTransition(aToB)
+        val aToB = transition(from = SceneA, to = SceneB, progress = { 0.5f })
+        state.startTransitionImmediately(animationScope = backgroundScope, aToB)
         assertThat(state.currentTransitions).containsExactly(aToB).inOrder()
 
         val bToC = transition(from = SceneB, to = SceneC, progress = { 0.8f })
-        state.startTransition(bToC)
+        state.startTransitionImmediately(animationScope = backgroundScope, bToC)
         assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
 
         // Ignore the request if the progress is not close to 0 or 1, using the threshold.
@@ -385,7 +375,8 @@
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, SceneTransitions.Empty)
         var progress by mutableStateOf(0f)
         var currentScene by mutableStateOf(SceneB)
-        state.startTransition(
+        state.startTransitionImmediately(
+            animationScope = backgroundScope,
             transition(
                 from = SceneA,
                 to = SceneB,
@@ -406,47 +397,51 @@
     }
 
     @Test
-    fun linkedTransition_fuzzyLinksAreMatchedAndStarted() {
+    fun linkedTransition_fuzzyLinksAreMatchedAndStarted() = runTest {
         val (parentState, childState) = setupLinkedStates(SceneC, SceneA, null, null, null, SceneD)
         val childTransition = transition(SceneA, SceneB)
 
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneA, SceneB)).isTrue()
         assertThat(parentState.isTransitioning(SceneC, SceneD)).isTrue()
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneB))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneD))
     }
 
     @Test
-    fun linkedTransition_fuzzyLinksAreMatchedAndResetToProperPreviousScene() {
+    fun linkedTransition_fuzzyLinksAreMatchedAndResetToProperPreviousScene() = runTest {
         val (parentState, childState) =
             setupLinkedStates(SceneC, SceneA, SceneA, null, null, SceneD)
 
         val childTransition = transition(SceneA, SceneB, current = { SceneA })
 
-        childState.startTransition(childTransition)
+        val job =
+            childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneA, SceneB)).isTrue()
         assertThat(parentState.isTransitioning(SceneC, SceneD)).isTrue()
 
-        childState.finishTransition(childTransition)
+        childTransition.finish()
+        job.join()
         assertThat(childState.transitionState).isEqualTo(TransitionState.Idle(SceneA))
         assertThat(parentState.transitionState).isEqualTo(TransitionState.Idle(SceneC))
     }
 
     @Test
-    fun linkedTransition_fuzzyLinksAreNotMatched() {
+    fun linkedTransition_fuzzyLinksAreNotMatched() = runTest {
         val (parentState, childState) =
             setupLinkedStates(SceneC, SceneA, SceneB, null, SceneC, SceneD)
         val childTransition = transition(SceneA, SceneB)
 
-        childState.startTransition(childTransition)
+        childState.startTransitionImmediately(animationScope = backgroundScope, childTransition)
         assertThat(childState.isTransitioning(SceneA, SceneB)).isTrue()
         assertThat(parentState.isTransitioning(SceneC, SceneD)).isFalse()
     }
 
-    private fun startOverscrollableTransistionFromAtoB(
+    private fun MonotonicClockTestScope.startOverscrollableTransistionFromAtoB(
         progress: () -> Float,
         sceneTransitions: SceneTransitions,
     ): MutableSceneTransitionLayoutStateImpl {
@@ -455,7 +450,8 @@
                 SceneA,
                 sceneTransitions,
             )
-        state.startTransition(
+        state.startTransitionImmediately(
+            animationScope = backgroundScope,
             transition(
                 from = SceneA,
                 to = SceneB,
@@ -560,54 +556,54 @@
 
     @Test
     fun multipleTransitions() = runTest {
-        val finishingTransitions = mutableSetOf<TransitionState.Transition>()
-        fun onFinish(transition: TransitionState.Transition): Job {
-            // Instead of letting the transition finish, we put the transition in the
-            // finishingTransitions set so that we can verify that finish() is called when expected
-            // and then we call state STLState.finishTransition() ourselves.
-            finishingTransitions.add(transition)
+        val frozenTransitions = mutableSetOf<TestTransition>()
+        fun onFreezeAndAnimate(transition: TestTransition): () -> Unit {
+            // Instead of letting the transition finish when it is frozen, we put the transition in
+            // the frozenTransitions set so that we can verify that freezeAndAnimateToCurrentState()
+            // is called when expected and then we call finish() ourselves to finish the
+            // transitions.
+            frozenTransitions.add(transition)
 
-            return backgroundScope.launch {
-                // Try to acquire a locked mutex so that this code never completes.
-                Mutex(locked = true).withLock {}
-            }
+            return { /* do nothing */ }
         }
 
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, EmptyTestTransitions)
-        val aToB = transition(SceneA, SceneB, onFinish = ::onFinish)
-        val bToC = transition(SceneB, SceneC, onFinish = ::onFinish)
-        val cToA = transition(SceneC, SceneA, onFinish = ::onFinish)
+        val aToB = transition(SceneA, SceneB, onFreezeAndAnimate = ::onFreezeAndAnimate)
+        val bToC = transition(SceneB, SceneC, onFreezeAndAnimate = ::onFreezeAndAnimate)
+        val cToA = transition(SceneC, SceneA, onFreezeAndAnimate = ::onFreezeAndAnimate)
 
         // Starting state.
-        assertThat(finishingTransitions).isEmpty()
+        assertThat(frozenTransitions).isEmpty()
         assertThat(state.currentTransitions).isEmpty()
 
         // A => B.
-        state.startTransition(aToB)
-        assertThat(finishingTransitions).isEmpty()
+        val aToBJob = state.startTransitionImmediately(animationScope = backgroundScope, aToB)
+        assertThat(frozenTransitions).isEmpty()
         assertThat(state.finishedTransitions).isEmpty()
         assertThat(state.currentTransitions).containsExactly(aToB).inOrder()
 
-        // B => C. This should automatically call finish() on aToB.
-        state.startTransition(bToC)
-        assertThat(finishingTransitions).containsExactly(aToB)
+        // B => C. This should automatically call freezeAndAnimateToCurrentState() on aToB.
+        val bToCJob = state.startTransitionImmediately(animationScope = backgroundScope, bToC)
+        assertThat(frozenTransitions).containsExactly(aToB)
         assertThat(state.finishedTransitions).isEmpty()
         assertThat(state.currentTransitions).containsExactly(aToB, bToC).inOrder()
 
-        // C => A. This should automatically call finish() on bToC.
-        state.startTransition(cToA)
-        assertThat(finishingTransitions).containsExactly(aToB, bToC)
+        // C => A. This should automatically call freezeAndAnimateToCurrentState() on bToC.
+        state.startTransitionImmediately(animationScope = backgroundScope, cToA)
+        assertThat(frozenTransitions).containsExactly(aToB, bToC)
         assertThat(state.finishedTransitions).isEmpty()
         assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
 
         // Mark bToC as finished. The list of current transitions does not change because aToB is
         // still not marked as finished.
-        state.finishTransition(bToC)
+        bToC.finish()
+        bToCJob.join()
         assertThat(state.finishedTransitions).containsExactly(bToC)
         assertThat(state.currentTransitions).containsExactly(aToB, bToC, cToA).inOrder()
 
         // Mark aToB as finished. This will remove both aToB and bToC from the list of transitions.
-        state.finishTransition(aToB)
+        aToB.finish()
+        aToBJob.join()
         assertThat(state.finishedTransitions).isEmpty()
         assertThat(state.currentTransitions).containsExactly(cToA).inOrder()
     }
@@ -617,8 +613,9 @@
         val state = MutableSceneTransitionLayoutStateImpl(SceneA, EmptyTestTransitions)
 
         fun startTransition() {
-            val transition = transition(SceneA, SceneB, onFinish = { launch { /* do nothing */ } })
-            state.startTransition(transition)
+            val transition =
+                transition(SceneA, SceneB, onFreezeAndAnimate = { launch { /* do nothing */ } })
+            state.startTransitionImmediately(animationScope = backgroundScope, transition)
         }
 
         var hasLoggedWtf = false
@@ -641,7 +638,7 @@
         val state = MutableSceneTransitionLayoutState(SceneA)
 
         // Transition to B.
-        state.setTargetScene(SceneB, coroutineScope = this)
+        state.setTargetScene(SceneB, animationScope = this)
         val transition = assertThat(state.transitionState).isSceneTransition()
         assertThat(transition).hasCurrentScene(SceneB)
 
@@ -650,4 +647,92 @@
         assertThat(state.transitionState).isIdle()
         assertThat(state.transitionState).hasCurrentScene(SceneC)
     }
+
+    @Test
+    fun snapToScene_freezesCurrentTransition() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutStateImpl(SceneA)
+
+        // Start a transition that is never finished. We don't use backgroundScope on purpose so
+        // that this test would fail if the transition was not frozen when snapping.
+        state.startTransitionImmediately(animationScope = this, transition(SceneA, SceneB))
+        val transition = assertThat(state.transitionState).isSceneTransition()
+        assertThat(transition).hasFromScene(SceneA)
+        assertThat(transition).hasToScene(SceneB)
+
+        // Snap to C.
+        state.snapToScene(SceneC)
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentScene(SceneC)
+    }
+
+    @Test
+    fun seekToScene() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(SceneA)
+        val progress = Channel<Float>()
+
+        val job =
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                state.seekToScene(SceneB, progress.consumeAsFlow())
+            }
+
+        val transition = assertThat(state.transitionState).isSceneTransition()
+        assertThat(transition).hasFromScene(SceneA)
+        assertThat(transition).hasToScene(SceneB)
+        assertThat(transition).hasProgress(0f)
+
+        // Change progress.
+        progress.send(0.4f)
+        assertThat(transition).hasProgress(0.4f)
+
+        // Close the channel normally to confirm the transition.
+        progress.close()
+        job.join()
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentScene(SceneB)
+    }
+
+    @Test
+    fun seekToScene_cancelled() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(SceneA)
+        val progress = Channel<Float>()
+
+        val job =
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                state.seekToScene(SceneB, progress.consumeAsFlow())
+            }
+
+        val transition = assertThat(state.transitionState).isSceneTransition()
+        assertThat(transition).hasFromScene(SceneA)
+        assertThat(transition).hasToScene(SceneB)
+        assertThat(transition).hasProgress(0f)
+
+        // Change progress.
+        progress.send(0.4f)
+        assertThat(transition).hasProgress(0.4f)
+
+        // Close the channel with a CancellationException to cancel the transition.
+        progress.close(CancellationException())
+        job.join()
+        assertThat(state.transitionState).isIdle()
+        assertThat(state.transitionState).hasCurrentScene(SceneA)
+    }
+
+    @Test
+    fun seekToScene_interrupted() = runMonotonicClockTest {
+        val state = MutableSceneTransitionLayoutState(SceneA)
+        val progress = Channel<Float>()
+
+        val job =
+            launch(start = CoroutineStart.UNDISPATCHED) {
+                state.seekToScene(SceneB, progress.consumeAsFlow())
+            }
+
+        assertThat(state.transitionState).isSceneTransition()
+
+        // Start a new transition, interrupting the seek transition.
+        state.setTargetScene(SceneB, animationScope = this)
+
+        // The previous job is cancelled and does not infinitely collect the progress.
+        job.join()
+    }
 }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
index b8e13da..63ab04f 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/SceneTransitionLayoutTest.kt
@@ -55,10 +55,13 @@
 import com.android.compose.animation.scene.TestScenes.SceneC
 import com.android.compose.animation.scene.subjects.assertThat
 import com.android.compose.test.assertSizeIsEqualTo
+import com.android.compose.test.setContentAndCreateMainScope
 import com.android.compose.test.subjects.DpOffsetSubject
 import com.android.compose.test.subjects.assertThat
+import com.android.compose.test.transition
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import org.junit.Assert.assertThrows
 import org.junit.Rule
 import org.junit.Test
@@ -327,17 +330,18 @@
             }
 
         val layoutTag = "layout"
-        rule.setContent {
-            SceneTransitionLayout(state, Modifier.testTag(layoutTag)) {
-                scene(SceneA) { Box(Modifier.size(50.dp)) }
-                scene(SceneB) { Box(Modifier.size(70.dp)) }
+        val scope =
+            rule.setContentAndCreateMainScope {
+                SceneTransitionLayout(state, Modifier.testTag(layoutTag)) {
+                    scene(SceneA) { Box(Modifier.size(50.dp)) }
+                    scene(SceneB) { Box(Modifier.size(70.dp)) }
+                }
             }
-        }
 
         // Overscroll on A at -100%: size should be interpolated given that there is no overscroll
         // defined for scene A.
         var progress by mutableStateOf(-1f)
-        rule.runOnIdle {
+        scope.launch {
             state.startTransition(transition(from = SceneA, to = SceneB, progress = { progress }))
         }
         rule.onNodeWithTag(layoutTag).assertSizeIsEqualTo(30.dp)
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
index bed6cef..f8068e6 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/TransitionDslTest.kt
@@ -232,6 +232,47 @@
     }
 
     @Test
+    fun defaultPredictiveBack() {
+        val transitions = transitions {
+            from(
+                TestScenes.SceneA,
+                to = TestScenes.SceneB,
+                preview = { fractionRange(start = 0.1f, end = 0.8f) { fade(TestElements.Foo) } }
+            ) {
+                spec = tween(500)
+                fractionRange(start = 0.1f, end = 0.8f) { fade(TestElements.Foo) }
+                timestampRange(startMillis = 100, endMillis = 300) { fade(TestElements.Foo) }
+            }
+        }
+
+        // Verify that fetching the transitionSpec with the PredictiveBack key defaults to the above
+        // transition despite it not having the PredictiveBack key set.
+        val transitionSpec =
+            transitions.transitionSpec(
+                from = TestScenes.SceneA,
+                to = TestScenes.SceneB,
+                key = TransitionKey.PredictiveBack
+            )
+
+        val transformations = transitionSpec.transformationSpec().transformations
+
+        assertThat(transformations)
+            .comparingElementsUsing(TRANSFORMATION_RANGE)
+            .containsExactly(
+                TransformationRange(start = 0.1f, end = 0.8f),
+                TransformationRange(start = 100 / 500f, end = 300 / 500f),
+            )
+
+        val previewTransformations = transitionSpec.previewTransformationSpec()?.transformations
+
+        assertThat(previewTransformations)
+            .comparingElementsUsing(TRANSFORMATION_RANGE)
+            .containsExactly(
+                TransformationRange(start = 0.1f, end = 0.8f),
+            )
+    }
+
+    @Test
     fun springSpec() {
         val defaultSpec = spring<Float>(stiffness = 1f)
         val specFromAToC = spring<Float>(stiffness = 2f)
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/subjects/TransitionStateSubject.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/subjects/TransitionStateSubject.kt
index 3fb5708..44e0ba5 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/subjects/TransitionStateSubject.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/subjects/TransitionStateSubject.kt
@@ -32,8 +32,8 @@
     return Truth.assertAbout(TransitionStateSubject.transitionStates()).that(state)
 }
 
-/** Assert on a [TransitionState.Transition.ChangeCurrentScene]. */
-fun assertThat(transition: TransitionState.Transition.ChangeCurrentScene): SceneTransitionSubject {
+/** Assert on a [TransitionState.Transition.ChangeScene]. */
+fun assertThat(transition: TransitionState.Transition.ChangeScene): SceneTransitionSubject {
     return Truth.assertAbout(SceneTransitionSubject.sceneTransitions()).that(transition)
 }
 
@@ -74,14 +74,14 @@
         return actual as TransitionState.Idle
     }
 
-    fun isSceneTransition(): TransitionState.Transition.ChangeCurrentScene {
-        if (actual !is TransitionState.Transition.ChangeCurrentScene) {
+    fun isSceneTransition(): TransitionState.Transition.ChangeScene {
+        if (actual !is TransitionState.Transition.ChangeScene) {
             failWithActual(
                 simpleFact("expected to be TransitionState.Transition.ChangeCurrentScene")
             )
         }
 
-        return actual as TransitionState.Transition.ChangeCurrentScene
+        return actual as TransitionState.Transition.ChangeScene
     }
 
     fun isShowOrHideOverlayTransition(): TransitionState.Transition.ShowOrHideOverlay {
@@ -183,8 +183,8 @@
 class SceneTransitionSubject
 private constructor(
     metadata: FailureMetadata,
-    actual: TransitionState.Transition.ChangeCurrentScene,
-) : BaseTransitionSubject<TransitionState.Transition.ChangeCurrentScene>(metadata, actual) {
+    actual: TransitionState.Transition.ChangeScene,
+) : BaseTransitionSubject<TransitionState.Transition.ChangeScene>(metadata, actual) {
     fun hasFromScene(sceneKey: SceneKey) {
         check("fromScene").that(actual.fromScene).isEqualTo(sceneKey)
     }
@@ -195,7 +195,7 @@
 
     companion object {
         fun sceneTransitions() =
-            Factory { metadata, actual: TransitionState.Transition.ChangeCurrentScene ->
+            Factory { metadata, actual: TransitionState.Transition.ChangeScene ->
                 SceneTransitionSubject(metadata, actual)
             }
     }
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
index 46075c3..5bfc947 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/transformation/AnchoredTranslateTest.kt
@@ -27,7 +27,6 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestElements
 import com.android.compose.animation.scene.testTransition
-import com.android.compose.animation.scene.transition
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/SetContentAndCreateScope.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/SetContentAndCreateScope.kt
new file mode 100644
index 0000000..28a864f
--- /dev/null
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/SetContentAndCreateScope.kt
@@ -0,0 +1,38 @@
+/*
+ * 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.compose.test
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.rememberCoroutineScope
+import androidx.compose.ui.test.junit4.ComposeContentTestRule
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+
+/**
+ * Set [content] as this rule's content and return a [CoroutineScope] bound to [Dispatchers.Main]
+ * and scoped to this rule.
+ */
+fun ComposeContentTestRule.setContentAndCreateMainScope(
+    content: @Composable () -> Unit,
+): CoroutineScope {
+    lateinit var coroutineScope: CoroutineScope
+    setContent {
+        coroutineScope = rememberCoroutineScope(getContext = { Dispatchers.Main })
+        content()
+    }
+    return coroutineScope
+}
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/Transition.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestTransition.kt
similarity index 60%
rename from packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/Transition.kt
rename to packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestTransition.kt
index 1f7fe37..a6a83ee 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/Transition.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/test/TestTransition.kt
@@ -14,17 +14,35 @@
  * limitations under the License.
  */
 
-package com.android.compose.animation.scene
+package com.android.compose.test
 
 import androidx.compose.foundation.gestures.Orientation
+import com.android.compose.animation.scene.ContentKey
+import com.android.compose.animation.scene.SceneKey
+import com.android.compose.animation.scene.SceneTransitionLayoutImpl
 import com.android.compose.animation.scene.content.state.TransitionState
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.sync.Mutex
-import kotlinx.coroutines.sync.withLock
-import kotlinx.coroutines.test.TestScope
+import com.android.compose.animation.scene.content.state.TransitionState.Transition
+import kotlinx.coroutines.CompletableDeferred
 
-/** A utility to easily create a [TransitionState.Transition] in tests. */
+/** A transition for tests that will be finished once [finish] is called. */
+abstract class TestTransition(
+    fromScene: SceneKey,
+    toScene: SceneKey,
+    replacedTransition: Transition?,
+) : Transition.ChangeScene(fromScene, toScene, replacedTransition) {
+    private val finishCompletable = CompletableDeferred<Unit>()
+
+    override suspend fun run() {
+        finishCompletable.await()
+    }
+
+    /** Finish this transition. */
+    fun finish() {
+        finishCompletable.complete(Unit)
+    }
+}
+
+/** A utility to easily create a [TestTransition] in tests. */
 fun transition(
     from: SceneKey,
     to: SceneKey,
@@ -40,12 +58,11 @@
     isUpOrLeft: Boolean = false,
     bouncingContent: ContentKey? = null,
     orientation: Orientation = Orientation.Horizontal,
-    onFinish: ((TransitionState.Transition) -> Job)? = null,
-    replacedTransition: TransitionState.Transition? = null,
-): TransitionState.Transition.ChangeCurrentScene {
+    onFreezeAndAnimate: ((TestTransition) -> Unit)? = null,
+    replacedTransition: Transition? = null,
+): TestTransition {
     return object :
-        TransitionState.Transition.ChangeCurrentScene(from, to, replacedTransition),
-        TransitionState.HasOverscrollProperties {
+        TestTransition(from, to, replacedTransition), TransitionState.HasOverscrollProperties {
         override val currentScene: SceneKey
             get() = current()
 
@@ -69,21 +86,14 @@
         override val isUpOrLeft: Boolean = isUpOrLeft
         override val bouncingContent: ContentKey? = bouncingContent
         override val orientation: Orientation = orientation
-        override val overscrollScope: OverscrollScope =
-            object : OverscrollScope {
-                override val density: Float = 1f
-                override val fontScale: Float = 1f
-                override val absoluteDistance = 0f
+        override val absoluteDistance = 0f
+
+        override fun freezeAndAnimateToCurrentState() {
+            if (onFreezeAndAnimate != null) {
+                onFreezeAndAnimate(this)
+            } else {
+                finish()
             }
-
-        override fun finish(): Job {
-            val onFinish =
-                onFinish
-                    ?: error(
-                        "onFinish() must be provided if finish() is called on test transitions"
-                    )
-
-            return onFinish(this)
         }
 
         override fun interruptionProgress(layoutImpl: SceneTransitionLayoutImpl): Float {
@@ -91,16 +101,3 @@
         }
     }
 }
-
-/**
- * Return a onFinish lambda that can be used with [transition] so that the transition never
- * finishes. This allows to keep the transition in the current transitions list.
- */
-fun TestScope.neverFinish(): (TransitionState.Transition) -> Job {
-    return {
-        backgroundScope.launch {
-            // Try to acquire a locked mutex so that this code never completes.
-            Mutex(locked = true).withLock {}
-        }
-    }
-}
diff --git a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
index c5a5173..25f9564 100644
--- a/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
+++ b/packages/SystemUI/compose/scene/tests/utils/src/com/android/compose/animation/scene/TestTransition.kt
@@ -246,7 +246,7 @@
         content = { play ->
             LaunchedEffect(play) {
                 if (play) {
-                    state.setTargetScene(toScene, coroutineScope = this)
+                    state.setTargetScene(toScene, animationScope = this)
                 }
             }
 
@@ -284,7 +284,7 @@
 
     testTransition(
         state = state,
-        changeState = { state -> state.setTargetScene(to, coroutineScope = this) },
+        changeState = { state -> state.setTargetScene(to, animationScope = this) },
         transitionLayout = transitionLayout,
         builder = builder,
     )
diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
index 362e23d..96d79df 100644
--- a/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/notifications/data/repository/NotificationSettingsRepository.kt
@@ -71,7 +71,7 @@
             .stateIn(
                 scope = backgroundScope,
                 started = SharingStarted.Eagerly,
-                initialValue = false,
+                initialValue = true,
             )
 
     /** The default duration for DND mode when enabled. See [Settings.Secure.ZEN_DURATION]. */
diff --git a/packages/SystemUI/docs/scene.md b/packages/SystemUI/docs/scene.md
index 2f50bbd..0ac15c5 100644
--- a/packages/SystemUI/docs/scene.md
+++ b/packages/SystemUI/docs/scene.md
@@ -121,11 +121,11 @@
 do so by defining their own scene. This section describes how to do that.
 
 Each scene is defined as an implementation of the
-[`ComposableScene`](https://cs.android.com/android/platform/superproject/main/+/main:frameworks/base/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/ComposableScene.kt)
+[`Scene`](https://cs.android.com/android/platform/superproject/main/+/main:frameworks/base/packages/SystemUI/compose/features/src/com/android/systemui/scene/ui/composable/Scene.kt)
 interface, which has three parts: 1. The `key` property returns the
 [`SceneKey`](https://cs.android.com/android/platform/superproject/main/+/main:frameworks/base/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneKey.kt)
-that uniquely identifies that scene 2. The `destinationScenes` `Flow` returns
-the (potentially ever-changing) set of navigation edges to other scenes, based
+that uniquely identifies that scene 2. The `userActions` `Flow` returns
+the (potentially ever-changing) set of navigation edges to other content, based
 on user-actions, which is how the navigation graph is defined (see
 [the Scene navigation](#Scene-navigation) section for more) 3. The `Content`
 function which uses
@@ -138,10 +138,10 @@
 For example:
 
 ```kotlin
-@SysUISingleton class YourScene @Inject constructor( /* your dependencies here */ ) : ComposableScene {
+@SysUISingleton class YourScene @Inject constructor( /* your dependencies here */ ) : Scene {
     override val key = SceneKey.YourScene
 
-    override val destinationScenes: StateFlow<Map<UserAction, SceneModel>> =
+    override val userActions: StateFlow<Map<UserAction, SceneModel>> =
         MutableStateFlow<Map<UserAction, SceneModel>>(
             mapOf(
                 // This is where scene navigation is defined, more on that below.
diff --git a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPasswordViewControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPasswordViewControllerTest.kt
index 7707a60..062d351 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPasswordViewControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPasswordViewControllerTest.kt
@@ -29,6 +29,7 @@
 import com.android.internal.widget.LockPatternUtils
 import com.android.internal.widget.LockscreenCredential
 import com.android.keyguard.domain.interactor.KeyguardKeyboardInteractor
+import com.android.systemui.Flags as AconfigFlags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.classifier.FalsingCollector
 import com.android.systemui.flags.FakeFeatureFlags
@@ -56,7 +57,6 @@
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
-import com.android.systemui.Flags as AconfigFlags
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -66,8 +66,7 @@
 class KeyguardPasswordViewControllerTest : SysuiTestCase() {
     @Mock private lateinit var keyguardPasswordView: KeyguardPasswordView
     @Mock private lateinit var passwordEntry: EditText
-    private var passwordEntryLayoutParams =
-        ViewGroup.LayoutParams(/* width = */ 0, /* height = */ 0)
+    private var passwordEntryLayoutParams = ViewGroup.LayoutParams(/* width= */ 0, /* height= */ 0)
     @Mock lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
     @Mock lateinit var securityMode: KeyguardSecurityModel.SecurityMode
     @Mock lateinit var lockPatternUtils: LockPatternUtils
@@ -85,6 +84,7 @@
     private lateinit var mKeyguardMessageAreaController:
         KeyguardMessageAreaController<BouncerKeyguardMessageArea>
     @Mock private lateinit var postureController: DevicePostureController
+    @Mock private lateinit var mUserActivityNotifier: UserActivityNotifier
     @Captor private lateinit var keyListenerArgumentCaptor: ArgumentCaptor<View.OnKeyListener>
 
     private lateinit var keyguardPasswordViewController: KeyguardPasswordViewController
@@ -106,6 +106,8 @@
         whenever(keyguardPasswordView.findViewById<ImageView>(R.id.switch_ime_button))
             .thenReturn(mock(ImageView::class.java))
         `when`(keyguardPasswordView.resources).thenReturn(context.resources)
+        // TODO(b/362362385): No need to mock keyguardPasswordView.context once this bug is fixed.
+        `when`(keyguardPasswordView.context).thenReturn(context)
         whenever(passwordEntry.layoutParams).thenReturn(passwordEntryLayoutParams)
         val keyguardKeyboardInteractor = KeyguardKeyboardInteractor(FakeKeyboardRepository())
         val fakeFeatureFlags = FakeFeatureFlags()
@@ -130,6 +132,8 @@
                 fakeFeatureFlags,
                 mSelectedUserInteractor,
                 keyguardKeyboardInteractor,
+                null,
+                mUserActivityNotifier
             )
     }
 
@@ -187,9 +191,11 @@
         verify(passwordEntry).setOnKeyListener(keyListenerArgumentCaptor.capture())
 
         val eventHandled =
-            keyListenerArgumentCaptor.value.onKey(keyguardPasswordView,
+            keyListenerArgumentCaptor.value.onKey(
+                keyguardPasswordView,
                 KeyEvent.KEYCODE_SPACE,
-                KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_SPACE))
+                KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_SPACE)
+            )
 
         assertFalse("Unlock attempted.", eventHandled)
     }
@@ -204,9 +210,11 @@
         verify(passwordEntry).setOnKeyListener(keyListenerArgumentCaptor.capture())
 
         val eventHandled =
-            keyListenerArgumentCaptor.value.onKey(keyguardPasswordView,
+            keyListenerArgumentCaptor.value.onKey(
+                keyguardPasswordView,
                 KeyEvent.KEYCODE_ENTER,
-                KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_ENTER))
+                KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_ENTER)
+            )
 
         assertTrue("Unlock not attempted.", eventHandled)
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPinBasedInputViewControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPinBasedInputViewControllerTest.java
index 0054d13..1076d90 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPinBasedInputViewControllerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardPinBasedInputViewControllerTest.java
@@ -87,6 +87,8 @@
     private View mOkButton;
     @Mock
     private SelectedUserInteractor mSelectedUserInteractor;
+    @Mock
+    private UserActivityNotifier mUserActivityNotifier;
     private NumPadKey[] mButtons = new NumPadKey[]{};
 
     private KeyguardPinBasedInputViewController mKeyguardPinViewController;
@@ -117,7 +119,7 @@
                 mKeyguardUpdateMonitor, mSecurityMode, mLockPatternUtils, mKeyguardSecurityCallback,
                 mKeyguardMessageAreaControllerFactory, mLatencyTracker, mLiftToactivateListener,
                 mEmergencyButtonController, mFalsingCollector, featureFlags,
-                mSelectedUserInteractor, keyguardKeyboardInteractor) {
+                mSelectedUserInteractor, keyguardKeyboardInteractor, null, mUserActivityNotifier) {
             @Override
             public void onResume(int reason) {
                 super.onResume(reason);
diff --git a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
index f7f69d3..18ed22a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/keyguard/KeyguardSecurityContainerControllerTest.kt
@@ -56,6 +56,7 @@
 import com.android.systemui.keyboard.data.repository.FakeKeyboardRepository
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintAuthRepository
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.domain.interactor.keyguardDismissTransitionInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.Kosmos
@@ -153,6 +154,7 @@
     @Mock private lateinit var deviceProvisionedController: DeviceProvisionedController
     @Mock private lateinit var postureController: DevicePostureController
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
+    @Mock private lateinit var mUserActivityNotifier: UserActivityNotifier
 
     @Captor
     private lateinit var swipeListenerArgumentCaptor:
@@ -238,6 +240,8 @@
                 featureFlags,
                 mSelectedUserInteractor,
                 keyguardKeyboardInteractor,
+                null,
+                mUserActivityNotifier
             )
 
         kosmos = testKosmos()
@@ -279,7 +283,7 @@
                 deviceProvisionedController,
                 faceAuthAccessibilityDelegate,
                 devicePolicyManager,
-                keyguardTransitionInteractor,
+                kosmos.keyguardDismissTransitionInteractor,
                 { primaryBouncerInteractor },
             ) {
                 deviceEntryInteractor
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/AccessibilityQsShortcutsRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/AccessibilityQsShortcutsRepositoryImplTest.kt
index 460461a..176c3ac 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/AccessibilityQsShortcutsRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/AccessibilityQsShortcutsRepositoryImplTest.kt
@@ -22,10 +22,11 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Rule
@@ -38,13 +39,15 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class AccessibilityQsShortcutsRepositoryImplTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val secureSettings = kosmos.fakeSettings
+
     @Rule @JvmField val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
     // mocks
     @Mock private lateinit var a11yManager: AccessibilityManager
-    private val testDispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
-    private val secureSettings = FakeSettings()
 
     private val userA11yQsShortcutsRepositoryFactory =
         object : UserA11yQsShortcutsRepository.Factory {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorCorrectionRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorCorrectionRepositoryImplTest.kt
index 4e1f82c..801d359 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorCorrectionRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorCorrectionRepositoryImplTest.kt
@@ -23,11 +23,12 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -41,9 +42,10 @@
 
     private val testUser1 = UserHandle.of(1)!!
     private val testUser2 = UserHandle.of(2)!!
-    private val testDispatcher = StandardTestDispatcher()
-    private val scope = TestScope(testDispatcher)
-    private val settings: FakeSettings = FakeSettings()
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val scope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
 
     private lateinit var underTest: ColorCorrectionRepository
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorInversionRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorInversionRepositoryImplTest.kt
index b99dec4..2f457be 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorInversionRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/ColorInversionRepositoryImplTest.kt
@@ -23,11 +23,12 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -41,9 +42,10 @@
 
     private val testUser1 = UserHandle.of(1)!!
     private val testUser2 = UserHandle.of(2)!!
-    private val testDispatcher = StandardTestDispatcher()
-    private val scope = TestScope(testDispatcher)
-    private val settings: FakeSettings = FakeSettings()
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val scope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
 
     private lateinit var underTest: ColorInversionRepository
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/NightDisplayRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/NightDisplayRepositoryTest.kt
index 5757f67..54dbed8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/NightDisplayRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/NightDisplayRepositoryTest.kt
@@ -26,7 +26,9 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.dagger.NightDisplayListenerModule
-import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
 import com.android.systemui.user.utils.UserScopedService
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.eq
@@ -38,8 +40,6 @@
 import com.google.common.truth.Truth.assertThat
 import java.time.LocalTime
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -51,7 +51,7 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class NightDisplayRepositoryTest : SysuiTestCase() {
-    private val kosmos = Kosmos()
+    private val kosmos = testKosmos()
     private val testUser = UserHandle.of(1)!!
     private val testStartTime = LocalTime.MIDNIGHT
     private val testEndTime = LocalTime.NOON
@@ -71,8 +71,8 @@
         }
     private val globalSettings = kosmos.fakeGlobalSettings
     private val secureSettings = kosmos.fakeSettings
-    private val testDispatcher = StandardTestDispatcher()
-    private val scope = TestScope(testDispatcher)
+    private val testDispatcher = kosmos.testDispatcher
+    private val scope = kosmos.testScope
     private val userScopedColorDisplayManager =
         mock<UserScopedService<ColorDisplayManager>> {
             whenever(forUser(eq(testUser))).thenReturn(colorDisplayManager)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/OneHandedModeRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/OneHandedModeRepositoryImplTest.kt
index 1378dac..729d356 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/OneHandedModeRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/OneHandedModeRepositoryImplTest.kt
@@ -22,11 +22,12 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -39,9 +40,10 @@
 
     private val testUser1 = UserHandle.of(1)!!
     private val testUser2 = UserHandle.of(2)!!
-    private val testDispatcher = StandardTestDispatcher()
-    private val scope = TestScope(testDispatcher)
-    private val settings: FakeSettings = FakeSettings()
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val scope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
 
     private val underTest: OneHandedModeRepository =
         OneHandedModeRepositoryImpl(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/UserA11yQsShortcutsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/UserA11yQsShortcutsRepositoryTest.kt
index ce22e28..62f13f8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/UserA11yQsShortcutsRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/data/repository/UserA11yQsShortcutsRepositoryTest.kt
@@ -21,10 +21,11 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -32,9 +33,10 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class UserA11yQsShortcutsRepositoryTest : SysuiTestCase() {
-    private val secureSettings = FakeSettings()
-    private val testDispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val secureSettings = kosmos.fakeSettings
 
     private val underTest =
         UserA11yQsShortcutsRepository(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationControllerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationControllerTest.java
index d7acaaf..80de087 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationControllerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationControllerTest.java
@@ -33,7 +33,7 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.accessibility.utils.TestUtils;
 import com.android.systemui.util.settings.SecureSettings;
-import com.android.wm.shell.common.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
 
 import org.junit.Before;
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/MenuListViewTouchHandlerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/MenuListViewTouchHandlerTest.java
index 4373c88..46f076a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/MenuListViewTouchHandlerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/accessibility/floatingmenu/MenuListViewTouchHandlerTest.java
@@ -46,7 +46,7 @@
 import com.android.systemui.accessibility.MotionEventHelper;
 import com.android.systemui.accessibility.utils.TestUtils;
 import com.android.systemui.util.settings.SecureSettings;
-import com.android.wm.shell.common.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 
 import org.junit.After;
 import org.junit.Before;
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerFullscreenSwipeTouchHandlerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerFullscreenSwipeTouchHandlerTest.java
index d244482..58c3fec 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerFullscreenSwipeTouchHandlerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerFullscreenSwipeTouchHandlerTest.java
@@ -18,6 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
+
 import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -35,15 +37,14 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.logging.UiEventLogger;
-import com.android.internal.widget.LockPatternUtils;
 import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.ambient.touch.scrim.ScrimController;
 import com.android.systemui.ambient.touch.scrim.ScrimManager;
 import com.android.systemui.communal.ui.viewmodel.CommunalViewModel;
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.plugins.ActivityStarter;
-import com.android.systemui.settings.FakeUserTracker;
 import com.android.systemui.shared.system.InputChannelCompat;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
 import com.android.systemui.statusbar.phone.CentralSurfaces;
@@ -57,7 +58,6 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Collections;
 import java.util.Optional;
 
 @SmallTest
@@ -106,15 +106,13 @@
     UiEventLogger mUiEventLogger;
 
     @Mock
-    LockPatternUtils mLockPatternUtils;
-
-    @Mock
     ActivityStarter mActivityStarter;
 
     @Mock
     CommunalViewModel mCommunalViewModel;
 
-    FakeUserTracker mUserTracker;
+    @Mock
+    KeyguardInteractor mKeyguardInteractor;
 
     private static final float TOUCH_REGION = .3f;
     private static final float MIN_BOUNCER_HEIGHT = .05f;
@@ -130,7 +128,6 @@
     public void setup() {
         mKosmos = new KosmosJavaAdapter(this);
         MockitoAnnotations.initMocks(this);
-        mUserTracker = new FakeUserTracker();
         mTouchHandler = new BouncerSwipeTouchHandler(
                 mKosmos.getTestScope(),
                 mScrimManager,
@@ -138,24 +135,21 @@
                 mNotificationShadeWindowController,
                 mValueAnimatorCreator,
                 mVelocityTrackerFactory,
-                mLockPatternUtils,
-                mUserTracker,
                 mCommunalViewModel,
                 mFlingAnimationUtils,
                 mFlingAnimationUtilsClosing,
                 TOUCH_REGION,
                 MIN_BOUNCER_HEIGHT,
                 mUiEventLogger,
-                mActivityStarter);
+                mActivityStarter,
+                mKeyguardInteractor);
 
         when(mScrimManager.getCurrentController()).thenReturn(mScrimController);
         when(mValueAnimatorCreator.create(anyFloat(), anyFloat())).thenReturn(mValueAnimator);
         when(mVelocityTrackerFactory.obtain()).thenReturn(mVelocityTracker);
         when(mFlingAnimationUtils.getMinVelocityPxPerSecond()).thenReturn(Float.MAX_VALUE);
         when(mTouchSession.getBounds()).thenReturn(SCREEN_BOUNDS);
-        when(mLockPatternUtils.isSecure(CURRENT_USER_INFO.id)).thenReturn(true);
-
-        mUserTracker.set(Collections.singletonList(CURRENT_USER_INFO), 0);
+        when(mKeyguardInteractor.isKeyguardDismissible()).thenReturn(MutableStateFlow(false));
     }
 
     /**
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandlerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandlerTest.java
index b85e32b..9568167 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandlerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandlerTest.java
@@ -18,6 +18,8 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static kotlinx.coroutines.flow.StateFlowKt.MutableStateFlow;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.ArgumentMatchers.eq;
@@ -44,16 +46,15 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.logging.UiEventLogger;
-import com.android.internal.widget.LockPatternUtils;
 import com.android.systemui.Flags;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.ambient.touch.scrim.ScrimController;
 import com.android.systemui.ambient.touch.scrim.ScrimManager;
 import com.android.systemui.bouncer.shared.constants.KeyguardBouncerConstants;
 import com.android.systemui.communal.ui.viewmodel.CommunalViewModel;
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.plugins.ActivityStarter;
-import com.android.systemui.settings.FakeUserTracker;
 import com.android.systemui.shade.ShadeExpansionChangeEvent;
 import com.android.systemui.shared.system.InputChannelCompat;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
@@ -69,7 +70,6 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Collections;
 import java.util.Optional;
 
 @SmallTest
@@ -116,9 +116,6 @@
     UiEventLogger mUiEventLogger;
 
     @Mock
-    LockPatternUtils mLockPatternUtils;
-
-    @Mock
     ActivityStarter mActivityStarter;
 
     @Mock
@@ -127,11 +124,12 @@
     @Mock
     CommunalViewModel mCommunalViewModel;
 
+    @Mock
+    KeyguardInteractor mKeyguardInteractor;
+
     @Captor
     ArgumentCaptor<Rect> mRectCaptor;
 
-    FakeUserTracker mUserTracker;
-
     private static final float TOUCH_REGION = .3f;
     private static final int SCREEN_WIDTH_PX = 1024;
     private static final int SCREEN_HEIGHT_PX = 100;
@@ -148,7 +146,6 @@
     public void setup() {
         mKosmos = new KosmosJavaAdapter(this);
         MockitoAnnotations.initMocks(this);
-        mUserTracker = new FakeUserTracker();
         mTouchHandler = new BouncerSwipeTouchHandler(
                 mKosmos.getTestScope(),
                 mScrimManager,
@@ -156,24 +153,21 @@
                 mNotificationShadeWindowController,
                 mValueAnimatorCreator,
                 mVelocityTrackerFactory,
-                mLockPatternUtils,
-                mUserTracker,
                 mCommunalViewModel,
                 mFlingAnimationUtils,
                 mFlingAnimationUtilsClosing,
                 TOUCH_REGION,
                 MIN_BOUNCER_HEIGHT,
                 mUiEventLogger,
-                mActivityStarter);
+                mActivityStarter,
+                mKeyguardInteractor);
 
         when(mScrimManager.getCurrentController()).thenReturn(mScrimController);
         when(mValueAnimatorCreator.create(anyFloat(), anyFloat())).thenReturn(mValueAnimator);
         when(mVelocityTrackerFactory.obtain()).thenReturn(mVelocityTracker);
         when(mFlingAnimationUtils.getMinVelocityPxPerSecond()).thenReturn(Float.MAX_VALUE);
         when(mTouchSession.getBounds()).thenReturn(SCREEN_BOUNDS);
-        when(mLockPatternUtils.isSecure(CURRENT_USER_INFO.id)).thenReturn(true);
-
-        mUserTracker.set(Collections.singletonList(CURRENT_USER_INFO), 0);
+        when(mKeyguardInteractor.isKeyguardDismissible()).thenReturn(MutableStateFlow(false));
     }
 
     /**
@@ -391,7 +385,7 @@
      */
     @Test
     public void testSwipeUp_keyguardNotSecure_doesNotExpand() {
-        when(mLockPatternUtils.isSecure(CURRENT_USER_INFO.id)).thenReturn(false);
+        when(mKeyguardInteractor.isKeyguardDismissible()).thenReturn(MutableStateFlow(true));
         mTouchHandler.onSessionStart(mTouchSession);
         ArgumentCaptor<GestureDetector.OnGestureListener> gestureListenerCaptor =
                 ArgumentCaptor.forClass(GestureDetector.OnGestureListener.class);
@@ -426,7 +420,7 @@
      */
     @Test
     public void testSwipeDown_keyguardNotSecure_doesNotExpand() {
-        when(mLockPatternUtils.isSecure(CURRENT_USER_INFO.id)).thenReturn(false);
+        when(mKeyguardInteractor.isKeyguardDismissible()).thenReturn(MutableStateFlow(true));
         mTouchHandler.onSessionStart(mTouchSession);
         ArgumentCaptor<GestureDetector.OnGestureListener> gestureListenerCaptor =
                 ArgumentCaptor.forClass(GestureDetector.OnGestureListener.class);
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
index 1cd9d76..72163e4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/authentication/data/repository/AuthenticationRepositoryTest.kt
@@ -31,9 +31,7 @@
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.log.table.TableLogBuffer
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.FakeUserRepository
 import com.android.systemui.util.mockito.whenever
@@ -58,14 +56,13 @@
 
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
     @Mock private lateinit var getSecurityMode: Function<Int, KeyguardSecurityModel.SecurityMode>
-    @Mock private lateinit var tableLogger: TableLogBuffer
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val clock = FakeSystemClock()
     private val userRepository = FakeUserRepository()
-    private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
+    private val mobileConnectionsRepository = kosmos.fakeMobileConnectionsRepository
 
     private lateinit var underTest: AuthenticationRepository
 
@@ -78,8 +75,6 @@
         userRepository.setUserInfos(USER_INFOS)
         runBlocking { userRepository.setSelectedUserInfo(USER_INFOS[0]) }
         whenever(getSecurityMode.apply(anyInt())).thenAnswer { currentSecurityMode }
-        mobileConnectionsRepository =
-            FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), tableLogger)
 
         underTest =
             AuthenticationRepositoryImpl(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDebouncerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDebouncerTest.kt
index baef620..a36b0bc 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDebouncerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDebouncerTest.kt
@@ -218,4 +218,60 @@
         assertThat(underTest.getMessageToShow(startWindow)?.msgId)
             .isEqualTo(BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE)
     }
+
+    @Test
+    fun messageMustMeetThreshold() {
+        underTest =
+            FaceHelpMessageDebouncer(
+                window = window,
+                startWindow = 0,
+                shownFaceMessageFrequencyBoost = 0,
+                threshold = .8f,
+            )
+
+        underTest.addMessage(
+            HelpFaceAuthenticationStatus(
+                BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE,
+                "tooClose",
+                0
+            )
+        )
+        underTest.addMessage(
+            HelpFaceAuthenticationStatus(
+                BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE,
+                "tooClose",
+                0
+            )
+        )
+        underTest.addMessage(
+            HelpFaceAuthenticationStatus(
+                BiometricFaceConstants.FACE_ACQUIRED_TOO_BRIGHT,
+                "tooBright",
+                0
+            )
+        )
+
+        // although tooClose message is the majority, it doesn't meet the 80% threshold
+        assertThat(underTest.getMessageToShow(startWindow)).isNull()
+
+        underTest.addMessage(
+            HelpFaceAuthenticationStatus(
+                BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE,
+                "tooClose",
+                0
+            )
+        )
+        underTest.addMessage(
+            HelpFaceAuthenticationStatus(
+                BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE,
+                "tooClose",
+                0
+            )
+        )
+
+        // message shows once it meets the threshold
+        assertThat(underTest.getMessageToShow(startWindow)?.msgId)
+            .isEqualTo(BiometricFaceConstants.FACE_ACQUIRED_TOO_CLOSE)
+        assertThat(underTest.getMessageToShow(startWindow)?.msg).isEqualTo("tooClose")
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDeferralTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDeferralTest.kt
index b31f6f5..add7a7f 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDeferralTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/biometrics/FaceHelpMessageDeferralTest.kt
@@ -16,11 +16,15 @@
 
 package com.android.systemui.biometrics
 
-import androidx.test.ext.junit.runners.AndroidJUnit4
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.FlagsParameterization
 import androidx.test.filters.SmallTest
 import com.android.keyguard.logging.BiometricMessageDeferralLogger
+import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.util.time.FakeSystemClock
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertFalse
 import org.junit.Assert.assertNull
@@ -31,14 +35,29 @@
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(ParameterizedAndroidJunit4::class)
 @android.platform.test.annotations.EnabledOnRavenwood
-class FaceHelpMessageDeferralTest : SysuiTestCase() {
+class FaceHelpMessageDeferralTest(flags: FlagsParameterization) : SysuiTestCase() {
     val threshold = .75f
     @Mock lateinit var logger: BiometricMessageDeferralLogger
     @Mock lateinit var dumpManager: DumpManager
+    val systemClock = FakeSystemClock()
+
+    companion object {
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams(): List<FlagsParameterization> {
+            return FlagsParameterization.allCombinationsOf(Flags.FLAG_FACE_MESSAGE_DEFER_UPDATE)
+        }
+    }
+
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags)
+    }
 
     @Before
     fun setUp() {
@@ -111,10 +130,11 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_FACE_MESSAGE_DEFER_UPDATE)
     fun testReturnsMostFrequentDeferredMessage() {
         val biometricMessageDeferral = createMsgDeferral(setOf(1, 2))
 
-        // WHEN there's 80%of the messages are msgId=1 and 20% is msgId=2
+        // WHEN there's 80% of the messages are msgId=1 and 20% is msgId=2
         biometricMessageDeferral.processFrame(1)
         biometricMessageDeferral.processFrame(1)
         biometricMessageDeferral.processFrame(1)
@@ -124,7 +144,41 @@
         biometricMessageDeferral.processFrame(2)
         biometricMessageDeferral.updateMessage(2, "msgId-2")
 
-        // THEN the most frequent deferred message is that meets the threshold is returned
+        // THEN the most frequent deferred message that meets the threshold is returned
+        assertEquals("msgId-1", biometricMessageDeferral.getDeferredMessage())
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_FACE_MESSAGE_DEFER_UPDATE)
+    fun testReturnsMostFrequentDeferredMessage_onlyAnalyzesLastNWindow() {
+        val biometricMessageDeferral = createMsgDeferral(setOf(1, 2))
+
+        // WHEN there's 80% of the messages are msgId=1 and 20% is msgId=2, but the last
+        // N window only contains messages with msgId=2
+        repeat(80) { biometricMessageDeferral.processFrame(1) }
+        biometricMessageDeferral.updateMessage(1, "msgId-1")
+        systemClock.setElapsedRealtime(systemClock.elapsedRealtime() + 501L)
+        repeat(20) { biometricMessageDeferral.processFrame(2) }
+        biometricMessageDeferral.updateMessage(2, "msgId-2")
+
+        // THEN the most frequent deferred message in the last N window (500L) is returned
+        assertEquals("msgId-2", biometricMessageDeferral.getDeferredMessage())
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_FACE_MESSAGE_DEFER_UPDATE)
+    fun testReturnsMostFrequentDeferredMessage_analyzesAllFrames() {
+        val biometricMessageDeferral = createMsgDeferral(setOf(1, 2))
+
+        // WHEN there's 80% of the messages are msgId=1 and 20% is msgId=2, but the last
+        // N window only contains messages with msgId=2
+        repeat(80) { biometricMessageDeferral.processFrame(1) }
+        biometricMessageDeferral.updateMessage(1, "msgId-1")
+        systemClock.setElapsedRealtime(systemClock.elapsedRealtime() + 501L)
+        repeat(20) { biometricMessageDeferral.processFrame(2) }
+        biometricMessageDeferral.updateMessage(2, "msgId-2")
+
+        // THEN the most frequent deferred message is returned
         assertEquals("msgId-1", biometricMessageDeferral.getDeferredMessage())
     }
 
@@ -213,14 +267,17 @@
     private fun createMsgDeferral(
         messagesToDefer: Set<Int>,
         acquiredInfoToIgnore: Set<Int> = emptySet(),
+        windowToAnalyzeLastNFrames: Long = 500L,
     ): BiometricMessageDeferral {
         return BiometricMessageDeferral(
-            messagesToDefer,
-            acquiredInfoToIgnore,
-            threshold,
-            logger,
-            dumpManager,
-            "0",
+            messagesToDefer = messagesToDefer,
+            acquiredInfoToIgnore = acquiredInfoToIgnore,
+            threshold = threshold,
+            windowToAnalyzeLastNFrames = windowToAnalyzeLastNFrames,
+            logBuffer = logger,
+            dumpManager = dumpManager,
+            id = "0",
+            systemClock = { systemClock },
         )
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelTest.kt
index c161525..8c8faee 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelTest.kt
@@ -19,9 +19,11 @@
 import android.content.pm.UserInfo
 import android.hardware.biometrics.BiometricFaceConstants
 import android.hardware.fingerprint.FingerprintManager
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.widget.LockPatternUtils
+import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.FakeAuthenticationRepository
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
@@ -35,7 +37,6 @@
 import com.android.systemui.biometrics.data.repository.fakeFingerprintPropertyRepository
 import com.android.systemui.biometrics.shared.model.SensorStrength
 import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
-import com.android.systemui.bouncer.shared.flag.fakeComposeBouncerFlags
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.domain.interactor.DeviceUnlockedInteractor
 import com.android.systemui.deviceentry.shared.model.ErrorFaceAuthenticationStatus
@@ -71,6 +72,7 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
+@EnableFlags(Flags.FLAG_COMPOSE_BOUNCER)
 class BouncerMessageViewModelTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
@@ -82,7 +84,6 @@
     @Before
     fun setUp() {
         kosmos.fakeUserRepository.setUserInfos(listOf(PRIMARY_USER))
-        kosmos.fakeComposeBouncerFlags.composeBouncerEnabled = true
         overrideResource(
             R.array.config_face_acquire_device_entry_ignorelist,
             intArrayOf(ignoreHelpMessageId)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModelTest.kt
similarity index 93%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModelTest.kt
index a86a0c0..f58bbc3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModelTest.kt
@@ -44,17 +44,17 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 @EnableSceneContainer
-class BouncerSceneActionsViewModelTest : SysuiTestCase() {
+class BouncerUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
-    private lateinit var underTest: BouncerSceneActionsViewModel
+    private lateinit var underTest: BouncerUserActionsViewModel
 
     @Before
     fun setUp() {
         kosmos.sceneContainerStartable.start()
-        underTest = kosmos.bouncerSceneActionsViewModel
+        underTest = kosmos.bouncerUserActionsViewModel
         underTest.activateIn(testScope)
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalDreamStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalDreamStartableTest.kt
index 76920e4..3b0057d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalDreamStartableTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalDreamStartableTest.kt
@@ -81,6 +81,7 @@
     @Test
     fun startDreamWhenTransitioningToHub() =
         testScope.runTest {
+            keyguardRepository.setKeyguardShowing(true)
             keyguardRepository.setDreaming(false)
             powerRepository.setScreenPowerState(ScreenPowerState.SCREEN_ON)
             whenever(dreamManager.canStartDreaming(/* isScreenOn= */ true)).thenReturn(true)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
index d251585..1a426d6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepositoryImplTest.kt
@@ -28,6 +28,7 @@
 import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.plugins.BcSmartspaceDataPlugin.SmartspaceTargetListener
 import com.android.systemui.testKosmos
 import com.android.systemui.util.time.fakeSystemClock
@@ -67,6 +68,7 @@
                 smartspaceController,
                 fakeExecutor,
                 systemClock,
+                logcatLogBuffer("CommunalSmartspaceRepositoryImplTest"),
             )
     }
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalTutorialRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalTutorialRepositoryImplTest.kt
index c37b33e..ae1c496 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalTutorialRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalTutorialRepositoryImplTest.kt
@@ -29,7 +29,7 @@
 import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.FakeUserRepository
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -45,8 +45,7 @@
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-
-    private lateinit var secureSettings: FakeSettings
+    private val secureSettings = kosmos.fakeSettings
     private lateinit var userRepository: FakeUserRepository
 
     private lateinit var underTest: CommunalTutorialRepositoryImpl
@@ -55,7 +54,6 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        secureSettings = FakeSettings()
         userRepository = FakeUserRepository()
         val listOfUserInfo = listOf(MAIN_USER_INFO)
         userRepository.setUserInfos(listOfUserInfo)
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 1d03ced..777ddab 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
@@ -68,15 +68,16 @@
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.settings.fakeUserTracker
+import com.android.systemui.statusbar.phone.fakeManagedProfileController
 import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.FakeUserRepository
 import com.android.systemui.user.data.repository.fakeUserRepository
 import com.android.systemui.util.mockito.any
 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.utils.leaks.FakeManagedProfileController
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -122,6 +123,7 @@
     private lateinit var userTracker: FakeUserTracker
     private lateinit var activityStarter: ActivityStarter
     private lateinit var userManager: UserManager
+    private lateinit var managedProfileController: FakeManagedProfileController
 
     private lateinit var underTest: CommunalInteractor
 
@@ -143,6 +145,7 @@
         userTracker = kosmos.fakeUserTracker
         activityStarter = kosmos.activityStarter
         userManager = kosmos.userManager
+        managedProfileController = kosmos.fakeManagedProfileController
 
         whenever(mainUser.isMain).thenReturn(true)
         whenever(secondaryUser.isMain).thenReturn(false)
@@ -352,7 +355,7 @@
 
             smartspaceRepository.setTimers(targets)
 
-            val smartspaceContent by collectLastValue(underTest.getOngoingContent(true))
+            val smartspaceContent by collectLastValue(underTest.ongoingContent)
             assertThat(smartspaceContent?.size).isEqualTo(totalTargets)
             for (index in 0 until totalTargets) {
                 assertThat(smartspaceContent?.get(index)?.size).isEqualTo(expectedSizes[index])
@@ -368,7 +371,7 @@
             // Media is playing.
             mediaRepository.mediaActive()
 
-            val umoContent by collectLastValue(underTest.getOngoingContent(true))
+            val umoContent by collectLastValue(underTest.ongoingContent)
 
             assertThat(umoContent?.size).isEqualTo(1)
             assertThat(umoContent?.get(0)).isInstanceOf(CommunalContentModel.Umo::class.java)
@@ -376,20 +379,6 @@
         }
 
     @Test
-    fun umo_mediaPlaying_doNotShowUmo() =
-        testScope.run {
-            // Tutorial completed.
-            tutorialRepository.setTutorialSettingState(HUB_MODE_TUTORIAL_COMPLETED)
-
-            // Media is playing.
-            mediaRepository.mediaActive()
-
-            val umoContent by collectLastValue(underTest.getOngoingContent(false))
-
-            assertThat(umoContent?.size).isEqualTo(0)
-        }
-
-    @Test
     fun ongoing_shouldOrderAndSizeByTimestamp() =
         testScope.runTest {
             // Keyguard showing, and tutorial completed.
@@ -412,7 +401,7 @@
             val timer3 = smartspaceTimer("timer3", timestamp = 4L)
             smartspaceRepository.setTimers(listOf(timer1, timer2, timer3))
 
-            val ongoingContent by collectLastValue(underTest.getOngoingContent(true))
+            val ongoingContent by collectLastValue(underTest.ongoingContent)
             assertThat(ongoingContent?.size).isEqualTo(4)
             assertThat(ongoingContent?.get(0)?.key)
                 .isEqualTo(CommunalContentModel.KEY.smartspace("timer3"))
@@ -1084,6 +1073,14 @@
         }
     }
 
+    @Test
+    fun unpauseWorkProfileEnablesWorkMode() =
+        testScope.runTest {
+            underTest.unpauseWorkProfile()
+
+            assertThat(managedProfileController.isWorkModeEnabled()).isTrue()
+        }
+
     private fun setKeyguardFeaturesDisabled(user: UserInfo, disabledFlags: Int) {
         whenever(kosmos.devicePolicyManager.getKeyguardDisabledFeatures(nullable(), eq(user.id)))
             .thenReturn(disabledFlags)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
index 8218178..179ba22 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalEditModeViewModelTest.kt
@@ -20,9 +20,7 @@
 import android.content.ActivityNotFoundException
 import android.content.ComponentName
 import android.content.Intent
-import android.content.pm.ActivityInfo
 import android.content.pm.PackageManager
-import android.content.pm.ResolveInfo
 import android.content.pm.UserInfo
 import android.provider.Settings
 import android.view.accessibility.AccessibilityEvent
@@ -72,7 +70,6 @@
 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
 import org.mockito.Mockito.never
@@ -141,6 +138,7 @@
                 context,
                 accessibilityManager,
                 packageManager,
+                WIDGET_PICKER_PACKAGE_NAME,
             )
     }
 
@@ -259,18 +257,8 @@
     @Test
     fun onOpenWidgetPicker_launchesWidgetPickerActivity() {
         testScope.runTest {
-            whenever(packageManager.resolveActivity(any(), anyInt())).then {
-                ResolveInfo().apply {
-                    activityInfo = ActivityInfo().apply { packageName = WIDGET_PICKER_PACKAGE_NAME }
-                }
-            }
-
             val success =
-                underTest.onOpenWidgetPicker(
-                    testableResources.resources,
-                    packageManager,
-                    activityResultLauncher
-                )
+                underTest.onOpenWidgetPicker(testableResources.resources, activityResultLauncher)
 
             verify(activityResultLauncher).launch(any())
             assertTrue(success)
@@ -278,38 +266,14 @@
     }
 
     @Test
-    fun onOpenWidgetPicker_launcherActivityNotResolved_doesNotLaunchWidgetPickerActivity() {
-        testScope.runTest {
-            whenever(packageManager.resolveActivity(any(), anyInt())).thenReturn(null)
-
-            val success =
-                underTest.onOpenWidgetPicker(
-                    testableResources.resources,
-                    packageManager,
-                    activityResultLauncher
-                )
-
-            verify(activityResultLauncher, never()).launch(any())
-            assertFalse(success)
-        }
-    }
-
-    @Test
     fun onOpenWidgetPicker_activityLaunchThrowsException_failure() {
         testScope.runTest {
-            whenever(packageManager.resolveActivity(any(), anyInt())).then {
-                ResolveInfo().apply {
-                    activityInfo = ActivityInfo().apply { packageName = WIDGET_PICKER_PACKAGE_NAME }
-                }
-            }
-
             whenever(activityResultLauncher.launch(any()))
                 .thenThrow(ActivityNotFoundException::class.java)
 
             val success =
                 underTest.onOpenWidgetPicker(
                     testableResources.resources,
-                    packageManager,
                     activityResultLauncher,
                 )
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
index f6f5bc0..780d357 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/view/viewmodel/CommunalViewModelTest.kt
@@ -97,6 +97,7 @@
 import org.mockito.Mockito
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.atLeastOnce
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
@@ -757,7 +758,7 @@
 
             // updateViewVisibility is called when the flow is collected.
             assertThat(communalContent).isNotNull()
-            verify(mediaHost).updateViewVisibility()
+            verify(mediaHost, atLeastOnce()).updateViewVisibility()
         }
 
     @Test
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 3e75ceb..71abed7 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
@@ -75,7 +75,7 @@
 import com.android.systemui.log.FaceAuthenticationLogger
 import com.android.systemui.log.SessionTracker
 import com.android.systemui.log.logcatLogBuffer
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
@@ -90,12 +90,12 @@
 import com.android.systemui.util.mockito.captureMany
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.io.PrintWriter
 import java.io.StringWriter
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
@@ -198,25 +198,8 @@
         fmOverride: FaceManager? = faceManager,
         bypassControllerOverride: KeyguardBypassController? = bypassController
     ): DeviceEntryFaceAuthRepositoryImpl {
-        val systemClock = FakeSystemClock()
-        val faceAuthBuffer =
-            TableLogBuffer(
-                10,
-                "face auth",
-                systemClock,
-                mock(),
-                testDispatcher,
-                testScope.backgroundScope
-            )
-        val faceDetectBuffer =
-            TableLogBuffer(
-                10,
-                "face detect",
-                systemClock,
-                mock(),
-                testDispatcher,
-                testScope.backgroundScope
-            )
+        val faceAuthBuffer = logcatTableLogBuffer(kosmos, "face auth")
+        val faceDetectBuffer = logcatTableLogBuffer(kosmos, "face detect")
 
         return DeviceEntryFaceAuthRepositoryImpl(
             mContext,
@@ -561,14 +544,29 @@
     fun withSceneContainerEnabled_authenticateDoesNotRunWhenKeyguardIsGoingAway() =
         testScope.runTest {
             testGatingCheckForFaceAuth(sceneContainerEnabled = true) {
-                keyguardTransitionRepository.sendTransitionStep(
-                    TransitionStep(
-                        KeyguardState.LOCKSCREEN,
-                        KeyguardState.UNDEFINED,
-                        value = 0.5f,
-                        transitionState = TransitionState.RUNNING
-                    ),
-                    validateStep = false
+                kosmos.sceneInteractor.setTransitionState(
+                    MutableStateFlow(
+                        ObservableTransitionState.Transition(
+                            fromScene = Scenes.Bouncer,
+                            toScene = Scenes.Gone,
+                            currentScene = flowOf(Scenes.Bouncer),
+                            progress = MutableStateFlow(0.2f),
+                            isInitiatedByUserInput = true,
+                            isUserInputOngoing = flowOf(false),
+                        )
+                    )
+                )
+                runCurrent()
+            }
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun withSceneContainerEnabled_authenticateDoesNotRunWhenLockscreenIsGone() =
+        testScope.runTest {
+            testGatingCheckForFaceAuth(sceneContainerEnabled = true) {
+                kosmos.sceneInteractor.setTransitionState(
+                    MutableStateFlow(ObservableTransitionState.Idle(Scenes.Gone))
                 )
                 runCurrent()
             }
@@ -898,15 +896,32 @@
     fun withSceneContainer_faceDetectDoesNotRunWhenKeyguardGoingAway() =
         testScope.runTest {
             testGatingCheckForDetect(sceneContainerEnabled = true) {
-                keyguardTransitionRepository.sendTransitionStep(
-                    TransitionStep(
-                        KeyguardState.LOCKSCREEN,
-                        KeyguardState.UNDEFINED,
-                        value = 0.5f,
-                        transitionState = TransitionState.RUNNING
-                    ),
-                    validateStep = false
+                kosmos.sceneInteractor.setTransitionState(
+                    MutableStateFlow(
+                        ObservableTransitionState.Transition(
+                            fromScene = Scenes.Bouncer,
+                            toScene = Scenes.Gone,
+                            currentScene = flowOf(Scenes.Bouncer),
+                            progress = MutableStateFlow(0.2f),
+                            isInitiatedByUserInput = true,
+                            isUserInputOngoing = flowOf(false),
+                        )
+                    )
                 )
+
+                runCurrent()
+            }
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun withSceneContainer_faceDetectDoesNotRunWhenLockscreenIsGone() =
+        testScope.runTest {
+            testGatingCheckForDetect(sceneContainerEnabled = true) {
+                kosmos.sceneInteractor.setTransitionState(
+                    MutableStateFlow(ObservableTransitionState.Idle(Scenes.Gone))
+                )
+
                 runCurrent()
             }
         }
@@ -1231,6 +1246,9 @@
                 TransitionStep(KeyguardState.OFF, KeyguardState.LOCKSCREEN, value = 1.0f),
                 validateStep = false
             )
+            kosmos.sceneInteractor.setTransitionState(
+                MutableStateFlow(ObservableTransitionState.Idle(Scenes.Lockscreen))
+            )
         } else {
             keyguardRepository.setKeyguardGoingAway(false)
         }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
index eda9039..5e6ff73 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/dreams/DreamOverlayServiceTest.kt
@@ -15,6 +15,7 @@
  */
 package com.android.systemui.dreams
 
+import android.app.WindowConfiguration
 import android.content.ComponentName
 import android.content.Intent
 import android.os.RemoteException
@@ -65,6 +66,8 @@
 import com.android.systemui.keyguard.gesture.domain.gestureInteractor
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.navigationbar.gestural.domain.GestureInteractor
+import com.android.systemui.navigationbar.gestural.domain.TaskInfo
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher
 import com.android.systemui.testKosmos
 import com.android.systemui.touch.TouchInsetManager
 import com.android.systemui.util.concurrency.FakeExecutor
@@ -83,13 +86,17 @@
 import org.mockito.Mockito
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.isNull
+import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.eq
+import org.mockito.kotlin.firstValue
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verifyZeroInteractions
 import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -315,6 +322,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -332,6 +340,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -351,6 +360,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -373,6 +383,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -391,6 +402,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -406,6 +418,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -413,6 +426,47 @@
     }
 
     @Test
+    fun testDeferredResetRespondsToAnimationEnd() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*isPreview*/,
+            true /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        whenever(mStateController.areExitAnimationsRunning()).thenReturn(true)
+        clearInvocations(mStateController, mTouchMonitor)
+
+        // Starting a dream will cause it to end first.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*isPreview*/,
+            true /*shouldShowComplication*/
+        )
+
+        mMainExecutor.runAllReady()
+
+        verifyZeroInteractions(mTouchMonitor)
+
+        val captor = ArgumentCaptor.forClass(DreamOverlayStateController.Callback::class.java)
+        verify(mStateController).addCallback(captor.capture())
+
+        whenever(mStateController.areExitAnimationsRunning()).thenReturn(false)
+
+        captor.firstValue.onStateChanged()
+
+        // Should only be called once since it should be null during the second reset.
+        verify(mTouchMonitor).destroy()
+    }
+
+    @Test
     fun testLowLightSetByStartDream() {
         val client = client
 
@@ -421,6 +475,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -437,6 +492,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             HOME_CONTROL_PANEL_DREAM_COMPONENT.flattenToString(),
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -453,6 +509,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -487,6 +544,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         // Immediately end the dream.
@@ -518,6 +576,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -537,6 +596,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             LOW_LIGHT_COMPONENT.flattenToString(),
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -588,6 +648,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -611,6 +672,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -631,6 +693,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -660,6 +723,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -683,6 +747,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -708,6 +773,7 @@
                 mWindowParams,
                 mDreamOverlayCallback,
                 DREAM_COMPONENT,
+                false /*isPreview*/,
                 false /*shouldShowComplication*/
             )
             mMainExecutor.runAllReady()
@@ -733,6 +799,7 @@
                 mWindowParams,
                 mDreamOverlayCallback,
                 DREAM_COMPONENT,
+                false /*isPreview*/,
                 false /*shouldShowComplication*/
             )
             // Set communal available, verify that overlay callback is informed.
@@ -761,6 +828,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -781,6 +849,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -800,6 +869,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             true /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -823,6 +893,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -853,6 +924,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         testScope.runCurrent()
@@ -869,6 +941,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -892,6 +965,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -923,6 +997,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -954,6 +1029,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -989,6 +1065,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
@@ -1015,7 +1092,7 @@
     }
 
     @Test
-    fun testDreamActivityGesturesBlockedOnStart() {
+    fun testDreamActivityGesturesBlockedWhenDreaming() {
         val client = client
 
         // Inform the overlay service of dream starting.
@@ -1023,36 +1100,75 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
-        val captor = argumentCaptor<ComponentName>()
+
+        val matcherCaptor = argumentCaptor<TaskMatcher>()
         verify(gestureInteractor)
-            .addGestureBlockedActivity(captor.capture(), eq(GestureInteractor.Scope.Global))
-        assertThat(captor.firstValue.packageName)
-            .isEqualTo(ComponentName.unflattenFromString(DREAM_COMPONENT)?.packageName)
-    }
+            .addGestureBlockedMatcher(matcherCaptor.capture(), eq(GestureInteractor.Scope.Global))
+        val matcher = matcherCaptor.firstValue
 
-    @Test
-    fun testDreamActivityGesturesUnblockedOnEnd() {
-        val client = client
-
-        // Inform the overlay service of dream starting.
-        client.startDream(
-            mWindowParams,
-            mDreamOverlayCallback,
-            DREAM_COMPONENT,
-            false /*shouldShowComplication*/
-        )
-        mMainExecutor.runAllReady()
+        val dreamTaskInfo = TaskInfo(mock<ComponentName>(), WindowConfiguration.ACTIVITY_TYPE_DREAM)
+        assertThat(matcher.matches(dreamTaskInfo)).isTrue()
 
         client.endDream()
         mMainExecutor.runAllReady()
-        val captor = argumentCaptor<ComponentName>()
+
         verify(gestureInteractor)
-            .removeGestureBlockedActivity(captor.capture(), eq(GestureInteractor.Scope.Global))
-        assertThat(captor.firstValue.packageName)
-            .isEqualTo(ComponentName.unflattenFromString(DREAM_COMPONENT)?.packageName)
+            .removeGestureBlockedMatcher(eq(matcher), eq(GestureInteractor.Scope.Global))
+    }
+
+    @Test
+    fun testDreamActivityGesturesNotBlockedWhenPreview() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            true /*isPreview*/,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        verify(gestureInteractor, never())
+            .addGestureBlockedMatcher(any(), eq(GestureInteractor.Scope.Global))
+    }
+
+    @Test
+    fun testDreamActivityGesturesNotBlockedWhenNotificationShadeShowing() {
+        val client = client
+
+        // Inform the overlay service of dream starting.
+        client.startDream(
+            mWindowParams,
+            mDreamOverlayCallback,
+            DREAM_COMPONENT,
+            false /*isPreview*/,
+            false /*shouldShowComplication*/
+        )
+        mMainExecutor.runAllReady()
+
+        val matcherCaptor = argumentCaptor<TaskMatcher>()
+        verify(gestureInteractor)
+            .addGestureBlockedMatcher(matcherCaptor.capture(), eq(GestureInteractor.Scope.Global))
+        val matcher = matcherCaptor.firstValue
+
+        val dreamTaskInfo = TaskInfo(mock<ComponentName>(), WindowConfiguration.ACTIVITY_TYPE_DREAM)
+        assertThat(matcher.matches(dreamTaskInfo)).isTrue()
+
+        val callbackCaptor = ArgumentCaptor.forClass(KeyguardUpdateMonitorCallback::class.java)
+        verify(mKeyguardUpdateMonitor).registerCallback(callbackCaptor.capture())
+
+        // Notification shade opens.
+        callbackCaptor.value.onShadeExpandedChanged(true)
+        mMainExecutor.runAllReady()
+
+        verify(gestureInteractor)
+            .removeGestureBlockedMatcher(eq(matcher), eq(GestureInteractor.Scope.Global))
     }
 
     @Test
@@ -1077,6 +1193,7 @@
             mWindowParams,
             mDreamOverlayCallback,
             DREAM_COMPONENT,
+            false /*isPreview*/,
             false /*shouldShowComplication*/
         )
         mMainExecutor.runAllReady()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
index 3aed79f..34d926a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
@@ -17,10 +17,14 @@
 package com.android.systemui.education.domain.interactor
 
 import android.content.pm.UserInfo
+import android.hardware.input.InputManager
+import android.hardware.input.KeyGestureEvent
+import android.view.KeyEvent
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.contextualeducation.GestureType.BACK
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.education.data.model.GestureEduModel
@@ -41,6 +45,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.kotlin.any
+import org.mockito.kotlin.verify
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -203,6 +210,33 @@
             assertThat(model.keyboardFirstConnectionTime).isEqualTo(newUserFirstConnectionTime)
         }
 
+    @Test
+    fun updateShortcutTimeOnKeyboardShortcutTriggered() =
+        testScope.runTest {
+            // runCurrent() to trigger inputManager#registerKeyGestureEventListener in the
+            // interactor
+            runCurrent()
+            val listenerCaptor =
+                ArgumentCaptor.forClass(InputManager.KeyGestureEventListener::class.java)
+            verify(kosmos.mockEduInputManager)
+                .registerKeyGestureEventListener(any(), listenerCaptor.capture())
+
+            val allAppsKeyGestureEvent =
+                KeyGestureEvent(
+                    /* deviceId= */ 1,
+                    IntArray(0),
+                    KeyEvent.META_META_ON,
+                    KeyGestureEvent.KEY_GESTURE_TYPE_ALL_APPS
+                )
+            listenerCaptor.value.onKeyGestureEvent(allAppsKeyGestureEvent)
+
+            val model by
+                collectLastValue(
+                    kosmos.contextualEducationRepository.readGestureEduModelFlow(ALL_APPS)
+                )
+            assertThat(model?.lastShortcutTriggeredTime).isEqualTo(eduClock.instant())
+        }
+
     private suspend fun triggerMaxEducationSignals(gestureType: GestureType) {
         // Increment max number of signal to try triggering education
         for (i in 1..KeyboardTouchpadEduInteractor.MAX_SIGNAL_COUNT) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
index cd0c58f..8b5f594 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
@@ -19,16 +19,23 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.contextualeducation.GestureType.BACK
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.education.data.repository.contextualEducationRepository
 import com.android.systemui.education.data.repository.fakeEduClock
+import com.android.systemui.inputdevice.data.model.UserDeviceConnectionStatus
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
+import kotlin.time.Duration.Companion.seconds
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -36,24 +43,190 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val underTest = kosmos.keyboardTouchpadEduStatsInteractor
+    private val repository = kosmos.contextualEducationRepository
+    private val fakeClock = kosmos.fakeEduClock
+    private val initialDelayElapsedDuration =
+        KeyboardTouchpadEduStatsInteractorImpl.initialDelayDuration + 1.seconds
 
     @Test
-    fun dataUpdatedOnIncrementSignalCount() =
+    fun dataUpdatedOnIncrementSignalCountWhenTouchpadConnected() =
         testScope.runTest {
-            val model by
-                collectLastValue(kosmos.contextualEducationRepository.readGestureEduModelFlow(BACK))
+            setUpForInitialDelayElapse()
+            whenever(mockUserInputDeviceRepository.isAnyTouchpadConnectedForUser)
+                .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = true, userId = 0)))
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
             val originalValue = model!!.signalCount
             underTest.incrementSignalCount(BACK)
+
             assertThat(model?.signalCount).isEqualTo(originalValue + 1)
         }
 
     @Test
+    fun dataUnchangedOnIncrementSignalCountWhenTouchpadDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            whenever(mockUserInputDeviceRepository.isAnyTouchpadConnectedForUser)
+                .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = false, userId = 0)))
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountWhenKeyboardConnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            whenever(mockUserInputDeviceRepository.isAnyKeyboardConnectedForUser)
+                .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = true, userId = 0)))
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(ALL_APPS)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenKeyboardDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            whenever(mockUserInputDeviceRepository.isAnyKeyboardConnectedForUser)
+                .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = false, userId = 0)))
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(ALL_APPS)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountAfterOobeLaunchInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            whenever(mockTutorialSchedulerRepository.launchTime(any<DeviceType>()))
+                .thenReturn(fakeClock.instant())
+            fakeClock.offset(initialDelayElapsedDuration)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountBeforeOobeLaunchInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            whenever(mockTutorialSchedulerRepository.launchTime(any<DeviceType>()))
+                .thenReturn(fakeClock.instant())
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountAfterTouchpadConnectionInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            repository.updateEduDeviceConnectionTime { model ->
+                model.copy(touchpadFirstConnectionTime = fakeClock.instant())
+            }
+            fakeClock.offset(initialDelayElapsedDuration)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountBeforeTouchpadConnectionInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            repository.updateEduDeviceConnectionTime { model ->
+                model.copy(touchpadFirstConnectionTime = fakeClock.instant())
+            }
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountAfterKeyboardConnectionInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            repository.updateEduDeviceConnectionTime { model ->
+                model.copy(keyboardFirstConnectionTime = fakeClock.instant())
+            }
+            fakeClock.offset(initialDelayElapsedDuration)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(ALL_APPS)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountBeforeKeyboardConnectionInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            repository.updateEduDeviceConnectionTime { model ->
+                model.copy(keyboardFirstConnectionTime = fakeClock.instant())
+            }
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(ALL_APPS)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenNoSetupTime() =
+        testScope.runTest {
+            whenever(mockUserInputDeviceRepository.isAnyTouchpadConnectedForUser)
+                .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = true, userId = 0)))
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
+            val originalValue = model!!.signalCount
+            underTest.incrementSignalCount(BACK)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
     fun dataAddedOnUpdateShortcutTriggerTime() =
         testScope.runTest {
-            val model by
-                collectLastValue(kosmos.contextualEducationRepository.readGestureEduModelFlow(BACK))
+            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
             assertThat(model?.lastShortcutTriggeredTime).isNull()
             underTest.updateShortcutTriggerTime(BACK)
             assertThat(model?.lastShortcutTriggeredTime).isEqualTo(kosmos.fakeEduClock.instant())
         }
+
+    private suspend fun setUpForInitialDelayElapse() {
+        whenever(mockTutorialSchedulerRepository.launchTime(any<DeviceType>()))
+            .thenReturn(fakeClock.instant())
+        fakeClock.offset(initialDelayElapsedDuration)
+    }
+
+    private fun setUpForDeviceConnection() {
+        whenever(mockUserInputDeviceRepository.isAnyTouchpadConnectedForUser)
+            .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = true, userId = 0)))
+        whenever(mockUserInputDeviceRepository.isAnyKeyboardConnectedForUser)
+            .thenReturn(flowOf(UserDeviceConnectionStatus(isConnected = true, userId = 0)))
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt
index 91d37cf..a764256 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/data/GestureRepositoryTest.kt
@@ -16,12 +16,14 @@
 
 package com.android.systemui.gesture.data
 
+import android.app.WindowConfiguration
 import android.content.ComponentName
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.navigationbar.gestural.data.respository.GestureRepositoryImpl
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
@@ -40,14 +42,36 @@
     @Test
     fun addRemoveComponentToBlock_updatesBlockedComponentSet() =
         testScope.runTest {
-            val component = mock<ComponentName>()
+            val matcher = TaskMatcher.TopActivityComponent(mock<ComponentName>())
 
-            underTest.addGestureBlockedActivity(component)
-            val addedBlockedComponents by collectLastValue(underTest.gestureBlockedActivities)
-            assertThat(addedBlockedComponents).contains(component)
+            kotlin.run {
+                underTest.addGestureBlockedMatcher(matcher)
+                val blockedMatchers by collectLastValue(underTest.gestureBlockedMatchers)
+                assertThat(blockedMatchers).contains(matcher)
+            }
 
-            underTest.removeGestureBlockedActivity(component)
-            val removedBlockedComponents by collectLastValue(underTest.gestureBlockedActivities)
-            assertThat(removedBlockedComponents).isEmpty()
+            kotlin.run {
+                underTest.removeGestureBlockedMatcher(matcher)
+                val blockedMatchers by collectLastValue(underTest.gestureBlockedMatchers)
+                assertThat(blockedMatchers).doesNotContain(matcher)
+            }
+        }
+
+    @Test
+    fun addRemoveActivityTypeToBlock_updatesBlockedActivityTypesSet() =
+        testScope.runTest {
+            val matcher = TaskMatcher.TopActivityType(WindowConfiguration.ACTIVITY_TYPE_STANDARD)
+
+            kotlin.run {
+                underTest.addGestureBlockedMatcher(matcher)
+                val blockedMatchers by collectLastValue(underTest.gestureBlockedMatchers)
+                assertThat(blockedMatchers).contains(matcher)
+            }
+
+            kotlin.run {
+                underTest.removeGestureBlockedMatcher(matcher)
+                val blockedMatchers by collectLastValue(underTest.gestureBlockedMatchers)
+                assertThat(blockedMatchers).doesNotContain(matcher)
+            }
         }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt
index 6395448..0ce0d93 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/GestureInteractorTest.kt
@@ -26,6 +26,7 @@
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.navigationbar.gestural.data.gestureRepository
 import com.android.systemui.navigationbar.gestural.domain.GestureInteractor
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher
 import com.android.systemui.shared.system.activityManagerWrapper
 import com.android.systemui.shared.system.taskStackChangeListeners
 import com.android.systemui.testKosmos
@@ -76,13 +77,16 @@
     fun addBlockedActivity_testCombination() =
         testScope.runTest {
             val globalComponent = mock<ComponentName>()
-            repository.addGestureBlockedActivity(globalComponent)
+            repository.addGestureBlockedMatcher(TaskMatcher.TopActivityComponent(globalComponent))
 
             val localComponent = mock<ComponentName>()
 
             val blocked by collectLastValue(underTest.topActivityBlocked)
 
-            underTest.addGestureBlockedActivity(localComponent, GestureInteractor.Scope.Local)
+            underTest.addGestureBlockedMatcher(
+                TaskMatcher.TopActivityComponent(localComponent),
+                GestureInteractor.Scope.Local
+            )
 
             assertThat(blocked).isFalse()
 
@@ -95,7 +99,7 @@
     fun initialization_testEmit() =
         testScope.runTest {
             val globalComponent = mock<ComponentName>()
-            repository.addGestureBlockedActivity(globalComponent)
+            repository.addGestureBlockedMatcher(TaskMatcher.TopActivityComponent(globalComponent))
             setTopActivity(globalComponent)
 
             val interactor = createInteractor()
@@ -114,10 +118,36 @@
 
             val localComponent = mock<ComponentName>()
 
-            interactor1.addGestureBlockedActivity(localComponent, GestureInteractor.Scope.Local)
+            interactor1.addGestureBlockedMatcher(
+                TaskMatcher.TopActivityComponent(localComponent),
+                GestureInteractor.Scope.Local
+            )
             setTopActivity(localComponent)
 
             assertThat(interactor1Blocked).isTrue()
             assertThat(interactor2Blocked).isFalse()
         }
+
+    @Test
+    fun matchingBlockers_separatelyManaged() =
+        testScope.runTest {
+            val interactor = createInteractor()
+            val interactorBlocked by collectLastValue(interactor.topActivityBlocked)
+
+            val localComponent = mock<ComponentName>()
+
+            val matcher1 = TaskMatcher.TopActivityComponent(localComponent)
+            val matcher2 = TaskMatcher.TopActivityComponent(localComponent)
+
+            interactor.addGestureBlockedMatcher(matcher1, GestureInteractor.Scope.Local)
+            interactor.addGestureBlockedMatcher(matcher2, GestureInteractor.Scope.Local)
+            setTopActivity(localComponent)
+            assertThat(interactorBlocked).isTrue()
+
+            interactor.removeGestureBlockedMatcher(matcher1, GestureInteractor.Scope.Local)
+            assertThat(interactorBlocked).isTrue()
+
+            interactor.removeGestureBlockedMatcher(matcher2, GestureInteractor.Scope.Local)
+            assertThat(interactorBlocked).isFalse()
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/TaskMatcherTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/TaskMatcherTest.kt
new file mode 100644
index 0000000..a246270
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/gesture/domain/TaskMatcherTest.kt
@@ -0,0 +1,91 @@
+/*
+ * 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.gesture.domain
+
+import android.app.WindowConfiguration
+import android.content.ComponentName
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.navigationbar.gestural.domain.TaskInfo
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class TaskMatcherTest : SysuiTestCase() {
+    @Test
+    fun activityMatcher_matchesComponentName() {
+        val componentName = ComponentName.unflattenFromString("com.foo/.bar")!!
+        val matcher = TaskMatcher.TopActivityComponent(componentName)
+
+        val taskInfo = TaskInfo(componentName, WindowConfiguration.ACTIVITY_TYPE_STANDARD)
+        assertThat(matcher.matches(taskInfo)).isTrue()
+    }
+
+    @Test
+    fun activityMatcher_doesNotMatchComponentName() {
+        val componentName = ComponentName.unflattenFromString("com.foo/.bar")!!
+        val matcher = TaskMatcher.TopActivityComponent(componentName)
+
+        val taskInfo =
+            TaskInfo(
+                ComponentName.unflattenFromString("com.bar/.baz"),
+                WindowConfiguration.ACTIVITY_TYPE_STANDARD
+            )
+        assertThat(matcher.matches(taskInfo)).isFalse()
+    }
+
+    @Test
+    fun activityMatcher_matchesActivityType() {
+        val activityType = WindowConfiguration.ACTIVITY_TYPE_HOME
+        val matcher = TaskMatcher.TopActivityType(activityType)
+
+        val taskInfo = TaskInfo(mock<ComponentName>(), activityType)
+        assertThat(matcher.matches(taskInfo)).isTrue()
+    }
+
+    @Test
+    fun activityMatcher_doesNotMatchEmptyActivityType() {
+        val activityType = WindowConfiguration.ACTIVITY_TYPE_HOME
+        val matcher = TaskMatcher.TopActivityType(activityType)
+
+        val taskInfo = TaskInfo(null, activityType)
+        assertThat(matcher.matches(taskInfo)).isFalse()
+    }
+
+    @Test
+    fun activityMatcher_doesNotMatchActivityType() {
+        val activityType = WindowConfiguration.ACTIVITY_TYPE_HOME
+        val matcher = TaskMatcher.TopActivityType(activityType)
+
+        val taskInfo = TaskInfo(mock<ComponentName>(), WindowConfiguration.ACTIVITY_TYPE_STANDARD)
+        assertThat(matcher.matches(taskInfo)).isFalse()
+    }
+
+    @Test
+    fun activityMatcher_equivalentMatchersAreNotEqual() {
+        val activityType = WindowConfiguration.ACTIVITY_TYPE_HOME
+        val matcher1 = TaskMatcher.TopActivityType(activityType)
+        val matcher2 = TaskMatcher.TopActivityType(activityType)
+
+        assertThat(matcher1).isNotEqualTo(matcher2)
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
index fd4ed38..686b518 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/haptics/qs/QSLongPressEffectTest.kt
@@ -23,7 +23,7 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.ActivityTransitionAnimator
-import com.android.systemui.haptics.vibratorHelper
+import com.android.systemui.haptics.fakeVibratorHelper
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.log.core.FakeLogBuffer
 import com.android.systemui.qs.qsTileFactory
@@ -50,7 +50,7 @@
 
     @Rule @JvmField val mMockitoRule: MockitoRule = MockitoJUnit.rule()
     private val kosmos = testKosmos()
-    private val vibratorHelper = kosmos.vibratorHelper
+    private val vibratorHelper = kosmos.fakeVibratorHelper
     private val qsTile = kosmos.qsTileFactory.createTile("Test Tile")
     @Mock private lateinit var callback: QSLongPressEffect.Callback
     @Mock private lateinit var controller: ActivityTransitionAnimator.Controller
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
index bbfaf6f..6c3c7ef 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/DoNotDisturbQuickAffordanceConfigTest.kt
@@ -31,19 +31,19 @@
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.statusbar.policy.ZenModeController
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.argumentCaptor
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
@@ -63,27 +63,24 @@
 @RunWith(AndroidJUnit4::class)
 class DoNotDisturbQuickAffordanceConfigTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
+
     @Mock private lateinit var zenModeController: ZenModeController
     @Mock private lateinit var userTracker: UserTracker
     @Mock private lateinit var conditionUri: Uri
     @Mock private lateinit var enableZenModeDialog: EnableZenModeDialog
     @Captor private lateinit var spyZenMode: ArgumentCaptor<Int>
     @Captor private lateinit var spyConditionId: ArgumentCaptor<Uri?>
-    private lateinit var settings: FakeSettings
 
     private lateinit var underTest: DoNotDisturbQuickAffordanceConfig
-    private lateinit var testDispatcher: TestDispatcher
-    private lateinit var testScope: TestScope
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        testDispatcher = StandardTestDispatcher()
-        testScope = TestScope(testDispatcher)
-
-        settings = FakeSettings()
-
         underTest =
             DoNotDisturbQuickAffordanceConfig(
                 context,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
index 26fcb23..0145f17 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/KeyguardQuickAffordanceLegacySettingSyncerTest.kt
@@ -23,19 +23,19 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.kosmos.unconfinedTestScope
 import com.android.systemui.res.R
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.systemui.testKosmos
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.unconfinedDispatcherFakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestDispatcher
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -51,14 +51,15 @@
 @RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceLegacySettingSyncerTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.unconfinedTestDispatcher
+    private val testScope = kosmos.unconfinedTestScope
+    private val settings = kosmos.unconfinedDispatcherFakeSettings
+
     @Mock private lateinit var sharedPrefs: FakeSharedPreferences
 
     private lateinit var underTest: KeyguardQuickAffordanceLegacySettingSyncer
-
-    private lateinit var testScope: TestScope
-    private lateinit var testDispatcher: TestDispatcher
     private lateinit var selectionManager: KeyguardQuickAffordanceLocalUserSelectionManager
-    private lateinit var settings: FakeSettings
 
     @Before
     fun setUp() {
@@ -73,8 +74,6 @@
         whenever(resources.getBoolean(R.bool.custom_lockscreen_shortcuts_enabled)).thenReturn(true)
         whenever(context.resources).thenReturn(resources)
 
-        testDispatcher = UnconfinedTestDispatcher()
-        testScope = TestScope(testDispatcher)
         selectionManager =
             KeyguardQuickAffordanceLocalUserSelectionManager(
                 context = context,
@@ -92,7 +91,6 @@
                 userTracker = FakeUserTracker(),
                 broadcastDispatcher = fakeBroadcastDispatcher,
             )
-        settings = FakeSettings()
         settings.putInt(Settings.Secure.LOCKSCREEN_SHOW_CONTROLS, 0)
         settings.putInt(Settings.Secure.LOCKSCREEN_SHOW_WALLET, 0)
         settings.putInt(Settings.Secure.LOCK_SCREEN_SHOW_QR_CODE_SCANNER, 0)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/BiometricSettingsRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/BiometricSettingsRepositoryTest.kt
index c5ba02d..4e429c3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/BiometricSettingsRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/BiometricSettingsRepositoryTest.kt
@@ -46,11 +46,10 @@
 import com.android.systemui.keyguard.data.repository.BiometricType.SIDE_FINGERPRINT
 import com.android.systemui.keyguard.data.repository.BiometricType.UNDER_DISPLAY_FINGERPRINT
 import com.android.systemui.keyguard.shared.model.DevicePosture
-import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.res.R
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.policy.DevicePostureController
+import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.FakeUserRepository
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
@@ -81,6 +80,8 @@
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 @RunWith(AndroidJUnit4::class)
 class BiometricSettingsRepositoryTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: BiometricSettingsRepository
 
     @Mock private lateinit var authController: AuthController
@@ -88,7 +89,6 @@
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
     @Mock private lateinit var dumpManager: DumpManager
     @Mock private lateinit var biometricManager: BiometricManager
-    @Mock private lateinit var tableLogger: TableLogBuffer
     @Captor
     private lateinit var strongAuthTracker: ArgumentCaptor<LockPatternUtils.StrongAuthTracker>
     @Captor private lateinit var authControllerCallback: ArgumentCaptor<AuthController.Callback>
@@ -99,7 +99,7 @@
     private lateinit var devicePostureRepository: FakeDevicePostureRepository
     private lateinit var facePropertyRepository: FakeFacePropertyRepository
     private lateinit var fingerprintPropertyRepository: FakeFingerprintPropertyRepository
-    private lateinit var mobileConnectionsRepository: FakeMobileConnectionsRepository
+    private val mobileConnectionsRepository = kosmos.fakeMobileConnectionsRepository
 
     private lateinit var testDispatcher: TestDispatcher
     private lateinit var testScope: TestScope
@@ -115,8 +115,6 @@
         devicePostureRepository = FakeDevicePostureRepository()
         facePropertyRepository = FakeFacePropertyRepository()
         fingerprintPropertyRepository = FakeFingerprintPropertyRepository()
-        mobileConnectionsRepository =
-            FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), tableLogger)
     }
 
     private suspend fun createBiometricSettingsRepository() {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
index c85cd66..1582e47 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/repository/KeyguardQuickAffordanceRepositoryTest.kt
@@ -31,20 +31,21 @@
 import com.android.systemui.keyguard.data.quickaffordance.KeyguardQuickAffordanceRemoteUserSelectionManager
 import com.android.systemui.keyguard.shared.model.KeyguardQuickAffordancePickerRepresentation
 import com.android.systemui.keyguard.shared.model.KeyguardSlotPickerRepresentation
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
 import com.android.systemui.settings.FakeUserTracker
 import com.android.systemui.settings.UserFileManager
 import com.android.systemui.shared.customization.data.content.FakeCustomizationProviderClient
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
+import com.android.systemui.testKosmos
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import java.util.Locale
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.After
 import org.junit.Before
@@ -58,6 +59,11 @@
 @RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceRepositoryTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
+
     private lateinit var underTest: KeyguardQuickAffordanceRepository
 
     private lateinit var config1: FakeKeyguardQuickAffordanceConfig
@@ -65,7 +71,6 @@
     private lateinit var userTracker: FakeUserTracker
     private lateinit var client1: FakeCustomizationProviderClient
     private lateinit var client2: FakeCustomizationProviderClient
-    private lateinit var testScope: TestScope
 
     @Before
     fun setUp() {
@@ -73,8 +78,6 @@
         context.resources.configuration.setLayoutDirection(Locale.US)
         config1 = FakeKeyguardQuickAffordanceConfig(FakeCustomizationProviderClient.AFFORDANCE_1)
         config2 = FakeKeyguardQuickAffordanceConfig(FakeCustomizationProviderClient.AFFORDANCE_2)
-        val testDispatcher = StandardTestDispatcher()
-        testScope = TestScope(testDispatcher)
         userTracker = FakeUserTracker()
         val localUserSelectionManager =
             KeyguardQuickAffordanceLocalUserSelectionManager(
@@ -128,7 +131,7 @@
                     KeyguardQuickAffordanceLegacySettingSyncer(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = testDispatcher,
-                        secureSettings = FakeSettings(),
+                        secureSettings = settings,
                         selectionsManager = localUserSelectionManager,
                     ),
                 configs = setOf(config1, config2),
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
index 3b8ffcd..17e3006 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractorTest.kt
@@ -116,7 +116,7 @@
             reset(transitionRepository)
 
             // Authentication results in calling startDismissKeyguardTransition.
-            kosmos.keyguardTransitionInteractor.startDismissKeyguardTransition()
+            kosmos.keyguardDismissTransitionInteractor.startDismissKeyguardTransition()
             runCurrent()
 
             assertThat(transitionRepository)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorTest.kt
index 6eb9862..33f3cd4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractorTest.kt
@@ -39,12 +39,15 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.keyguardOcclusionRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.keyguard.util.KeyguardTransitionRepositorySpySubject.Companion.assertThat
@@ -53,6 +56,7 @@
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.domain.interactor.keyguardOcclusionInteractor
 import com.android.systemui.testKosmos
 import junit.framework.Assert.assertEquals
@@ -166,6 +170,10 @@
     @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToGone_onWakeUp_ifPowerButtonGestureDetected_fromGone() =
         testScope.runTest {
+            val isGone by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Gone, GONE)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.AOD,
@@ -175,7 +183,7 @@
             runCurrent()
 
             // Make sure we're GONE.
-            assertEquals(KeyguardState.GONE, kosmos.keyguardTransitionInteractor.getFinishedState())
+            assertEquals(true, isGone)
 
             // Get part way to AOD.
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
@@ -204,6 +212,10 @@
     @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToOccluded_onWakeUp_ifPowerButtonGestureDetectedAfterFinishedInAod_fromGone() =
         testScope.runTest {
+            val isGone by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Gone, GONE)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.AOD,
@@ -213,7 +225,7 @@
             runCurrent()
 
             // Make sure we're GONE.
-            assertEquals(KeyguardState.GONE, kosmos.keyguardTransitionInteractor.getFinishedState())
+            assertEquals(true, isGone)
 
             // Get all the way to AOD
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
@@ -239,6 +251,10 @@
     @EnableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToOccluded_onWakeUp_ifPowerButtonGestureDetected_fromLockscreen() =
         testScope.runTest {
+            val isLockscreen by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Lockscreen, LOCKSCREEN)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.AOD,
@@ -248,10 +264,7 @@
             runCurrent()
 
             // Make sure we're in LOCKSCREEN.
-            assertEquals(
-                KeyguardState.LOCKSCREEN,
-                kosmos.keyguardTransitionInteractor.getFinishedState()
-            )
+            assertEquals(true, isLockscreen)
 
             // Get part way to AOD.
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
@@ -314,7 +327,7 @@
         testScope.runTest {
             kosmos.fakeKeyguardRepository.setKeyguardShowing(false)
             kosmos.fakeKeyguardRepository.setKeyguardDismissible(true)
-            kosmos.keyguardTransitionInteractor.startDismissKeyguardTransition()
+            kosmos.keyguardDismissTransitionInteractor.startDismissKeyguardTransition()
             powerInteractor.setAwakeForTest()
             advanceTimeBy(100) // account for debouncing
 
@@ -327,6 +340,10 @@
     @DisableFlags(Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToGone_onWakeUp_ifPowerButtonGestureDetectedInAod_fromGone() =
         testScope.runTest {
+            val isGone by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Gone, GONE)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.AOD,
@@ -336,7 +353,7 @@
             runCurrent()
 
             // Make sure we're GONE.
-            assertEquals(KeyguardState.GONE, kosmos.keyguardTransitionInteractor.getFinishedState())
+            assertEquals(true, isGone)
 
             // Start going to AOD on first button push
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorTest.kt
index ee4a0d2d..c18deb1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/FromDozingTransitionInteractorTest.kt
@@ -48,12 +48,15 @@
 import com.android.systemui.communal.domain.interactor.setCommunalAvailable
 import com.android.systemui.communal.shared.model.CommunalScenes
 import com.android.systemui.communal.shared.model.CommunalTransitionKeys
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.keyguardOcclusionRepository
 import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.util.KeyguardTransitionRepositorySpySubject.Companion.assertThat
 import com.android.systemui.kosmos.applicationCoroutineScope
@@ -62,6 +65,7 @@
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.testKosmos
 import junit.framework.Assert.assertEquals
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -316,6 +320,10 @@
     @EnableFlags(FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToGone_onWakeUp_ifPowerButtonGestureDetected_fromGone() =
         testScope.runTest {
+            val isGone by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Gone, GONE)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.DOZING,
@@ -325,7 +333,7 @@
             runCurrent()
 
             // Make sure we're GONE.
-            assertEquals(KeyguardState.GONE, kosmos.keyguardTransitionInteractor.getFinishedState())
+            assertEquals(true, isGone)
 
             // Get part way to AOD.
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
@@ -355,6 +363,10 @@
     @Suppress("ktlint:standard:max-line-length")
     fun testTransitionToOccluded_onWakeUp_ifPowerButtonGestureDetectedAfterFinishedInAod_fromGone() =
         testScope.runTest {
+            val isGone by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Gone, GONE)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.DOZING,
@@ -364,7 +376,7 @@
             runCurrent()
 
             // Make sure we're GONE.
-            assertEquals(KeyguardState.GONE, kosmos.keyguardTransitionInteractor.getFinishedState())
+            assertEquals(true, isGone)
 
             // Get all the way to AOD
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
@@ -390,6 +402,10 @@
     @EnableFlags(FLAG_KEYGUARD_WM_STATE_REFACTOR)
     fun testTransitionToOccluded_onWakeUp_ifPowerButtonGestureDetected_fromLockscreen() =
         testScope.runTest {
+            val isLockscreen by
+                collectLastValue(
+                    kosmos.keyguardTransitionInteractor.isFinishedIn(Scenes.Lockscreen, LOCKSCREEN)
+                )
             powerInteractor.setAwakeForTest()
             transitionRepository.sendTransitionSteps(
                 from = KeyguardState.DOZING,
@@ -399,10 +415,7 @@
             runCurrent()
 
             // Make sure we're in LOCKSCREEN.
-            assertEquals(
-                KeyguardState.LOCKSCREEN,
-                kosmos.keyguardTransitionInteractor.getFinishedState()
-            )
+            assertEquals(true, isLockscreen)
 
             // Get part way to AOD.
             powerInteractor.onStartedGoingToSleep(PowerManager.GO_TO_SLEEP_REASON_MIN)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
index ad07c1c..a8bb2b0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractorTest.kt
@@ -61,7 +61,7 @@
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -80,6 +80,10 @@
 @RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordanceInteractorTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
+
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
     @Mock private lateinit var keyguardStateController: KeyguardStateController
     @Mock private lateinit var userTracker: UserTracker
@@ -90,11 +94,8 @@
     @Mock private lateinit var logger: KeyguardQuickAffordancesLogger
     @Mock private lateinit var metricsLogger: KeyguardQuickAffordancesMetricsLogger
 
-    private val kosmos = testKosmos()
-
     private lateinit var underTest: KeyguardQuickAffordanceInteractor
 
-    private val testScope = kosmos.testScope
     private lateinit var repository: FakeKeyguardRepository
     private lateinit var homeControls: FakeKeyguardQuickAffordanceConfig
     private lateinit var quickAccessWallet: FakeKeyguardQuickAffordanceConfig
@@ -170,7 +171,7 @@
                     KeyguardQuickAffordanceLegacySettingSyncer(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = kosmos.testDispatcher,
-                        secureSettings = FakeSettings(),
+                        secureSettings = settings,
                         selectionsManager = localUserSelectionManager,
                     ),
                 configs = setOf(homeControls, quickAccessWallet, qrCodeScanner),
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
index 6e76cbc..12039c1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorTest.kt
@@ -90,52 +90,6 @@
         }
 
     @Test
-    fun finishedKeyguardStateTests() =
-        testScope.runTest {
-            val finishedSteps by collectValues(underTest.finishedKeyguardState)
-            runCurrent()
-            val steps = mutableListOf<TransitionStep>()
-
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 0f, STARTED))
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 0.5f, RUNNING))
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 1f, FINISHED))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 0f, STARTED))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 0.9f, RUNNING))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 1f, FINISHED))
-            steps.add(TransitionStep(AOD, GONE, 1f, STARTED))
-
-            steps.forEach {
-                repository.sendTransitionStep(it)
-                runCurrent()
-            }
-
-            assertThat(finishedSteps).isEqualTo(listOf(LOCKSCREEN, PRIMARY_BOUNCER, AOD))
-        }
-
-    @Test
-    fun startedKeyguardStateTests() =
-        testScope.runTest {
-            val startedStates by collectValues(underTest.startedKeyguardState)
-            runCurrent()
-            val steps = mutableListOf<TransitionStep>()
-
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 0f, STARTED))
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 0.5f, RUNNING))
-            steps.add(TransitionStep(AOD, PRIMARY_BOUNCER, 1f, FINISHED))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 0f, STARTED))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 0.9f, RUNNING))
-            steps.add(TransitionStep(PRIMARY_BOUNCER, AOD, 1f, FINISHED))
-            steps.add(TransitionStep(AOD, GONE, 1f, STARTED))
-
-            steps.forEach {
-                repository.sendTransitionStep(it)
-                runCurrent()
-            }
-
-            assertThat(startedStates).isEqualTo(listOf(LOCKSCREEN, PRIMARY_BOUNCER, AOD, GONE))
-        }
-
-    @Test
     fun startedKeyguardTransitionStepTests() =
         testScope.runTest {
             val startedSteps by collectValues(underTest.startedKeyguardTransitionStep)
@@ -1206,95 +1160,6 @@
         }
 
     @Test
-    fun finishedKeyguardState_emitsAgainIfCancelledAndReversed() =
-        testScope.runTest {
-            val finishedStates by collectValues(underTest.finishedKeyguardState)
-
-            // We default FINISHED in LOCKSCREEN.
-            assertEquals(listOf(LOCKSCREEN), finishedStates)
-
-            sendSteps(
-                TransitionStep(LOCKSCREEN, AOD, 0f, STARTED),
-                TransitionStep(LOCKSCREEN, AOD, 0.5f, RUNNING),
-                TransitionStep(LOCKSCREEN, AOD, 1f, FINISHED),
-            )
-
-            // We're FINISHED in AOD.
-            assertEquals(
-                listOf(
-                    LOCKSCREEN,
-                    AOD,
-                ),
-                finishedStates
-            )
-
-            // Transition back to LOCKSCREEN.
-            sendSteps(
-                TransitionStep(AOD, LOCKSCREEN, 0f, STARTED),
-                TransitionStep(AOD, LOCKSCREEN, 0.5f, RUNNING),
-                TransitionStep(AOD, LOCKSCREEN, 1f, FINISHED),
-            )
-
-            // We're FINISHED in LOCKSCREEN.
-            assertEquals(
-                listOf(
-                    LOCKSCREEN,
-                    AOD,
-                    LOCKSCREEN,
-                ),
-                finishedStates
-            )
-
-            sendSteps(
-                TransitionStep(LOCKSCREEN, GONE, 0f, STARTED),
-                TransitionStep(LOCKSCREEN, GONE, 0.5f, RUNNING),
-            )
-
-            // We've STARTED a transition to GONE but not yet finished it so we're still FINISHED in
-            // LOCKSCREEN.
-            assertEquals(
-                listOf(
-                    LOCKSCREEN,
-                    AOD,
-                    LOCKSCREEN,
-                ),
-                finishedStates
-            )
-
-            sendSteps(
-                TransitionStep(LOCKSCREEN, GONE, 0.6f, CANCELED),
-            )
-
-            // We've CANCELED a transition to GONE, we're still FINISHED in LOCKSCREEN.
-            assertEquals(
-                listOf(
-                    LOCKSCREEN,
-                    AOD,
-                    LOCKSCREEN,
-                ),
-                finishedStates
-            )
-
-            sendSteps(
-                TransitionStep(GONE, LOCKSCREEN, 0.6f, STARTED),
-                TransitionStep(GONE, LOCKSCREEN, 0.9f, RUNNING),
-                TransitionStep(GONE, LOCKSCREEN, 1f, FINISHED),
-            )
-
-            // Expect another emission of LOCKSCREEN, as we have FINISHED a second transition to
-            // LOCKSCREEN after the cancellation.
-            assertEquals(
-                listOf(
-                    LOCKSCREEN,
-                    AOD,
-                    LOCKSCREEN,
-                    LOCKSCREEN,
-                ),
-                finishedStates
-            )
-        }
-
-    @Test
     fun testCurrentState() =
         testScope.runTest {
             val currentStates by collectValues(underTest.currentKeyguardState)
@@ -1529,9 +1394,11 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Lockscreen, Scenes.Gone))
             val sendStep1 = TransitionStep(LOCKSCREEN, UNDEFINED, 0f, STARTED)
+            sendSteps(sendStep1)
+            kosmos.setSceneTransition(Idle(Scenes.Gone))
             val sendStep2 = TransitionStep(LOCKSCREEN, UNDEFINED, 1f, FINISHED)
             val sendStep3 = TransitionStep(LOCKSCREEN, AOD, 0f, STARTED)
-            sendSteps(sendStep1, sendStep2, sendStep3)
+            sendSteps(sendStep2, sendStep3)
 
             assertEquals(listOf(sendStep1, sendStep2), currentStates)
             assertEquals(listOf(sendStep1, sendStep2), currentStatesConverted)
@@ -1545,6 +1412,7 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Gone, Scenes.Lockscreen))
             val sendStep1 = TransitionStep(LOCKSCREEN, UNDEFINED, 0f, STARTED)
+            kosmos.setSceneTransition(Idle(Scenes.Lockscreen))
             val sendStep2 = TransitionStep(LOCKSCREEN, UNDEFINED, 1f, FINISHED)
             val sendStep3 = TransitionStep(LOCKSCREEN, AOD, 0f, STARTED)
             sendSteps(sendStep1, sendStep2, sendStep3)
@@ -1561,6 +1429,7 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Gone, Scenes.Lockscreen))
             val sendStep1 = TransitionStep(LOCKSCREEN, DOZING, 0f, STARTED)
+            kosmos.setSceneTransition(Idle(Scenes.Lockscreen))
             val sendStep2 = TransitionStep(LOCKSCREEN, DOZING, 1f, FINISHED)
             val sendStep3 = TransitionStep(LOCKSCREEN, AOD, 0f, STARTED)
             sendSteps(sendStep1, sendStep2, sendStep3)
@@ -1578,6 +1447,7 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Gone, Scenes.Lockscreen))
             val sendStep1 = TransitionStep(LOCKSCREEN, DOZING, 0f, STARTED)
+            kosmos.setSceneTransition(Idle(Scenes.Lockscreen))
             val sendStep2 = TransitionStep(LOCKSCREEN, DOZING, 1f, FINISHED)
             val sendStep3 = TransitionStep(LOCKSCREEN, AOD, 0f, STARTED)
             val sendStep4 = TransitionStep(AOD, LOCKSCREEN, 0f, STARTED)
@@ -1596,10 +1466,12 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Lockscreen, Scenes.Gone))
             val sendStep1 = TransitionStep(LOCKSCREEN, UNDEFINED, 0f, STARTED)
+            sendSteps(sendStep1)
+            kosmos.setSceneTransition(Idle(Scenes.Gone))
             val sendStep2 = TransitionStep(LOCKSCREEN, UNDEFINED, 1f, FINISHED)
             val sendStep3 = TransitionStep(UNDEFINED, AOD, 0f, STARTED)
             val sendStep4 = TransitionStep(AOD, LOCKSCREEN, 0f, STARTED)
-            sendSteps(sendStep1, sendStep2, sendStep3, sendStep4)
+            sendSteps(sendStep2, sendStep3, sendStep4)
 
             assertEquals(listOf(sendStep1, sendStep2), currentStates)
             assertEquals(listOf(sendStep1, sendStep2), currentStatesMapped)
@@ -1613,10 +1485,12 @@
 
             kosmos.setSceneTransition(Transition(Scenes.Gone, Scenes.Lockscreen))
             val sendStep1 = TransitionStep(LOCKSCREEN, UNDEFINED, 0f, STARTED)
+            sendSteps(sendStep1)
+            kosmos.setSceneTransition(Idle(Scenes.Gone))
             val sendStep2 = TransitionStep(LOCKSCREEN, UNDEFINED, 1f, FINISHED)
             val sendStep3 = TransitionStep(UNDEFINED, AOD, 0f, STARTED)
             val sendStep4 = TransitionStep(AOD, LOCKSCREEN, 0f, STARTED)
-            sendSteps(sendStep1, sendStep2, sendStep3, sendStep4)
+            sendSteps(sendStep2, sendStep3, sendStep4)
 
             assertEquals(listOf<TransitionStep>(), currentStatesMapped)
         }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
index 3e0a1f3..073ed61 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/domain/interactor/LightRevealScrimInteractorTest.kt
@@ -19,18 +19,20 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectValues
 import com.android.systemui.keyguard.data.fakeLightRevealScrimRepository
+import com.android.systemui.keyguard.data.repository.DEFAULT_REVEAL_EFFECT
 import com.android.systemui.keyguard.data.repository.FakeLightRevealScrimRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+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.testScope
 import com.android.systemui.statusbar.LightRevealEffect
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.testKosmos
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
 import org.junit.Test
@@ -64,35 +66,45 @@
 
     @Test
     fun lightRevealEffect_doesNotChangeDuringKeyguardTransition() =
-        runTest(UnconfinedTestDispatcher()) {
-            val values = mutableListOf<LightRevealEffect>()
-            val job = underTest.lightRevealEffect.onEach(values::add).launchIn(this)
+        kosmos.testScope.runTest {
+            val values by collectValues(underTest.lightRevealEffect)
+            runCurrent()
+            assertEquals(listOf(DEFAULT_REVEAL_EFFECT), values)
 
             fakeLightRevealScrimRepository.setRevealEffect(reveal1)
-
+            runCurrent()
             // The reveal effect shouldn't emit anything until a keyguard transition starts.
-            assertEquals(values.size, 0)
+            assertEquals(listOf(DEFAULT_REVEAL_EFFECT), values)
 
             // Once it starts, it should emit reveal1.
             fakeKeyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(transitionState = TransitionState.STARTED)
+                TransitionStep(to = KeyguardState.AOD, transitionState = TransitionState.STARTED)
             )
-            assertEquals(values, listOf(reveal1))
+            runCurrent()
+            assertEquals(listOf(DEFAULT_REVEAL_EFFECT, reveal1), values)
 
             // Until the next transition starts, reveal2 should not be emitted.
             fakeLightRevealScrimRepository.setRevealEffect(reveal2)
+            runCurrent()
             fakeKeyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(transitionState = TransitionState.RUNNING)
+                TransitionStep(
+                    to = KeyguardState.LOCKSCREEN,
+                    transitionState = TransitionState.RUNNING
+                )
             )
+            runCurrent()
             fakeKeyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(transitionState = TransitionState.FINISHED)
+                TransitionStep(to = KeyguardState.AOD, transitionState = TransitionState.FINISHED)
             )
-            assertEquals(values, listOf(reveal1))
+            runCurrent()
+            assertEquals(listOf(DEFAULT_REVEAL_EFFECT, reveal1), values)
             fakeKeyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(transitionState = TransitionState.STARTED)
+                TransitionStep(
+                    to = KeyguardState.LOCKSCREEN,
+                    transitionState = TransitionState.STARTED
+                )
             )
-            assertEquals(values, listOf(reveal1, reveal2))
-
-            job.cancel()
+            runCurrent()
+            assertEquals(listOf(DEFAULT_REVEAL_EFFECT, reveal1, reveal2), values)
         }
 }
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 3e1f4f6..3b2b12c 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
@@ -360,6 +360,7 @@
         }
 
     @Test
+    @DisableSceneContainer
     fun alpha_transitionBetweenHubAndDream_isZero() =
         testScope.runTest {
             val alpha by collectLastValue(underTest.alpha(viewState))
@@ -388,8 +389,8 @@
                 ObservableTransitionState.Transition(
                     fromScene = Scenes.Lockscreen,
                     toScene = Scenes.Communal,
-                    emptyFlow(),
-                    emptyFlow(),
+                    flowOf(Scenes.Communal),
+                    flowOf(0.5f),
                     false,
                     emptyFlow()
                 )
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelTest.kt
similarity index 91%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelTest.kt
index c66ebf3..4253c29 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModelTest.kt
@@ -63,7 +63,7 @@
 @RunWith(ParameterizedAndroidJunit4::class)
 @RunWithLooper
 @EnableSceneContainer
-class LockscreenSceneActionsViewModelTest : SysuiTestCase() {
+class LockscreenUserActionsViewModelTest : SysuiTestCase() {
 
     companion object {
         private const val parameterCount = 6
@@ -170,7 +170,7 @@
 
     @Test
     @EnableFlags(Flags.FLAG_COMMUNAL_HUB)
-    fun destinationScenes() =
+    fun userActions() =
         testScope.runTest {
             underTest.activateIn(this)
             kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
@@ -193,9 +193,9 @@
                     },
             )
 
-            val destinationScenes by collectLastValue(underTest.actions)
+            val userActions by collectLastValue(underTest.actions)
             val downDestination =
-                destinationScenes?.get(
+                userActions?.get(
                     Swipe(
                         SwipeDirection.Down,
                         fromSource = Edge.Top.takeIf { downFromEdge },
@@ -227,11 +227,10 @@
 
             val upScene by
                 collectLastValue(
-                    (destinationScenes?.get(Swipe(SwipeDirection.Up))
-                            as? UserActionResult.ChangeScene)
-                        ?.toScene
-                        ?.let { scene -> kosmos.sceneInteractor.resolveSceneFamily(scene) }
-                        ?: flowOf(null)
+                    (userActions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene?.let {
+                        scene ->
+                        kosmos.sceneInteractor.resolveSceneFamily(scene)
+                    } ?: flowOf(null)
                 )
 
             assertThat(upScene)
@@ -244,11 +243,10 @@
 
             val leftScene by
                 collectLastValue(
-                    (destinationScenes?.get(Swipe(SwipeDirection.Left))
-                            as? UserActionResult.ChangeScene)
-                        ?.toScene
-                        ?.let { scene -> kosmos.sceneInteractor.resolveSceneFamily(scene) }
-                        ?: flowOf(null)
+                    (userActions?.get(Swipe.Left) as? UserActionResult.ChangeScene)?.toScene?.let {
+                        scene ->
+                        kosmos.sceneInteractor.resolveSceneFamily(scene)
+                    } ?: flowOf(null)
                 )
 
             assertThat(leftScene)
@@ -260,8 +258,8 @@
                 )
         }
 
-    private fun createLockscreenSceneViewModel(): LockscreenSceneActionsViewModel {
-        return LockscreenSceneActionsViewModel(
+    private fun createLockscreenSceneViewModel(): LockscreenUserActionsViewModel {
+        return LockscreenUserActionsViewModel(
             deviceEntryInteractor = kosmos.deviceEntryInteractor,
             communalInteractor = kosmos.communalInteractor,
             shadeInteractor = kosmos.shadeInteractor,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/HydratorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/HydratorTest.kt
index 8c9c527..ec6045c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/HydratorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/lifecycle/HydratorTest.kt
@@ -48,12 +48,13 @@
         composeRule.setContent {
             val keepAlive by keepAliveMutable
             if (keepAlive) {
-                val viewModel = rememberViewModel {
-                    FakeSysUiViewModel(
-                        upstreamFlow = upstreamFlow,
-                        upstreamStateFlow = upstreamStateFlow,
-                    )
-                }
+                val viewModel =
+                    rememberViewModel("test") {
+                        FakeSysUiViewModel(
+                            upstreamFlow = upstreamFlow,
+                            upstreamStateFlow = upstreamStateFlow,
+                        )
+                    }
 
                 Column {
                     Text(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt
new file mode 100644
index 0000000..f6865f13
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModelTest.kt
@@ -0,0 +1,69 @@
+/*
+ * 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.notifications.ui.viewmodel
+
+import android.testing.TestableLooper
+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.UserActionResult
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.shade.ui.viewmodel.notificationsShadeOverlayActionsViewModel
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+@EnableSceneContainer
+class NotificationsShadeOverlayActionsViewModelTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+
+    private val underTest = kosmos.notificationsShadeOverlayActionsViewModel
+
+    @Test
+    fun upTransitionSceneKey_hidesShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat((actions?.get(Swipe.Up) as? UserActionResult.HideOverlay)?.overlay)
+                .isEqualTo(Overlays.NotificationsShade)
+            assertThat(actions?.get(Swipe.Down)).isNull()
+        }
+
+    @Test
+    fun back_hidesShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat((actions?.get(Back) as? UserActionResult.HideOverlay)?.overlay)
+                .isEqualTo(Overlays.NotificationsShade)
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModelTest.kt
new file mode 100644
index 0000000..88a1df1
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModelTest.kt
@@ -0,0 +1,61 @@
+/*
+ * 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.notifications.ui.viewmodel
+
+import android.testing.TestableLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.shade.ui.viewmodel.notificationsShadeOverlayContentViewModel
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+@EnableSceneContainer
+class NotificationsShadeOverlayContentViewModelTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val sceneInteractor = kosmos.sceneInteractor
+
+    private val underTest = kosmos.notificationsShadeOverlayContentViewModel
+
+    @Test
+    fun onScrimClicked_hidesShade() =
+        testScope.runTest {
+            val currentOverlays by collectLastValue(sceneInteractor.currentOverlays)
+            sceneInteractor.showOverlay(
+                overlay = Overlays.NotificationsShade,
+                loggingReason = "test",
+            )
+            assertThat(currentOverlays).contains(Overlays.NotificationsShade)
+
+            underTest.onScrimClicked()
+
+            assertThat(currentOverlays).doesNotContain(Overlays.NotificationsShade)
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModelTest.kt
similarity index 81%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModelTest.kt
index 0505e19..46b02e92 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModelTest.kt
@@ -37,8 +37,7 @@
 import com.android.systemui.scene.domain.resolver.homeSceneFamilyResolver
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.data.repository.fakeShadeRepository
-import com.android.systemui.shade.ui.viewmodel.notificationsShadeSceneActionsViewModel
+import com.android.systemui.shade.ui.viewmodel.notificationsShadeUserActionsViewModel
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -53,14 +52,14 @@
 @RunWith(AndroidJUnit4::class)
 @TestableLooper.RunWithLooper
 @EnableSceneContainer
-class NotificationsShadeSceneActionsViewModelTest : SysuiTestCase() {
+class NotificationsShadeUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val sceneInteractor by lazy { kosmos.sceneInteractor }
     private val deviceUnlockedInteractor by lazy { kosmos.deviceUnlockedInteractor }
 
-    private val underTest by lazy { kosmos.notificationsShadeSceneActionsViewModel }
+    private val underTest by lazy { kosmos.notificationsShadeUserActionsViewModel }
 
     @Test
     fun upTransitionSceneKey_deviceLocked_lockscreen() =
@@ -104,36 +103,6 @@
         }
 
     @Test
-    fun downTransitionSceneKey_deviceLocked_bottomAligned_lockscreen() =
-        testScope.runTest {
-            kosmos.fakeShadeRepository.setDualShadeAlignedToBottom(true)
-            val actions by collectLastValue(underTest.actions)
-            lockDevice()
-            underTest.activateIn(this)
-
-            assertThat((actions?.get(Swipe.Down) as? UserActionResult.ChangeScene)?.toScene)
-                .isEqualTo(SceneFamilies.Home)
-            assertThat(actions?.get(Swipe.Up)).isNull()
-            assertThat(kosmos.homeSceneFamilyResolver.resolvedScene.value)
-                .isEqualTo(Scenes.Lockscreen)
-        }
-
-    @Test
-    fun downTransitionSceneKey_deviceUnlocked_bottomAligned_gone() =
-        testScope.runTest {
-            kosmos.fakeShadeRepository.setDualShadeAlignedToBottom(true)
-            val actions by collectLastValue(underTest.actions)
-            lockDevice()
-            unlockDevice()
-            underTest.activateIn(this)
-
-            assertThat((actions?.get(Swipe.Down) as? UserActionResult.ChangeScene)?.toScene)
-                .isEqualTo(SceneFamilies.Home)
-            assertThat(actions?.get(Swipe.Up)).isNull()
-            assertThat(sceneInteractor.currentScene.value).isEqualTo(Scenes.Gone)
-        }
-
-    @Test
     fun upTransitionSceneKey_authMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
         testScope.runTest {
             val actions by collectLastValue(underTest.actions)
@@ -153,11 +122,13 @@
     @Test
     fun upTransitionSceneKey_authMethodSwipe_lockscreenDismissed_goesToGone() =
         testScope.runTest {
+            val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
             val actions by collectLastValue(underTest.actions)
             kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
             kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
                 AuthenticationMethodModel.None
             )
+            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
             sceneInteractor // force the lazy; this will kick off StateFlows
             runCurrent()
             sceneInteractor.changeScene(Scenes.Gone, "reason")
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
index 768fbca..7203b61 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.qs.composefragment.viewmodel
 
+import android.app.StatusBarManager
 import android.content.testableContext
 import android.testing.TestableLooper.RunWithLooper
 import androidx.lifecycle.Lifecycle
@@ -32,6 +33,8 @@
 import com.android.systemui.res.R
 import com.android.systemui.shade.largeScreenHeaderHelper
 import com.android.systemui.statusbar.StatusBarState
+import com.android.systemui.statusbar.disableflags.data.model.DisableFlagsModel
+import com.android.systemui.statusbar.disableflags.data.repository.fakeDisableFlagsRepository
 import com.android.systemui.statusbar.sysuiStatusBarStateController
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
@@ -178,6 +181,23 @@
             }
         }
 
+    @Test
+    fun qsEnabled_followsRepository() =
+        with(kosmos) {
+            testScope.testWithinLifecycle {
+                val qsEnabled by collectLastValue(underTest.qsEnabled)
+
+                fakeDisableFlagsRepository.disableFlags.value =
+                    DisableFlagsModel(disable2 = QS_DISABLE_FLAG)
+
+                assertThat(qsEnabled).isFalse()
+
+                fakeDisableFlagsRepository.disableFlags.value = DisableFlagsModel()
+
+                assertThat(qsEnabled).isTrue()
+            }
+        }
+
     private inline fun TestScope.testWithinLifecycle(
         crossinline block: suspend TestScope.() -> TestResult
     ): TestResult {
@@ -186,4 +206,8 @@
             block().also { lifecycleOwner.setCurrentState(Lifecycle.State.DESTROYED) }
         }
     }
+
+    companion object {
+        private const val QS_DISABLE_FLAG = StatusBarManager.DISABLE2_QUICK_SETTINGS
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/autoaddable/AutoAddableSettingTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/autoaddable/AutoAddableSettingTest.kt
index d153e9d..5619022 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/autoaddable/AutoAddableSettingTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/autoaddable/AutoAddableSettingTest.kt
@@ -21,14 +21,15 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.qs.pipeline.domain.model.AutoAddSignal
 import com.android.systemui.qs.pipeline.domain.model.AutoAddTracking
 import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -38,10 +39,11 @@
 @RunWith(AndroidJUnit4::class)
 class AutoAddableSettingTest : SysuiTestCase() {
 
-    private val testDispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val secureSettings = kosmos.fakeSettings
 
-    private val secureSettings = FakeSettings()
     private val underTest =
         AutoAddableSetting(
             secureSettings,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/airplane/domain/interactor/AirplaneModeTileUserActionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/airplane/domain/interactor/AirplaneModeTileUserActionInteractorTest.kt
index 79fcc92..d27e810 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/airplane/domain/interactor/AirplaneModeTileUserActionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/airplane/domain/interactor/AirplaneModeTileUserActionInteractorTest.kt
@@ -29,8 +29,9 @@
 import com.android.systemui.qs.tiles.impl.airplane.domain.model.AirplaneModeTileModel
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -40,8 +41,9 @@
 @EnabledOnRavenwood
 @RunWith(AndroidJUnit4::class)
 class AirplaneModeTileUserActionInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
-    private val mobileConnectionsRepository = FakeMobileConnectionsRepository()
+    private val mobileConnectionsRepository = kosmos.fakeMobileConnectionsRepository
     private val connectivityRepository = FakeConnectivityRepository()
     private val airplaneModeRepository = FakeAirplaneModeRepository()
     private val inputHandler = FakeQSTileIntentUserInputHandler()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/internet/domain/interactor/InternetTileDataInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/internet/domain/interactor/InternetTileDataInteractorTest.kt
index 1ea8abc..6f11b2a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/internet/domain/interactor/InternetTileDataInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/internet/domain/interactor/InternetTileDataInteractorTest.kt
@@ -33,7 +33,7 @@
 import com.android.systemui.flags.Flags
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.qs.tiles.base.interactor.DataUpdateTrigger
 import com.android.systemui.qs.tiles.impl.internet.domain.model.InternetTileModel
 import com.android.systemui.res.R
@@ -86,7 +86,7 @@
     private val wifiInteractor =
         WifiInteractorImpl(connectivityRepository, wifiRepository, testScope.backgroundScope)
 
-    private val tableLogBuffer: TableLogBuffer = mock()
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "InternetTileDataInteractorTest")
     private val carrierConfigTracker: CarrierConfigTracker = mock()
 
     private val mobileConnectionsRepository =
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
index 5925819..91d8e2a 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
@@ -24,9 +24,10 @@
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.settingslib.notification.modes.ZenIconLoader
+import com.android.internal.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.asIcon
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
@@ -36,7 +37,6 @@
 import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
-import com.google.common.util.concurrent.MoreExecutors
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.toCollection
@@ -60,10 +60,10 @@
     @Before
     fun setUp() {
         context.orCreateTestableResources.apply {
-            addOverride(com.android.internal.R.drawable.ic_zen_mode_type_bedtime, BEDTIME_DRAWABLE)
-            addOverride(com.android.internal.R.drawable.ic_zen_mode_type_driving, DRIVING_DRAWABLE)
+            addOverride(MODES_DRAWABLE_ID, MODES_DRAWABLE)
+            addOverride(R.drawable.ic_zen_mode_type_bedtime, BEDTIME_DRAWABLE)
+            addOverride(R.drawable.ic_zen_mode_type_driving, DRIVING_DRAWABLE)
         }
-        ZenIconLoader.setInstance(ZenIconLoader(MoreExecutors.newDirectExecutorService()))
     }
 
     @EnableFlags(Flags.FLAG_MODES_UI)
@@ -128,28 +128,34 @@
 
     @Test
     @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
-    fun changesIconWhenActiveModesChange() =
+    fun tileData_iconsFlagEnabled_changesIconWhenActiveModesChange() =
         testScope.runTest {
-            val dataList: List<ModesTileModel> by
-                collectValues(
+            val tileData by
+                collectLastValue(
                     underTest.tileData(TEST_USER, flowOf(DataUpdateTrigger.InitialRequest))
                 )
-            runCurrent()
-            assertThat(dataList.map { it.icon }).containsExactly(null).inOrder()
 
-            // Add an inactive mode: state hasn't changed, so this shouldn't cause another emission
+            // Tile starts with the generic Modes icon.
+            runCurrent()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
+
+            // Add an inactive mode -> Still modes icon
             zenModeRepository.addMode(id = "Mode", active = false)
             runCurrent()
-            assertThat(dataList.map { it.icon }).containsExactly(null).inOrder()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
-            // Add an active mode: icon should be the mode icon
+            // Add an active mode: icon should be the mode icon. No iconResId, because we don't
+            // really know that it's a system icon.
             zenModeRepository.addMode(
                 id = "Bedtime",
                 type = AutomaticZenRule.TYPE_BEDTIME,
                 active = true
             )
             runCurrent()
-            assertThat(dataList.map { it.icon }).containsExactly(null, BEDTIME_ICON).inOrder()
+            assertThat(tileData?.icon).isEqualTo(BEDTIME_ICON)
+            assertThat(tileData?.iconResId).isNull()
 
             // Add another, less-prioritized mode: icon should remain the first mode icon
             zenModeRepository.addMode(
@@ -158,29 +164,91 @@
                 active = true
             )
             runCurrent()
-            assertThat(dataList.map { it.icon })
-                .containsExactly(null, BEDTIME_ICON, BEDTIME_ICON)
-                .inOrder()
+            assertThat(tileData?.icon).isEqualTo(BEDTIME_ICON)
+            assertThat(tileData?.iconResId).isNull()
 
-            // Deactivate more important mode: icon should be the less important, still active mode.
+            // Deactivate more important mode: icon should be the less important, still active mode
             zenModeRepository.deactivateMode("Bedtime")
             runCurrent()
-            assertThat(dataList.map { it.icon })
-                .containsExactly(null, BEDTIME_ICON, BEDTIME_ICON, DRIVING_ICON)
-                .inOrder()
+            assertThat(tileData?.icon).isEqualTo(DRIVING_ICON)
+            assertThat(tileData?.iconResId).isNull()
 
-            // Deactivate remaining mode: no icon
+            // Deactivate remaining mode: back to the default modes icon
             zenModeRepository.deactivateMode("Driving")
             runCurrent()
-            assertThat(dataList.map { it.icon })
-                .containsExactly(null, BEDTIME_ICON, BEDTIME_ICON, DRIVING_ICON, null)
-                .inOrder()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
         }
 
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    @DisableFlags(Flags.FLAG_MODES_UI_ICONS)
+    fun tileData_iconsFlagDisabled_hasPriorityModesIcon() =
+        testScope.runTest {
+            val tileData by
+                collectLastValue(
+                    underTest.tileData(TEST_USER, flowOf(DataUpdateTrigger.InitialRequest))
+                )
+
+            runCurrent()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
+
+            // Activate a Mode -> Icon doesn't change.
+            zenModeRepository.addMode(id = "Mode", active = true)
+            runCurrent()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
+
+            zenModeRepository.deactivateMode(id = "Mode")
+            runCurrent()
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
+            assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
+        }
+
+    @EnableFlags(Flags.FLAG_MODES_UI)
+    @Test
+    fun getCurrentTileModel_returnsActiveModes() = runTest {
+        var tileData = underTest.getCurrentTileModel()
+        assertThat(tileData.isActivated).isFalse()
+        assertThat(tileData.activeModes).isEmpty()
+
+        // Add active mode
+        zenModeRepository.addMode(id = "One", active = true)
+        tileData = underTest.getCurrentTileModel()
+        assertThat(tileData.isActivated).isTrue()
+        assertThat(tileData.activeModes).containsExactly("Mode One")
+
+        // Add an inactive mode: state hasn't changed
+        zenModeRepository.addMode(id = "Two", active = false)
+        tileData = underTest.getCurrentTileModel()
+        assertThat(tileData.isActivated).isTrue()
+        assertThat(tileData.activeModes).containsExactly("Mode One")
+
+        // Add another active mode
+        zenModeRepository.addMode(id = "Three", active = true)
+        tileData = underTest.getCurrentTileModel()
+        assertThat(tileData.isActivated).isTrue()
+        assertThat(tileData.activeModes).containsExactly("Mode One", "Mode Three").inOrder()
+
+        // Remove a mode and deactivate the other
+        zenModeRepository.removeMode("One")
+        zenModeRepository.deactivateMode("Three")
+        tileData = underTest.getCurrentTileModel()
+        assertThat(tileData.isActivated).isFalse()
+        assertThat(tileData.activeModes).isEmpty()
+    }
+
     private companion object {
         val TEST_USER = UserHandle.of(1)!!
+
+        val MODES_DRAWABLE_ID = R.drawable.ic_zen_priority_modes
+
+        val MODES_DRAWABLE = TestStubDrawable("modes_icon")
         val BEDTIME_DRAWABLE = TestStubDrawable("bedtime")
         val DRIVING_DRAWABLE = TestStubDrawable("driving")
+
+        val MODES_ICON = MODES_DRAWABLE.asIcon()
         val BEDTIME_ICON = BEDTIME_DRAWABLE.asIcon()
         val DRIVING_ICON = DRIVING_DRAWABLE.asIcon()
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
index 4b75649..cd58127 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileUserActionInteractorTest.kt
@@ -16,12 +16,14 @@
 
 package com.android.systemui.qs.tiles.impl.modes.domain.interactor
 
+import android.graphics.drawable.TestStubDrawable
 import android.platform.test.annotations.EnableFlags
 import android.provider.Settings
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.animation.Expandable
+import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.qs.tiles.base.actions.QSTileIntentUserInputHandlerSubject
 import com.android.systemui.qs.tiles.base.actions.qsTileIntentUserInputHandler
 import com.android.systemui.qs.tiles.base.interactor.QSTileInputTestKtx
@@ -54,10 +56,7 @@
     fun handleClick_active() = runTest {
         val expandable = mock<Expandable>()
         underTest.handleInput(
-            QSTileInputTestKtx.click(
-                data = ModesTileModel(true, listOf("DND")),
-                expandable = expandable
-            )
+            QSTileInputTestKtx.click(data = modelOf(true, listOf("DND")), expandable = expandable)
         )
 
         verify(mockDialogDelegate).showDialog(eq(expandable))
@@ -67,10 +66,7 @@
     fun handleClick_inactive() = runTest {
         val expandable = mock<Expandable>()
         underTest.handleInput(
-            QSTileInputTestKtx.click(
-                data = ModesTileModel(false, emptyList()),
-                expandable = expandable
-            )
+            QSTileInputTestKtx.click(data = modelOf(false, emptyList()), expandable = expandable)
         )
 
         verify(mockDialogDelegate).showDialog(eq(expandable))
@@ -78,7 +74,7 @@
 
     @Test
     fun handleLongClick_active() = runTest {
-        underTest.handleInput(QSTileInputTestKtx.longClick(ModesTileModel(true, listOf("DND"))))
+        underTest.handleInput(QSTileInputTestKtx.longClick(modelOf(true, listOf("DND"))))
 
         QSTileIntentUserInputHandlerSubject.assertThat(inputHandler).handledOneIntentInput {
             assertThat(it.intent.action).isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
@@ -87,10 +83,14 @@
 
     @Test
     fun handleLongClick_inactive() = runTest {
-        underTest.handleInput(QSTileInputTestKtx.longClick(ModesTileModel(false, emptyList())))
+        underTest.handleInput(QSTileInputTestKtx.longClick(modelOf(false, emptyList())))
 
         QSTileIntentUserInputHandlerSubject.assertThat(inputHandler).handledOneIntentInput {
             assertThat(it.intent.action).isEqualTo(Settings.ACTION_ZEN_MODE_SETTINGS)
         }
     }
+
+    private fun modelOf(isActivated: Boolean, activeModeNames: List<String>): ModesTileModel {
+        return ModesTileModel(isActivated, activeModeNames, TestStubDrawable("icon").asIcon(), 123)
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
index a41f15d..f7bdcb8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
@@ -18,11 +18,12 @@
 
 import android.app.Flags
 import android.graphics.drawable.TestStubDrawable
+import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfigTestBuilder
 import com.android.systemui.qs.tiles.viewmodel.QSTileState
@@ -58,47 +59,88 @@
 
     @Test
     fun inactiveState() {
-        val model = ModesTileModel(isActivated = false, activeModes = emptyList())
+        val icon = TestStubDrawable("res123").asIcon()
+        val model =
+            ModesTileModel(
+                isActivated = false,
+                activeModes = emptyList(),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
         assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.INACTIVE)
-        assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_off)
+        assertThat(state.icon()).isEqualTo(icon)
         assertThat(state.secondaryLabel).isEqualTo("No active modes")
     }
 
     @Test
     fun activeState_oneMode() {
-        val model = ModesTileModel(isActivated = true, activeModes = listOf("DND"))
+        val icon = TestStubDrawable("res123").asIcon()
+        val model =
+            ModesTileModel(
+                isActivated = true,
+                activeModes = listOf("DND"),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
         assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.ACTIVE)
-        assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_on)
+        assertThat(state.icon()).isEqualTo(icon)
         assertThat(state.secondaryLabel).isEqualTo("DND is active")
     }
 
     @Test
     fun activeState_multipleModes() {
+        val icon = TestStubDrawable("res123").asIcon()
         val model =
-            ModesTileModel(isActivated = true, activeModes = listOf("Mode 1", "Mode 2", "Mode 3"))
+            ModesTileModel(
+                isActivated = true,
+                activeModes = listOf("Mode 1", "Mode 2", "Mode 3"),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
         assertThat(state.activationState).isEqualTo(QSTileState.ActivationState.ACTIVE)
-        assertThat(state.iconRes).isEqualTo(R.drawable.qs_dnd_icon_on)
+        assertThat(state.icon()).isEqualTo(icon)
         assertThat(state.secondaryLabel).isEqualTo("3 modes are active")
     }
 
     @Test
     @EnableFlags(Flags.FLAG_MODES_UI_ICONS)
-    fun activeState_withIcon() {
-        val icon = Icon.Resource(1234, contentDescription = null)
-        val model = ModesTileModel(isActivated = true, activeModes = listOf("DND"), icon = icon)
+    fun state_withEnabledFlag_noIconResId() {
+        val icon = TestStubDrawable("res123").asIcon()
+        val model =
+            ModesTileModel(
+                isActivated = false,
+                activeModes = emptyList(),
+                icon = icon,
+                iconResId = 123 // Should not be populated, but is ignored even if present
+            )
 
         val state = underTest.map(config, model)
 
-        assertThat(state.iconRes).isNull()
         assertThat(state.icon()).isEqualTo(icon)
+        assertThat(state.iconRes).isNull()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI_ICONS)
+    fun state_withDisabledFlag_includesIconResId() {
+        val icon = TestStubDrawable("res123").asIcon()
+        val model =
+            ModesTileModel(
+                isActivated = false,
+                activeModes = emptyList(),
+                icon = icon,
+                iconResId = 123
+            )
+
+        val state = underTest.map(config, model)
+
+        assertThat(state.icon()).isEqualTo(icon)
+        assertThat(state.iconRes).isEqualTo(123)
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModelTest.kt
index e2149d9..424afe1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModelTest.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.qs.ui.viewmodel
 
+import android.platform.test.annotations.DisableFlags
 import android.testing.TestableLooper.RunWithLooper
 import androidx.lifecycle.LifecycleOwner
 import androidx.test.ext.junit.runners.AndroidJUnit4
@@ -26,20 +27,26 @@
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.fakeFeatureFlagsClassic
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.media.controls.data.repository.mediaFilterRepository
 import com.android.systemui.media.controls.domain.pipeline.interactor.mediaCarouselInteractor
 import com.android.systemui.media.controls.shared.model.MediaData
 import com.android.systemui.qs.FooterActionsController
 import com.android.systemui.qs.footer.ui.viewmodel.FooterActionsViewModel
 import com.android.systemui.qs.ui.adapter.FakeQSSceneAdapter
+import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.domain.startable.sceneContainerStartable
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.settings.brightness.ui.viewmodel.brightnessMirrorViewModelFactory
+import com.android.systemui.shade.data.repository.shadeRepository
+import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.shade.ui.viewmodel.shadeHeaderViewModelFactory
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -51,6 +58,7 @@
 @RunWith(AndroidJUnit4::class)
 @RunWithLooper
 @EnableSceneContainer
+@DisableFlags(com.android.systemui.Flags.FLAG_DUAL_SHADE)
 class QuickSettingsSceneContentViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
@@ -64,6 +72,8 @@
     private val footerActionsController = mock<FooterActionsController>()
 
     private val sceneContainerStartable = kosmos.sceneContainerStartable
+    private val sceneInteractor by lazy { kosmos.sceneInteractor }
+    private val shadeInteractor by lazy { kosmos.shadeInteractor }
 
     private lateinit var underTest: QuickSettingsSceneContentViewModel
 
@@ -80,7 +90,10 @@
                 footerActionsViewModelFactory = footerActionsViewModelFactory,
                 footerActionsController = footerActionsController,
                 mediaCarouselInteractor = kosmos.mediaCarouselInteractor,
+                shadeInteractor = shadeInteractor,
+                sceneInteractor = sceneInteractor,
             )
+        underTest.activateIn(testScope)
     }
 
     @Test
@@ -122,4 +135,16 @@
 
             assertThat(isMediaVisible).isTrue()
         }
+
+    @Test
+    fun shadeModeChange_switchToShadeScene() =
+        testScope.runTest {
+            val scene by collectLastValue(sceneInteractor.currentScene)
+
+            // switch to split shade
+            kosmos.shadeRepository.setShadeLayoutWide(true)
+            runCurrent()
+
+            assertThat(scene).isEqualTo(Scenes.Shade)
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt
new file mode 100644
index 0000000..762941d
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelTest.kt
@@ -0,0 +1,68 @@
+/*
+ * 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.qs.ui.viewmodel
+
+import android.testing.TestableLooper
+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.UserActionResult
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.lifecycle.activateIn
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+@EnableSceneContainer
+class QuickSettingsShadeOverlayActionsViewModelTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+
+    private val underTest = kosmos.quickSettingsShadeOverlayActionsViewModel
+
+    @Test
+    fun upTransitionSceneKey_hidesShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat((actions?.get(Swipe.Up) as? UserActionResult.HideOverlay)?.overlay)
+                .isEqualTo(Overlays.QuickSettingsShade)
+            assertThat(actions?.get(Swipe.Down)).isNull()
+        }
+
+    @Test
+    fun back_hidesShade() =
+        testScope.runTest {
+            val actions by collectLastValue(underTest.actions)
+            underTest.activateIn(this)
+
+            assertThat((actions?.get(Back) as? UserActionResult.HideOverlay)?.overlay)
+                .isEqualTo(Overlays.QuickSettingsShade)
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelTest.kt
new file mode 100644
index 0000000..abd1e2c
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelTest.kt
@@ -0,0 +1,60 @@
+/*
+ * 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.qs.ui.viewmodel
+
+import android.testing.TestableLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+@EnableSceneContainer
+class QuickSettingsShadeOverlayContentViewModelTest : SysuiTestCase() {
+
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val sceneInteractor = kosmos.sceneInteractor
+
+    private val underTest = kosmos.quickSettingsShadeOverlayContentViewModel
+
+    @Test
+    fun onScrimClicked_hidesShade() =
+        testScope.runTest {
+            val currentOverlays by collectLastValue(sceneInteractor.currentOverlays)
+            sceneInteractor.showOverlay(
+                overlay = Overlays.QuickSettingsShade,
+                loggingReason = "test",
+            )
+            assertThat(currentOverlays).contains(Overlays.QuickSettingsShade)
+
+            underTest.onScrimClicked()
+
+            assertThat(currentOverlays).doesNotContain(Overlays.QuickSettingsShade)
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelTest.kt
similarity index 82%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelTest.kt
index 45a8b3e..32772d2 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelTest.kt
@@ -39,14 +39,12 @@
 import com.android.systemui.scene.domain.resolver.homeSceneFamilyResolver
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.data.repository.fakeShadeRepository
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 
@@ -55,23 +53,19 @@
 @RunWith(AndroidJUnit4::class)
 @TestableLooper.RunWithLooper
 @EnableSceneContainer
-class QuickSettingsShadeSceneActionsViewModelTest : SysuiTestCase() {
+class QuickSettingsShadeUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val sceneInteractor = kosmos.sceneInteractor
     private val deviceUnlockedInteractor = kosmos.deviceUnlockedInteractor
 
-    private val underTest by lazy { kosmos.quickSettingsShadeSceneActionsViewModel }
-
-    @Before
-    fun setUp() {
-        underTest.activateIn(testScope)
-    }
+    private val underTest by lazy { kosmos.quickSettingsShadeUserActionsViewModel }
 
     @Test
     fun upTransitionSceneKey_deviceLocked_lockscreen() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             lockDevice()
@@ -85,6 +79,7 @@
     @Test
     fun upTransitionSceneKey_deviceLocked_keyguardDisabled_gone() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             lockDevice()
@@ -98,6 +93,7 @@
     @Test
     fun upTransitionSceneKey_deviceUnlocked_gone() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             lockDevice()
@@ -110,37 +106,9 @@
         }
 
     @Test
-    fun downTransitionSceneKey_deviceLocked_bottomAligned_lockscreen() =
-        testScope.runTest {
-            kosmos.fakeShadeRepository.setDualShadeAlignedToBottom(true)
-            val actions by collectLastValue(underTest.actions)
-            val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
-            lockDevice()
-
-            assertThat((actions?.get(Swipe.Down) as? UserActionResult.ChangeScene)?.toScene)
-                .isEqualTo(SceneFamilies.Home)
-            assertThat(actions?.get(Swipe.Up)).isNull()
-            assertThat(homeScene).isEqualTo(Scenes.Lockscreen)
-        }
-
-    @Test
-    fun downTransitionSceneKey_deviceUnlocked_bottomAligned_gone() =
-        testScope.runTest {
-            kosmos.fakeShadeRepository.setDualShadeAlignedToBottom(true)
-            val actions by collectLastValue(underTest.actions)
-            val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
-            lockDevice()
-            unlockDevice()
-
-            assertThat((actions?.get(Swipe.Down) as? UserActionResult.ChangeScene)?.toScene)
-                .isEqualTo(SceneFamilies.Home)
-            assertThat(actions?.get(Swipe.Up)).isNull()
-            assertThat(homeScene).isEqualTo(Scenes.Gone)
-        }
-
-    @Test
     fun upTransitionSceneKey_authMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
@@ -157,6 +125,7 @@
     @Test
     fun upTransitionSceneKey_authMethodSwipe_lockscreenDismissed_goesToGone() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
@@ -174,6 +143,7 @@
     @Test
     fun backTransitionSceneKey_notEditing_Home() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
 
             assertThat((actions?.get(Back) as? UserActionResult.ChangeScene)?.toScene)
@@ -183,6 +153,7 @@
     @Test
     fun backTransition_editing_noDestination() =
         testScope.runTest {
+            underTest.activateIn(this)
             val actions by collectLastValue(underTest.actions)
             kosmos.editModeViewModel.startEditing()
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
similarity index 98%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
index f26a9db..6986cf8e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModelTest.kt
@@ -57,7 +57,7 @@
 @RunWith(AndroidJUnit4::class)
 @RunWithLooper
 @EnableSceneContainer
-class QuickSettingsSceneActionsViewModelTest : SysuiTestCase() {
+class QuickSettingsUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
@@ -67,7 +67,7 @@
     private val sceneBackInteractor = kosmos.sceneBackInteractor
     private val sceneContainerStartable = kosmos.sceneContainerStartable
 
-    private lateinit var underTest: QuickSettingsSceneActionsViewModel
+    private lateinit var underTest: QuickSettingsUserActionsViewModel
 
     @Before
     fun setUp() {
@@ -75,7 +75,7 @@
 
         sceneContainerStartable.start()
         underTest =
-            QuickSettingsSceneActionsViewModel(
+            QuickSettingsUserActionsViewModel(
                 qsSceneAdapter = qsFlexiglassAdapter,
                 sceneBackInteractor = sceneBackInteractor,
             )
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
index 4d3909c..4f7c013 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/SceneFrameworkIntegrationTest.kt
@@ -49,7 +49,7 @@
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.fakeFeatureFlagsClassic
-import com.android.systemui.keyguard.ui.viewmodel.LockscreenSceneActionsViewModel
+import com.android.systemui.keyguard.ui.viewmodel.LockscreenUserActionsViewModel
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.lifecycle.activateIn
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
@@ -65,10 +65,10 @@
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.shade.domain.interactor.shadeInteractor
-import com.android.systemui.shade.ui.viewmodel.ShadeSceneActionsViewModel
 import com.android.systemui.shade.ui.viewmodel.ShadeSceneContentViewModel
-import com.android.systemui.shade.ui.viewmodel.shadeSceneActionsViewModel
+import com.android.systemui.shade.ui.viewmodel.ShadeUserActionsViewModel
 import com.android.systemui.shade.ui.viewmodel.shadeSceneContentViewModel
+import com.android.systemui.shade.ui.viewmodel.shadeUserActionsViewModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.telephony.data.repository.fakeTelephonyRepository
@@ -145,8 +145,8 @@
     private lateinit var bouncerActionButtonInteractor: BouncerActionButtonInteractor
     private lateinit var bouncerSceneContentViewModel: BouncerSceneContentViewModel
 
-    private val lockscreenSceneActionsViewModel by lazy {
-        LockscreenSceneActionsViewModel(
+    private val mLockscreenUserActionsViewModel by lazy {
+        LockscreenUserActionsViewModel(
             deviceEntryInteractor = deviceEntryInteractor,
             communalInteractor = communalInteractor,
             shadeInteractor = kosmos.shadeInteractor,
@@ -154,7 +154,7 @@
     }
 
     private lateinit var shadeSceneContentViewModel: ShadeSceneContentViewModel
-    private lateinit var shadeSceneActionsViewModel: ShadeSceneActionsViewModel
+    private lateinit var mShadeUserActionsViewModel: ShadeUserActionsViewModel
 
     private val powerInteractor by lazy { kosmos.powerInteractor }
 
@@ -191,14 +191,14 @@
         bouncerSceneContentViewModel = kosmos.bouncerSceneContentViewModel
 
         shadeSceneContentViewModel = kosmos.shadeSceneContentViewModel
-        shadeSceneActionsViewModel = kosmos.shadeSceneActionsViewModel
+        mShadeUserActionsViewModel = kosmos.shadeUserActionsViewModel
 
         val startable = kosmos.sceneContainerStartable
         startable.start()
 
-        lockscreenSceneActionsViewModel.activateIn(testScope)
+        mLockscreenUserActionsViewModel.activateIn(testScope)
         shadeSceneContentViewModel.activateIn(testScope)
-        shadeSceneActionsViewModel.activateIn(testScope)
+        mShadeUserActionsViewModel.activateIn(testScope)
         bouncerSceneContentViewModel.activateIn(testScope)
         sceneContainerViewModel.activateIn(testScope)
 
@@ -229,7 +229,7 @@
     @Test
     fun swipeUpOnLockscreen_enterCorrectPin_unlocksDevice() =
         testScope.runTest {
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
@@ -250,7 +250,7 @@
         testScope.runTest {
             setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
 
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
@@ -262,7 +262,7 @@
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
         testScope.runTest {
-            val actions by collectLastValue(shadeSceneActionsViewModel.actions)
+            val actions by collectLastValue(mShadeUserActionsViewModel.actions)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
             setAuthMethod(AuthenticationMethodModel.None, enableLockscreen = true)
             assertCurrentScene(Scenes.Lockscreen)
@@ -283,7 +283,7 @@
     @Test
     fun swipeUpOnShadeScene_withAuthMethodSwipe_lockscreenDismissed_goesToGone() =
         testScope.runTest {
-            val actions by collectLastValue(shadeSceneActionsViewModel.actions)
+            val actions by collectLastValue(mShadeUserActionsViewModel.actions)
             val canSwipeToEnter by collectLastValue(deviceEntryInteractor.canSwipeToEnter)
             val homeScene by collectLastValue(kosmos.homeSceneFamilyResolver.resolvedScene)
 
@@ -369,7 +369,7 @@
     fun swipeUpOnLockscreenWhileUnlocked_dismissesLockscreen() =
         testScope.runTest {
             unlockDevice()
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Gone)
@@ -392,7 +392,7 @@
     fun dismissingIme_whileOnPasswordBouncer_navigatesToLockscreen() =
         testScope.runTest {
             setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
@@ -411,7 +411,7 @@
     fun bouncerActionButtonClick_opensEmergencyServicesDialer() =
         testScope.runTest {
             setAuthMethod(AuthenticationMethodModel.Password)
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
@@ -432,7 +432,7 @@
         testScope.runTest {
             setAuthMethod(AuthenticationMethodModel.Password)
             startPhoneCall()
-            val actions by collectLastValue(lockscreenSceneActionsViewModel.actions)
+            val actions by collectLastValue(mLockscreenUserActionsViewModel.actions)
             val upDestinationSceneKey =
                 (actions?.get(Swipe.Up) as? UserActionResult.ChangeScene)?.toScene
             assertThat(upDestinationSceneKey).isEqualTo(Scenes.Bouncer)
@@ -501,7 +501,7 @@
     private fun getCurrentSceneInUi(): SceneKey {
         return when (val state = transitionState.value) {
             is ObservableTransitionState.Idle -> state.currentScene
-            is ObservableTransitionState.Transition.ChangeCurrentScene -> state.fromScene
+            is ObservableTransitionState.Transition.ChangeScene -> state.fromScene
             is ObservableTransitionState.Transition.ShowOrHideOverlay -> state.currentScene
             is ObservableTransitionState.Transition.ReplaceOverlay -> state.currentScene
         }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
index df30c4b..227b3a6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.scene.overlayKeys
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.model.Scenes
@@ -46,19 +47,9 @@
     private val testScope = kosmos.testScope
 
     @Test
-    fun allSceneKeys() {
+    fun allContentKeys() {
         val underTest = kosmos.sceneContainerRepository
-        assertThat(underTest.allSceneKeys())
-            .isEqualTo(
-                listOf(
-                    Scenes.QuickSettings,
-                    Scenes.Shade,
-                    Scenes.Lockscreen,
-                    Scenes.Bouncer,
-                    Scenes.Gone,
-                    Scenes.Communal,
-                )
-            )
+        assertThat(underTest.allContentKeys).isEqualTo(kosmos.sceneKeys + kosmos.overlayKeys)
     }
 
     @Test
@@ -75,6 +66,18 @@
             assertThat(currentScene).isEqualTo(Scenes.QuickSettings)
         }
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
+    @Test
+    fun currentOverlays() =
+        testScope.runTest {
+            val underTest = kosmos.sceneContainerRepository
+            val currentOverlays by collectLastValue(underTest.currentOverlays)
+            assertThat(currentOverlays).isEmpty()
+
+            // TODO(b/356596436): When we have a first overlay, add it here and assert contains.
+        }
+
     @Test(expected = IllegalStateException::class)
     fun changeScene_noSuchSceneInContainer_throws() {
         kosmos.sceneKeys = listOf(Scenes.QuickSettings, Scenes.Lockscreen)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
index 35cefa6..4a7d8b0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/domain/interactor/SceneInteractorTest.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.scene.data.repository.sceneContainerRepository
 import com.android.systemui.scene.data.repository.setSceneTransition
 import com.android.systemui.scene.domain.resolver.homeSceneFamilyResolver
+import com.android.systemui.scene.overlayKeys
 import com.android.systemui.scene.sceneContainerConfig
 import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.model.SceneFamilies
@@ -72,9 +73,11 @@
         kosmos.keyguardEnabledInteractor
     }
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
     @Test
-    fun allSceneKeys() {
-        assertThat(underTest.allSceneKeys()).isEqualTo(kosmos.sceneKeys)
+    fun allContentKeys() {
+        assertThat(underTest.allContentKeys).isEqualTo(kosmos.sceneKeys + kosmos.overlayKeys)
     }
 
     @Test
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 8f8d2e2..ec79cc6 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
@@ -19,6 +19,7 @@
 package com.android.systemui.scene.domain.startable
 
 import android.app.StatusBarManager
+import android.hardware.face.FaceManager
 import android.os.PowerManager
 import android.view.Display
 import androidx.test.ext.junit.runners.AndroidJUnit4
@@ -30,6 +31,9 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.biometrics.data.repository.fingerprintPropertyRepository
+import com.android.systemui.biometrics.shared.model.FingerprintSensorType
+import com.android.systemui.biometrics.shared.model.SensorStrength
 import com.android.systemui.bouncer.data.repository.fakeKeyguardBouncerRepository
 import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
 import com.android.systemui.bouncer.shared.logging.BouncerUiEvent
@@ -39,8 +43,14 @@
 import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryHapticsInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
+import com.android.systemui.deviceentry.shared.model.FailedFaceAuthenticationStatus
+import com.android.systemui.deviceentry.shared.model.SuccessFaceAuthenticationStatus
 import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.haptics.vibratorHelper
+import com.android.systemui.keyevent.data.repository.fakeKeyEventRepository
+import com.android.systemui.keyguard.data.repository.biometricSettingsRepository
 import com.android.systemui.keyguard.data.repository.deviceEntryFingerprintAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeBiometricSettingsRepository
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
@@ -51,13 +61,17 @@
 import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.keyguard.domain.interactor.keyguardEnabledInteractor
-import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
+import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.scenetransition.lockscreenSceneTransitionInteractor
+import com.android.systemui.keyguard.shared.model.BiometricUnlockMode
+import com.android.systemui.keyguard.shared.model.BiometricUnlockSource
+import com.android.systemui.keyguard.shared.model.FailFingerprintAuthenticationStatus
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.model.sysUiState
 import com.android.systemui.power.data.repository.fakePowerRepository
+import com.android.systemui.power.data.repository.powerRepository
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAsleepForTest
 import com.android.systemui.power.domain.interactor.PowerInteractor.Companion.setAwakeForTest
 import com.android.systemui.power.domain.interactor.powerInteractor
@@ -69,6 +83,7 @@
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.shared.system.QuickStepContract
+import com.android.systemui.statusbar.VibratorHelper
 import com.android.systemui.statusbar.domain.interactor.keyguardOcclusionInteractor
 import com.android.systemui.statusbar.notification.data.repository.FakeHeadsUpRowRepository
 import com.android.systemui.statusbar.notification.data.repository.HeadsUpRowRepository
@@ -85,6 +100,7 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -92,6 +108,8 @@
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.ArgumentMatchers.anyString
+import org.mockito.Mockito
 import org.mockito.Mockito.clearInvocations
 import org.mockito.Mockito.never
 import org.mockito.Mockito.times
@@ -105,12 +123,14 @@
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
+    private val deviceEntryHapticsInteractor by lazy { kosmos.deviceEntryHapticsInteractor }
     private val sceneInteractor by lazy { kosmos.sceneInteractor }
     private val bouncerInteractor by lazy { kosmos.bouncerInteractor }
     private val faceAuthRepository by lazy { kosmos.fakeDeviceEntryFaceAuthRepository }
     private val bouncerRepository by lazy { kosmos.fakeKeyguardBouncerRepository }
     private val sysUiState = kosmos.sysUiState
     private val falsingCollector = mock<FalsingCollector>().also { kosmos.falsingCollector = it }
+    private val vibratorHelper = mock<VibratorHelper>().also { kosmos.vibratorHelper = it }
     private val fakeSceneDataSource = kosmos.fakeSceneDataSource
     private val windowController = kosmos.notificationShadeWindowController
     private val centralSurfaces = kosmos.centralSurfaces
@@ -551,8 +571,7 @@
     fun switchToAOD_whenAvailable_whenDeviceSleepsLocked() =
         testScope.runTest {
             kosmos.lockscreenSceneTransitionInteractor.start()
-            val asleepState by
-                collectLastValue(kosmos.keyguardTransitionInteractor.asleepKeyguardState)
+            val asleepState by collectLastValue(kosmos.keyguardInteractor.asleepKeyguardState)
             val currentTransitionInfo by
                 collectLastValue(kosmos.keyguardTransitionRepository.currentTransitionInfoInternal)
             val transitionState =
@@ -584,8 +603,7 @@
     fun switchToDozing_whenAodUnavailable_whenDeviceSleepsLocked() =
         testScope.runTest {
             kosmos.lockscreenSceneTransitionInteractor.start()
-            val asleepState by
-                collectLastValue(kosmos.keyguardTransitionInteractor.asleepKeyguardState)
+            val asleepState by collectLastValue(kosmos.keyguardInteractor.asleepKeyguardState)
             val currentTransitionInfo by
                 collectLastValue(kosmos.keyguardTransitionRepository.currentTransitionInfoInternal)
             val transitionState =
@@ -636,6 +654,194 @@
         }
 
     @Test
+    fun playSuccessHaptics_onSuccessfulLockscreenAuth_udfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playSuccessHaptic by
+                collectLastValue(deviceEntryHapticsInteractor.playSuccessHaptic)
+
+            setupBiometricAuth(hasUdfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            unlockWithFingerprintAuth()
+
+            assertThat(playSuccessHaptic).isNotNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper)
+                .vibrateAuthSuccess(
+                    "SceneContainerStartable, $currentSceneKey device-entry::success"
+                )
+            verify(vibratorHelper, never()).vibrateAuthError(anyString())
+
+            updateFingerprintAuthStatus(isSuccess = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun playSuccessHaptics_onSuccessfulLockscreenAuth_sfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playSuccessHaptic by
+                collectLastValue(deviceEntryHapticsInteractor.playSuccessHaptic)
+
+            setupBiometricAuth(hasSfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            allowHapticsOnSfps()
+            unlockWithFingerprintAuth()
+
+            assertThat(playSuccessHaptic).isNotNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper)
+                .vibrateAuthSuccess(
+                    "SceneContainerStartable, $currentSceneKey device-entry::success"
+                )
+            verify(vibratorHelper, never()).vibrateAuthError(anyString())
+
+            updateFingerprintAuthStatus(isSuccess = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun playErrorHaptics_onFailedLockscreenAuth_udfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playErrorHaptic by collectLastValue(deviceEntryHapticsInteractor.playErrorHaptic)
+
+            setupBiometricAuth(hasUdfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            updateFingerprintAuthStatus(isSuccess = false)
+
+            assertThat(playErrorHaptic).isNotNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper)
+                .vibrateAuthError("SceneContainerStartable, $currentSceneKey device-entry::error")
+            verify(vibratorHelper, never()).vibrateAuthSuccess(anyString())
+        }
+
+    @Test
+    fun playErrorHaptics_onFailedLockscreenAuth_sfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playErrorHaptic by collectLastValue(deviceEntryHapticsInteractor.playErrorHaptic)
+
+            setupBiometricAuth(hasSfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            updateFingerprintAuthStatus(isSuccess = false)
+
+            assertThat(playErrorHaptic).isNotNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper)
+                .vibrateAuthError("SceneContainerStartable, $currentSceneKey device-entry::error")
+            verify(vibratorHelper, never()).vibrateAuthSuccess(anyString())
+        }
+
+    @Test
+    fun skipsSuccessHaptics_whenPowerButtonDown_sfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playSuccessHaptic by
+                collectLastValue(deviceEntryHapticsInteractor.playSuccessHaptic)
+
+            setupBiometricAuth(hasSfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            allowHapticsOnSfps(isPowerButtonDown = true)
+            unlockWithFingerprintAuth()
+
+            assertThat(playSuccessHaptic).isNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper, never())
+                .vibrateAuthSuccess(
+                    "SceneContainerStartable, $currentSceneKey device-entry::success"
+                )
+            verify(vibratorHelper, never()).vibrateAuthError(anyString())
+
+            updateFingerprintAuthStatus(isSuccess = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun skipsSuccessHaptics_whenPowerButtonRecentlyPressed_sfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playSuccessHaptic by
+                collectLastValue(deviceEntryHapticsInteractor.playSuccessHaptic)
+
+            setupBiometricAuth(hasSfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            allowHapticsOnSfps(lastPowerPress = 50)
+            unlockWithFingerprintAuth()
+
+            assertThat(playSuccessHaptic).isNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper, never())
+                .vibrateAuthSuccess(
+                    "SceneContainerStartable, $currentSceneKey device-entry::success"
+                )
+            verify(vibratorHelper, never()).vibrateAuthError(anyString())
+
+            updateFingerprintAuthStatus(isSuccess = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Gone)
+        }
+
+    @Test
+    fun skipsErrorHaptics_whenPowerButtonDown_sfps() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playErrorHaptic by collectLastValue(deviceEntryHapticsInteractor.playErrorHaptic)
+
+            setupBiometricAuth(hasSfps = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            kosmos.fakeKeyEventRepository.setPowerButtonDown(true)
+            updateFingerprintAuthStatus(isSuccess = false)
+
+            assertThat(playErrorHaptic).isNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper, never())
+                .vibrateAuthError("SceneContainerStartable, $currentSceneKey device-entry::error")
+            verify(vibratorHelper, never()).vibrateAuthSuccess(anyString())
+        }
+
+    @Test
+    fun skipsFaceErrorHaptics_nonSfps_coEx() =
+        testScope.runTest {
+            val currentSceneKey by collectLastValue(sceneInteractor.currentScene)
+            val playErrorHaptic by collectLastValue(deviceEntryHapticsInteractor.playErrorHaptic)
+
+            setupBiometricAuth(hasUdfps = true, hasFace = true)
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isFalse()
+
+            underTest.start()
+            updateFaceAuthStatus(isSuccess = false)
+
+            assertThat(playErrorHaptic).isNull()
+            assertThat(currentSceneKey).isEqualTo(Scenes.Lockscreen)
+            verify(vibratorHelper, never())
+                .vibrateAuthError("SceneContainerStartable, $currentSceneKey device-entry::error")
+            verify(vibratorHelper, never()).vibrateAuthSuccess(anyString())
+        }
+
+    @Test
     fun hydrateSystemUiState() =
         testScope.runTest {
             val transitionStateFlow = prepareState()
@@ -1194,41 +1400,6 @@
         }
 
     @Test
-    fun hydrateWindowController_setBouncerShowing() =
-        testScope.runTest {
-            underTest.start()
-            val notificationShadeWindowController = kosmos.notificationShadeWindowController
-            val transitionStateFlow = prepareState(initialSceneKey = Scenes.Lockscreen)
-            val currentScene by collectLastValue(sceneInteractor.currentScene)
-            assertThat(currentScene).isEqualTo(Scenes.Lockscreen)
-            verify(notificationShadeWindowController, never()).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(false)
-
-            emulateSceneTransition(transitionStateFlow, Scenes.Bouncer)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(false)
-
-            emulateSceneTransition(transitionStateFlow, Scenes.Lockscreen)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(2)).setBouncerShowing(false)
-
-            kosmos.deviceEntryFingerprintAuthRepository.setAuthenticationStatus(
-                SuccessFingerprintAuthenticationStatus(0, true)
-            )
-            assertThat(currentScene).isEqualTo(Scenes.Gone)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(2)).setBouncerShowing(false)
-
-            emulateSceneTransition(transitionStateFlow, Scenes.Lockscreen)
-            verify(notificationShadeWindowController, times(1)).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(2)).setBouncerShowing(false)
-
-            emulateSceneTransition(transitionStateFlow, Scenes.Bouncer)
-            verify(notificationShadeWindowController, times(2)).setBouncerShowing(true)
-            verify(notificationShadeWindowController, times(2)).setBouncerShowing(false)
-        }
-
-    @Test
     fun hydrateWindowController_setKeyguardOccluded() =
         testScope.runTest {
             underTest.start()
@@ -1878,4 +2049,92 @@
         FakeHeadsUpRowRepository(key = key, elementKey = Any()).apply {
             this.isPinned.value = isPinned
         }
+
+    private fun setFingerprintSensorType(fingerprintSensorType: FingerprintSensorType) {
+        kosmos.fingerprintPropertyRepository.setProperties(
+            sensorId = 0,
+            strength = SensorStrength.STRONG,
+            sensorType = fingerprintSensorType,
+            sensorLocations = mapOf(),
+        )
+        kosmos.biometricSettingsRepository.setIsFingerprintAuthEnrolledAndEnabled(true)
+    }
+
+    private fun setFaceEnrolled() {
+        kosmos.biometricSettingsRepository.setIsFaceAuthEnrolledAndEnabled(true)
+    }
+
+    private fun TestScope.allowHapticsOnSfps(
+        isPowerButtonDown: Boolean = false,
+        lastPowerPress: Long = 10000
+    ) {
+        kosmos.fakeKeyEventRepository.setPowerButtonDown(isPowerButtonDown)
+
+        kosmos.powerRepository.updateWakefulness(
+            WakefulnessState.AWAKE,
+            WakeSleepReason.POWER_BUTTON,
+            WakeSleepReason.POWER_BUTTON,
+            powerButtonLaunchGestureTriggered = false,
+        )
+
+        advanceTimeBy(lastPowerPress)
+        runCurrent()
+    }
+
+    private fun unlockWithFingerprintAuth() {
+        kosmos.fakeKeyguardRepository.setBiometricUnlockSource(
+            BiometricUnlockSource.FINGERPRINT_SENSOR
+        )
+        kosmos.fakeKeyguardRepository.setBiometricUnlockState(BiometricUnlockMode.UNLOCK_COLLAPSING)
+    }
+
+    private fun TestScope.setupBiometricAuth(
+        hasSfps: Boolean = false,
+        hasUdfps: Boolean = false,
+        hasFace: Boolean = false
+    ) {
+        if (hasSfps) {
+            setFingerprintSensorType(FingerprintSensorType.POWER_BUTTON)
+        }
+
+        if (hasUdfps) {
+            setFingerprintSensorType(FingerprintSensorType.UDFPS_ULTRASONIC)
+        }
+
+        if (hasFace) {
+            setFaceEnrolled()
+        }
+
+        prepareState(
+            authenticationMethod = AuthenticationMethodModel.Pin,
+            isDeviceUnlocked = false,
+            initialSceneKey = Scenes.Lockscreen,
+        )
+    }
+
+    private fun updateFingerprintAuthStatus(isSuccess: Boolean) {
+        if (isSuccess) {
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                SuccessFingerprintAuthenticationStatus(0, true)
+            )
+        } else {
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                FailFingerprintAuthenticationStatus
+            )
+        }
+    }
+
+    private fun updateFaceAuthStatus(isSuccess: Boolean) {
+        if (isSuccess) {
+            kosmos.fakeDeviceEntryFaceAuthRepository.setAuthenticationStatus(
+                SuccessFaceAuthenticationStatus(
+                    successResult = Mockito.mock(FaceManager.AuthenticationResult::class.java)
+                )
+            )
+        } else {
+            kosmos.fakeDeviceEntryFaceAuthRepository.setAuthenticationStatus(
+                FailedFaceAuthenticationStatus()
+            )
+        }
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
index 32c0172..2720c57 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegatorTest.kt
@@ -37,6 +37,8 @@
     private val initialSceneKey = kosmos.initialSceneKey
     private val fakeSceneDataSource = kosmos.fakeSceneDataSource
 
+    // TODO(b/356596436): Add tests for showing, hiding, and replacing overlays after we've defined
+    //  them.
     private val underTest = kosmos.sceneDataSourceDelegator
 
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModelTest.kt
similarity index 83%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModelTest.kt
index b526275..03106ec 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModelTest.kt
@@ -43,17 +43,17 @@
 @RunWith(AndroidJUnit4::class)
 @TestableLooper.RunWithLooper
 @EnableSceneContainer
-class GoneSceneActionsViewModelTest : SysuiTestCase() {
+class GoneUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val shadeRepository by lazy { kosmos.shadeRepository }
-    private lateinit var underTest: GoneSceneActionsViewModel
+    private lateinit var underTest: GoneUserActionsViewModel
 
     @Before
     fun setUp() {
         underTest =
-            GoneSceneActionsViewModel(
+            GoneUserActionsViewModel(
                 shadeInteractor = kosmos.shadeInteractor,
             )
         underTest.activateIn(testScope)
@@ -62,21 +62,21 @@
     @Test
     fun downTransitionKey_splitShadeEnabled_isGoneToSplitShade() =
         testScope.runTest {
-            val destinationScenes by collectLastValue(underTest.actions)
+            val userActions by collectLastValue(underTest.actions)
             shadeRepository.setShadeLayoutWide(true)
             runCurrent()
 
-            assertThat(destinationScenes?.get(Swipe(SwipeDirection.Down))?.transitionKey)
+            assertThat(userActions?.get(Swipe(SwipeDirection.Down))?.transitionKey)
                 .isEqualTo(ToSplitShade)
         }
 
     @Test
     fun downTransitionKey_splitShadeDisabled_isNull() =
         testScope.runTest {
-            val destinationScenes by collectLastValue(underTest.actions)
+            val userActions by collectLastValue(underTest.actions)
             shadeRepository.setShadeLayoutWide(false)
             runCurrent()
 
-            assertThat(destinationScenes?.get(Swipe(SwipeDirection.Down))?.transitionKey).isNull()
+            assertThat(userActions?.get(Swipe(SwipeDirection.Down))?.transitionKey).isNull()
         }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
index 832e7b1..3558f17 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModelTest.kt
@@ -32,7 +32,6 @@
 import com.android.systemui.power.domain.interactor.powerInteractor
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.scene.sceneContainerConfig
-import com.android.systemui.scene.sceneKeys
 import com.android.systemui.scene.shared.logger.sceneLogger
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.fakeSceneDataSource
@@ -49,6 +48,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 @EnableSceneContainer
@@ -113,11 +113,6 @@
         }
 
     @Test
-    fun allSceneKeys() {
-        assertThat(underTest.allSceneKeys).isEqualTo(kosmos.sceneKeys)
-    }
-
-    @Test
     fun sceneTransition() =
         testScope.runTest {
             val currentScene by collectLastValue(underTest.currentScene)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModelTest.kt
similarity index 93%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModelTest.kt
index dd4432d..972afb5 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModelTest.kt
@@ -35,7 +35,6 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -43,12 +42,12 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-class SceneActionsViewModelTest : SysuiTestCase() {
+class UserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
 
-    private val underTest = FakeSceneActionsViewModel()
+    private val underTest = FakeUserActionsViewModel()
 
     @Test
     fun actions_emptyBeforeActivation() =
@@ -116,14 +115,14 @@
             assertThat(actions).isEmpty()
         }
 
-    private class FakeSceneActionsViewModel : SceneActionsViewModel() {
+    private class FakeUserActionsViewModel : UserActionsViewModel() {
 
         val upstream = MutableStateFlow<Map<UserAction, UserActionResult>>(emptyMap())
 
         override suspend fun hydrateActions(
             setActions: (Map<UserAction, UserActionResult>) -> Unit,
         ) {
-            upstream.collectLatest { setActions(it) }
+            upstream.collect { setActions(it) }
         }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModelTest.kt
index 8b97739..f5022b9 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModelTest.kt
@@ -16,18 +16,28 @@
 
 package com.android.systemui.shade.ui.viewmodel
 
+import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.bouncer.data.repository.fakeKeyguardBouncerRepository
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 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.testScope
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -150,4 +160,90 @@
             )
             assertThat(isKeyguardOccluded).isTrue()
         }
+
+    @Test
+    @EnableSceneContainer
+    fun withSceneContainer_bouncerShowing_providesTheCorrectState() =
+        testScope.runTest {
+            val bouncerShowing by collectLastValue(underTest.isBouncerShowing)
+
+            val transitionState =
+                MutableStateFlow<ObservableTransitionState>(
+                    ObservableTransitionState.Idle(Scenes.Lockscreen)
+                )
+            kosmos.sceneInteractor.setTransitionState(transitionState)
+            runCurrent()
+            assertThat(bouncerShowing).isFalse()
+
+            transitionState.value = ObservableTransitionState.Idle(Scenes.Bouncer)
+            runCurrent()
+            assertThat(bouncerShowing).isTrue()
+        }
+
+    @Test
+    @EnableFlags(com.android.systemui.Flags.FLAG_COMPOSE_BOUNCER)
+    fun withComposeBouncer_bouncerShowing_providesTheCorrectState() =
+        testScope.runTest {
+            val bouncerShowing by collectLastValue(underTest.isBouncerShowing)
+
+            kosmos.fakeKeyguardBouncerRepository.setPrimaryShow(isShowing = false)
+            runCurrent()
+            assertThat(bouncerShowing).isFalse()
+
+            kosmos.fakeKeyguardBouncerRepository.setPrimaryShow(isShowing = true)
+            runCurrent()
+            assertThat(bouncerShowing).isTrue()
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun withSceneContainer_doesBouncerRequireIme_providesTheCorrectState() =
+        testScope.runTest {
+            val bouncerRequiresIme by collectLastValue(underTest.doesBouncerRequireIme)
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Pin
+            )
+
+            val transitionState =
+                MutableStateFlow<ObservableTransitionState>(
+                    ObservableTransitionState.Idle(Scenes.Bouncer)
+                )
+            kosmos.sceneInteractor.setTransitionState(transitionState)
+            runCurrent()
+            assertThat(bouncerRequiresIme).isFalse()
+
+            // go back to lockscreen
+            transitionState.value = ObservableTransitionState.Idle(Scenes.Lockscreen)
+            runCurrent()
+
+            // change auth method
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Password
+            )
+            // go back to bouncer
+            transitionState.value = ObservableTransitionState.Idle(Scenes.Bouncer)
+            runCurrent()
+            assertThat(bouncerRequiresIme).isTrue()
+        }
+
+    @Test
+    @EnableFlags(com.android.systemui.Flags.FLAG_COMPOSE_BOUNCER)
+    fun withComposeBouncer_doesBouncerRequireIme_providesTheCorrectState() =
+        testScope.runTest {
+            val bouncerRequiresIme by collectLastValue(underTest.doesBouncerRequireIme)
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Pin
+            )
+
+            kosmos.fakeKeyguardBouncerRepository.setPrimaryShow(isShowing = true)
+            runCurrent()
+            assertThat(bouncerRequiresIme).isFalse()
+
+            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
+                AuthenticationMethodModel.Password
+            )
+            kosmos.fakeKeyguardBouncerRepository.setPrimaryShow(isShowing = true)
+            runCurrent()
+            assertThat(bouncerRequiresIme).isFalse()
+        }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt
deleted file mode 100644
index 3f087b4..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelTest.kt
+++ /dev/null
@@ -1,170 +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.shade.ui.viewmodel
-
-import android.testing.TestableLooper
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.authentication.data.repository.fakeAuthenticationRepository
-import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.deviceentry.data.repository.fakeDeviceEntryRepository
-import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
-import com.android.systemui.flags.EnableSceneContainer
-import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFingerprintAuthRepository
-import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.lifecycle.activateIn
-import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.testKosmos
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.runCurrent
-import kotlinx.coroutines.test.runTest
-import org.junit.Before
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@OptIn(ExperimentalCoroutinesApi::class)
-@SmallTest
-@RunWith(AndroidJUnit4::class)
-@TestableLooper.RunWithLooper
-@EnableSceneContainer
-class OverlayShadeViewModelTest : SysuiTestCase() {
-
-    private val kosmos = testKosmos()
-    private val testScope = kosmos.testScope
-    private val sceneInteractor = kosmos.sceneInteractor
-    private val deviceUnlockedInteractor by lazy { kosmos.deviceUnlockedInteractor }
-
-    private val underTest = kosmos.overlayShadeViewModel
-
-    @Before
-    fun setUp() {
-        underTest.activateIn(testScope)
-    }
-
-    @Test
-    fun backgroundScene_deviceLocked_lockscreen() =
-        testScope.runTest {
-            val backgroundScene by collectLastValue(underTest.backgroundScene)
-
-            lockDevice()
-
-            assertThat(backgroundScene).isEqualTo(Scenes.Lockscreen)
-        }
-
-    @Test
-    fun backgroundScene_deviceUnlocked_gone() =
-        testScope.runTest {
-            val backgroundScene by collectLastValue(underTest.backgroundScene)
-
-            lockDevice()
-            unlockDevice()
-
-            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
-        }
-
-    @Test
-    fun backgroundScene_authMethodSwipe_lockscreenNotDismissed_goesToLockscreen() =
-        testScope.runTest {
-            val backgroundScene by collectLastValue(underTest.backgroundScene)
-            val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
-
-            kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
-            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.None
-            )
-            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
-            sceneInteractor.changeScene(Scenes.Lockscreen, "reason")
-            runCurrent()
-
-            assertThat(backgroundScene).isEqualTo(Scenes.Lockscreen)
-        }
-
-    @Test
-    fun backgroundScene_authMethodSwipe_lockscreenDismissed_goesToGone() =
-        testScope.runTest {
-            val backgroundScene by collectLastValue(underTest.backgroundScene)
-            val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
-
-            kosmos.fakeDeviceEntryRepository.setLockscreenEnabled(true)
-            kosmos.fakeAuthenticationRepository.setAuthenticationMethod(
-                AuthenticationMethodModel.None
-            )
-            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
-            sceneInteractor.changeScene(Scenes.Gone, "reason")
-            runCurrent()
-
-            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
-        }
-
-    @Test
-    fun onScrimClicked_onLockscreen_goesToLockscreen() =
-        testScope.runTest {
-            val currentScene by collectLastValue(sceneInteractor.currentScene)
-            lockDevice()
-            sceneInteractor.changeScene(Scenes.Bouncer, "reason")
-            runCurrent()
-            assertThat(currentScene).isNotEqualTo(Scenes.Lockscreen)
-
-            underTest.onScrimClicked()
-
-            assertThat(currentScene).isEqualTo(Scenes.Lockscreen)
-        }
-
-    @Test
-    fun onScrimClicked_deviceWasEntered_goesToGone() =
-        testScope.runTest {
-            val currentScene by collectLastValue(sceneInteractor.currentScene)
-            val backgroundScene by collectLastValue(underTest.backgroundScene)
-
-            lockDevice()
-            unlockDevice()
-            sceneInteractor.changeScene(Scenes.QuickSettings, "reason")
-            runCurrent()
-            assertThat(backgroundScene).isEqualTo(Scenes.Gone)
-            assertThat(currentScene).isNotEqualTo(Scenes.Gone)
-
-            underTest.onScrimClicked()
-
-            assertThat(currentScene).isEqualTo(Scenes.Gone)
-        }
-
-    private fun TestScope.lockDevice() {
-        val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
-
-        kosmos.fakeAuthenticationRepository.setAuthenticationMethod(AuthenticationMethodModel.Pin)
-        assertThat(deviceUnlockStatus?.isUnlocked).isFalse()
-        sceneInteractor.changeScene(Scenes.Lockscreen, "reason")
-        runCurrent()
-    }
-
-    private fun TestScope.unlockDevice() {
-        val deviceUnlockStatus by collectLastValue(deviceUnlockedInteractor.deviceUnlockStatus)
-
-        kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
-            SuccessFingerprintAuthenticationStatus(0, true)
-        )
-        assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
-        sceneInteractor.changeScene(Scenes.Gone, "reason")
-        runCurrent()
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelTest.kt
similarity index 98%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelTest.kt
rename to packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelTest.kt
index a931e65..9f3e126e 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelTest.kt
@@ -64,7 +64,7 @@
 @TestableLooper.RunWithLooper
 @EnableSceneContainer
 @DisableFlags(DualShade.FLAG_NAME)
-class ShadeSceneActionsViewModelTest : SysuiTestCase() {
+class ShadeUserActionsViewModelTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
@@ -72,7 +72,7 @@
     private val shadeRepository by lazy { kosmos.shadeRepository }
     private val qsSceneAdapter by lazy { kosmos.fakeQSSceneAdapter }
 
-    private val underTest: ShadeSceneActionsViewModel by lazy { kosmos.shadeSceneActionsViewModel }
+    private val underTest: ShadeUserActionsViewModel by lazy { kosmos.shadeUserActionsViewModel }
 
     @Before
     fun setUp() {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractorTest.kt
index 14134cc..cea8857 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractorTest.kt
@@ -18,11 +18,11 @@
 
 package com.android.systemui.statusbar.notification.domain.interactor
 
-import android.platform.test.annotations.EnableFlags
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.keyguard.data.repository.FakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
@@ -31,7 +31,6 @@
 import com.android.systemui.shade.shadeTestUtil
 import com.android.systemui.statusbar.notification.data.repository.FakeHeadsUpRowRepository
 import com.android.systemui.statusbar.notification.data.repository.notificationsKeyguardViewStateRepository
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.stack.data.repository.headsUpNotificationRepository
 import com.android.systemui.statusbar.notification.stack.domain.interactor.headsUpNotificationInteractor
 import com.android.systemui.testKosmos
@@ -44,7 +43,7 @@
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
-@EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+@EnableSceneContainer
 class HeadsUpNotificationInteractorTest : SysuiTestCase() {
     private val kosmos =
         testKosmos().apply {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModelTest.kt
index f96cf10..840aa92 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModelTest.kt
@@ -26,6 +26,7 @@
 import com.android.settingslib.notification.data.repository.updateNotificationPolicy
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.flags.fakeFeatureFlagsClassic
@@ -41,7 +42,6 @@
 import com.android.systemui.statusbar.notification.data.repository.activeNotificationListRepository
 import com.android.systemui.statusbar.notification.data.repository.setActiveNotifs
 import com.android.systemui.statusbar.notification.footer.shared.FooterViewRefactor
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.stack.data.repository.headsUpNotificationRepository
 import com.android.systemui.statusbar.policy.data.repository.fakeUserSetupRepository
 import com.android.systemui.statusbar.policy.data.repository.zenModeRepository
@@ -535,7 +535,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun pinnedHeadsUpRows_filtersForPinnedItems() =
         testScope.runTest {
             val pinnedHeadsUpRows by collectLastValue(underTest.pinnedHeadsUpRows)
@@ -576,7 +576,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun hasPinnedHeadsUpRows_true() =
         testScope.runTest {
             val hasPinnedHeadsUpRow by collectLastValue(underTest.hasPinnedHeadsUpRow)
@@ -591,7 +591,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun hasPinnedHeadsUpRows_false() =
         testScope.runTest {
             val hasPinnedHeadsUpRow by collectLastValue(underTest.hasPinnedHeadsUpRow)
@@ -606,7 +606,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun topHeadsUpRow_emptyList_null() =
         testScope.runTest {
             val topHeadsUpRow by collectLastValue(underTest.topHeadsUpRow)
@@ -618,7 +618,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun headsUpAnimationsEnabled_true() =
         testScope.runTest {
             val animationsEnabled by collectLastValue(underTest.headsUpAnimationsEnabled)
@@ -631,7 +631,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun headsUpAnimationsEnabled_keyguardShowing_true() =
         testScope.runTest {
             val animationsEnabled by collectLastValue(underTest.headsUpAnimationsEnabled)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
index 733cac9..3f97f0b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModelTest.kt
@@ -42,10 +42,14 @@
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
 import com.android.systemui.keyguard.shared.model.BurnInModel
-import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
+import com.android.systemui.keyguard.shared.model.KeyguardState.DREAMING
+import com.android.systemui.keyguard.shared.model.KeyguardState.GLANCEABLE_HUB
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
+import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
+import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
@@ -56,6 +60,10 @@
 import com.android.systemui.keyguard.ui.viewmodel.keyguardRootViewModel
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
+import com.android.systemui.scene.data.repository.Idle
+import com.android.systemui.scene.data.repository.Transition
+import com.android.systemui.scene.data.repository.setTransition
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.data.repository.fakeShadeRepository
 import com.android.systemui.shade.mockLargeScreenHeaderHelper
 import com.android.systemui.shade.shadeTestUtil
@@ -66,6 +74,7 @@
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runCurrent
@@ -295,34 +304,47 @@
 
             // Start transitioning to glanceable hub
             val progress = 0.6f
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.STARTED,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 0f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = 0f,
+                    )
             )
+
             runCurrent()
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.RUNNING,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = progress,
-                )
+            kosmos.setTransition(
+                sceneTransition =
+                    Transition(
+                        from = Scenes.Lockscreen,
+                        to = Scenes.Communal,
+                        progress = flowOf(progress)
+                    ),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.RUNNING,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = progress,
+                    )
             )
+
             runCurrent()
             assertThat(alpha).isIn(Range.closed(0f, 1f))
 
             // Finish transition to glanceable hub
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.FINISHED,
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 1f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.FINISHED,
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                        value = 1f,
+                    )
             )
             assertThat(alpha).isEqualTo(0f)
 
@@ -348,35 +370,46 @@
 
             // Start transitioning to glanceable hub
             val progress = 0.6f
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.STARTED,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 0f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = 0f,
+                    )
             )
             runCurrent()
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.RUNNING,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = progress,
-                )
+            kosmos.setTransition(
+                sceneTransition =
+                    Transition(
+                        from = Scenes.Lockscreen,
+                        to = Scenes.Communal,
+                        progress = flowOf(progress)
+                    ),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.RUNNING,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = progress,
+                    )
             )
             runCurrent()
             // Keep notifications hidden during the transition from dream to hub
             assertThat(alpha).isEqualTo(0)
 
             // Finish transition to glanceable hub
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    transitionState = TransitionState.FINISHED,
-                    from = KeyguardState.DREAMING,
-                    to = KeyguardState.GLANCEABLE_HUB,
-                    value = 1f,
-                )
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        transitionState = TransitionState.FINISHED,
+                        from = DREAMING,
+                        to = GLANCEABLE_HUB,
+                        value = 1f,
+                    )
             )
             assertThat(alpha).isEqualTo(0f)
         }
@@ -400,35 +433,47 @@
         testScope.runTest {
             val isOnLockscreen by collectLastValue(underTest.isOnLockscreen)
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GONE,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(from = LOCKSCREEN, to = GONE)
             )
             assertThat(isOnLockscreen).isFalse()
 
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(from = GONE, to = LOCKSCREEN)
+            )
+            assertThat(isOnLockscreen).isTrue()
             // While progressing from lockscreen, should still be true
-            keyguardTransitionRepository.sendTransitionStep(
-                TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
-                    value = 0.8f,
-                    transitionState = TransitionState.RUNNING
-                )
+            kosmos.setTransition(
+                sceneTransition = Transition(from = Scenes.Lockscreen, to = Scenes.Gone),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GONE,
+                        value = 0.8f,
+                        transitionState = TransitionState.RUNNING
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.GONE,
-                to = KeyguardState.LOCKSCREEN,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition =
+                    TransitionStep(
+                        from = GONE,
+                        to = LOCKSCREEN,
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
 
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.PRIMARY_BOUNCER,
-                testScope,
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Bouncer),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = PRIMARY_BOUNCER,
+                    )
             )
             assertThat(isOnLockscreen).isTrue()
         }
@@ -442,8 +487,8 @@
             shadeTestUtil.setLockscreenShadeExpansion(0f)
             shadeTestUtil.setQsExpansion(0f)
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.OCCLUDED,
+                from = LOCKSCREEN,
+                to = OCCLUDED,
                 testScope,
             )
             assertThat(isOnLockscreenWithoutShade).isFalse()
@@ -480,11 +525,15 @@
             assertThat(isOnGlanceableHubWithoutShade).isFalse()
 
             // Move to glanceable hub
-            keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GLANCEABLE_HUB,
-                testScope = this
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                    )
             )
+
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
 
             // While state is GLANCEABLE_HUB, validate variations of both shade and qs expansion
@@ -502,6 +551,14 @@
 
             shadeTestUtil.setQsExpansion(0f)
             shadeTestUtil.setLockscreenShadeExpansion(0f)
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Communal),
+                stateTransition =
+                    TransitionStep(
+                        from = LOCKSCREEN,
+                        to = GLANCEABLE_HUB,
+                    )
+            )
             assertThat(isOnGlanceableHubWithoutShade).isTrue()
         }
 
@@ -808,8 +865,8 @@
             // GONE transition gets to 90% complete
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
+                    from = LOCKSCREEN,
+                    to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
                 )
@@ -817,8 +874,8 @@
             runCurrent()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.GONE,
+                    from = LOCKSCREEN,
+                    to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
                 )
@@ -843,8 +900,8 @@
             // OCCLUDED transition gets to 90% complete
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.OCCLUDED,
+                    from = LOCKSCREEN,
+                    to = OCCLUDED,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
                 )
@@ -852,8 +909,8 @@
             runCurrent()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.LOCKSCREEN,
-                    to = KeyguardState.OCCLUDED,
+                    from = LOCKSCREEN,
+                    to = OCCLUDED,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
                 )
@@ -877,8 +934,8 @@
             showLockscreen()
 
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.GONE,
+                from = LOCKSCREEN,
+                to = GONE,
                 testScope
             )
             keyguardRepository.setStatusBarState(StatusBarState.SHADE)
@@ -922,8 +979,8 @@
 
             // ... then user hits power to go to AOD
             keyguardTransitionRepository.sendTransitionSteps(
-                from = KeyguardState.LOCKSCREEN,
-                to = KeyguardState.AOD,
+                from = LOCKSCREEN,
+                to = AOD,
                 testScope,
             )
             // ... followed by a shade collapse
@@ -945,7 +1002,7 @@
             // PRIMARY_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
@@ -956,7 +1013,7 @@
             // PRIMARY_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -967,7 +1024,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -979,7 +1036,7 @@
             hideCommunalScene()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1003,7 +1060,7 @@
             // PRIMARY_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                 )
@@ -1013,7 +1070,7 @@
             // PRIMARY_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1024,7 +1081,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1035,7 +1092,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.PRIMARY_BOUNCER,
+                    from = PRIMARY_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1058,7 +1115,7 @@
             // ALTERNATE_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                     value = 0f,
@@ -1069,7 +1126,7 @@
             // ALTERNATE_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1080,7 +1137,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1092,7 +1149,7 @@
             hideCommunalScene()
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1116,7 +1173,7 @@
             // ALTERNATE_BOUNCER->GONE transition is started
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.STARTED,
                 )
@@ -1126,7 +1183,7 @@
             // ALTERNATE_BOUNCER->GONE transition running
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.1f,
@@ -1137,7 +1194,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.RUNNING,
                     value = 0.9f,
@@ -1148,7 +1205,7 @@
 
             keyguardTransitionRepository.sendTransitionStep(
                 TransitionStep(
-                    from = KeyguardState.ALTERNATE_BOUNCER,
+                    from = ALTERNATE_BOUNCER,
                     to = GONE,
                     transitionState = TransitionState.FINISHED,
                     value = 1f
@@ -1165,8 +1222,8 @@
         keyguardRepository.setStatusBarState(StatusBarState.KEYGUARD)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1178,8 +1235,8 @@
         keyguardRepository.setDreaming(true)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.LOCKSCREEN,
-            to = KeyguardState.DREAMING,
+            from = LOCKSCREEN,
+            to = DREAMING,
             testScope,
         )
     }
@@ -1191,8 +1248,8 @@
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1204,8 +1261,8 @@
         keyguardRepository.setStatusBarState(StatusBarState.SHADE_LOCKED)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.AOD,
-            to = KeyguardState.LOCKSCREEN,
+            from = AOD,
+            to = LOCKSCREEN,
             testScope,
         )
     }
@@ -1219,8 +1276,8 @@
         kosmos.keyguardBouncerRepository.setPrimaryShow(true)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.GLANCEABLE_HUB,
-            to = KeyguardState.PRIMARY_BOUNCER,
+            from = GLANCEABLE_HUB,
+            to = PRIMARY_BOUNCER,
             testScope,
         )
     }
@@ -1234,8 +1291,8 @@
         kosmos.keyguardBouncerRepository.setPrimaryShow(false)
         runCurrent()
         keyguardTransitionRepository.sendTransitionSteps(
-            from = KeyguardState.GLANCEABLE_HUB,
-            to = KeyguardState.ALTERNATE_BOUNCER,
+            from = GLANCEABLE_HUB,
+            to = ALTERNATE_BOUNCER,
             testScope,
         )
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
index 1237347..dc24cf7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
@@ -25,14 +25,14 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.ContentDescription.Companion.loadContentDescription
 import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.statusbar.connectivity.WifiIcons
 import com.android.systemui.statusbar.phone.StatusBarLocation
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModel
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModelImpl
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.model.DataActivityModel
@@ -44,6 +44,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiNetworkModel
 import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel.Companion.viewModelForLocation
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.TestScope
@@ -58,10 +59,11 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class WifiViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var underTest: WifiViewModel
 
-    @Mock private lateinit var tableLogBuffer: TableLogBuffer
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "WifiViewModelTest")
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -86,7 +88,7 @@
                 AirplaneModeInteractor(
                     airplaneModeRepository,
                     connectivityRepository,
-                    FakeMobileConnectionsRepository(),
+                    kosmos.fakeMobileConnectionsRepository,
                 ),
                 tableLogBuffer,
                 testScope.backgroundScope,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HeadsUpManagerPhoneTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HeadsUpManagerPhoneTest.kt
index ca106fa..469a7bc 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HeadsUpManagerPhoneTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/HeadsUpManagerPhoneTest.kt
@@ -23,6 +23,7 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.plugins.statusbar.StatusBarStateController
@@ -36,7 +37,6 @@
 import com.android.systemui.statusbar.notification.collection.render.GroupMembershipManager
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.phone.ConfigurationControllerImpl
 import com.android.systemui.statusbar.phone.HeadsUpManagerPhone
 import com.android.systemui.statusbar.phone.KeyguardBypassController
@@ -469,9 +469,9 @@
         @get:Parameters(name = "{0}")
         val flags: List<FlagsParameterization>
             get() = buildList {
-                addAll(FlagsParameterization.allCombinationsOf(NotificationThrottleHun.FLAG_NAME))
                 addAll(
-                    FlagsParameterization.allCombinationsOf(NotificationsHeadsUpRefactor.FLAG_NAME)
+                    FlagsParameterization.allCombinationsOf(NotificationThrottleHun.FLAG_NAME)
+                        .andSceneContainer()
                 )
             }
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
index 20d3a7b..fb32855 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
@@ -22,8 +22,10 @@
 import android.provider.Settings.Secure.ZEN_DURATION
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
 import android.provider.Settings.Secure.ZEN_DURATION_PROMPT
+import android.service.notification.SystemZenRules
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.R
 import com.android.settingslib.notification.data.repository.updateNotificationPolicy
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
@@ -220,30 +222,116 @@
         }
 
     @Test
-    fun mainActiveMode_returnsMainActiveMode() =
+    fun activeModes_computesMainActiveMode() =
         testScope.runTest {
-            val mainActiveMode by collectLastValue(underTest.mainActiveMode)
+            val activeModes by collectLastValue(underTest.activeModes)
 
             zenModeRepository.addMode(id = "Bedtime", type = AutomaticZenRule.TYPE_BEDTIME)
             zenModeRepository.addMode(id = "Other", type = AutomaticZenRule.TYPE_OTHER)
 
             runCurrent()
+            assertThat(activeModes?.modeNames).hasSize(0)
+            assertThat(activeModes?.mainMode).isNull()
+
+            zenModeRepository.activateMode("Other")
+            runCurrent()
+            assertThat(activeModes?.modeNames).containsExactly("Mode Other")
+            assertThat(activeModes?.mainMode?.name).isEqualTo("Mode Other")
+
+            zenModeRepository.activateMode("Bedtime")
+            runCurrent()
+            assertThat(activeModes?.modeNames)
+                .containsExactly("Mode Bedtime", "Mode Other")
+                .inOrder()
+            assertThat(activeModes?.mainMode?.name).isEqualTo("Mode Bedtime")
+
+            zenModeRepository.deactivateMode("Other")
+            runCurrent()
+            assertThat(activeModes?.modeNames).containsExactly("Mode Bedtime")
+            assertThat(activeModes?.mainMode?.name).isEqualTo("Mode Bedtime")
+
+            zenModeRepository.deactivateMode("Bedtime")
+            runCurrent()
+            assertThat(activeModes?.modeNames).hasSize(0)
+            assertThat(activeModes?.mainMode).isNull()
+        }
+
+    @Test
+    fun getActiveModes_computesMainActiveMode() = runTest {
+        zenModeRepository.addMode(id = "Bedtime", type = AutomaticZenRule.TYPE_BEDTIME)
+        zenModeRepository.addMode(id = "Other", type = AutomaticZenRule.TYPE_OTHER)
+
+        var activeModes = underTest.getActiveModes()
+        assertThat(activeModes.modeNames).hasSize(0)
+        assertThat(activeModes.mainMode).isNull()
+
+        zenModeRepository.activateMode("Other")
+        activeModes = underTest.getActiveModes()
+        assertThat(activeModes.modeNames).containsExactly("Mode Other")
+        assertThat(activeModes.mainMode?.name).isEqualTo("Mode Other")
+
+        zenModeRepository.activateMode("Bedtime")
+        activeModes = underTest.getActiveModes()
+        assertThat(activeModes.modeNames).containsExactly("Mode Bedtime", "Mode Other").inOrder()
+        assertThat(activeModes.mainMode?.name).isEqualTo("Mode Bedtime")
+
+        zenModeRepository.deactivateMode("Other")
+        activeModes = underTest.getActiveModes()
+        assertThat(activeModes.modeNames).containsExactly("Mode Bedtime")
+        assertThat(activeModes.mainMode?.name).isEqualTo("Mode Bedtime")
+
+        zenModeRepository.deactivateMode("Bedtime")
+        activeModes = underTest.getActiveModes()
+        assertThat(activeModes.modeNames).hasSize(0)
+        assertThat(activeModes.mainMode).isNull()
+    }
+
+    @Test
+    fun mainActiveMode_flows() =
+        testScope.runTest {
+            val mainActiveMode by collectLastValue(underTest.mainActiveMode)
+
+            zenModeRepository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId("Bedtime")
+                        .setName("Mode Bedtime")
+                        .setType(AutomaticZenRule.TYPE_BEDTIME)
+                        .setActive(false)
+                        .setPackage(mContext.packageName)
+                        .setIconResId(R.drawable.ic_zen_mode_type_bedtime)
+                        .build(),
+                    TestModeBuilder()
+                        .setId("Other")
+                        .setName("Mode Other")
+                        .setType(AutomaticZenRule.TYPE_OTHER)
+                        .setActive(false)
+                        .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                        .setIconResId(R.drawable.ic_zen_mode_type_other)
+                        .build(),
+                )
+            )
+
+            runCurrent()
             assertThat(mainActiveMode).isNull()
 
             zenModeRepository.activateMode("Other")
             runCurrent()
-            assertThat(mainActiveMode).isNotNull()
-            assertThat(mainActiveMode!!.id).isEqualTo("Other")
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Other")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_other)
 
             zenModeRepository.activateMode("Bedtime")
             runCurrent()
-            assertThat(mainActiveMode).isNotNull()
-            assertThat(mainActiveMode!!.id).isEqualTo("Bedtime")
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
 
             zenModeRepository.deactivateMode("Other")
             runCurrent()
-            assertThat(mainActiveMode).isNotNull()
-            assertThat(mainActiveMode!!.id).isEqualTo("Bedtime")
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
 
             zenModeRepository.deactivateMode("Bedtime")
             runCurrent()
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
index 33f379d..a0f6431 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModelTest.kt
@@ -23,7 +23,6 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.settingslib.notification.modes.TestModeBuilder
-import com.android.settingslib.notification.modes.ZenIconLoader
 import com.android.settingslib.notification.modes.ZenMode
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
@@ -35,12 +34,10 @@
 import com.android.systemui.statusbar.policy.ui.dialog.mockModesDialogEventLogger
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
-import com.google.common.util.concurrent.MoreExecutors
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito.clearInvocations
@@ -67,11 +64,6 @@
             mockDialogEventLogger
         )
 
-    @Before
-    fun setUp() {
-        ZenIconLoader.setInstance(ZenIconLoader(MoreExecutors.newDirectExecutorService()))
-    }
-
     @Test
     fun tiles_filtersOutUserDisabledModes() =
         testScope.runTest {
@@ -338,6 +330,83 @@
         }
 
     @Test
+    fun tiles_populatesFieldsForAccessibility() =
+        testScope.runTest {
+            val tiles by collectLastValue(underTest.tiles)
+
+            repository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setName("With description, inactive")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription("When the going gets tough")
+                        .setActive(false)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("With description, active")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription("When in Rome")
+                        .setActive(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("With description, needs setup")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription("When you find yourself in a hole")
+                        .setEnabled(false, /* byUser= */ false)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Without description, inactive")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription(null)
+                        .setActive(false)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Without description, active")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription(null)
+                        .setActive(true)
+                        .build(),
+                    TestModeBuilder()
+                        .setName("Without description, needs setup")
+                        .setManualInvocationAllowed(true)
+                        .setTriggerDescription(null)
+                        .setEnabled(false, /* byUser= */ false)
+                        .build(),
+                )
+            )
+            runCurrent()
+
+            assertThat(tiles!!).hasSize(6)
+            with(tiles?.elementAt(0)!!) {
+                assertThat(this.stateDescription).isEqualTo("Off")
+                assertThat(this.subtextDescription).isEqualTo("When the going gets tough")
+            }
+            with(tiles?.elementAt(1)!!) {
+                assertThat(this.stateDescription).isEqualTo("On")
+                assertThat(this.subtextDescription).isEqualTo("When in Rome")
+            }
+            with(tiles?.elementAt(2)!!) {
+                assertThat(this.stateDescription).isEqualTo("Off")
+                assertThat(this.subtextDescription).isEqualTo("Set up")
+            }
+            with(tiles?.elementAt(3)!!) {
+                assertThat(this.stateDescription).isEqualTo("Off")
+                assertThat(this.subtextDescription).isEmpty()
+            }
+            with(tiles?.elementAt(4)!!) {
+                assertThat(this.stateDescription).isEqualTo("On")
+                assertThat(this.subtextDescription).isEmpty()
+            }
+            with(tiles?.elementAt(5)!!) {
+                assertThat(this.stateDescription).isEqualTo("Off")
+                assertThat(this.subtextDescription).isEqualTo("Set up")
+            }
+
+            // All tiles have the same long click info
+            tiles!!.forEach { assertThat(it.onLongClickLabel).isEqualTo("Open settings") }
+        }
+
+    @Test
     fun onClick_togglesTileState() =
         testScope.runTest {
             val tiles by collectLastValue(underTest.tiles)
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractorTest.kt
index 0f56d0b..fa7f37c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractorTest.kt
@@ -90,8 +90,9 @@
                 assertThat(model)
                     .isEqualTo(
                         MediaOutputComponentModel.Calling(
-                            AudioOutputDevice.BuiltIn(builtInDeviceName, testIcon),
-                            false,
+                            device = AudioOutputDevice.BuiltIn(builtInDeviceName, testIcon),
+                            isInAudioSharing = false,
+                            canOpenAudioSwitcher = false,
                         )
                     )
             }
@@ -101,6 +102,9 @@
     fun hasSession_stateIs_MediaSession() =
         with(kosmos) {
             testScope.runTest {
+                localMediaRepository.updateCurrentConnectedDevice(
+                    TestMediaDevicesFactory.builtInMediaDevice()
+                )
                 mediaControllerRepository.setActiveSessions(listOf(localMediaController))
 
                 val model by collectLastValue(underTest.mediaOutputModel.filterData())
@@ -113,6 +117,7 @@
                     assertThat(device)
                         .isEqualTo(AudioOutputDevice.BuiltIn("built_in_media", testIcon))
                     assertThat(isInAudioSharing).isFalse()
+                    assertThat(canOpenAudioSwitcher).isTrue()
                 }
             }
         }
@@ -129,8 +134,9 @@
                 assertThat(model)
                     .isEqualTo(
                         MediaOutputComponentModel.Idle(
-                            AudioOutputDevice.BuiltIn("built_in_media", testIcon),
-                            true,
+                            device = AudioOutputDevice.BuiltIn("built_in_media", testIcon),
+                            isInAudioSharing = true,
+                            canOpenAudioSwitcher = false,
                         )
                     )
             }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
index 51a70bd..fe6c741 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/panel/ui/viewmodel/VolumePanelViewModelTest.kt
@@ -28,6 +28,7 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.configurationController
 import com.android.systemui.statusbar.policy.fakeConfigurationController
 import com.android.systemui.testKosmos
@@ -157,6 +158,10 @@
     @Test
     fun testDumpingState() =
         test({
+            testableResources.addOverride(R.bool.volume_panel_is_large_screen, false)
+            testableResources.overrideConfiguration(
+                Configuration().apply { orientation = Configuration.ORIENTATION_PORTRAIT }
+            )
             componentByKey =
                 mapOf(
                     COMPONENT_1 to mockVolumePanelUiComponentProvider,
diff --git a/packages/SystemUI/res/values/config.xml b/packages/SystemUI/res/values/config.xml
index 823ff9f..38ef0e9 100644
--- a/packages/SystemUI/res/values/config.xml
+++ b/packages/SystemUI/res/values/config.xml
@@ -727,6 +727,10 @@
         .75
     </item>
 
+    <!-- The last x ms of face acquired info messages to analyze to determine
+         whether to show a deferred face auth help message. -->
+    <integer name="config_face_help_msgs_defer_analyze_timeframe">500</integer>
+
     <!-- Which face help messages to surface when fingerprint is also enrolled.
          Message ids correspond with the acquired ids in BiometricFaceConstants -->
     <integer-array name="config_face_help_msgs_when_fingerprint_enrolled">
@@ -1048,9 +1052,6 @@
     <!-- The width of the shortcut helper container, as a fraction of the screen's width. -->
     <item name="shortcut_helper_screen_width_fraction" format="float" type="dimen">1.0</item>
 
-    <!-- Only applicable for dual shade - Allow Notifications/QS shade to anchor to the bottom. -->
-    <bool name="config_dualShadeAlignedToBottom">false</bool>
-
     <!-- List of packages for which we want to use activity info (instead of application info) for biometric prompt logo. Empty for AOSP. [DO NOT TRANSLATE] -->
     <string-array name="config_useActivityLogoForBiometricPrompt" translatable="false"/>
 
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index fd943d0..d3d757b 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -270,6 +270,7 @@
     <!-- Add to note button used in App Clips flow to return the saved screenshot image to notes app. [CHAR LIMIT=NONE] -->
     <string name="app_clips_save_add_to_note">Add to note</string>
     <string name="backlinks_include_link">Include link</string>
+    <string name="backlinks_duplicate_label_format"><xliff:g id="appName" example="Google Chrome">%1$s</xliff:g> <xliff:g id="frequencyCount" example="(1)">(%2$d)</xliff:g></string>
 
     <!-- Notification title displayed for screen recording [CHAR LIMIT=50]-->
     <string name="screenrecord_title">Screen Recorder</string>
@@ -720,8 +721,8 @@
     <!-- QuickSettings: Do not disturb - Priority only [CHAR LIMIT=NONE] -->
     <!-- QuickSettings: Do not disturb - Alarms only [CHAR LIMIT=NONE] -->
     <!-- QuickSettings: Do not disturb - Total silence [CHAR LIMIT=NONE] -->
-    <!-- QuickSettings: Priority modes [CHAR LIMIT=NONE] -->
-    <string name="quick_settings_modes_label">Priority modes</string>
+    <!-- QuickSettings: Modes [CHAR LIMIT=NONE] -->
+    <string name="quick_settings_modes_label">Modes</string>
     <!-- QuickSettings: Bluetooth [CHAR LIMIT=NONE] -->
     <string name="quick_settings_bluetooth_label">Bluetooth</string>
     <!-- QuickSettings: Bluetooth (Multiple) [CHAR LIMIT=NONE] -->
@@ -740,6 +741,8 @@
     <string name="quick_settings_bluetooth_device_connected">Connected</string>
     <!-- QuickSettings: Bluetooth dialog device in audio sharing default summary [CHAR LIMIT=50]-->
     <string name="quick_settings_bluetooth_device_audio_sharing">Audio Sharing</string>
+    <!-- QuickSettings: Bluetooth dialog device summary for devices that are capable of audio sharing and switching to active[CHAR LIMIT=NONE]-->
+    <string name="quick_settings_bluetooth_device_audio_sharing_or_switch_active">Tap to switch or share audio</string>
     <!-- QuickSettings: Bluetooth dialog device saved default summary [CHAR LIMIT=NONE]-->
     <string name="quick_settings_bluetooth_device_saved">Saved</string>
     <!-- QuickSettings: Accessibility label to disconnect a device [CHAR LIMIT=NONE]-->
@@ -1094,28 +1097,28 @@
     <!-- QuickStep: Accessibility to toggle overview [CHAR LIMIT=40] -->
     <string name="quick_step_accessibility_toggle_overview">Toggle Overview</string>
 
-    <!-- Priority modes dialog title [CHAR LIMIT=35] -->
-    <string name="zen_modes_dialog_title">Priority modes</string>
+    <!-- Modes dialog title [CHAR LIMIT=35] -->
+    <string name="zen_modes_dialog_title">Modes</string>
 
-    <!-- Priority modes dialog confirmation button [CHAR LIMIT=15] -->
+    <!-- Modes dialog confirmation button [CHAR LIMIT=15] -->
     <string name="zen_modes_dialog_done">Done</string>
 
-    <!-- Priority modes dialog settings shortcut button [CHAR LIMIT=15] -->
+    <!-- Modes dialog settings shortcut button [CHAR LIMIT=15] -->
     <string name="zen_modes_dialog_settings">Settings</string>
 
-    <!-- Priority modes: label for an active mode [CHAR LIMIT=35] -->
+    <!-- Modes: label for an active mode [CHAR LIMIT=35] -->
     <string name="zen_mode_on">On</string>
 
-    <!-- Priority modes: label for an active mode, with details [CHAR LIMIT=10] -->
+    <!-- Modes: label for an active mode, with details [CHAR LIMIT=10] -->
     <string name="zen_mode_on_with_details">On • <xliff:g id="trigger_description" example="Mon-Fri, 23:00-7:00">%1$s</xliff:g></string>
 
-    <!-- Priority modes: label for an inactive mode [CHAR LIMIT=35] -->
+    <!-- Modes: label for an inactive mode [CHAR LIMIT=35] -->
     <string name="zen_mode_off">Off</string>
 
-    <!-- Priority modes: label for a mode that needs to be set up [CHAR LIMIT=35] -->
+    <!-- Modes: label for a mode that needs to be set up [CHAR LIMIT=35] -->
     <string name="zen_mode_set_up">Set up</string>
 
-    <!-- Priority modes: label for a mode that cannot be manually turned on [CHAR LIMIT=35] -->
+    <!-- Modes: label for a mode that cannot be manually turned on [CHAR LIMIT=35] -->
     <string name="zen_mode_no_manual_invocation">Manage in settings</string>
 
     <string name="zen_mode_active_modes">
@@ -1377,9 +1380,13 @@
     <string name="media_projection_entry_app_permission_dialog_title">Share your screen with <xliff:g id="app_seeking_permission" example="Meet">%s</xliff:g>?</string>
 
     <!-- 1P/3P app media projection permission option for capturing just a single app [CHAR LIMIT=50] -->
-    <string name="media_projection_entry_app_permission_dialog_option_text_single_app">Share one app</string>
+    <string name="screen_share_permission_dialog_option_single_app">Share one app</string>
+    <!-- CTS tests rely on the `screen_share_permission_dialog_option_single_app` resource name, so just point the updated resource name to the old resource name. -->
+    <string name="media_projection_entry_app_permission_dialog_option_text_single_app">@string/screen_share_permission_dialog_option_single_app</string>
     <!-- 1P/3P app media projection permission option for capturing the whole screen [CHAR LIMIT=50] -->
-    <string name="media_projection_entry_app_permission_dialog_option_text_entire_screen">Share entire screen</string>
+    <string name="screen_share_permission_dialog_option_entire_screen">Share entire screen</string>
+    <!-- CTS tests rely on the `screen_share_permission_dialog_option_entire_screen` resource name, so just point the updated resource name to the old resource name. -->
+    <string name="media_projection_entry_app_permission_dialog_option_text_entire_screen">@string/screen_share_permission_dialog_option_entire_screen</string>
     <!-- 1P/3P app media projection permission warning for capturing the whole screen. [CHAR LIMIT=350] -->
     <string name="media_projection_entry_app_permission_dialog_warning_entire_screen">When you’re sharing your entire screen, anything on your screen is visible to <xliff:g id="app_seeking_permission" example="Meet">%s</xliff:g>. So be careful with things like passwords, payment details, messages, photos, and audio and video.</string>
     <!-- 1P/3P app media projection permission warning for capturing an app. [CHAR LIMIT=350] -->
diff --git a/packages/SystemUI/shared/Android.bp b/packages/SystemUI/shared/Android.bp
index e68da09..8f55961 100644
--- a/packages/SystemUI/shared/Android.bp
+++ b/packages/SystemUI/shared/Android.bp
@@ -49,7 +49,6 @@
         "src/**/*.aidl",
         ":wm_shell-aidls",
         ":wm_shell-shared-aidls",
-        ":wm_shell_util-sources",
     ],
     static_libs: [
         "BiometricsSharedLib",
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/Task.java b/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/Task.java
index 64fe78d..7ec977a 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/Task.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/Task.java
@@ -218,6 +218,7 @@
     @ViewDebug.ExportedProperty(category="recents")
     public String title;
     @ViewDebug.ExportedProperty(category="recents")
+    @Nullable
     public String titleDescription;
     @ViewDebug.ExportedProperty(category="recents")
     public int colorPrimary;
diff --git a/packages/SystemUI/src/com/android/keyguard/AuthInteractionProperties.kt b/packages/SystemUI/src/com/android/keyguard/AuthInteractionProperties.kt
new file mode 100644
index 0000000..efa13c6
--- /dev/null
+++ b/packages/SystemUI/src/com/android/keyguard/AuthInteractionProperties.kt
@@ -0,0 +1,29 @@
+/*
+ * 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.keyguard
+
+import android.os.VibrationAttributes
+import com.google.android.msdl.domain.InteractionProperties
+
+/**
+ * This class represents the set of [InteractionProperties] that only hold [VibrationAttributes] for
+ * the case of user authentication.
+ */
+data class AuthInteractionProperties(
+    override val vibrationAttributes: VibrationAttributes =
+        VibrationAttributes.createForUsage(VibrationAttributes.USAGE_COMMUNICATION_REQUEST)
+) : InteractionProperties
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardAbsKeyInputViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardAbsKeyInputViewController.java
index dad4400..28f1381 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardAbsKeyInputViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardAbsKeyInputViewController.java
@@ -19,7 +19,10 @@
 import static com.android.internal.util.LatencyTracker.ACTION_CHECK_CREDENTIAL;
 import static com.android.internal.util.LatencyTracker.ACTION_CHECK_CREDENTIAL_UNLOCKED;
 import static com.android.keyguard.KeyguardAbsKeyInputView.MINIMUM_PASSWORD_LENGTH_BEFORE_REPORT;
+import static com.android.systemui.Flags.msdlFeedback;
+import static com.android.systemui.Flags.notifyPasswordTextViewUserActivityInBackground;
 
+import android.annotation.Nullable;
 import android.content.res.ColorStateList;
 import android.os.AsyncTask;
 import android.os.CountDownTimer;
@@ -40,6 +43,9 @@
 import com.android.systemui.res.R;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.data.model.MSDLToken;
+import com.google.android.msdl.domain.MSDLPlayer;
+
 import java.util.HashMap;
 import java.util.Map;
 
@@ -50,11 +56,14 @@
     private final LatencyTracker mLatencyTracker;
     private final FalsingCollector mFalsingCollector;
     private final EmergencyButtonController mEmergencyButtonController;
+    private final UserActivityNotifier mUserActivityNotifier;
     private CountDownTimer mCountdownTimer;
     private boolean mDismissing;
     protected AsyncTask<?, ?, ?> mPendingLockCheck;
     protected boolean mResumed;
     protected boolean mLockedOut;
+    @Nullable
+    protected MSDLPlayer mMSDLPlayer;
 
     private final KeyDownListener mKeyDownListener = (keyCode, keyEvent) -> {
         // Fingerprint sensor sends a KeyEvent.KEYCODE_UNKNOWN.
@@ -81,7 +90,9 @@
             KeyguardMessageAreaController.Factory messageAreaControllerFactory,
             LatencyTracker latencyTracker, FalsingCollector falsingCollector,
             EmergencyButtonController emergencyButtonController,
-            FeatureFlags featureFlags, SelectedUserInteractor selectedUserInteractor) {
+            FeatureFlags featureFlags, SelectedUserInteractor selectedUserInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, securityMode, keyguardSecurityCallback, emergencyButtonController,
                 messageAreaControllerFactory, featureFlags, selectedUserInteractor);
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
@@ -89,6 +100,8 @@
         mLatencyTracker = latencyTracker;
         mFalsingCollector = falsingCollector;
         mEmergencyButtonController = emergencyButtonController;
+        mMSDLPlayer = msdlPlayer;
+        mUserActivityNotifier = userActivityNotifier;
     }
 
     abstract void resetState();
@@ -178,6 +191,7 @@
     void onPasswordChecked(int userId, boolean matched, int timeoutMs, boolean isValidPassword) {
         boolean dismissKeyguard = mSelectedUserInteractor.getSelectedUserId() == userId;
         if (matched) {
+            playAuthenticationHaptics(/* unlock= */true);
             getKeyguardSecurityCallback().reportUnlockAttempt(userId, true, 0);
             if (dismissKeyguard) {
                 mDismissing = true;
@@ -185,6 +199,7 @@
                 getKeyguardSecurityCallback().dismiss(true, userId, getSecurityMode());
             }
         } else {
+            playAuthenticationHaptics(/* unlock= */false);
             mView.resetPasswordText(true /* animate */, false /* announce deletion if no match */);
             if (isValidPassword) {
                 getKeyguardSecurityCallback().reportUnlockAttempt(userId, false, timeoutMs);
@@ -201,6 +216,18 @@
         }
     }
 
+    private void playAuthenticationHaptics(boolean unlock) {
+        if (!msdlFeedback() || mMSDLPlayer == null) return;
+
+        MSDLToken token;
+        if (unlock) {
+            token = MSDLToken.UNLOCK;
+        } else {
+            token = MSDLToken.FAILURE;
+        }
+        mMSDLPlayer.playToken(token, mAuthInteractionProperties);
+    }
+
     protected void startErrorAnimation() { /* no-op */ }
 
     protected void verifyPasswordAndUnlock() {
@@ -280,6 +307,9 @@
         getKeyguardSecurityCallback().userActivity();
         getKeyguardSecurityCallback().onUserInput();
         mMessageAreaController.setMessage("");
+        if (notifyPasswordTextViewUserActivityInBackground()) {
+            mUserActivityNotifier.notifyUserActivity();
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
index db14a0f..dd84bc6 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardInputViewController.java
@@ -45,6 +45,9 @@
 import com.android.systemui.util.ViewController;
 import com.android.systemui.util.concurrency.DelayableExecutor;
 
+import com.google.android.msdl.domain.InteractionProperties;
+import com.google.android.msdl.domain.MSDLPlayer;
+
 import javax.inject.Inject;
 
 /** Controller for a {@link KeyguardSecurityView}. */
@@ -63,6 +66,8 @@
     private KeyguardSecurityCallback mNullCallback = new KeyguardSecurityCallback() {};
     private final FeatureFlags mFeatureFlags;
     protected final SelectedUserInteractor mSelectedUserInteractor;
+    protected final InteractionProperties mAuthInteractionProperties =
+            new AuthInteractionProperties();
 
     protected KeyguardInputViewController(T view, SecurityMode securityMode,
             KeyguardSecurityCallback keyguardSecurityCallback,
@@ -214,6 +219,8 @@
         private final SelectedUserInteractor mSelectedUserInteractor;
         private final UiEventLogger mUiEventLogger;
         private final KeyguardKeyboardInteractor mKeyguardKeyboardInteractor;
+        private final MSDLPlayer mMSDLPlayer;
+        private final UserActivityNotifier mUserActivityNotifier;
 
         @Inject
         public Factory(KeyguardUpdateMonitor keyguardUpdateMonitor,
@@ -228,7 +235,9 @@
                 KeyguardViewController keyguardViewController,
                 FeatureFlags featureFlags, SelectedUserInteractor selectedUserInteractor,
                 UiEventLogger uiEventLogger,
-                KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+                KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+                MSDLPlayer msdlPlayer,
+                UserActivityNotifier userActivityNotifier) {
             mKeyguardUpdateMonitor = keyguardUpdateMonitor;
             mLockPatternUtils = lockPatternUtils;
             mLatencyTracker = latencyTracker;
@@ -246,6 +255,8 @@
             mSelectedUserInteractor = selectedUserInteractor;
             mUiEventLogger = uiEventLogger;
             mKeyguardKeyboardInteractor = keyguardKeyboardInteractor;
+            mMSDLPlayer = msdlPlayer;
+            mUserActivityNotifier = userActivityNotifier;
         }
 
         /** Create a new {@link KeyguardInputViewController}. */
@@ -268,29 +279,29 @@
                         mInputMethodManager, emergencyButtonController, mMainExecutor, mResources,
                         mFalsingCollector, mKeyguardViewController,
                         mDevicePostureController, mFeatureFlags, mSelectedUserInteractor,
-                        mKeyguardKeyboardInteractor);
+                        mKeyguardKeyboardInteractor, mMSDLPlayer, mUserActivityNotifier);
             } else if (keyguardInputView instanceof KeyguardPINView) {
                 return new KeyguardPinViewController((KeyguardPINView) keyguardInputView,
                         mKeyguardUpdateMonitor, securityMode, mLockPatternUtils,
                         keyguardSecurityCallback, mMessageAreaControllerFactory, mLatencyTracker,
                         mLiftToActivateListener, emergencyButtonController, mFalsingCollector,
                         mDevicePostureController, mFeatureFlags, mSelectedUserInteractor,
-                        mUiEventLogger, mKeyguardKeyboardInteractor
-                );
+                        mUiEventLogger, mKeyguardKeyboardInteractor, mMSDLPlayer,
+                        mUserActivityNotifier);
             } else if (keyguardInputView instanceof KeyguardSimPinView) {
                 return new KeyguardSimPinViewController((KeyguardSimPinView) keyguardInputView,
                         mKeyguardUpdateMonitor, securityMode, mLockPatternUtils,
                         keyguardSecurityCallback, mMessageAreaControllerFactory, mLatencyTracker,
                         mLiftToActivateListener, mTelephonyManager, mFalsingCollector,
                         emergencyButtonController, mFeatureFlags, mSelectedUserInteractor,
-                        mKeyguardKeyboardInteractor);
+                        mKeyguardKeyboardInteractor, mMSDLPlayer, mUserActivityNotifier);
             } else if (keyguardInputView instanceof KeyguardSimPukView) {
                 return new KeyguardSimPukViewController((KeyguardSimPukView) keyguardInputView,
                         mKeyguardUpdateMonitor, securityMode, mLockPatternUtils,
                         keyguardSecurityCallback, mMessageAreaControllerFactory, mLatencyTracker,
                         mLiftToActivateListener, mTelephonyManager, mFalsingCollector,
                         emergencyButtonController, mFeatureFlags, mSelectedUserInteractor,
-                        mKeyguardKeyboardInteractor
+                        mKeyguardKeyboardInteractor, mMSDLPlayer, mUserActivityNotifier
                 );
             }
 
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardPasswordViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardPasswordViewController.java
index 490ad5c..905fa09 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardPasswordViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardPasswordViewController.java
@@ -19,9 +19,13 @@
 import static com.android.systemui.flags.Flags.LOCKSCREEN_ENABLE_LANDSCAPE;
 import static com.android.systemui.util.kotlin.JavaAdapterKt.collectFlow;
 
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
 import android.content.res.Resources;
 import android.graphics.drawable.Drawable;
 import android.os.UserHandle;
+import android.os.UserManager;
 import android.text.Editable;
 import android.text.InputType;
 import android.text.TextUtils;
@@ -52,6 +56,8 @@
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 import com.android.systemui.util.concurrency.DelayableExecutor;
 
+import com.google.android.msdl.domain.MSDLPlayer;
+
 import java.util.List;
 
 public class KeyguardPasswordViewController
@@ -131,10 +137,13 @@
             DevicePostureController postureController,
             FeatureFlags featureFlags,
             SelectedUserInteractor selectedUserInteractor,
-            KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+            KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, keyguardUpdateMonitor, securityMode, lockPatternUtils, keyguardSecurityCallback,
                 messageAreaControllerFactory, latencyTracker, falsingCollector,
-                emergencyButtonController, featureFlags, selectedUserInteractor);
+                emergencyButtonController, featureFlags, selectedUserInteractor, msdlPlayer,
+                userActivityNotifier);
         mKeyguardSecurityCallback = keyguardSecurityCallback;
         mInputMethodManager = inputMethodManager;
         mPostureController = postureController;
@@ -170,8 +179,33 @@
         mPasswordEntry.setOnEditorActionListener(mOnEditorActionListener);
         mPasswordEntry.setOnKeyListener(mKeyListener);
         mPasswordEntry.addTextChangedListener(mTextWatcher);
+
         // Poke the wakelock any time the text is selected or modified
-        mPasswordEntry.setOnClickListener(v -> mKeyguardSecurityCallback.userActivity());
+        // TODO(b/362362385): Revert to the previous onClickListener implementation once this bug is
+        //  fixed.
+        mPasswordEntry.setOnClickListener(new View.OnClickListener() {
+
+            private final boolean mAutomotiveAndVisibleBackgroundUsers =
+                    isAutomotiveAndVisibleBackgroundUsers();
+
+            @Override
+            public void onClick(View v) {
+                if (mAutomotiveAndVisibleBackgroundUsers) {
+                    mInputMethodManager.restartInput(v);
+                }
+                mKeyguardSecurityCallback.userActivity();
+            }
+
+            private boolean isAutomotiveAndVisibleBackgroundUsers() {
+                final Context context = getContext();
+                return context.getPackageManager().hasSystemFeature(
+                        PackageManager.FEATURE_AUTOMOTIVE)
+                        && UserManager.isVisibleBackgroundUsersEnabled()
+                        && context.getResources().getBoolean(
+                        android.R.bool.config_perDisplayFocusEnabled);
+            }
+        });
+
         mSwitchImeButton.setOnClickListener(v -> {
             mKeyguardSecurityCallback.userActivity(); // Leave the screen on a bit longer
             // Do not show auxiliary subtypes in password lock screen.
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardPinBasedInputViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardPinBasedInputViewController.java
index 0f61233..f575cf2 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardPinBasedInputViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardPinBasedInputViewController.java
@@ -16,9 +16,11 @@
 
 package com.android.keyguard;
 
+import static com.android.systemui.Flags.msdlFeedback;
 import static com.android.systemui.Flags.pinInputFieldStyledFocusState;
 import static com.android.systemui.util.kotlin.JavaAdapterKt.collectFlow;
 
+import android.annotation.Nullable;
 import android.graphics.drawable.Drawable;
 import android.graphics.drawable.GradientDrawable;
 import android.graphics.drawable.StateListDrawable;
@@ -40,6 +42,9 @@
 import com.android.systemui.res.R;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.data.model.MSDLToken;
+import com.google.android.msdl.domain.MSDLPlayer;
+
 public abstract class KeyguardPinBasedInputViewController<T extends KeyguardPinBasedInputView>
         extends KeyguardAbsKeyInputViewController<T> {
 
@@ -77,10 +82,13 @@
             FalsingCollector falsingCollector,
             FeatureFlags featureFlags,
             SelectedUserInteractor selectedUserInteractor,
-            KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+            KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, keyguardUpdateMonitor, securityMode, lockPatternUtils, keyguardSecurityCallback,
                 messageAreaControllerFactory, latencyTracker, falsingCollector,
-                emergencyButtonController, featureFlags, selectedUserInteractor);
+                emergencyButtonController, featureFlags, selectedUserInteractor, msdlPlayer,
+                userActivityNotifier);
         mLiftToActivateListener = liftToActivateListener;
         mFalsingCollector = falsingCollector;
         mKeyguardKeyboardInteractor = keyguardKeyboardInteractor;
@@ -102,12 +110,22 @@
                 return false;
             });
             button.setAnimationEnabled(showAnimations);
+            button.setMSDLPlayer(mMSDLPlayer);
         }
         mPasswordEntry.setOnKeyListener(mOnKeyListener);
         mPasswordEntry.setUserActivityListener(this::onUserInput);
 
         View deleteButton = mView.findViewById(R.id.delete_button);
-        deleteButton.setOnTouchListener(mActionButtonTouchListener);
+        if (msdlFeedback()) {
+            deleteButton.setOnTouchListener((View view, MotionEvent event) -> {
+                if (event.getActionMasked() == MotionEvent.ACTION_DOWN && mMSDLPlayer != null) {
+                    mMSDLPlayer.playToken(MSDLToken.KEYPRESS_DELETE, null);
+                }
+                return false;
+            });
+        } else {
+            deleteButton.setOnTouchListener(mActionButtonTouchListener);
+        }
         deleteButton.setOnClickListener(v -> {
             // check for time-based lockouts
             if (mPasswordEntry.isEnabled()) {
@@ -119,13 +137,19 @@
             if (mPasswordEntry.isEnabled()) {
                 mView.resetPasswordText(true /* animate */, true /* announce */);
             }
-            mView.doHapticKeyClick();
+            if (msdlFeedback() && mMSDLPlayer != null) {
+                mMSDLPlayer.playToken(MSDLToken.LONG_PRESS, null);
+            } else {
+                mView.doHapticKeyClick();
+            }
             return true;
         });
 
         View okButton = mView.findViewById(R.id.key_enter);
         if (okButton != null) {
-            okButton.setOnTouchListener(mActionButtonTouchListener);
+            if (!msdlFeedback()) {
+                okButton.setOnTouchListener(mActionButtonTouchListener);
+            }
             okButton.setOnClickListener(v -> {
                 if (mPasswordEntry.isEnabled()) {
                     verifyPasswordAndUnlock();
@@ -177,6 +201,7 @@
 
         for (NumPadKey button : mView.getButtons()) {
             button.setOnTouchListener(null);
+            button.setMSDLPlayer(null);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardPinViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardPinViewController.java
index f4cda02..3b5bf1a 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardPinViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardPinViewController.java
@@ -18,6 +18,7 @@
 
 import static com.android.systemui.flags.Flags.LOCKSCREEN_ENABLE_LANDSCAPE;
 
+import android.annotation.Nullable;
 import android.view.View;
 
 import com.android.internal.logging.UiEvent;
@@ -33,6 +34,8 @@
 import com.android.systemui.statusbar.policy.DevicePostureController;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.domain.MSDLPlayer;
+
 public class KeyguardPinViewController
         extends KeyguardPinBasedInputViewController<KeyguardPINView> {
     private final KeyguardUpdateMonitor mKeyguardUpdateMonitor;
@@ -61,11 +64,13 @@
             FalsingCollector falsingCollector,
             DevicePostureController postureController, FeatureFlags featureFlags,
             SelectedUserInteractor selectedUserInteractor, UiEventLogger uiEventLogger,
-            KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+            KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, keyguardUpdateMonitor, securityMode, lockPatternUtils, keyguardSecurityCallback,
                 messageAreaControllerFactory, latencyTracker, liftToActivateListener,
                 emergencyButtonController, falsingCollector, featureFlags, selectedUserInteractor,
-                keyguardKeyboardInteractor);
+                keyguardKeyboardInteractor, msdlPlayer, userActivityNotifier);
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mPostureController = postureController;
         mLockPatternUtils = lockPatternUtils;
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
index afd42cb..61f9800 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSecurityContainerController.java
@@ -81,7 +81,7 @@
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.keyguard.KeyguardWmStateRefactor;
-import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
+import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor;
 import com.android.systemui.log.SessionTracker;
 import com.android.systemui.plugins.ActivityStarter;
 import com.android.systemui.plugins.FalsingManager;
@@ -134,7 +134,7 @@
     private final DeviceEntryFaceAuthInteractor mDeviceEntryFaceAuthInteractor;
     private final BouncerMessageInteractor mBouncerMessageInteractor;
     private int mTranslationY;
-    private final KeyguardTransitionInteractor mKeyguardTransitionInteractor;
+    private final KeyguardDismissTransitionInteractor mKeyguardDismissTransitionInteractor;
     private final DevicePolicyManager mDevicePolicyManager;
     // Whether the volume keys should be handled by keyguard. If true, then
     // they will be handled here for specific media types such as music, otherwise
@@ -321,7 +321,7 @@
             }
 
             if (KeyguardWmStateRefactor.isEnabled()) {
-                mKeyguardTransitionInteractor.startDismissKeyguardTransition(
+                mKeyguardDismissTransitionInteractor.startDismissKeyguardTransition(
                         "KeyguardSecurityContainerController#finish");
             }
         }
@@ -458,7 +458,7 @@
             DeviceProvisionedController deviceProvisionedController,
             FaceAuthAccessibilityDelegate faceAuthAccessibilityDelegate,
             DevicePolicyManager devicePolicyManager,
-            KeyguardTransitionInteractor keyguardTransitionInteractor,
+            KeyguardDismissTransitionInteractor keyguardDismissTransitionInteractor,
             Lazy<PrimaryBouncerInteractor> primaryBouncerInteractor,
             Provider<DeviceEntryInteractor> deviceEntryInteractor
     ) {
@@ -490,7 +490,7 @@
         mSelectedUserInteractor = selectedUserInteractor;
         mDeviceEntryInteractor = deviceEntryInteractor;
         mJavaAdapter = javaAdapter;
-        mKeyguardTransitionInteractor = keyguardTransitionInteractor;
+        mKeyguardDismissTransitionInteractor = keyguardDismissTransitionInteractor;
         mDeviceProvisionedController = deviceProvisionedController;
         mPrimaryBouncerInteractor = primaryBouncerInteractor;
         mDevicePolicyManager = devicePolicyManager;
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
index 3ef3418..47fe2b2 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPinViewController.java
@@ -21,6 +21,7 @@
 import static com.android.systemui.util.PluralMessageFormaterKt.icuMessageFormat;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.AlertDialog;
 import android.app.AlertDialog.Builder;
 import android.app.Dialog;
@@ -48,6 +49,8 @@
 import com.android.systemui.res.R;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.domain.MSDLPlayer;
+
 public class KeyguardSimPinViewController
         extends KeyguardPinBasedInputViewController<KeyguardSimPinView> {
     public static final String TAG = "KeyguardSimPinView";
@@ -95,11 +98,13 @@
             TelephonyManager telephonyManager, FalsingCollector falsingCollector,
             EmergencyButtonController emergencyButtonController, FeatureFlags featureFlags,
             SelectedUserInteractor selectedUserInteractor,
-            KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+            KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, keyguardUpdateMonitor, securityMode, lockPatternUtils, keyguardSecurityCallback,
                 messageAreaControllerFactory, latencyTracker, liftToActivateListener,
                 emergencyButtonController, falsingCollector, featureFlags, selectedUserInteractor,
-                keyguardKeyboardInteractor);
+                keyguardKeyboardInteractor, msdlPlayer, userActivityNotifier);
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mTelephonyManager = telephonyManager;
         mSimImageView = mView.findViewById(R.id.keyguard_sim);
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
index 46225c7..c688acb 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardSimPukViewController.java
@@ -17,6 +17,7 @@
 package com.android.keyguard;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.Activity;
 import android.app.AlertDialog;
 import android.app.Dialog;
@@ -43,6 +44,8 @@
 import com.android.systemui.res.R;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.domain.MSDLPlayer;
+
 public class KeyguardSimPukViewController
         extends KeyguardPinBasedInputViewController<KeyguardSimPukView> {
     private static final boolean DEBUG = KeyguardConstants.DEBUG;
@@ -92,11 +95,13 @@
             TelephonyManager telephonyManager, FalsingCollector falsingCollector,
             EmergencyButtonController emergencyButtonController, FeatureFlags featureFlags,
             SelectedUserInteractor selectedUserInteractor,
-            KeyguardKeyboardInteractor keyguardKeyboardInteractor) {
+            KeyguardKeyboardInteractor keyguardKeyboardInteractor,
+            @Nullable MSDLPlayer msdlPlayer,
+            UserActivityNotifier userActivityNotifier) {
         super(view, keyguardUpdateMonitor, securityMode, lockPatternUtils, keyguardSecurityCallback,
                 messageAreaControllerFactory, latencyTracker, liftToActivateListener,
                 emergencyButtonController, falsingCollector, featureFlags, selectedUserInteractor,
-                keyguardKeyboardInteractor);
+                keyguardKeyboardInteractor, msdlPlayer, userActivityNotifier);
         mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mTelephonyManager = telephonyManager;
         mSimImageView = mView.findViewById(R.id.keyguard_sim);
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index 9b45fa4..f731186 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -3811,7 +3811,8 @@
         if (!mSimDatas.containsKey(subId)) {
             refreshSimState(subId, SubscriptionManager.getSlotIndex(subId));
         }
-        return mSimDatas.get(subId).slotId;
+        SimData simData = mSimDatas.get(subId);
+        return simData != null ? simData.slotId : SubscriptionManager.INVALID_SUBSCRIPTION_ID;
     }
 
     private final TaskStackChangeListener mTaskStackListener = new TaskStackChangeListener() {
diff --git a/packages/SystemUI/src/com/android/keyguard/NumPadKey.java b/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
index dcfa775..4fb80de 100644
--- a/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
+++ b/packages/SystemUI/src/com/android/keyguard/NumPadKey.java
@@ -15,6 +15,7 @@
  */
 package com.android.keyguard;
 
+import static com.android.systemui.Flags.msdlFeedback;
 import static com.android.systemui.bouncer.shared.constants.KeyguardBouncerConstants.ColorId.NUM_PAD_KEY;
 
 import android.content.Context;
@@ -38,6 +39,9 @@
 import com.android.settingslib.Utils;
 import com.android.systemui.res.R;
 
+import com.google.android.msdl.data.model.MSDLToken;
+import com.google.android.msdl.domain.MSDLPlayer;
+
 /**
  * Viewgroup for the bouncer numpad button, specifically for digits.
  */
@@ -57,6 +61,8 @@
     @Nullable
     private NumPadAnimator mAnimator;
     private int mOrientation;
+    @Nullable
+    private MSDLPlayer mMSDLPlayer;
 
     private View.OnClickListener mListener = new View.OnClickListener() {
         @Override
@@ -221,8 +227,12 @@
 
     // Cause a VIRTUAL_KEY vibration
     public void doHapticKeyClick() {
-        performHapticFeedback(HapticFeedbackConstants.VIRTUAL_KEY,
-                HapticFeedbackConstants.FLAG_IGNORE_VIEW_SETTING);
+        if (msdlFeedback() && mMSDLPlayer != null) {
+            mMSDLPlayer.playToken(MSDLToken.KEYPRESS_STANDARD, null);
+        } else {
+            performHapticFeedback(HapticFeedbackConstants.VIRTUAL_KEY,
+                    HapticFeedbackConstants.FLAG_IGNORE_VIEW_SETTING);
+        }
     }
 
     @Override
@@ -244,4 +254,8 @@
         super.onInitializeAccessibilityNodeInfo(info);
         info.setTextEntryKey(true);
     }
+
+    public void setMSDLPlayer(@Nullable MSDLPlayer player) {
+        mMSDLPlayer = player;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/keyguard/PasswordTextView.java b/packages/SystemUI/src/com/android/keyguard/PasswordTextView.java
index 85f8b48..0c4bc0e 100644
--- a/packages/SystemUI/src/com/android/keyguard/PasswordTextView.java
+++ b/packages/SystemUI/src/com/android/keyguard/PasswordTextView.java
@@ -16,6 +16,8 @@
 
 package com.android.keyguard;
 
+import static com.android.systemui.Flags.notifyPasswordTextViewUserActivityInBackground;
+
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
@@ -257,7 +259,9 @@
 
     @Override
     protected void onUserActivity() {
-        mPM.userActivity(SystemClock.uptimeMillis(), false);
+        if (!notifyPasswordTextViewUserActivityInBackground()) {
+            mPM.userActivity(SystemClock.uptimeMillis(), false);
+        }
         super.onUserActivity();
     }
 
diff --git a/packages/SystemUI/src/com/android/keyguard/UserActivityNotifier.kt b/packages/SystemUI/src/com/android/keyguard/UserActivityNotifier.kt
new file mode 100644
index 0000000..9b1ddb7
--- /dev/null
+++ b/packages/SystemUI/src/com/android/keyguard/UserActivityNotifier.kt
@@ -0,0 +1,41 @@
+/*
+ * 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.keyguard
+
+import android.os.PowerManager
+import android.os.SystemClock
+import com.android.systemui.dagger.qualifiers.UiBackground
+import java.util.concurrent.Executor
+import javax.inject.Inject
+
+/** Wrapper class for notifying the system about user activity in the background. */
+class UserActivityNotifier
+@Inject
+constructor(
+    @UiBackground private val uiBgExecutor: Executor,
+    private val powerManager: PowerManager
+) {
+
+    fun notifyUserActivity() {
+        uiBgExecutor.execute {
+            powerManager.userActivity(
+                SystemClock.uptimeMillis(),
+                PowerManager.USER_ACTIVITY_EVENT_OTHER,
+                0
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java b/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
index 394f8dd..04afd86 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/FullscreenMagnificationController.java
@@ -408,6 +408,10 @@
         if (!isActivated()) {
             return;
         }
+        if (!(mFullscreenBorder.getBackground() instanceof GradientDrawable)) {
+            // Wear doesn't use the same magnification border background. So early return here.
+            return;
+        }
 
         float cornerRadius = ScreenDecorationsUtils.getWindowCornerRadius(mContext);
         GradientDrawable backgroundDrawable = (GradientDrawable) mFullscreenBorder.getBackground();
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/WindowMagnificationSettings.java b/packages/SystemUI/src/com/android/systemui/accessibility/WindowMagnificationSettings.java
index 9b6501e..2f0ca6e 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/WindowMagnificationSettings.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/WindowMagnificationSettings.java
@@ -482,6 +482,10 @@
                 } else { // mode = ACCESSIBILITY_MAGNIFICATION_MODE_WINDOW
                     mEditButton.setVisibility(View.VISIBLE);
                     mAllowDiagonalScrollingView.setVisibility(View.VISIBLE);
+                    if (Flags.saveAndRestoreMagnificationSettingsButtons()) {
+                        selectedButtonIndex =
+                                windowMagnificationFrameSizePrefs.getIndexForCurrentDensity();
+                    }
                 }
                 break;
 
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuController.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuController.java
index e4b7b7e..275147e 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuController.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuController.java
@@ -21,11 +21,13 @@
 
 import android.content.Context;
 import android.hardware.display.DisplayManager;
+import android.os.Handler;
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.view.Display;
 import android.view.WindowManager;
 import android.view.accessibility.AccessibilityManager;
+import android.view.accessibility.IUserInitializationCompleteCallback;
 
 import androidx.annotation.MainThread;
 
@@ -68,6 +70,9 @@
     private int mBtnMode;
     private String mBtnTargets;
     private boolean mIsKeyguardVisible;
+    private boolean mIsUserInInitialization;
+    @VisibleForTesting
+    Handler mHandler;
 
     @VisibleForTesting
     final KeyguardUpdateMonitorCallback mKeyguardCallback = new KeyguardUpdateMonitorCallback() {
@@ -86,18 +91,14 @@
         @Override
         public void onUserSwitching(int userId) {
             destroyFloatingMenu();
-        }
-
-        @Override
-        public void onUserSwitchComplete(int userId) {
-            mContext = mContext.createContextAsUser(UserHandle.of(userId), /* flags= */ 0);
-            mBtnMode = mAccessibilityButtonModeObserver.getCurrentAccessibilityButtonMode();
-            mBtnTargets =
-                    mAccessibilityButtonTargetsObserver.getCurrentAccessibilityButtonTargets();
-            handleFloatingMenuVisibility(mIsKeyguardVisible, mBtnMode, mBtnTargets);
+            mIsUserInInitialization = true;
         }
     };
 
+    @VisibleForTesting
+    final UserInitializationCompleteCallback mUserInitializationCompleteCallback =
+            new UserInitializationCompleteCallback();
+
     @Inject
     public AccessibilityFloatingMenuController(Context context,
             WindowManager windowManager,
@@ -109,7 +110,8 @@
             KeyguardUpdateMonitor keyguardUpdateMonitor,
             SecureSettings secureSettings,
             DisplayTracker displayTracker,
-            NavigationModeController navigationModeController) {
+            NavigationModeController navigationModeController,
+            Handler handler) {
         mContext = context;
         mWindowManager = windowManager;
         mViewCaptureAwareWindowManager = viewCaptureAwareWindowManager;
@@ -121,6 +123,7 @@
         mSecureSettings = secureSettings;
         mDisplayTracker = displayTracker;
         mNavigationModeController = navigationModeController;
+        mHandler = handler;
 
         mIsKeyguardVisible = false;
     }
@@ -159,6 +162,8 @@
         mAccessibilityButtonModeObserver.addListener(this);
         mAccessibilityButtonTargetsObserver.addListener(this);
         mKeyguardUpdateMonitor.registerCallback(mKeyguardCallback);
+        mAccessibilityManager.registerUserInitializationCompleteCallback(
+                mUserInitializationCompleteCallback);
     }
 
     /**
@@ -172,7 +177,7 @@
      */
     private void handleFloatingMenuVisibility(boolean keyguardVisible,
             @AccessibilityButtonMode int mode, String targets) {
-        if (keyguardVisible) {
+        if (keyguardVisible || mIsUserInInitialization) {
             destroyFloatingMenu();
             return;
         }
@@ -210,4 +215,18 @@
         mFloatingMenu.hide();
         mFloatingMenu = null;
     }
+
+    class UserInitializationCompleteCallback
+            extends IUserInitializationCompleteCallback.Stub {
+        @Override
+        public void onUserInitializationComplete(int userId) {
+            mIsUserInInitialization = false;
+            mContext = mContext.createContextAsUser(UserHandle.of(userId), /* flags= */ 0);
+            mBtnMode = mAccessibilityButtonModeObserver.getCurrentAccessibilityButtonMode();
+            mBtnTargets =
+                    mAccessibilityButtonTargetsObserver.getCurrentAccessibilityButtonTargets();
+            mHandler.post(
+                    () -> handleFloatingMenuVisibility(mIsKeyguardVisible, mBtnMode, mBtnTargets));
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationController.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationController.java
index d718ae3..708f7f1 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationController.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractAnimationController.java
@@ -30,8 +30,8 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.systemui.Flags;
-import com.android.wm.shell.common.bubbles.DismissCircleView;
-import com.android.wm.shell.common.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.DismissCircleView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
 
 import java.util.Map;
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractView.kt b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractView.kt
index c1b3962..13c1a45 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractView.kt
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/DragToInteractView.kt
@@ -39,9 +39,9 @@
 import androidx.dynamicanimation.animation.SpringForce.DAMPING_RATIO_LOW_BOUNCY
 import androidx.dynamicanimation.animation.SpringForce.STIFFNESS_LOW
 import com.android.wm.shell.R
-import com.android.wm.shell.common.bubbles.DismissCircleView
-import com.android.wm.shell.common.bubbles.DismissView
 import com.android.wm.shell.shared.animation.PhysicsAnimator
+import com.android.wm.shell.shared.bubbles.DismissCircleView
+import com.android.wm.shell.shared.bubbles.DismissView
 
 /**
  * View that handles interactions between DismissCircleView and BubbleStackView.
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayer.java b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayer.java
index d62162b..7a674e2 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayer.java
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/floatingmenu/MenuViewLayer.java
@@ -81,7 +81,7 @@
 import com.android.systemui.res.R;
 import com.android.systemui.util.settings.SecureSettings;
 import com.android.wm.shell.bubbles.DismissViewUtils;
-import com.android.wm.shell.common.bubbles.DismissView;
+import com.android.wm.shell.shared.bubbles.DismissView;
 import com.android.wm.shell.shared.magnetictarget.MagnetizedObject;
 
 import java.lang.annotation.Retention;
diff --git a/packages/SystemUI/src/com/android/systemui/accessibility/qs/QSAccessibilityModule.kt b/packages/SystemUI/src/com/android/systemui/accessibility/qs/QSAccessibilityModule.kt
index 03f282e..bb80396 100644
--- a/packages/SystemUI/src/com/android/systemui/accessibility/qs/QSAccessibilityModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/accessibility/qs/QSAccessibilityModule.kt
@@ -120,42 +120,42 @@
     @IntoMap
     @StringKey(COLOR_CORRECTION_TILE_SPEC)
     fun provideColorCorrectionAvailabilityInteractor(
-            impl: ColorCorrectionTileDataInteractor
+        impl: ColorCorrectionTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     @Binds
     @IntoMap
     @StringKey(COLOR_INVERSION_TILE_SPEC)
     fun provideColorInversionAvailabilityInteractor(
-            impl: ColorCorrectionTileDataInteractor
+        impl: ColorCorrectionTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     @Binds
     @IntoMap
     @StringKey(FONT_SCALING_TILE_SPEC)
     fun provideFontScalingAvailabilityInteractor(
-            impl: FontScalingTileDataInteractor
+        impl: FontScalingTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     @Binds
     @IntoMap
     @StringKey(REDUCE_BRIGHTNESS_TILE_SPEC)
     fun provideReduceBrightnessAvailabilityInteractor(
-            impl: ReduceBrightColorsTileDataInteractor
+        impl: ReduceBrightColorsTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     @Binds
     @IntoMap
     @StringKey(ONE_HANDED_TILE_SPEC)
     fun provideOneHandedAvailabilityInteractor(
-            impl: OneHandedModeTileDataInteractor
+        impl: OneHandedModeTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     @Binds
     @IntoMap
     @StringKey(NIGHT_DISPLAY_TILE_SPEC)
     fun provideNightDisplayAvailabilityInteractor(
-            impl: NightDisplayTileDataInteractor
+        impl: NightDisplayTileDataInteractor
     ): QSTileAvailabilityInteractor
 
     companion object {
@@ -165,6 +165,7 @@
         const val REDUCE_BRIGHTNESS_TILE_SPEC = "reduce_brightness"
         const val ONE_HANDED_TILE_SPEC = "onehanded"
         const val NIGHT_DISPLAY_TILE_SPEC = "night"
+        const val HEARING_DEVICES_TILE_SPEC = "hearing_devices"
 
         @Provides
         @IntoMap
@@ -273,6 +274,20 @@
                 instanceId = uiEventLogger.getNewInstanceId(),
             )
 
+        @Provides
+        @IntoMap
+        @StringKey(HEARING_DEVICES_TILE_SPEC)
+        fun provideHearingDevicesTileConfig(uiEventLogger: QsEventLogger): QSTileConfig =
+            QSTileConfig(
+                tileSpec = TileSpec.create(HEARING_DEVICES_TILE_SPEC),
+                uiConfig =
+                    QSTileUIConfig.Resource(
+                        iconRes = R.drawable.qs_hearing_devices_icon,
+                        labelRes = R.string.quick_settings_hearing_devices_label,
+                    ),
+                instanceId = uiEventLogger.getNewInstanceId(),
+            )
+
         /**
          * Inject Reduce Bright Colors Tile into tileViewModelMap in QSModule. The tile is hidden
          * behind a flag.
diff --git a/packages/SystemUI/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandler.kt b/packages/SystemUI/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandler.kt
index a093f58..fd06bb1 100644
--- a/packages/SystemUI/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandler.kt
+++ b/packages/SystemUI/src/com/android/systemui/ambient/touch/BouncerSwipeTouchHandler.kt
@@ -29,7 +29,6 @@
 import androidx.annotation.VisibleForTesting
 import com.android.internal.logging.UiEvent
 import com.android.internal.logging.UiEventLogger
-import com.android.internal.widget.LockPatternUtils
 import com.android.systemui.Flags
 import com.android.systemui.ambient.touch.TouchHandler.TouchSession
 import com.android.systemui.ambient.touch.dagger.BouncerSwipeModule
@@ -37,8 +36,8 @@
 import com.android.systemui.ambient.touch.scrim.ScrimManager
 import com.android.systemui.bouncer.shared.constants.KeyguardBouncerConstants
 import com.android.systemui.communal.ui.viewmodel.CommunalViewModel
+import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.plugins.ActivityStarter
-import com.android.systemui.settings.UserTracker
 import com.android.systemui.shade.ShadeExpansionChangeEvent
 import com.android.systemui.statusbar.NotificationShadeWindowController
 import com.android.systemui.statusbar.phone.CentralSurfaces
@@ -63,8 +62,6 @@
     private val notificationShadeWindowController: NotificationShadeWindowController,
     private val valueAnimatorCreator: ValueAnimatorCreator,
     private val velocityTrackerFactory: VelocityTrackerFactory,
-    private val lockPatternUtils: LockPatternUtils,
-    private val userTracker: UserTracker,
     private val communalViewModel: CommunalViewModel,
     @param:Named(BouncerSwipeModule.SWIPE_TO_BOUNCER_FLING_ANIMATION_UTILS_OPENING)
     private val flingAnimationUtils: FlingAnimationUtils,
@@ -75,7 +72,8 @@
     @param:Named(BouncerSwipeModule.MIN_BOUNCER_ZONE_SCREEN_PERCENTAGE)
     private val minBouncerZoneScreenPercentage: Float,
     private val uiEventLogger: UiEventLogger,
-    private val activityStarter: ActivityStarter
+    private val activityStarter: ActivityStarter,
+    private val keyguardInteractor: KeyguardInteractor,
 ) : TouchHandler {
     /** An interface for creating ValueAnimators. */
     interface ValueAnimatorCreator {
@@ -148,7 +146,7 @@
 
                     // If scrolling up and keyguard is not locked, dismiss both keyguard and the
                     // dream since there's no bouncer to show.
-                    if (y > e2.y && !lockPatternUtils.isSecure(userTracker.userId)) {
+                    if (y > e2.y && keyguardInteractor.isKeyguardDismissible.value) {
                         activityStarter.executeRunnableDismissingKeyguard(
                             { centralSurfaces.get().awakenDreams() },
                             /* cancelAction= */ null,
@@ -331,8 +329,8 @@
             return
         }
 
-        // Don't set expansion if the user doesn't have a pin/password set.
-        if (!lockPatternUtils.isSecure(userTracker.userId)) {
+        // Don't set expansion if keyguard is dismissible (i.e. unlocked).
+        if (keyguardInteractor.isKeyguardDismissible.value) {
             return
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
index 468737d..732a90d 100644
--- a/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/authentication/data/repository/AuthenticationRepository.kt
@@ -32,11 +32,11 @@
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Pin
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel.Sim
 import com.android.systemui.authentication.shared.model.AuthenticationResultModel
+import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlags
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
 import com.android.systemui.user.data.repository.UserRepository
 import com.android.systemui.util.kotlin.onSubscriberAdded
@@ -254,7 +254,7 @@
     override val hasLockoutOccurred: StateFlow<Boolean> = _hasLockoutOccurred.asStateFlow()
 
     init {
-        if (SceneContainerFlag.isEnabled) {
+        if (ComposeBouncerFlags.isComposeBouncerOrSceneContainerEnabled()) {
             // Hydrate failedAuthenticationAttempts initially and whenever the selected user
             // changes.
             applicationScope.launch {
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDebouncer.kt b/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDebouncer.kt
index 1685f49..4731ebb 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDebouncer.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDebouncer.kt
@@ -25,11 +25,13 @@
  * - startWindow: Window of time on start required before showing the first help message
  * - shownFaceMessageFrequencyBoost: Frequency boost given to messages that are currently shown to
  *   the user
+ * - threshold: minimum percentage of frames a message must appear in order to show it
  */
 class FaceHelpMessageDebouncer(
     private val window: Long = DEFAULT_WINDOW_MS,
     private val startWindow: Long = window,
     private val shownFaceMessageFrequencyBoost: Int = 4,
+    private val threshold: Float = 0f,
 ) {
     private val TAG = "FaceHelpMessageDebouncer"
     private var startTime = 0L
@@ -56,7 +58,7 @@
         }
     }
 
-    private fun getMostFrequentHelpMessage(): HelpFaceAuthenticationStatus? {
+    private fun getMostFrequentHelpMessageSurpassingThreshold(): HelpFaceAuthenticationStatus? {
         // freqMap: msgId => frequency
         val freqMap = helpFaceAuthStatuses.groupingBy { it.msgId }.eachCount().toMutableMap()
 
@@ -83,7 +85,25 @@
                     }
                 }
                 ?.key
-        return helpFaceAuthStatuses.findLast { it.msgId == msgIdWithHighestFrequency }
+
+        if (msgIdWithHighestFrequency == null) {
+            return null
+        }
+
+        val freq =
+            if (msgIdWithHighestFrequency == lastMessageIdShown) {
+                    freqMap[msgIdWithHighestFrequency]!! - shownFaceMessageFrequencyBoost
+                } else {
+                    freqMap[msgIdWithHighestFrequency]!!
+                }
+                .toFloat()
+
+        return if ((freq / helpFaceAuthStatuses.size.toFloat()) >= threshold) {
+            helpFaceAuthStatuses.findLast { it.msgId == msgIdWithHighestFrequency }
+        } else {
+            Log.v(TAG, "most frequent helpFaceAuthStatus didn't make the threshold: $threshold")
+            null
+        }
     }
 
     fun addMessage(helpFaceAuthStatus: HelpFaceAuthenticationStatus) {
@@ -98,14 +118,15 @@
             return null
         }
         removeOldMessages(atTimestamp)
-        val messageToShow = getMostFrequentHelpMessage()
+        val messageToShow = getMostFrequentHelpMessageSurpassingThreshold()
         if (lastMessageIdShown != messageToShow?.msgId) {
             Log.v(
                 TAG,
                 "showMessage previousLastMessageId=$lastMessageIdShown" +
                     "\n\tmessageToShow=$messageToShow " +
                     "\n\thelpFaceAuthStatusesSize=${helpFaceAuthStatuses.size}" +
-                    "\n\thelpFaceAuthStatuses=$helpFaceAuthStatuses"
+                    "\n\thelpFaceAuthStatuses=$helpFaceAuthStatuses" +
+                    "\n\tthreshold=$threshold"
             )
             lastMessageIdShown = messageToShow?.msgId
         }
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDeferral.kt b/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDeferral.kt
index 90d06fb..d382ada 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDeferral.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/FaceHelpMessageDeferral.kt
@@ -17,14 +17,19 @@
 package com.android.systemui.biometrics
 
 import android.content.res.Resources
+import android.os.SystemClock.elapsedRealtime
 import com.android.keyguard.logging.BiometricMessageDeferralLogger
 import com.android.systemui.Dumpable
+import com.android.systemui.Flags.faceMessageDeferUpdate
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.deviceentry.shared.model.HelpFaceAuthenticationStatus
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.dagger.BiometricLog
 import com.android.systemui.res.R
+import com.android.systemui.util.time.SystemClock
+import dagger.Lazy
 import java.io.PrintWriter
 import java.util.Objects
 import java.util.UUID
@@ -36,7 +41,8 @@
 constructor(
     @Main private val resources: Resources,
     @BiometricLog private val logBuffer: LogBuffer,
-    private val dumpManager: DumpManager
+    private val dumpManager: DumpManager,
+    private val systemClock: Lazy<SystemClock>,
 ) {
     fun create(): FaceHelpMessageDeferral {
         val id = UUID.randomUUID().toString()
@@ -45,6 +51,7 @@
             logBuffer = BiometricMessageDeferralLogger(logBuffer, "FaceHelpMessageDeferral[$id]"),
             dumpManager = dumpManager,
             id = id,
+            systemClock,
         )
     }
 }
@@ -58,14 +65,17 @@
     logBuffer: BiometricMessageDeferralLogger,
     dumpManager: DumpManager,
     val id: String,
+    val systemClock: Lazy<SystemClock>,
 ) :
     BiometricMessageDeferral(
         resources.getIntArray(R.array.config_face_help_msgs_defer_until_timeout).toHashSet(),
         resources.getIntArray(R.array.config_face_help_msgs_ignore).toHashSet(),
         resources.getFloat(R.dimen.config_face_help_msgs_defer_until_timeout_threshold),
+        resources.getInteger(R.integer.config_face_help_msgs_defer_analyze_timeframe).toLong(),
         logBuffer,
         dumpManager,
         id,
+        systemClock,
     )
 
 /**
@@ -77,10 +87,24 @@
     private val messagesToDefer: Set<Int>,
     private val acquiredInfoToIgnore: Set<Int>,
     private val threshold: Float,
+    private val windowToAnalyzeLastNFrames: Long,
     private val logBuffer: BiometricMessageDeferralLogger,
     dumpManager: DumpManager,
     id: String,
+    private val systemClock: Lazy<SystemClock>,
 ) : Dumpable {
+
+    private val faceHelpMessageDebouncer: FaceHelpMessageDebouncer? =
+        if (faceMessageDeferUpdate()) {
+            FaceHelpMessageDebouncer(
+                window = windowToAnalyzeLastNFrames,
+                startWindow = 0L,
+                shownFaceMessageFrequencyBoost = 0,
+                threshold = threshold,
+            )
+        } else {
+            null
+        }
     private val acquiredInfoToFrequency: MutableMap<Int, Int> = HashMap()
     private val acquiredInfoToHelpString: MutableMap<Int, String> = HashMap()
     private var mostFrequentAcquiredInfoToDefer: Int? = null
@@ -97,13 +121,20 @@
         pw.println("messagesToDefer=$messagesToDefer")
         pw.println("totalFrames=$totalFrames")
         pw.println("threshold=$threshold")
+        pw.println("faceMessageDeferUpdateFlagEnabled=${faceMessageDeferUpdate()}")
+        if (faceMessageDeferUpdate()) {
+            pw.println("windowToAnalyzeLastNFrames(ms)=$windowToAnalyzeLastNFrames")
+        }
     }
 
     /** Reset all saved counts. */
     fun reset() {
         totalFrames = 0
-        mostFrequentAcquiredInfoToDefer = null
-        acquiredInfoToFrequency.clear()
+        if (!faceMessageDeferUpdate()) {
+            mostFrequentAcquiredInfoToDefer = null
+            acquiredInfoToFrequency.clear()
+        }
+
         acquiredInfoToHelpString.clear()
         logBuffer.reset()
     }
@@ -137,24 +168,48 @@
             logBuffer.logFrameIgnored(acquiredInfo)
             return
         }
-
         totalFrames++
 
-        val newAcquiredInfoCount = acquiredInfoToFrequency.getOrDefault(acquiredInfo, 0) + 1
-        acquiredInfoToFrequency[acquiredInfo] = newAcquiredInfoCount
-        if (
-            messagesToDefer.contains(acquiredInfo) &&
-                (mostFrequentAcquiredInfoToDefer == null ||
-                    newAcquiredInfoCount >
-                        acquiredInfoToFrequency.getOrDefault(mostFrequentAcquiredInfoToDefer!!, 0))
-        ) {
-            mostFrequentAcquiredInfoToDefer = acquiredInfo
+        if (faceMessageDeferUpdate()) {
+            faceHelpMessageDebouncer?.let {
+                val helpFaceAuthStatus =
+                    HelpFaceAuthenticationStatus(
+                        msgId = acquiredInfo,
+                        msg = null,
+                        systemClock.get().elapsedRealtime()
+                    )
+                if (totalFrames == 1) { // first frame
+                    it.startNewFaceAuthSession(helpFaceAuthStatus.createdAt)
+                }
+                it.addMessage(helpFaceAuthStatus)
+            }
+        } else {
+            val newAcquiredInfoCount = acquiredInfoToFrequency.getOrDefault(acquiredInfo, 0) + 1
+            acquiredInfoToFrequency[acquiredInfo] = newAcquiredInfoCount
+            if (
+                messagesToDefer.contains(acquiredInfo) &&
+                    (mostFrequentAcquiredInfoToDefer == null ||
+                        newAcquiredInfoCount >
+                            acquiredInfoToFrequency.getOrDefault(
+                                mostFrequentAcquiredInfoToDefer!!,
+                                0
+                            ))
+            ) {
+                mostFrequentAcquiredInfoToDefer = acquiredInfo
+            }
         }
 
         logBuffer.logFrameProcessed(
             acquiredInfo,
             totalFrames,
-            mostFrequentAcquiredInfoToDefer?.toString()
+            if (faceMessageDeferUpdate()) {
+                faceHelpMessageDebouncer
+                    ?.getMessageToShow(systemClock.get().elapsedRealtime())
+                    ?.msgId
+                    .toString()
+            } else {
+                mostFrequentAcquiredInfoToDefer?.toString()
+            }
         )
     }
 
@@ -166,9 +221,16 @@
      *   [threshold] percentage.
      */
     fun getDeferredMessage(): CharSequence? {
-        mostFrequentAcquiredInfoToDefer?.let {
-            if (acquiredInfoToFrequency.getOrDefault(it, 0) > (threshold * totalFrames)) {
-                return acquiredInfoToHelpString[it]
+        if (faceMessageDeferUpdate()) {
+            faceHelpMessageDebouncer?.let {
+                val helpFaceAuthStatus = it.getMessageToShow(systemClock.get().elapsedRealtime())
+                return acquiredInfoToHelpString[helpFaceAuthStatus?.msgId]
+            }
+        } else {
+            mostFrequentAcquiredInfoToDefer?.let {
+                if (acquiredInfoToFrequency.getOrDefault(it, 0) > (threshold * totalFrames)) {
+                    return acquiredInfoToHelpString[it]
+                }
             }
         }
         return null
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/ui/binder/BiometricViewBinder.kt b/packages/SystemUI/src/com/android/systemui/biometrics/ui/binder/BiometricViewBinder.kt
index cd9b9bc..0b440ad 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/ui/binder/BiometricViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/ui/binder/BiometricViewBinder.kt
@@ -45,6 +45,7 @@
 import androidx.lifecycle.repeatOnLifecycle
 import com.airbnb.lottie.LottieAnimationView
 import com.airbnb.lottie.LottieCompositionFactory
+import com.android.systemui.Flags.bpIconA11y
 import com.android.systemui.biometrics.Utils.ellipsize
 import com.android.systemui.biometrics.shared.model.BiometricModalities
 import com.android.systemui.biometrics.shared.model.BiometricModality
@@ -54,6 +55,7 @@
 import com.android.systemui.biometrics.ui.viewmodel.PromptMessage
 import com.android.systemui.biometrics.ui.viewmodel.PromptSize
 import com.android.systemui.biometrics.ui.viewmodel.PromptViewModel
+import com.android.systemui.common.ui.view.onTouchListener
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.VibratorHelper
@@ -330,17 +332,31 @@
 
                 // reuse the icon as a confirm button
                 launch {
-                    viewModel.isIconConfirmButton
-                        .map { isPending ->
-                            when {
-                                isPending && modalities.hasFaceAndFingerprint ->
-                                    View.OnTouchListener { _: View, event: MotionEvent ->
-                                        viewModel.onOverlayTouch(event)
-                                    }
-                                else -> null
+                    if (bpIconA11y()) {
+                        viewModel.isIconConfirmButton.collect { isButton ->
+                            if (isButton) {
+                                iconView.onTouchListener { _: View, event: MotionEvent ->
+                                    viewModel.onOverlayTouch(event)
+                                }
+                                iconView.setOnClickListener { viewModel.confirmAuthenticated() }
+                            } else {
+                                iconView.setOnTouchListener(null)
+                                iconView.setOnClickListener(null)
                             }
                         }
-                        .collect { onTouch -> iconView.setOnTouchListener(onTouch) }
+                    } else {
+                        viewModel.isIconConfirmButton
+                            .map { isPending ->
+                                when {
+                                    isPending && modalities.hasFaceAndFingerprint ->
+                                        View.OnTouchListener { _: View, event: MotionEvent ->
+                                            viewModel.onOverlayTouch(event)
+                                        }
+                                    else -> null
+                                }
+                            }
+                            .collect { onTouch -> iconView.setOnTouchListener(onTouch) }
+                    }
                 }
 
                 // dismiss prompt when authenticated and confirmed
@@ -358,7 +374,8 @@
                             // Allow icon to be used as confirmation button with udfps and a11y
                             // enabled
                             if (
-                                accessibilityManager.isTouchExplorationEnabled &&
+                                !bpIconA11y() &&
+                                    accessibilityManager.isTouchExplorationEnabled &&
                                     modalities.hasUdfps
                             ) {
                                 iconView.setOnClickListener { viewModel.confirmAuthenticated() }
diff --git a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
index 25d43d9..4c2fe07 100644
--- a/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModel.kt
@@ -941,16 +941,16 @@
     private fun vibrateOnSuccess() {
         _hapticsToPlay.value =
             HapticsToPlay(
-                HapticFeedbackConstants.CONFIRM,
-                HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING,
+                HapticFeedbackConstants.BIOMETRIC_CONFIRM,
+                null,
             )
     }
 
     private fun vibrateOnError() {
         _hapticsToPlay.value =
             HapticsToPlay(
-                HapticFeedbackConstants.REJECT,
-                HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING,
+                HapticFeedbackConstants.BIOMETRIC_REJECT,
+                null,
             )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/BluetoothTileDialogUiEvent.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/BluetoothTileDialogUiEvent.kt
index 4a358c0..bdd4c16 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/BluetoothTileDialogUiEvent.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/BluetoothTileDialogUiEvent.kt
@@ -31,6 +31,8 @@
     @UiEvent(doc = "Saved clicked to connect") SAVED_DEVICE_CONNECT(1500),
     @UiEvent(doc = "Active device clicked to disconnect") ACTIVE_DEVICE_DISCONNECT(1507),
     @UiEvent(doc = "Audio sharing device clicked, do nothing") AUDIO_SHARING_DEVICE_CLICKED(1699),
+    @UiEvent(doc = "Available audio sharing device clicked, do nothing")
+    AVAILABLE_AUDIO_SHARING_DEVICE_CLICKED(1880),
     @UiEvent(doc = "Connected other device clicked to disconnect")
     CONNECTED_OTHER_DEVICE_DISCONNECT(1508),
     @UiEvent(doc = "The auto on toggle is clicked") BLUETOOTH_AUTO_ON_TOGGLE_CLICKED(1617),
@@ -44,8 +46,14 @@
         doc = "Not broadcasting, having one connected, another saved LE audio device is clicked"
     )
     LAUNCH_SETTINGS_NOT_SHARING_SAVED_LE_DEVICE_CLICKED(1719),
+    @Deprecated(
+        "Use case no longer needed",
+        ReplaceWith("LAUNCH_SETTINGS_NOT_SHARING_ACTIVE_LE_DEVICE_CLICKED")
+    )
     @UiEvent(doc = "Not broadcasting, one of the two connected LE audio devices is clicked")
-    LAUNCH_SETTINGS_NOT_SHARING_CONNECTED_LE_DEVICE_CLICKED(1720);
+    LAUNCH_SETTINGS_NOT_SHARING_CONNECTED_LE_DEVICE_CLICKED(1720),
+    @UiEvent(doc = "Not broadcasting, having two connected, the active LE audio devices is clicked")
+    LAUNCH_SETTINGS_NOT_SHARING_ACTIVE_LE_DEVICE_CLICKED(1881);
 
     override fun getId() = metricId
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItem.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItem.kt
index a78130f..2ba4c73 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItem.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItem.kt
@@ -38,6 +38,7 @@
 enum class DeviceItemType {
     ACTIVE_MEDIA_BLUETOOTH_DEVICE,
     AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
+    AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
     AVAILABLE_MEDIA_BLUETOOTH_DEVICE,
     CONNECTED_BLUETOOTH_DEVICE,
     SAVED_BLUETOOTH_DEVICE,
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractor.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractor.kt
index 9d82e76..f1894d3 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractor.kt
@@ -67,7 +67,7 @@
                     backgroundDispatcher,
                     logger
                 ),
-                NotSharingClickedConnected(
+                NotSharingClickedActive(
                     leAudioProfile,
                     assistantProfile,
                     backgroundDispatcher,
@@ -106,6 +106,12 @@
                     DeviceItemType.AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE -> {
                         uiEventLogger.log(BluetoothTileDialogUiEvent.AUDIO_SHARING_DEVICE_CLICKED)
                     }
+                    DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE -> {
+                        // TODO(b/360759048): pop up dialog
+                        uiEventLogger.log(
+                            BluetoothTileDialogUiEvent.AVAILABLE_AUDIO_SHARING_DEVICE_CLICKED
+                        )
+                    }
                     DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE -> {
                         setActive()
                         uiEventLogger.log(BluetoothTileDialogUiEvent.CONNECTED_DEVICE_SET_ACTIVE)
@@ -238,14 +244,14 @@
             BluetoothTileDialogUiEvent.LAUNCH_SETTINGS_NOT_SHARING_SAVED_LE_DEVICE_CLICKED
     }
 
-    private class NotSharingClickedConnected(
+    private class NotSharingClickedActive(
         private val leAudioProfile: LeAudioProfile?,
         private val assistantProfile: LocalBluetoothLeBroadcastAssistant?,
         @Background private val backgroundDispatcher: CoroutineDispatcher,
         private val logger: BluetoothTileDialogLogger,
     ) : LaunchSettingsCriteria {
-        // If not broadcasting, having two device connected, clicked on any connected LE audio
-        // devices
+        // If not broadcasting, having two device connected, clicked on the active LE audio
+        // device
         override suspend fun matched(inAudioSharing: Boolean, deviceItem: DeviceItem): Boolean {
             return withContext(backgroundDispatcher) {
                 val matched =
@@ -259,7 +265,7 @@
                                         logger
                                     )
                                     .size == 2 &&
-                                deviceItem.isActiveOrConnectedLeAudioSupported
+                                deviceItem.isActiveLeAudioSupported
                         }
                     } ?: false
 
@@ -275,7 +281,7 @@
         }
 
         override suspend fun getClickUiEvent(deviceItem: DeviceItem) =
-            BluetoothTileDialogUiEvent.LAUNCH_SETTINGS_NOT_SHARING_CONNECTED_LE_DEVICE_CLICKED
+            BluetoothTileDialogUiEvent.LAUNCH_SETTINGS_NOT_SHARING_ACTIVE_LE_DEVICE_CLICKED
     }
 
     private companion object {
@@ -290,10 +296,8 @@
         val DeviceItem.isNotConnectedLeAudioSupported: Boolean
             get() = type == DeviceItemType.SAVED_BLUETOOTH_DEVICE && isLeAudioSupported
 
-        val DeviceItem.isActiveOrConnectedLeAudioSupported: Boolean
-            get() =
-                (type == DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE ||
-                    type == DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE) && isLeAudioSupported
+        val DeviceItem.isActiveLeAudioSupported: Boolean
+            get() = type == DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE && isLeAudioSupported
 
         val DeviceItem.isMediaDevice: Boolean
             get() =
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactory.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactory.kt
index e846bf7..7280489 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactory.kt
@@ -125,6 +125,37 @@
     }
 }
 
+internal class AvailableAudioSharingMediaDeviceItemFactory(
+    private val localBluetoothManager: LocalBluetoothManager?
+) : AvailableMediaDeviceItemFactory() {
+    override fun isFilterMatched(
+        context: Context,
+        cachedDevice: CachedBluetoothDevice,
+        audioManager: AudioManager
+    ): Boolean {
+        return BluetoothUtils.isAudioSharingEnabled() &&
+            super.isFilterMatched(context, cachedDevice, audioManager) &&
+            BluetoothUtils.isAvailableAudioSharingMediaBluetoothDevice(
+                cachedDevice,
+                localBluetoothManager
+            )
+    }
+
+    override fun create(context: Context, cachedDevice: CachedBluetoothDevice): DeviceItem {
+        return createDeviceItem(
+            context,
+            cachedDevice,
+            DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
+            context.getString(
+                R.string.quick_settings_bluetooth_device_audio_sharing_or_switch_active
+            ),
+            if (cachedDevice.isBusy) backgroundOffBusy else backgroundOff,
+            "",
+            isActive = false
+        )
+    }
+}
+
 internal class ActiveHearingDeviceItemFactory : ActiveMediaDeviceItemFactory() {
     override fun isFilterMatched(
         context: Context,
diff --git a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemInteractor.kt b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemInteractor.kt
index 9524496..9114eca 100644
--- a/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/bluetooth/qsdialog/DeviceItemInteractor.kt
@@ -118,6 +118,7 @@
         listOf(
             ActiveMediaDeviceItemFactory(),
             AudioSharingMediaDeviceItemFactory(localBluetoothManager),
+            AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager),
             AvailableMediaDeviceItemFactory(),
             ConnectedDeviceItemFactory(),
             SavedDeviceItemFactory()
@@ -127,6 +128,7 @@
         listOf(
             DeviceItemType.ACTIVE_MEDIA_BLUETOOTH_DEVICE,
             DeviceItemType.AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
+            DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
             DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE,
             DeviceItemType.CONNECTED_BLUETOOTH_DEVICE,
             DeviceItemType.SAVED_BLUETOOTH_DEVICE,
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlags.kt b/packages/SystemUI/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlags.kt
index 62ef365..a1111f6 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlags.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlags.kt
@@ -17,18 +17,17 @@
 package com.android.systemui.bouncer.shared.flag
 
 import com.android.systemui.Flags
-import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import dagger.Module
-import dagger.Provides
 
-interface ComposeBouncerFlags {
+object ComposeBouncerFlags {
 
     /**
      * Returns `true` if the Compose bouncer is enabled or if the scene container framework is
      * enabled; `false` otherwise.
      */
-    fun isComposeBouncerOrSceneContainerEnabled(): Boolean
+    fun isComposeBouncerOrSceneContainerEnabled(): Boolean {
+        return SceneContainerFlag.isEnabled || Flags.composeBouncer()
+    }
 
     /**
      * Returns `true` if only compose bouncer is enabled and scene container framework is not
@@ -39,30 +38,7 @@
             "that includes compose bouncer in legacy keyguard.",
         replaceWith = ReplaceWith("isComposeBouncerOrSceneContainerEnabled()")
     )
-    fun isOnlyComposeBouncerEnabled(): Boolean
-}
-
-class ComposeBouncerFlagsImpl() : ComposeBouncerFlags {
-
-    override fun isComposeBouncerOrSceneContainerEnabled(): Boolean {
-        return SceneContainerFlag.isEnabled || Flags.composeBouncer()
-    }
-
-    @Deprecated(
-        "Avoid using this, this is meant to be used only by the glue code " +
-            "that includes compose bouncer in legacy keyguard.",
-        replaceWith = ReplaceWith("isComposeBouncerOrSceneContainerEnabled()")
-    )
-    override fun isOnlyComposeBouncerEnabled(): Boolean {
+    fun isOnlyComposeBouncerEnabled(): Boolean {
         return !SceneContainerFlag.isEnabled && Flags.composeBouncer()
     }
 }
-
-@Module
-object ComposeBouncerFlagsModule {
-    @Provides
-    @SysUISingleton
-    fun impl(): ComposeBouncerFlags {
-        return ComposeBouncerFlagsImpl()
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/BouncerViewBinder.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/BouncerViewBinder.kt
index ad93a25..cc8dce79 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/BouncerViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/BouncerViewBinder.kt
@@ -55,12 +55,11 @@
 class BouncerViewBinder
 @Inject
 constructor(
-    private val composeBouncerFlags: ComposeBouncerFlags,
     private val legacyBouncerDependencies: Lazy<LegacyBouncerDependencies>,
     private val composeBouncerDependencies: Lazy<ComposeBouncerDependencies>,
 ) {
     fun bind(view: ViewGroup) {
-        if (composeBouncerFlags.isOnlyComposeBouncerEnabled()) {
+        if (ComposeBouncerFlags.isOnlyComposeBouncerEnabled()) {
             val deps = composeBouncerDependencies.get()
             ComposeBouncerViewBinder.bind(
                 view,
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/ComposeBouncerViewBinder.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/ComposeBouncerViewBinder.kt
index c1f7d59..c4bbd9c 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/ComposeBouncerViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/binder/ComposeBouncerViewBinder.kt
@@ -8,14 +8,12 @@
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.repeatOnLifecycle
-import com.android.compose.theme.PlatformTheme
 import com.android.keyguard.ViewMediatorCallback
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor
 import com.android.systemui.bouncer.ui.BouncerDialogFactory
-import com.android.systemui.bouncer.ui.composable.BouncerContent
+import com.android.systemui.bouncer.ui.composable.BouncerContainer
 import com.android.systemui.bouncer.ui.viewmodel.BouncerSceneContentViewModel
-import com.android.systemui.lifecycle.rememberViewModel
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import kotlinx.coroutines.flow.collectLatest
@@ -49,14 +47,7 @@
                                     this@repeatWhenAttached.lifecycle
                             }
                         )
-                        setContent {
-                            PlatformTheme {
-                                BouncerContent(
-                                    rememberViewModel { viewModelFactory.create() },
-                                    dialogFactory,
-                                )
-                            }
-                        }
+                        setContent { BouncerContainer(viewModelFactory, dialogFactory) }
                     }
                 }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/composable/BouncerContainer.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/composable/BouncerContainer.kt
new file mode 100644
index 0000000..c05dcd5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/composable/BouncerContainer.kt
@@ -0,0 +1,54 @@
+/*
+ * 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.bouncer.ui.composable
+
+import androidx.compose.foundation.Canvas
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.material3.MaterialTheme
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.theme.PlatformTheme
+import com.android.systemui.bouncer.ui.BouncerDialogFactory
+import com.android.systemui.bouncer.ui.viewmodel.BouncerSceneContentViewModel
+import com.android.systemui.compose.modifiers.sysuiResTag
+import com.android.systemui.lifecycle.rememberViewModel
+
+/** Container that includes the compose bouncer and is meant to be included in legacy keyguard. */
+@Composable
+fun BouncerContainer(
+    viewModelFactory: BouncerSceneContentViewModel.Factory,
+    dialogFactory: BouncerDialogFactory,
+) {
+    PlatformTheme {
+        val backgroundColor = MaterialTheme.colorScheme.surface
+
+        val bouncerViewModel = rememberViewModel("BouncerContainer") { viewModelFactory.create() }
+        Box {
+            Canvas(Modifier.fillMaxSize()) { drawRect(color = backgroundColor) }
+
+            // Separate the bouncer content into a reusable composable that
+            // doesn't have any SceneScope
+            // dependencies
+            BouncerContent(
+                bouncerViewModel,
+                dialogFactory,
+                Modifier.sysuiResTag(Bouncer.TestTags.Root).fillMaxSize()
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModel.kt
index abca518..c67b354 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/AuthMethodBouncerViewModel.kt
@@ -22,7 +22,6 @@
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -40,7 +39,10 @@
      * being able to attempt to unlock the device.
      */
     val isInputEnabled: StateFlow<Boolean>,
-) : SysUiViewModel, ExclusiveActivatable() {
+
+    /** Name to use for performance tracing purposes. */
+    val traceName: String,
+) : ExclusiveActivatable() {
 
     private val _animateFailure = MutableStateFlow(false)
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModel.kt
index d21eccd..c383b8d 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModel.kt
@@ -39,7 +39,6 @@
 import com.android.systemui.deviceentry.shared.model.FingerprintFailureMessage
 import com.android.systemui.deviceentry.shared.model.FingerprintLockoutMessage
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.res.R.string.kg_too_many_failed_attempts_countdown
 import com.android.systemui.user.ui.viewmodel.UserSwitcherViewModel
 import com.android.systemui.util.kotlin.Utils.Companion.sample
@@ -79,8 +78,7 @@
     private val faceAuthInteractor: DeviceEntryFaceAuthInteractor,
     private val deviceUnlockedInteractor: DeviceUnlockedInteractor,
     private val deviceEntryBiometricsAllowedInteractor: DeviceEntryBiometricsAllowedInteractor,
-    private val flags: ComposeBouncerFlags,
-) : SysUiViewModel, ExclusiveActivatable() {
+) : ExclusiveActivatable() {
     /**
      * A message shown when the user has attempted the wrong credential too many times and now must
      * wait a while before attempting to authenticate again.
@@ -97,7 +95,7 @@
     val message: MutableStateFlow<MessageViewModel?> = MutableStateFlow(null)
 
     override suspend fun onActivated(): Nothing {
-        if (!flags.isComposeBouncerOrSceneContainerEnabled()) {
+        if (!ComposeBouncerFlags.isComposeBouncerOrSceneContainerEnabled()) {
             return awaitCancellation()
         }
 
@@ -156,7 +154,7 @@
                     emptyFlow()
                 }
             }
-            .collectLatest { messageViewModel -> message.value = messageViewModel }
+            .collect { messageViewModel -> message.value = messageViewModel }
     }
 
     private suspend fun listenForSimBouncerEvents() {
@@ -171,7 +169,7 @@
                     emptyFlow()
                 }
             }
-            .collectLatest {
+            .collect {
                 if (it != null) {
                     message.value = it
                 } else {
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
index 79e5f8d..0aada06 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneContentViewModel.kt
@@ -23,18 +23,17 @@
 import androidx.compose.ui.input.key.KeyEvent
 import androidx.compose.ui.input.key.type
 import androidx.core.graphics.drawable.toBitmap
+import com.android.app.tracing.coroutines.traceCoroutine
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.authentication.shared.model.AuthenticationWipeModel
 import com.android.systemui.bouncer.domain.interactor.BouncerActionButtonInteractor
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
-import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlags
 import com.android.systemui.bouncer.shared.model.BouncerActionButtonModel
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.shared.model.Text
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.user.ui.viewmodel.UserSwitcherViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
@@ -57,13 +56,12 @@
     private val authenticationInteractor: AuthenticationInteractor,
     private val devicePolicyManager: DevicePolicyManager,
     private val bouncerMessageViewModelFactory: BouncerMessageViewModel.Factory,
-    private val flags: ComposeBouncerFlags,
     private val userSwitcher: UserSwitcherViewModel,
     private val actionButtonInteractor: BouncerActionButtonInteractor,
     private val pinViewModelFactory: PinBouncerViewModel.Factory,
     private val patternViewModelFactory: PatternBouncerViewModel.Factory,
     private val passwordViewModelFactory: PasswordBouncerViewModel.Factory,
-) : SysUiViewModel, ExclusiveActivatable() {
+) : ExclusiveActivatable() {
     private val _selectedUserImage = MutableStateFlow<Bitmap?>(null)
     val selectedUserImage: StateFlow<Bitmap?> = _selectedUserImage.asStateFlow()
 
@@ -147,7 +145,7 @@
                     .map(::getChildViewModel)
                     .collectLatest { childViewModelOrNull ->
                         _authMethodViewModel.value = childViewModelOrNull
-                        childViewModelOrNull?.activate()
+                        childViewModelOrNull?.let { traceCoroutine(it.traceName) { it.activate() } }
                     }
             }
 
@@ -160,7 +158,7 @@
             launch {
                 userSwitcher.selectedUser
                     .map { it.image.toBitmap() }
-                    .collectLatest { _selectedUserImage.value = it }
+                    .collect { _selectedUserImage.value = it }
             }
 
             launch {
@@ -187,34 +185,32 @@
                                 )
                             }
                     }
-                    .collectLatest { _userSwitcherDropdown.value = it }
+                    .collect { _userSwitcherDropdown.value = it }
             }
 
             launch {
                 combine(wipeDialogMessage, lockoutDialogMessage) { _, _ -> createDialogViewModel() }
-                    .collectLatest { _dialogViewModel.value = it }
+                    .collect { _dialogViewModel.value = it }
             }
 
-            launch {
-                actionButtonInteractor.actionButton.collectLatest { _actionButton.value = it }
-            }
+            launch { actionButtonInteractor.actionButton.collect { _actionButton.value = it } }
 
             launch {
                 authMethodViewModel
                     .map { authMethod -> isSideBySideSupported(authMethod) }
-                    .collectLatest { _isSideBySideSupported.value = it }
+                    .collect { _isSideBySideSupported.value = it }
             }
 
             launch {
                 authMethodViewModel
                     .map { authMethod -> isFoldSplitRequired(authMethod) }
-                    .collectLatest { _isFoldSplitRequired.value = it }
+                    .collect { _isFoldSplitRequired.value = it }
             }
 
             launch {
                 message.isLockoutMessagePresent
                     .map { lockoutMessagePresent -> !lockoutMessagePresent }
-                    .collectLatest { _isInputEnabled.value = it }
+                    .collect { _isInputEnabled.value = it }
             }
 
             awaitCancellation()
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModel.kt
similarity index 85%
rename from packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModel.kt
index 2a27271..4fe6fc6 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/BouncerUserActionsViewModel.kt
@@ -22,21 +22,20 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.bouncer.domain.interactor.BouncerInteractor
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.map
 
 /**
  * Models UI state for user actions that can lead to navigation to other scenes when showing the
  * bouncer scene.
  */
-class BouncerSceneActionsViewModel
+class BouncerUserActionsViewModel
 @AssistedInject
 constructor(
     private val bouncerInteractor: BouncerInteractor,
-) : SceneActionsViewModel() {
+) : UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         bouncerInteractor.dismissDestination
@@ -46,11 +45,11 @@
                     Swipe(SwipeDirection.Down) to UserActionResult(prevScene),
                 )
             }
-            .collectLatest { actions -> setActions(actions) }
+            .collect { actions -> setActions(actions) }
     }
 
     @AssistedFactory
     interface Factory {
-        fun create(): BouncerSceneActionsViewModel
+        fun create(): BouncerUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
index c91fd6a..2493cf1 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PasswordBouncerViewModel.kt
@@ -34,7 +34,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
@@ -53,6 +52,7 @@
     AuthMethodBouncerViewModel(
         interactor = interactor,
         isInputEnabled = isInputEnabled,
+        traceName = "PasswordBouncerViewModel",
     ) {
 
     private val _password = MutableStateFlow("")
@@ -104,11 +104,9 @@
                 combine(isInputEnabled, isTextFieldFocused) { hasInput, hasFocus ->
                         hasInput && !hasFocus
                     }
-                    .collectLatest { _isTextFieldFocusRequested.value = it }
+                    .collect { _isTextFieldFocusRequested.value = it }
             }
-            launch {
-                selectedUserInteractor.selectedUser.collectLatest { _selectedUserId.value = it }
-            }
+            launch { selectedUserInteractor.selectedUser.collect { _selectedUserId.value = it } }
             launch {
                 // Re-fetch the currently-enabled IMEs whenever the selected user changes, and
                 // whenever
@@ -124,7 +122,7 @@
                     ) { selectedUserId, _ ->
                         inputMethodInteractor.hasMultipleEnabledImesOrSubtypes(selectedUserId)
                     }
-                    .collectLatest { _isImeSwitcherButtonVisible.value = it }
+                    .collect { _isImeSwitcherButtonVisible.value = it }
             }
             awaitCancellation()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModel.kt
index 4c02929..0a866b4 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PatternBouncerViewModel.kt
@@ -34,7 +34,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.launch
@@ -51,6 +50,7 @@
     AuthMethodBouncerViewModel(
         interactor = interactor,
         isInputEnabled = isInputEnabled,
+        traceName = "PatternBouncerViewModel",
     ) {
 
     /** The number of columns in the dot grid. */
@@ -85,9 +85,7 @@
         coroutineScope {
             launch { super.onActivated() }
             launch {
-                selectedDotSet
-                    .map { it.toList() }
-                    .collectLatest { selectedDotList.value = it.toList() }
+                selectedDotSet.map { it.toList() }.collect { selectedDotList.value = it.toList() }
             }
             awaitCancellation()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModel.kt
index c611954..df6ca9b 100644
--- a/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/bouncer/ui/viewmodel/PinBouncerViewModel.kt
@@ -42,7 +42,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
@@ -63,6 +62,7 @@
     AuthMethodBouncerViewModel(
         interactor = interactor,
         isInputEnabled = isInputEnabled,
+        traceName = "PinBouncerViewModel",
     ) {
     /**
      * Whether the sim-related UI in the pin view is showing.
@@ -122,7 +122,7 @@
                     } else {
                         interactor.hintedPinLength
                     }
-                    .collectLatest { _hintedPinLength.value = it }
+                    .collect { _hintedPinLength.value = it }
             }
             launch {
                 combine(
@@ -134,17 +134,17 @@
                             isAutoConfirmEnabled = isAutoConfirmEnabled,
                         )
                     }
-                    .collectLatest { _backspaceButtonAppearance.value = it }
+                    .collect { _backspaceButtonAppearance.value = it }
             }
             launch {
                 interactor.isAutoConfirmEnabled
                     .map { if (it) ActionButtonAppearance.Hidden else ActionButtonAppearance.Shown }
-                    .collectLatest { _confirmButtonAppearance.value = it }
+                    .collect { _confirmButtonAppearance.value = it }
             }
             launch {
                 interactor.isPinEnhancedPrivacyEnabled
                     .map { !it }
-                    .collectLatest { _isDigitButtonAnimationEnabled.value = it }
+                    .collect { _isDigitButtonAnimationEnabled.value = it }
             }
             awaitCancellation()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayWindow.java b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayWindow.java
index 0dc6fda..dc3b50c 100644
--- a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayWindow.java
+++ b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayWindow.java
@@ -27,6 +27,7 @@
 import android.view.WindowInsets;
 import android.view.WindowManager;
 
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.internal.policy.PhoneWindow;
 import com.android.systemui.clipboardoverlay.dagger.ClipboardOverlayModule.OverlayWindowContext;
 import com.android.systemui.screenshot.FloatingWindowUtil;
@@ -44,6 +45,7 @@
 
     private final Context mContext;
     private final WindowManager mWindowManager;
+    private final ViewCaptureAwareWindowManager mViewCaptureAwareWindowManager;
     private final WindowManager.LayoutParams mWindowLayoutParams;
 
     private boolean mKeyboardVisible;
@@ -52,7 +54,9 @@
     private Runnable mOnOrientationChangeListener;
 
     @Inject
-    ClipboardOverlayWindow(@OverlayWindowContext Context context) {
+    ClipboardOverlayWindow(@OverlayWindowContext Context context,
+            @OverlayWindowContext ViewCaptureAwareWindowManager viewCaptureAwareWindowManager,
+            @OverlayWindowContext WindowManager windowManager) {
         super(context);
         mContext = context;
         mOrientation = mContext.getResources().getConfiguration().orientation;
@@ -61,10 +65,11 @@
         requestFeature(Window.FEATURE_NO_TITLE);
         requestFeature(Window.FEATURE_ACTIVITY_TRANSITIONS);
         setBackgroundDrawableResource(android.R.color.transparent);
-        mWindowManager = mContext.getSystemService(WindowManager.class);
+        mWindowManager = windowManager;
+        mViewCaptureAwareWindowManager = viewCaptureAwareWindowManager;
         mWindowLayoutParams = FloatingWindowUtil.getFloatingWindowParams();
         mWindowLayoutParams.setTitle("ClipboardOverlay");
-        setWindowManager(mWindowManager, null, null);
+        setWindowManager(windowManager, null, null);
         setWindowFocusable(false);
     }
 
@@ -81,10 +86,12 @@
 
         attach();
         withWindowAttached(() -> {
-            WindowInsets currentInsets = mWindowManager.getCurrentWindowMetrics().getWindowInsets();
+            WindowInsets currentInsets = mWindowManager.getCurrentWindowMetrics()
+                    .getWindowInsets();
             mKeyboardVisible = currentInsets.isVisible(WindowInsets.Type.ime());
             peekDecorView().getViewTreeObserver().addOnGlobalLayoutListener(() -> {
-                WindowInsets insets = mWindowManager.getCurrentWindowMetrics().getWindowInsets();
+                WindowInsets insets = mWindowManager.getCurrentWindowMetrics()
+                        .getWindowInsets();
                 boolean keyboardVisible = insets.isVisible(WindowInsets.Type.ime());
                 if (keyboardVisible != mKeyboardVisible) {
                     mKeyboardVisible = keyboardVisible;
@@ -105,7 +112,7 @@
     void remove() {
         final View decorView = peekDecorView();
         if (decorView != null && decorView.isAttachedToWindow()) {
-            mWindowManager.removeViewImmediate(decorView);
+            mViewCaptureAwareWindowManager.removeViewImmediate(decorView);
         }
     }
 
@@ -139,7 +146,7 @@
         if (decorView.isAttachedToWindow()) {
             return;
         }
-        mWindowManager.addView(decorView, mWindowLayoutParams);
+        mViewCaptureAwareWindowManager.addView(decorView, mWindowLayoutParams);
         decorView.requestApplyInsets();
     }
 
@@ -160,7 +167,7 @@
         }
         final View decorView = peekDecorView();
         if (decorView != null && decorView.isAttachedToWindow()) {
-            mWindowManager.updateViewLayout(decorView, mWindowLayoutParams);
+            mViewCaptureAwareWindowManager.updateViewLayout(decorView, mWindowLayoutParams);
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/dagger/ClipboardOverlayModule.java b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/dagger/ClipboardOverlayModule.java
index ff9fba4..307a07f 100644
--- a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/dagger/ClipboardOverlayModule.java
+++ b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/dagger/ClipboardOverlayModule.java
@@ -18,17 +18,24 @@
 
 import static android.view.WindowManager.LayoutParams.TYPE_SCREENSHOT;
 
+import static com.android.systemui.Flags.enableViewCaptureTracing;
+import static com.android.systemui.util.ConvenienceExtensionsKt.toKotlinLazy;
+
 import static java.lang.annotation.RetentionPolicy.RUNTIME;
 
 import android.content.Context;
 import android.hardware.display.DisplayManager;
 import android.view.Display;
 import android.view.LayoutInflater;
+import android.view.WindowManager;
 
+import com.android.app.viewcapture.ViewCapture;
+import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.systemui.clipboardoverlay.ClipboardOverlayView;
 import com.android.systemui.res.R;
 import com.android.systemui.settings.DisplayTracker;
 
+import dagger.Lazy;
 import dagger.Module;
 import dagger.Provides;
 
@@ -61,6 +68,28 @@
                 R.layout.clipboard_overlay, null);
     }
 
+    /**
+     *
+     */
+    @Provides
+    @OverlayWindowContext
+    static WindowManager provideWindowManager(@OverlayWindowContext Context context) {
+        return context.getSystemService(WindowManager.class);
+    }
+
+    /**
+     *
+     */
+    @Provides
+    @OverlayWindowContext
+    static ViewCaptureAwareWindowManager provideViewCaptureAwareWindowManager(
+            @OverlayWindowContext WindowManager windowManager,
+            Lazy<ViewCapture> daggerLazyViewCapture) {
+        return new ViewCaptureAwareWindowManager(windowManager,
+                /* lazyViewCapture= */ toKotlinLazy(daggerLazyViewCapture),
+                /* isViewCaptureEnabled= */ enableViewCaptureTracing());
+    }
+
     @Qualifier
     @Documented
     @Retention(RUNTIME)
diff --git a/packages/SystemUI/src/com/android/systemui/common/shared/model/Icon.kt b/packages/SystemUI/src/com/android/systemui/common/shared/model/Icon.kt
index 3cdb573..aef5f1f 100644
--- a/packages/SystemUI/src/com/android/systemui/common/shared/model/Icon.kt
+++ b/packages/SystemUI/src/com/android/systemui/common/shared/model/Icon.kt
@@ -38,5 +38,5 @@
 }
 
 /** Creates [Icon.Loaded] for a given drawable with an optional [contentDescription]. */
-fun Drawable.asIcon(contentDescription: ContentDescription? = null): Icon =
+fun Drawable.asIcon(contentDescription: ContentDescription? = null): Icon.Loaded =
     Icon.Loaded(this, contentDescription)
diff --git a/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationState.kt b/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationState.kt
index 578389b..13f6bba 100644
--- a/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationState.kt
+++ b/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationState.kt
@@ -23,35 +23,25 @@
 import androidx.annotation.DimenRes
 import androidx.annotation.LayoutRes
 import com.android.settingslib.Utils
-import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.statusbar.policy.onDensityOrFontScaleChanged
 import com.android.systemui.statusbar.policy.onThemeChanged
 import com.android.systemui.util.kotlin.emitOnStart
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
 
-/** Configuration-aware-state-tracking utilities. */
-class ConfigurationState
-@Inject
-constructor(
-    private val configurationController: ConfigurationController,
-    @Application private val context: Context,
-    private val layoutInflater: LayoutInflater,
-) {
+interface ConfigurationState {
     /**
      * Returns a [Flow] that emits a dimension pixel size that is kept in sync with the device
      * configuration.
      *
      * @see android.content.res.Resources.getDimensionPixelSize
      */
-    fun getDimensionPixelSize(@DimenRes id: Int): Flow<Int> {
-        return configurationController.onDensityOrFontScaleChanged.emitOnStart().map {
-            context.resources.getDimensionPixelSize(id)
-        }
-    }
+    fun getDimensionPixelSize(@DimenRes id: Int): Flow<Int>
 
     /**
      * Returns a [Flow] that emits a dimension pixel size that is kept in sync with the device
@@ -59,22 +49,14 @@
      *
      * @see android.content.res.Resources.getDimensionPixelSize
      */
-    fun getDimensionPixelOffset(@DimenRes id: Int): Flow<Int> {
-        return configurationController.onDensityOrFontScaleChanged.emitOnStart().map {
-            context.resources.getDimensionPixelOffset(id)
-        }
-    }
+    fun getDimensionPixelOffset(@DimenRes id: Int): Flow<Int>
 
     /**
      * Returns a [Flow] that emits a color that is kept in sync with the device theme.
      *
      * @see Utils.getColorAttrDefaultColor
      */
-    fun getColorAttr(@AttrRes id: Int, @ColorInt defaultValue: Int): Flow<Int> {
-        return configurationController.onThemeChanged.emitOnStart().map {
-            Utils.getColorAttrDefaultColor(context, id, defaultValue)
-        }
-    }
+    fun getColorAttr(@AttrRes id: Int, @ColorInt defaultValue: Int): Flow<Int>
 
     /**
      * Returns a [Flow] that emits a [View] that is re-inflated as necessary to remain in sync with
@@ -87,6 +69,65 @@
         @LayoutRes id: Int,
         root: ViewGroup?,
         attachToRoot: Boolean,
+    ): Flow<T>
+}
+
+/** Configuration-aware-state-tracking utilities. */
+class ConfigurationStateImpl
+@AssistedInject
+constructor(
+    @Assisted private val configurationController: ConfigurationController,
+    @Assisted private val context: Context,
+) : ConfigurationState {
+
+    private val layoutInflater = LayoutInflater.from(context)
+
+    /**
+     * Returns a [Flow] that emits a dimension pixel size that is kept in sync with the device
+     * configuration.
+     *
+     * @see android.content.res.Resources.getDimensionPixelSize
+     */
+    override fun getDimensionPixelSize(@DimenRes id: Int): Flow<Int> {
+        return configurationController.onDensityOrFontScaleChanged.emitOnStart().map {
+            context.resources.getDimensionPixelSize(id)
+        }
+    }
+
+    /**
+     * Returns a [Flow] that emits a dimension pixel size that is kept in sync with the device
+     * configuration.
+     *
+     * @see android.content.res.Resources.getDimensionPixelSize
+     */
+    override fun getDimensionPixelOffset(@DimenRes id: Int): Flow<Int> {
+        return configurationController.onDensityOrFontScaleChanged.emitOnStart().map {
+            context.resources.getDimensionPixelOffset(id)
+        }
+    }
+
+    /**
+     * Returns a [Flow] that emits a color that is kept in sync with the device theme.
+     *
+     * @see Utils.getColorAttrDefaultColor
+     */
+    override fun getColorAttr(@AttrRes id: Int, @ColorInt defaultValue: Int): Flow<Int> {
+        return configurationController.onThemeChanged.emitOnStart().map {
+            Utils.getColorAttrDefaultColor(context, id, defaultValue)
+        }
+    }
+
+    /**
+     * Returns a [Flow] that emits a [View] that is re-inflated as necessary to remain in sync with
+     * the device configuration.
+     *
+     * @see LayoutInflater.inflate
+     */
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : View> inflateLayout(
+        @LayoutRes id: Int,
+        root: ViewGroup?,
+        attachToRoot: Boolean,
     ): Flow<T> {
         // TODO(b/305930747): This may lead to duplicate invocations if both flows emit, find a
         //  solution to only emit one event.
@@ -97,4 +138,16 @@
             .emitOnStart()
             .map { layoutInflater.inflate(id, root, attachToRoot) as T }
     }
+
+    @AssistedFactory
+    interface Factory {
+        /**
+         * Creates a configurationState for a given context. The [configurationController] is
+         * supposed to give config events specific for that context.
+         */
+        fun create(
+            context: Context,
+            configurationController: ConfigurationController
+        ): ConfigurationStateImpl
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationStateModule.kt b/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationStateModule.kt
new file mode 100644
index 0000000..b36da3b
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/common/ui/ConfigurationStateModule.kt
@@ -0,0 +1,61 @@
+/*
+ * 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.common.ui
+
+import android.content.Context
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.policy.ConfigurationController
+import dagger.Binds
+import dagger.Module
+import dagger.Provides
+import javax.inject.Qualifier
+
+/**
+ * Annotates elements that provide information from the global configuration.
+ *
+ * The global configuration is the one associted with the main display. Secondary displays will
+ * apply override to the global configuration. Elements annotated with this shouldn't be used for
+ * secondary displays.
+ */
+@Qualifier @Retention(AnnotationRetention.RUNTIME) annotation class GlobalConfig
+
+@Module
+interface ConfigurationStateModule {
+
+    /**
+     * Deprecated: [ConfigurationState] should be injected only with the correct annotation. For
+     * now, without annotation the global config associated state is provided.
+     */
+    @Binds
+    fun provideGlobalConfigurationState(
+        @GlobalConfig configurationState: ConfigurationState
+    ): ConfigurationState
+
+    companion object {
+        @SysUISingleton
+        @Provides
+        @GlobalConfig
+        fun provideGlobalConfigurationState(
+            configStateFactory: ConfigurationStateImpl.Factory,
+            configurationController: ConfigurationController,
+            @Application context: Context,
+        ): ConfigurationState {
+            return configStateFactory.create(context, configurationController)
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/CommunalDreamStartable.kt b/packages/SystemUI/src/com/android/systemui/communal/CommunalDreamStartable.kt
index 0582cc2..c69cea4 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/CommunalDreamStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/CommunalDreamStartable.kt
@@ -32,11 +32,14 @@
 import com.android.systemui.keyguard.shared.model.filterState
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.util.kotlin.BooleanFlowOperators.allOf
+import com.android.systemui.util.kotlin.BooleanFlowOperators.not
 import com.android.systemui.util.kotlin.Utils.Companion.sampleFilter
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 
 /**
@@ -54,6 +57,18 @@
     private val dreamManager: DreamManager,
     @Background private val bgScope: CoroutineScope,
 ) : CoreStartable {
+    /** Flow that emits when the dream should be started underneath the glanceable hub. */
+    val startDream =
+        allOf(
+                keyguardTransitionInteractor
+                    .transitionValue(Scenes.Communal, KeyguardState.GLANCEABLE_HUB)
+                    .map { it == 1f },
+                not(keyguardInteractor.isDreaming),
+                // TODO(b/362830856): Remove this workaround.
+                keyguardInteractor.isKeyguardShowing,
+            )
+            .filter { it }
+
     @SuppressLint("MissingPermission")
     override fun start() {
         if (!communalSettingsInteractor.isCommunalFlagEnabled()) {
@@ -72,17 +87,10 @@
 
         // Restart the dream underneath the hub in order to support the ability to swipe
         // away the hub to enter the dream.
-        keyguardTransitionInteractor
-            .transition(
-                edge = Edge.create(to = Scenes.Communal),
-                edgeWithoutSceneContainer = Edge.create(to = KeyguardState.GLANCEABLE_HUB)
-            )
-            .filterState(TransitionState.FINISHED)
+        startDream
             .sampleFilter(powerInteractor.isAwake) { isAwake ->
-                dreamManager.canStartDreaming(isAwake)
+                !glanceableHubAllowKeyguardWhenDreaming() && dreamManager.canStartDreaming(isAwake)
             }
-            .sampleFilter(keyguardInteractor.isDreaming) { isDreaming -> !isDreaming }
-            .filter { !glanceableHubAllowKeyguardWhenDreaming() }
             .onEach { dreamManager.startDream() }
             .launchIn(bgScope)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
index ba2b7bf..a33e0ac 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.communal.dagger
 
 import android.content.Context
+import android.content.res.Resources
 import com.android.systemui.CoreStartable
 import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalDatabaseModule
@@ -38,6 +39,8 @@
 import com.android.systemui.communal.widgets.EditWidgetsActivityStarterImpl
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.res.R
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSource
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
@@ -90,6 +93,7 @@
 
     companion object {
         const val LOGGABLE_PREFIXES = "loggable_prefixes"
+        const val LAUNCHER_PACKAGE = "launcher_package"
 
         @Provides
         @Communal
@@ -126,5 +130,12 @@
                 .getStringArray(com.android.internal.R.array.config_loggable_dream_prefixes)
                 .toList()
         }
+
+        /** The package name of the launcher */
+        @Provides
+        @Named(LAUNCHER_PACKAGE)
+        fun provideLauncherPackage(@Main resources: Resources): String {
+            return resources.getString(R.string.launcher_overlayable_package)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
index 86241a5..f77dd58 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalSmartspaceRepository.kt
@@ -24,6 +24,9 @@
 import com.android.systemui.communal.smartspace.CommunalSmartspaceController
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.Logger
+import com.android.systemui.log.dagger.CommunalLog
 import com.android.systemui.plugins.BcSmartspaceDataPlugin
 import com.android.systemui.util.time.SystemClock
 import java.util.concurrent.Executor
@@ -49,8 +52,11 @@
     private val communalSmartspaceController: CommunalSmartspaceController,
     @Main private val uiExecutor: Executor,
     private val systemClock: SystemClock,
+    @CommunalLog logBuffer: LogBuffer,
 ) : CommunalSmartspaceRepository, BcSmartspaceDataPlugin.SmartspaceTargetListener {
 
+    private val logger = Logger(logBuffer, "CommunalSmartspaceRepository")
+
     private val _timers: MutableStateFlow<List<CommunalSmartspaceTimer>> =
         MutableStateFlow(emptyList())
     override val timers: Flow<List<CommunalSmartspaceTimer>> = _timers
@@ -87,6 +93,8 @@
                     remoteViews = target.remoteViews!!,
                 )
             }
+
+        logger.d({ "Smartspace timers updated: $str1" }) { str1 = _timers.value.toString() }
     }
 
     override fun startListening() {
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 98abbeb..b570e14 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
@@ -61,6 +61,7 @@
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.settings.UserTracker
+import com.android.systemui.statusbar.phone.ManagedProfileController
 import com.android.systemui.util.kotlin.BooleanFlowOperators.allOf
 import com.android.systemui.util.kotlin.BooleanFlowOperators.not
 import com.android.systemui.util.kotlin.emitOnStart
@@ -116,6 +117,7 @@
     sceneInteractor: SceneInteractor,
     @CommunalLog logBuffer: LogBuffer,
     @CommunalTableLog tableLogBuffer: TableLogBuffer,
+    private val managedProfileController: ManagedProfileController
 ) {
     private val logger = Logger(logBuffer, "CommunalInteractor")
 
@@ -401,12 +403,7 @@
 
     /** Request to unpause work profile that is currently in quiet mode. */
     fun unpauseWorkProfile() {
-        userTracker.userProfiles
-            .find { it.isManagedProfile }
-            ?.userHandle
-            ?.let { userHandle ->
-                userManager.requestQuietModeEnabled(/* enableQuietMode */ false, userHandle)
-            }
+        managedProfileController.setWorkModeEnabled(true)
     }
 
     /** Returns true if work profile is in quiet mode (disabled) for user handle. */
@@ -510,7 +507,7 @@
      * A flow of ongoing content, including smartspace timers and umo, ordered by creation time and
      * sized dynamically.
      */
-    fun getOngoingContent(mediaHostVisible: Boolean): Flow<List<CommunalContentModel.Ongoing>> =
+    val ongoingContent: Flow<List<CommunalContentModel.Ongoing>> =
         combine(smartspaceRepository.timers, mediaRepository.mediaModel) { timers, media ->
                 val ongoingContent = mutableListOf<CommunalContentModel.Ongoing>()
 
@@ -526,7 +523,7 @@
                 )
 
                 // Add UMO
-                if (mediaHostVisible && media.hasAnyMediaOrRecommendation) {
+                if (media.hasAnyMediaOrRecommendation) {
                     ongoingContent.add(
                         CommunalContentModel.Umo(
                             createdTimestampMillis = media.createdTimestampMillis,
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneInteractor.kt
index 8f756a2..ac496f0 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneInteractor.kt
@@ -195,7 +195,7 @@
                     is ObservableTransitionState.Idle ->
                         flowOf(CommunalTransitionProgressModel.Idle(state.currentScene))
                     is ObservableTransitionState.Transition ->
-                        if (state.toScene == targetScene) {
+                        if (state.toContent == targetScene) {
                             state.progress.map {
                                 CommunalTransitionProgressModel.Transition(
                                     // Clamp the progress values between 0 and 1 as actual progress
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
index 5bbb46d..c7538bb4 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalSceneTransitionInteractor.kt
@@ -161,7 +161,7 @@
         if (
             prevTransition is ObservableTransitionState.Transition &&
                 currentTransitionId != null &&
-                idle.currentScene == prevTransition.toScene
+                idle.currentScene == prevTransition.toContent
         ) {
             finishCurrentTransition()
         } else {
@@ -182,6 +182,7 @@
     }
 
     private suspend fun finishCurrentTransition() {
+        if (currentTransitionId == null) return
         internalTransitionInteractor.updateTransition(
             currentTransitionId!!,
             1f,
@@ -218,17 +219,19 @@
         prevTransition: ObservableTransitionState,
         transition: ObservableTransitionState.Transition
     ) {
-        if (prevTransition.isTransitioning(from = transition.fromScene, to = transition.toScene)) {
+        if (
+            prevTransition.isTransitioning(from = transition.fromContent, to = transition.toContent)
+        ) {
             // This is a new transition, but exactly the same as the previous state. Skip resetting
             // KTF for this case and just collect the new progress instead.
             collectProgress(transition)
-        } else if (transition.toScene == CommunalScenes.Communal) {
+        } else if (transition.toContent == CommunalScenes.Communal) {
             if (currentToState == KeyguardState.GLANCEABLE_HUB) {
-                transitionKtfTo(transitionInteractor.getStartedFromState())
+                transitionKtfTo(transitionInteractor.startedKeyguardTransitionStep.value.from)
             }
             startTransitionToGlanceableHub()
             collectProgress(transition)
-        } else if (transition.toScene == CommunalScenes.Blank) {
+        } else if (transition.toContent == CommunalScenes.Blank) {
             // Another transition started before this one is completed. Transition to the
             // GLANCEABLE_HUB state so that we can properly transition away from it.
             transitionKtfTo(KeyguardState.GLANCEABLE_HUB)
diff --git a/packages/SystemUI/src/com/android/systemui/communal/log/CommunalLoggerStartable.kt b/packages/SystemUI/src/com/android/systemui/communal/log/CommunalLoggerStartable.kt
index 2352841f..1def5a3 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/log/CommunalLoggerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/log/CommunalLoggerStartable.kt
@@ -126,13 +126,13 @@
 /** Whether currently transitioning from another scene to communal. */
 private fun ObservableTransitionState.isSwipingToCommunal(): Boolean {
     return this is ObservableTransitionState.Transition &&
-        toScene == CommunalScenes.Communal &&
+        toContent == CommunalScenes.Communal &&
         isInitiatedByUserInput
 }
 
 /** Whether currently transitioning from communal to another scene. */
 private fun ObservableTransitionState.isSwipingFromCommunal(): Boolean {
     return this is ObservableTransitionState.Transition &&
-        fromScene == CommunalScenes.Communal &&
+        fromContent == CommunalScenes.Communal &&
         isInitiatedByUserInput
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/shared/log/CommunalSceneLogger.kt b/packages/SystemUI/src/com/android/systemui/communal/shared/log/CommunalSceneLogger.kt
index aed9215..83f31e5 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/shared/log/CommunalSceneLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/shared/log/CommunalSceneLogger.kt
@@ -74,8 +74,8 @@
                     tag = TAG,
                     level = LogLevel.INFO,
                     messageInitializer = {
-                        str1 = transitionState.fromScene.toString()
-                        str2 = transitionState.toScene.toString()
+                        str1 = transitionState.fromContent.toString()
+                        str2 = transitionState.toContent.toString()
                     },
                     messagePrinter = { "Scene transition started: $str1 → $str2" },
                 )
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
index 16788d1..65f0679 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalEditModeViewModel.kt
@@ -29,6 +29,7 @@
 import android.view.accessibility.AccessibilityManager
 import androidx.activity.result.ActivityResultLauncher
 import com.android.internal.logging.UiEventLogger
+import com.android.systemui.communal.dagger.CommunalModule.Companion.LAUNCHER_PACKAGE
 import com.android.systemui.communal.data.model.CommunalWidgetCategories
 import com.android.systemui.communal.domain.interactor.CommunalInteractor
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
@@ -81,6 +82,7 @@
     @Application private val context: Context,
     private val accessibilityManager: AccessibilityManager,
     private val packageManager: PackageManager,
+    @Named(LAUNCHER_PACKAGE) private val launcherPackage: String,
 ) : BaseCommunalViewModel(communalSceneInteractor, communalInteractor, mediaHost) {
 
     private val logger = Logger(logBuffer, "CommunalEditModeViewModel")
@@ -185,7 +187,6 @@
     /** Launch the widget picker activity using the given {@link ActivityResultLauncher}. */
     suspend fun onOpenWidgetPicker(
         resources: Resources,
-        packageManager: PackageManager,
         activityLauncher: ActivityResultLauncher<Intent>
     ): Boolean =
         withContext(backgroundDispatcher) {
@@ -196,7 +197,7 @@
                 ) {
                     it.providerInfo
                 }
-            getWidgetPickerActivityIntent(resources, packageManager, excludeList)?.let {
+            getWidgetPickerActivityIntent(resources, excludeList)?.let {
                 try {
                     activityLauncher.launch(it)
                     return@withContext true
@@ -209,18 +210,10 @@
 
     private fun getWidgetPickerActivityIntent(
         resources: Resources,
-        packageManager: PackageManager,
         excludeList: ArrayList<AppWidgetProviderInfo>
     ): Intent? {
-        val packageName =
-            getLauncherPackageName(packageManager)
-                ?: run {
-                    Log.e(TAG, "Couldn't resolve launcher package name")
-                    return@getWidgetPickerActivityIntent null
-                }
-
         return Intent(Intent.ACTION_PICK).apply {
-            setPackage(packageName)
+            setPackage(launcherPackage)
             putExtra(
                 EXTRA_DESIRED_WIDGET_WIDTH,
                 resources.getDimensionPixelSize(R.dimen.communal_widget_picker_desired_width)
@@ -247,16 +240,6 @@
         }
     }
 
-    private fun getLauncherPackageName(packageManager: PackageManager): String? {
-        return packageManager
-            .resolveActivity(
-                Intent(Intent.ACTION_MAIN).also { it.addCategory(Intent.CATEGORY_HOME) },
-                PackageManager.MATCH_DEFAULT_ONLY
-            )
-            ?.activityInfo
-            ?.packageName
-    }
-
     /** Sets whether edit mode is currently open */
     fun setEditModeOpen(isOpen: Boolean) = communalInteractor.setEditModeOpen(isOpen)
 
diff --git a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
index 5a39a62..d69ba1b 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/ui/viewmodel/CommunalViewModel.kt
@@ -99,7 +99,7 @@
 
     private val logger = Logger(logBuffer, "CommunalViewModel")
 
-    private val _isMediaHostVisible =
+    private val isMediaHostVisible =
         conflatedCallbackFlow {
                 val callback = { visible: Boolean ->
                     trySend(visible)
@@ -117,12 +117,26 @@
                 mediaHost.updateViewVisibility()
                 emit(mediaHost.visible)
             }
+            .distinctUntilChanged()
             .onEach { logger.d({ "_isMediaHostVisible: $bool1" }) { bool1 = it } }
             .flowOn(mainDispatcher)
 
     /** Communal content saved from the previous emission when the flow is active (not "frozen"). */
     private var frozenCommunalContent: List<CommunalContentModel>? = null
 
+    private val ongoingContent =
+        combine(
+            isMediaHostVisible,
+            communalInteractor.ongoingContent.onEach { mediaHost.updateViewVisibility() }
+        ) { mediaVisible, ongoingContent ->
+            if (mediaVisible) {
+                ongoingContent
+            } else {
+                // Media is not visible, don't show UMO
+                ongoingContent.filterNot { it is CommunalContentModel.Umo }
+            }
+        }
+
     @OptIn(ExperimentalCoroutinesApi::class)
     private val latestCommunalContent: Flow<List<CommunalContentModel>> =
         tutorialInteractor.isTutorialAvailable
@@ -130,8 +144,6 @@
                 if (isTutorialMode) {
                     return@flatMapLatest flowOf(communalInteractor.tutorialContent)
                 }
-                val ongoingContent =
-                    _isMediaHostVisible.flatMapLatest { communalInteractor.getOngoingContent(it) }
                 combine(
                     ongoingContent,
                     communalInteractor.widgetContent,
@@ -254,6 +266,7 @@
             expandedMatchesParentHeight = true
             showsOnlyActiveMedia = false
             falsingProtectionNeeded = false
+            disablePagination = true
             init(MediaHierarchyManager.LOCATION_COMMUNAL_HUB)
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
index 55a24d0..d84dc20 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/EditWidgetsActivity.kt
@@ -270,11 +270,7 @@
 
     private fun onOpenWidgetPicker() {
         lifecycleScope.launch {
-            communalViewModel.onOpenWidgetPicker(
-                resources,
-                packageManager,
-                addWidgetActivityLauncher
-            )
+            communalViewModel.onOpenWidgetPicker(resources, addWidgetActivityLauncher)
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
index 1dd3722..3fe6669 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
@@ -21,6 +21,7 @@
 import com.android.systemui.Dependency;
 import com.android.systemui.InitController;
 import com.android.systemui.SystemUIAppComponentFactoryBase;
+import com.android.systemui.common.ui.GlobalConfig;
 import com.android.systemui.dagger.qualifiers.PerUser;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.KeyguardSliceProvider;
@@ -127,6 +128,7 @@
      * Creates a ContextComponentHelper.
      */
     @SysUISingleton
+    @GlobalConfig
     ConfigurationController getConfigurationController();
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index 411cbd5..b55108d 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -48,6 +48,7 @@
 import com.android.systemui.classifier.FalsingModule;
 import com.android.systemui.clipboardoverlay.dagger.ClipboardOverlayModule;
 import com.android.systemui.common.data.CommonDataLayerModule;
+import com.android.systemui.common.ui.ConfigurationStateModule;
 import com.android.systemui.common.usagestats.data.CommonUsageStatsDataLayerModule;
 import com.android.systemui.communal.dagger.CommunalModule;
 import com.android.systemui.complication.dagger.ComplicationComponent;
@@ -207,6 +208,7 @@
         ClockRegistryModule.class,
         CommunalModule.class,
         CommonDataLayerModule.class,
+        ConfigurationStateModule.class,
         CommonUsageStatsDataLayerModule.class,
         ConfigurationControllerModule.class,
         ConnectivityModule.class,
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 d288cce..37c6e17 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
@@ -391,8 +391,10 @@
             ),
             Pair(
                 if (SceneContainerFlag.isEnabled) {
-                    keyguardTransitionInteractor
-                        .isInTransitionWhere(toStatePredicate = { it == KeyguardState.UNDEFINED })
+                    sceneInteractor
+                        .get()
+                        .transitionState
+                        .map { it.isTransitioning(to = Scenes.Gone) || it.isIdle(Scenes.Gone) }
                         .isFalse()
                 } else {
                     keyguardRepository.isKeyguardGoingAway.isFalse()
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractor.kt
index dff391a..cdd2b05 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractor.kt
@@ -60,14 +60,23 @@
     fun unregisterListener(listener: FaceAuthenticationListener)
 
     fun onUdfpsSensorTouched()
+
     fun onAssistantTriggeredOnLockScreen()
+
     fun onDeviceLifted()
-    fun onQsExpansionStared()
+
+    fun onShadeExpansionStarted()
+
     fun onNotificationPanelClicked()
+
     fun onSwipeUpOnBouncer()
+
     fun onPrimaryBouncerUserInput()
+
     fun onAccessibilityAction()
+
     fun onWalletLaunched()
+
     fun onDeviceUnfolded()
 
     /** Whether face auth is considered class 3 */
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
index 39f4e31..7018f9d 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractor.kt
@@ -16,12 +16,14 @@
 
 package com.android.systemui.deviceentry.domain.interactor
 
+import com.android.internal.policy.IKeyguardDismissCallback
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
 import com.android.systemui.bouncer.domain.interactor.AlternateBouncerInteractor
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.deviceentry.data.repository.DeviceEntryRepository
+import com.android.systemui.keyguard.DismissCallbackRegistry
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.utils.coroutines.flow.mapLatestConflated
@@ -56,6 +58,7 @@
     private val sceneInteractor: SceneInteractor,
     private val deviceUnlockedInteractor: DeviceUnlockedInteractor,
     private val alternateBouncerInteractor: AlternateBouncerInteractor,
+    private val dismissCallbackRegistry: DismissCallbackRegistry,
 ) {
     /**
      * Whether the device is unlocked.
@@ -126,17 +129,14 @@
                 },
                 isLockscreenEnabled,
                 deviceUnlockedInteractor.deviceUnlockStatus,
-                isDeviceEntered) {
-                    isNoneAuthMethod,
-                    isLockscreenEnabled,
-                    deviceUnlockStatus,
-                    isDeviceEntered ->
-                    val isSwipeAuthMethod = isNoneAuthMethod && isLockscreenEnabled
-                    (isSwipeAuthMethod ||
-                        (deviceUnlockStatus.isUnlocked &&
-                            deviceUnlockStatus.deviceUnlockSource?.dismissesLockscreen == false)) &&
-                        !isDeviceEntered
-                }
+                isDeviceEntered
+            ) { isNoneAuthMethod, isLockscreenEnabled, deviceUnlockStatus, isDeviceEntered ->
+                val isSwipeAuthMethod = isNoneAuthMethod && isLockscreenEnabled
+                (isSwipeAuthMethod ||
+                    (deviceUnlockStatus.isUnlocked &&
+                        deviceUnlockStatus.deviceUnlockSource?.dismissesLockscreen == false)) &&
+                    !isDeviceEntered
+            }
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.Eagerly,
@@ -150,8 +150,16 @@
     /**
      * Attempt to enter the device and dismiss the lockscreen. If authentication is required to
      * unlock the device it will transition to bouncer.
+     *
+     * @param callback An optional callback to invoke when the attempt succeeds, fails, or is
+     *   canceled
      */
-    fun attemptDeviceEntry() {
+    @JvmOverloads
+    fun attemptDeviceEntry(
+        callback: IKeyguardDismissCallback? = null,
+    ) {
+        callback?.let { dismissCallbackRegistry.addCallback(it) }
+
         // TODO (b/307768356),
         //       1. Check if the device is already authenticated by trust agent/passive biometrics
         //       2. Show SPFS/UDFPS bouncer if it is available AlternateBouncerInteractor.show
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/NoopDeviceEntryFaceAuthInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/NoopDeviceEntryFaceAuthInteractor.kt
index de5d0aa..9b8c2b1 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/NoopDeviceEntryFaceAuthInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/NoopDeviceEntryFaceAuthInteractor.kt
@@ -47,6 +47,7 @@
     override fun isFaceAuthEnabledAndEnrolled(): Boolean = false
 
     override fun isFaceAuthStrong(): Boolean = false
+
     override fun start() = Unit
 
     override fun registerListener(listener: FaceAuthenticationListener) {}
@@ -59,13 +60,17 @@
 
     override fun onDeviceLifted() {}
 
-    override fun onQsExpansionStared() {}
+    override fun onShadeExpansionStarted() {}
 
     override fun onNotificationPanelClicked() {}
 
     override fun onSwipeUpOnBouncer() {}
+
     override fun onPrimaryBouncerUserInput() {}
+
     override fun onAccessibilityAction() {}
+
     override fun onWalletLaunched() = Unit
+
     override fun onDeviceUnfolded() {}
 }
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/SystemUIDeviceEntryFaceAuthInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/SystemUIDeviceEntryFaceAuthInteractor.kt
index 183e0e9..3b5d5a8 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/SystemUIDeviceEntryFaceAuthInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/SystemUIDeviceEntryFaceAuthInteractor.kt
@@ -60,6 +60,7 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
@@ -214,6 +215,16 @@
                 }
             }
             .launchIn(applicationScope)
+
+        if (SceneContainerFlag.isEnabled) {
+            sceneInteractor
+                .get()
+                .transitionState
+                .filter { it.isTransitioning(from = Scenes.Lockscreen, to = Scenes.Shade) }
+                .distinctUntilChanged()
+                .onEach { onShadeExpansionStarted() }
+                .launchIn(applicationScope)
+        }
     }
 
     private val isBouncerVisible: Flow<Boolean> by lazy {
@@ -239,8 +250,8 @@
         runFaceAuth(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_NOTIFICATION_PANEL_CLICKED, true)
     }
 
-    override fun onQsExpansionStared() {
-        runFaceAuth(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, true)
+    override fun onShadeExpansionStarted() {
+        runFaceAuth(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, false)
     }
 
     override fun onDeviceLifted() {
diff --git a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
index 40e2f17..1f5878b 100644
--- a/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/display/data/repository/DisplayRepository.kt
@@ -48,6 +48,7 @@
 import kotlinx.coroutines.flow.filterIsInstance
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.scan
 import kotlinx.coroutines.flow.shareIn
@@ -184,6 +185,9 @@
         if (Flags.enableEfficientDisplayRepository()) {
             enabledDisplayIds
                 .mapElementsLazily { displayId -> getDisplay(displayId) }
+                .onEach {
+                    if (it.isEmpty()) Log.wtf(TAG, "No enabled displays. This should never happen.")
+                }
                 .flowOn(backgroundCoroutineDispatcher)
                 .debugLog("enabledDisplays")
                 .stateIn(
@@ -194,7 +198,8 @@
                     // performance concerns.
                     // Ultimately, this is a trade-off between a one-time UI thread binder call and
                     // the constant overhead of sharedFlows.
-                    initialValue = getDisplays())
+                    initialValue = getDisplays()
+                )
         } else {
             oldEnabledDisplays
         }
@@ -380,9 +385,8 @@
             val resultSet: Set<V>
         )
 
-        val initialState = State(emptySet<T>(), emptyMap(), emptySet<V>())
-
-        return this.scan(initialState) { state, currentSet ->
+        val emptyInitialState = State(emptySet<T>(), emptyMap(), emptySet<V>())
+        return this.scan(emptyInitialState) { state, currentSet ->
                 if (currentSet == state.previousSet) {
                     state
                 } else {
@@ -397,6 +401,7 @@
                     State(currentSet, newMap, resultSet)
                 }
             }
+            .filter { it != emptyInitialState }
             .map { it.resultSet }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
index caf5b01..1c263ae 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/DreamOverlayService.java
@@ -24,12 +24,11 @@
 import static com.android.systemui.dreams.dagger.DreamModule.HOME_CONTROL_PANEL_DREAM_COMPONENT;
 import static com.android.systemui.util.kotlin.JavaAdapterKt.collectFlow;
 
+import android.app.WindowConfiguration;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
-import android.content.pm.ActivityInfo;
 import android.graphics.drawable.ColorDrawable;
-import android.service.dreams.DreamActivity;
 import android.util.Log;
 import android.view.View;
 import android.view.ViewGroup;
@@ -65,6 +64,7 @@
 import com.android.systemui.dreams.dagger.DreamOverlayComponent;
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor;
 import com.android.systemui.navigationbar.gestural.domain.GestureInteractor;
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher;
 import com.android.systemui.shade.ShadeExpansionChangeEvent;
 import com.android.systemui.touch.TouchInsetManager;
 import com.android.systemui.util.concurrency.DelayableExecutor;
@@ -89,6 +89,8 @@
         LifecycleOwner {
     private static final String TAG = "DreamOverlayService";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
+    private static final TaskMatcher DREAM_TYPE_MATCHER =
+            new TaskMatcher.TopActivityType(WindowConfiguration.ACTIVITY_TYPE_DREAM);
 
     // The Context is used to construct the hosting constraint layout and child overlay views.
     private final Context mContext;
@@ -141,10 +143,6 @@
     private final TouchInsetManager mTouchInsetManager;
     private final LifecycleOwner mLifecycleOwner;
 
-
-
-    private ComponentName mCurrentBlockedGestureDreamActivityComponent;
-
     private final ArrayList<Job> mFlows = new ArrayList<>();
 
     /**
@@ -187,6 +185,7 @@
                         mShadeExpanded = expanded;
 
                         updateLifecycleStateLocked();
+                        updateGestureBlockingLocked();
                     });
                 }
             };
@@ -217,20 +216,127 @@
                 mBouncerShowing = bouncerShowing;
 
                 updateLifecycleStateLocked();
+                updateGestureBlockingLocked();
             });
         }
     };
 
-    private final DreamOverlayStateController.Callback mExitAnimationFinishedCallback =
-            new DreamOverlayStateController.Callback() {
-                @Override
-                public void onStateChanged() {
-                    if (!mStateController.areExitAnimationsRunning()) {
-                        mStateController.removeCallback(mExitAnimationFinishedCallback);
-                        resetCurrentDreamOverlayLocked();
+    /**
+     * {@link ResetHandler} protects resetting {@link DreamOverlayService} by making sure reset
+     * requests are processed before subsequent actions proceed. Requests themselves are also
+     * ordered between each other as well to ensure actions are correctly sequenced.
+     */
+    private final class ResetHandler {
+        @FunctionalInterface
+        interface Callback {
+            void onComplete();
+        }
+
+        private record Info(Callback callback, String source) {}
+
+        private final ArrayList<Info> mPendingCallbacks = new ArrayList<>();
+
+        DreamOverlayStateController.Callback mStateCallback =
+                new DreamOverlayStateController.Callback() {
+                    @Override
+                    public void onStateChanged() {
+                        process(true);
                     }
+                };
+
+        /**
+         * Called from places where there is no need to wait for the reset to complete. This still
+         * will defer the reset until it is okay to reset and also sequences the request with
+         * others.
+         */
+        public void reset(String source) {
+            reset(() -> {}, source);
+        }
+
+        /**
+         * Invoked to request a reset with a callback that will fire after reset if it is deferred.
+         *
+         * @return {@code true} if the reset happened immediately, {@code false} if it was deferred
+         * and will fire later, invoking the callback.
+         */
+        public boolean reset(Callback callback, String source) {
+            // Always add listener pre-emptively
+            if (mPendingCallbacks.isEmpty()) {
+                mStateController.addCallback(mStateCallback);
+            }
+
+            final Info info = new Info(callback, source);
+            mPendingCallbacks.add(info);
+            process(false);
+
+            boolean processed = !mPendingCallbacks.contains(info);
+
+            if (!processed) {
+                Log.d(TAG, "delayed resetting from: " + source);
+            }
+
+            return processed;
+        }
+
+        private void resetInternal() {
+            // This ensures the container view of the current dream is removed before
+            // the controller is potentially reset.
+            removeContainerViewFromParentLocked();
+
+            if (mStarted && mWindow != null) {
+                try {
+                    mWindow.clearContentView();
+                    mWindowManager.removeView(mWindow.getDecorView());
+                } catch (IllegalArgumentException e) {
+                    Log.e(TAG, "Error removing decor view when resetting overlay", e);
                 }
-            };
+            }
+
+            mStateController.setOverlayActive(false);
+            mStateController.setLowLightActive(false);
+            mStateController.setEntryAnimationsFinished(false);
+
+            if (mDreamOverlayContainerViewController != null) {
+                mDreamOverlayContainerViewController.destroy();
+                mDreamOverlayContainerViewController = null;
+            }
+
+            if (mTouchMonitor != null) {
+                mTouchMonitor.destroy();
+                mTouchMonitor = null;
+            }
+
+            mWindow = null;
+
+            // Always unregister the any set DreamActivity from being blocked from gestures.
+            mGestureInteractor.removeGestureBlockedMatcher(DREAM_TYPE_MATCHER,
+                    GestureInteractor.Scope.Global);
+
+            mStarted = false;
+        }
+
+        private boolean canReset() {
+            return !mStateController.areExitAnimationsRunning();
+        }
+
+        private void process(boolean fromDelayedCallback) {
+            while (canReset() && !mPendingCallbacks.isEmpty()) {
+                final Info callbackInfo = mPendingCallbacks.removeFirst();
+                resetInternal();
+                callbackInfo.callback.onComplete();
+
+                if (fromDelayedCallback) {
+                    Log.d(TAG, "reset overlay (delayed) for " + callbackInfo.source);
+                }
+            }
+
+            if (mPendingCallbacks.isEmpty()) {
+                mStateController.removeCallback(mStateCallback);
+            }
+        }
+    }
+
+    private final ResetHandler mResetHandler = new ResetHandler();
 
     private final DreamOverlayStateController mStateController;
 
@@ -344,10 +450,8 @@
 
         mExecutor.execute(() -> {
             setLifecycleStateLocked(Lifecycle.State.DESTROYED);
-
-            resetCurrentDreamOverlayLocked();
-
             mDestroyed = true;
+            mResetHandler.reset("destroying");
         });
 
         mDispatcher.onServicePreSuperOnDestroy();
@@ -387,7 +491,10 @@
             // Reset the current dream overlay before starting a new one. This can happen
             // when two dreams overlap (briefly, for a smoother dream transition) and both
             // dreams are bound to the dream overlay service.
-            resetCurrentDreamOverlayLocked();
+            if (!mResetHandler.reset(() -> onStartDream(layoutParams),
+                    "starting with dream already started")) {
+                return;
+            }
         }
 
         mDreamOverlayContainerViewController =
@@ -399,7 +506,7 @@
 
         // If we are not able to add the overlay window, reset the overlay.
         if (!addOverlayWindowLocked(layoutParams)) {
-            resetCurrentDreamOverlayLocked();
+            mResetHandler.reset("couldn't add window while starting");
             return;
         }
 
@@ -420,7 +527,7 @@
         mStarted = true;
 
         updateRedirectWakeup();
-        updateBlockedGestureDreamActivityComponent();
+        updateGestureBlockingLocked();
     }
 
     private void updateRedirectWakeup() {
@@ -431,21 +538,9 @@
         redirectWake(mCommunalAvailable && !glanceableHubAllowKeyguardWhenDreaming());
     }
 
-    private void updateBlockedGestureDreamActivityComponent() {
-        // TODO(b/343815446): We should not be crafting this ActivityInfo ourselves. It should be
-        // in a common place, Such as DreamActivity itself.
-        final ActivityInfo info = new ActivityInfo();
-        info.name = DreamActivity.class.getName();
-        info.packageName = getDreamComponent().getPackageName();
-        mCurrentBlockedGestureDreamActivityComponent = info.getComponentName();
-
-        mGestureInteractor.addGestureBlockedActivity(mCurrentBlockedGestureDreamActivityComponent,
-                GestureInteractor.Scope.Global);
-    }
-
     @Override
     public void onEndDream() {
-        resetCurrentDreamOverlayLocked();
+        mResetHandler.reset("ending dream");
     }
 
     @Override
@@ -456,6 +551,18 @@
                 null);
     }
 
+    private void updateGestureBlockingLocked() {
+        final boolean shouldBlock = !isDreamInPreviewMode() && !mShadeExpanded && !mBouncerShowing;
+
+        if (shouldBlock) {
+            mGestureInteractor.addGestureBlockedMatcher(DREAM_TYPE_MATCHER,
+                    GestureInteractor.Scope.Global);
+        } else {
+            mGestureInteractor.removeGestureBlockedMatcher(DREAM_TYPE_MATCHER,
+                    GestureInteractor.Scope.Global);
+        }
+    }
+
     private Lifecycle.State getLifecycleStateLocked() {
         return mLifecycleRegistry.getCurrentState();
     }
@@ -576,49 +683,4 @@
         Log.w(TAG, "Removing dream overlay container view parent!");
         parentView.removeView(containerView);
     }
-
-    private void resetCurrentDreamOverlayLocked() {
-        if (mStateController.areExitAnimationsRunning()) {
-            mStateController.addCallback(mExitAnimationFinishedCallback);
-            return;
-        }
-
-        // This ensures the container view of the current dream is removed before
-        // the controller is potentially reset.
-        removeContainerViewFromParentLocked();
-
-        if (mStarted && mWindow != null) {
-            try {
-                mWindow.clearContentView();
-                mWindowManager.removeView(mWindow.getDecorView());
-            } catch (IllegalArgumentException e) {
-                Log.e(TAG, "Error removing decor view when resetting overlay", e);
-            }
-        }
-
-        mStateController.setOverlayActive(false);
-        mStateController.setLowLightActive(false);
-        mStateController.setEntryAnimationsFinished(false);
-
-        if (mDreamOverlayContainerViewController != null) {
-            mDreamOverlayContainerViewController.destroy();
-            mDreamOverlayContainerViewController = null;
-        }
-
-        if (mTouchMonitor != null) {
-            mTouchMonitor.destroy();
-            mTouchMonitor = null;
-        }
-
-        mWindow = null;
-
-        // Always unregister the any set DreamActivity from being blocked from gestures.
-        if (mCurrentBlockedGestureDreamActivityComponent != null) {
-            mGestureInteractor.removeGestureBlockedActivity(
-                    mCurrentBlockedGestureDreamActivityComponent, GestureInteractor.Scope.Global);
-            mCurrentBlockedGestureDreamActivityComponent = null;
-        }
-
-        mStarted = false;
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
index e88349b2..3b2d771 100644
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
@@ -16,7 +16,13 @@
 
 package com.android.systemui.education.domain.interactor
 
+import android.hardware.input.InputManager
+import android.hardware.input.InputManager.KeyGestureEventListener
+import android.hardware.input.KeyGestureEvent
 import com.android.systemui.CoreStartable
+import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
+import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.contextualeducation.GestureType.BACK
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
@@ -25,10 +31,14 @@
 import com.android.systemui.education.shared.model.EducationInfo
 import com.android.systemui.education.shared.model.EducationUiType
 import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import java.time.Clock
+import java.util.concurrent.Executor
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.hours
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
@@ -41,10 +51,12 @@
     @Background private val backgroundScope: CoroutineScope,
     private val contextualEducationInteractor: ContextualEducationInteractor,
     private val userInputDeviceRepository: UserInputDeviceRepository,
+    private val inputManager: InputManager,
     @EduClock private val clock: Clock,
 ) : CoreStartable {
 
     companion object {
+        const val TAG = "KeyboardTouchpadEduInteractor"
         const val MAX_SIGNAL_COUNT: Int = 2
         val usageSessionDuration = 72.hours
     }
@@ -52,7 +64,26 @@
     private val _educationTriggered = MutableStateFlow<EducationInfo?>(null)
     val educationTriggered = _educationTriggered.asStateFlow()
 
+    private val keyboardShortcutTriggered: Flow<GestureType> = conflatedCallbackFlow {
+        val listener = KeyGestureEventListener { event ->
+            // Only store keyboard shortcut time for gestures providing keyboard education
+            val shortcutType =
+                when (event.keyGestureType) {
+                    KeyGestureEvent.KEY_GESTURE_TYPE_ALL_APPS -> ALL_APPS
+                    else -> null
+                }
+
+            if (shortcutType != null) {
+                trySendWithFailureLogging(shortcutType, TAG)
+            }
+        }
+
+        inputManager.registerKeyGestureEventListener(Executor(Runnable::run), listener)
+        awaitClose { inputManager.unregisterKeyGestureEventListener(listener) }
+    }
+
     override fun start() {
+        // Listen to back gesture model changes and trigger education if needed
         backgroundScope.launch {
             contextualEducationInteractor.backGestureModelFlow.collect {
                 if (isUsageSessionExpired(it)) {
@@ -64,6 +95,7 @@
             }
         }
 
+        // Listen to touchpad connection changes and update the first connection time
         backgroundScope.launch {
             userInputDeviceRepository.isAnyTouchpadConnectedForUser.collect {
                 if (
@@ -77,6 +109,7 @@
             }
         }
 
+        // Listen to keyboard connection changes and update the first connection time
         backgroundScope.launch {
             userInputDeviceRepository.isAnyKeyboardConnectedForUser.collect {
                 if (
@@ -89,6 +122,13 @@
                 }
             }
         }
+
+        // Listen to keyboard shortcut triggered and update the last trigger time
+        backgroundScope.launch {
+            keyboardShortcutTriggered.collect {
+                contextualEducationInteractor.updateShortcutTriggerTime(it)
+            }
+        }
     }
 
     private fun isEducationNeeded(model: GestureEduModel): Boolean {
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
index 3223433..7821f69 100644
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
@@ -16,11 +16,25 @@
 
 package com.android.systemui.education.domain.interactor
 
+import android.os.SystemProperties
+import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.education.dagger.ContextualEducationModule.EduClock
+import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.KEYBOARD
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.TOUCHPAD
+import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
+import java.time.Clock
 import javax.inject.Inject
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.hours
+import kotlin.time.DurationUnit
+import kotlin.time.toDuration
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 
 /**
@@ -39,12 +53,29 @@
 @Inject
 constructor(
     @Background private val backgroundScope: CoroutineScope,
-    private val contextualEducationInteractor: ContextualEducationInteractor
+    private val contextualEducationInteractor: ContextualEducationInteractor,
+    private val inputDeviceRepository: UserInputDeviceRepository,
+    private val tutorialRepository: TutorialSchedulerRepository,
+    @EduClock private val clock: Clock,
 ) : KeyboardTouchpadEduStatsInteractor {
 
+    companion object {
+        val initialDelayDuration: Duration
+            get() =
+                SystemProperties.getLong(
+                        "persist.contextual_edu.initial_delay_sec",
+                        /* defaultValue= */ 72.hours.inWholeSeconds
+                    )
+                    .toDuration(DurationUnit.SECONDS)
+    }
+
     override fun incrementSignalCount(gestureType: GestureType) {
-        // Todo: check if keyboard/touchpad is connected before update
-        backgroundScope.launch { contextualEducationInteractor.incrementSignalCount(gestureType) }
+        backgroundScope.launch {
+            val targetDevice = getTargetDevice(gestureType)
+            if (isTargetDeviceConnected(targetDevice) && hasInitialDelayElapsed(targetDevice)) {
+                contextualEducationInteractor.incrementSignalCount(gestureType)
+            }
+        }
     }
 
     override fun updateShortcutTriggerTime(gestureType: GestureType) {
@@ -52,4 +83,31 @@
             contextualEducationInteractor.updateShortcutTriggerTime(gestureType)
         }
     }
+
+    private suspend fun isTargetDeviceConnected(deviceType: DeviceType): Boolean {
+        if (deviceType == KEYBOARD) {
+            return inputDeviceRepository.isAnyKeyboardConnectedForUser.first().isConnected
+        } else if (deviceType == TOUCHPAD) {
+            return inputDeviceRepository.isAnyTouchpadConnectedForUser.first().isConnected
+        }
+        return false
+    }
+
+    /**
+     * Keyboard shortcut education would be provided for All Apps. Touchpad gesture education would
+     * be provided for the rest of the gesture types (i.e. Home, Overview, Back). This method maps
+     * gesture to its target education device.
+     */
+    private fun getTargetDevice(gestureType: GestureType) =
+        when (gestureType) {
+            ALL_APPS -> KEYBOARD
+            else -> TOUCHPAD
+        }
+
+    private suspend fun hasInitialDelayElapsed(deviceType: DeviceType): Boolean {
+        val oobeLaunchTime = tutorialRepository.launchTime(deviceType) ?: return false
+        return clock
+            .instant()
+            .isAfter(oobeLaunchTime.plusSeconds(initialDelayDuration.inWholeSeconds))
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt b/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
index cd0b3f9..6318dc0 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/FlagDependencies.kt
@@ -41,7 +41,6 @@
 import com.android.systemui.statusbar.notification.shared.NotificationAvalancheSuppression
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalismPrototype
 import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.shared.PriorityPeopleSection
 import javax.inject.Inject
 
@@ -59,7 +58,6 @@
         NotificationAvalancheSuppression.token dependsOn VisualInterruptionRefactor.token
         PriorityPeopleSection.token dependsOn SortBySectionTimeFlag.token
         NotificationMinimalismPrototype.token dependsOn NotificationThrottleHun.token
-        NotificationsHeadsUpRefactor.token dependsOn NotificationThrottleHun.token
 
         // SceneContainer dependencies
         SceneContainerFlag.getFlagDependencies().forEach { (alpha, beta) -> alpha dependsOn beta }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
index b654307..a20dfa5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/data/repository/KeyboardRepository.kt
@@ -25,7 +25,6 @@
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.inputdevice.data.repository.InputDeviceRepository
 import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceAdded
-import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceChange
 import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.DeviceRemoved
 import com.android.systemui.inputdevice.data.repository.InputDeviceRepository.FreshStart
 import com.android.systemui.keyboard.data.model.Keyboard
@@ -78,24 +77,16 @@
     inputDeviceRepository: InputDeviceRepository
 ) : KeyboardRepository {
 
-    private val keyboardsChange: Flow<Pair<Collection<Int>, DeviceChange>> =
-        inputDeviceRepository.deviceChange
-            .map { (ids, change) -> ids.filter { id -> isPhysicalFullKeyboard(id) } to change }
-            .filter { (_, change) ->
-                when (change) {
-                    FreshStart -> true
-                    is DeviceAdded -> isPhysicalFullKeyboard(change.deviceId)
-                    is DeviceRemoved -> isPhysicalFullKeyboard(change.deviceId)
-                }
-            }
-
     @FlowPreview
     override val newlyConnectedKeyboard: Flow<Keyboard> =
-        keyboardsChange
+        inputDeviceRepository.deviceChange
             .flatMapConcat { (devices, operation) ->
                 when (operation) {
-                    FreshStart -> devices.asFlow()
-                    is DeviceAdded -> flowOf(operation.deviceId)
+                    FreshStart -> devices.filter { id -> isPhysicalFullKeyboard(id) }.asFlow()
+                    is DeviceAdded -> {
+                        if (isPhysicalFullKeyboard(operation.deviceId)) flowOf(operation.deviceId)
+                        else emptyFlow()
+                    }
                     is DeviceRemoved -> emptyFlow()
                 }
             }
@@ -103,8 +94,8 @@
             .flowOn(backgroundDispatcher)
 
     override val isAnyKeyboardConnected: Flow<Boolean> =
-        keyboardsChange
-            .map { (devices, _) -> devices.isNotEmpty() }
+        inputDeviceRepository.deviceChange
+            .map { (ids, _) -> ids.any { id -> isPhysicalFullKeyboard(id) } }
             .distinctUntilChanged()
             .flowOn(backgroundDispatcher)
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index 0feb5ec..1bc91ca 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -76,6 +76,7 @@
 import com.android.systemui.SystemUIApplication;
 import com.android.systemui.dagger.qualifiers.Application;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardEnabledInteractor;
@@ -122,6 +123,7 @@
     private final PowerInteractor mPowerInteractor;
     private final KeyguardInteractor mKeyguardInteractor;
     private final Lazy<SceneInteractor> mSceneInteractorLazy;
+    private final Lazy<DeviceEntryInteractor> mDeviceEntryInteractorLazy;
     private final Executor mMainExecutor;
     private final Lazy<KeyguardStateCallbackStartable> mKeyguardStateCallbackStartableLazy;
 
@@ -347,7 +349,8 @@
             KeyguardEnabledInteractor keyguardEnabledInteractor,
             Lazy<KeyguardStateCallbackStartable> keyguardStateCallbackStartableLazy,
             KeyguardWakeDirectlyToGoneInteractor keyguardWakeDirectlyToGoneInteractor,
-            KeyguardDismissInteractor keyguardDismissInteractor) {
+            KeyguardDismissInteractor keyguardDismissInteractor,
+            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy) {
         super();
         mKeyguardViewMediator = keyguardViewMediator;
         mKeyguardLifecyclesDispatcher = keyguardLifecyclesDispatcher;
@@ -360,6 +363,7 @@
         mSceneInteractorLazy = sceneInteractorLazy;
         mMainExecutor = mainExecutor;
         mKeyguardStateCallbackStartableLazy = keyguardStateCallbackStartableLazy;
+        mDeviceEntryInteractorLazy = deviceEntryInteractorLazy;
 
         if (KeyguardWmStateRefactor.isEnabled()) {
             WindowManagerLockscreenVisibilityViewBinder.bind(
@@ -484,7 +488,9 @@
         public void dismiss(IKeyguardDismissCallback callback, CharSequence message) {
             trace("dismiss message=" + message);
             checkPermission();
-            if (KeyguardWmStateRefactor.isEnabled()) {
+            if (SceneContainerFlag.isEnabled()) {
+                mDeviceEntryInteractorLazy.get().attemptDeviceEntry(callback);
+            } else if (KeyguardWmStateRefactor.isEnabled()) {
                 mKeyguardDismissInteractor.dismissKeyguardWithCallback(callback);
             } else {
                 mKeyguardViewMediator.dismiss(callback, message);
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 8c82900..d38c952 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -3568,12 +3568,16 @@
                     }
                     return;
                 }
-                try {
-                    mStatusBarService.disableForUser(flags, mStatusBarDisableToken,
-                            mContext.getPackageName(),
-                            mSelectedUserInteractor.getSelectedUserId());
-                } catch (RemoteException e) {
-                    Log.d(TAG, "Failed to set disable flags: " + flags, e);
+
+                // Handled in StatusBarDisableFlagsInteractor.
+                if (!KeyguardWmStateRefactor.isEnabled()) {
+                    try {
+                        mStatusBarService.disableForUser(flags, mStatusBarDisableToken,
+                                mContext.getPackageName(),
+                                mSelectedUserInteractor.getSelectedUserId());
+                    } catch (RemoteException e) {
+                        Log.d(TAG, "Failed to set disable flags: " + flags, e);
+                    }
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt b/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
index 180afb2..e89594e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/WindowManagerLockscreenVisibilityManager.kt
@@ -23,7 +23,7 @@
 import android.view.WindowManager
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
-import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor
 import com.android.systemui.keyguard.ui.binder.KeyguardSurfaceBehindParamsApplier
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import java.util.concurrent.Executor
@@ -41,7 +41,7 @@
     private val activityTaskManagerService: IActivityTaskManager,
     private val keyguardStateController: KeyguardStateController,
     private val keyguardSurfaceBehindAnimator: KeyguardSurfaceBehindParamsApplier,
-    private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
+    private val keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor,
 ) {
 
     /**
@@ -148,7 +148,7 @@
         // a transition to GONE. This transition needs to start even if we're not provided an app
         // animation target - it's possible the app is destroyed on creation, etc. but we'll still
         // be unlocking.
-        keyguardTransitionInteractor.startDismissKeyguardTransition(
+        keyguardDismissTransitionInteractor.startDismissKeyguardTransition(
             reason = "Going away remote animation started"
         )
 
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 7087752..ec52055 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
@@ -101,7 +101,7 @@
         secureSettings
             .observerFlow(
                 names = arrayOf(Settings.Secure.LOCKSCREEN_USE_DOUBLE_LINE_CLOCK),
-                userId = UserHandle.USER_SYSTEM,
+                userId = UserHandle.USER_ALL,
             )
             .onStart { emit(Unit) } // Forces an initial update.
             .map { withContext(backgroundDispatcher) { getClockSize() } }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepository.kt
index a1e4af5..b67fd4b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepository.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.keyguard.data.repository
 
 import android.content.Context
+import android.os.UserHandle
 import android.provider.Settings
 import android.view.View
 import com.android.systemui.dagger.SysUISingleton
@@ -37,6 +38,7 @@
 interface KeyguardSmartspaceRepository {
     val bcSmartspaceVisibility: StateFlow<Int>
     val isWeatherEnabled: StateFlow<Boolean>
+
     fun setBcSmartspaceVisibility(visibility: Int)
 }
 
@@ -55,7 +57,7 @@
         secureSettings
             .observerFlow(
                 names = arrayOf(Settings.Secure.LOCK_SCREEN_WEATHER_ENABLED),
-                userId = userTracker.userId,
+                userId = UserHandle.USER_ALL,
             )
             .onStart { emit(Unit) }
             .map { getLockscreenWeatherEnabled() }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
index 1042ae3..e4b0f6e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAlternateBouncerTransitionInteractor.kt
@@ -35,7 +35,6 @@
 import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.distinctUntilChanged
@@ -49,7 +48,6 @@
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.launch
 
-@ExperimentalCoroutinesApi
 @SysUISingleton
 class FromAlternateBouncerTransitionInteractor
 @Inject
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
index 6e04133..4cf9ec8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromAodTransitionInteractor.kt
@@ -89,7 +89,7 @@
                 .filterRelevantKeyguardStateAnd { wakefulness -> wakefulness.isAwake() }
                 .debounce(50L)
                 .sample(
-                    startedKeyguardTransitionStep,
+                    transitionInteractor.startedKeyguardTransitionStep,
                     wakeToGoneInteractor.canWakeDirectlyToGone,
                 )
                 .collect {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
index 4666430..2434b29 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromDreamingTransitionInteractor.kt
@@ -117,7 +117,7 @@
         if (SceneContainerFlag.isEnabled) return
         scope.launch {
             keyguardInteractor.primaryBouncerShowing
-                .sample(startedKeyguardTransitionStep, ::Pair)
+                .sample(transitionInteractor.startedKeyguardTransitionStep, ::Pair)
                 .collect { pair ->
                     val (isBouncerShowing, lastStartedTransitionStep) = pair
                     if (
@@ -132,7 +132,7 @@
     fun startToLockscreenOrGlanceableHubTransition(openHub: Boolean) {
         scope.launch {
             if (
-                transitionInteractor.startedKeyguardState.replayCache.last() ==
+                transitionInteractor.startedKeyguardTransitionStep.value.to ==
                     KeyguardState.DREAMING
             ) {
                 if (powerInteractor.detailedWakefulness.value.isAwake()) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
index cd3df07..228e01e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/FromLockscreenTransitionInteractor.kt
@@ -134,16 +134,12 @@
                 .filterRelevantKeyguardState()
                 .sampleCombine(
                     internalTransitionInteractor.currentTransitionInfoInternal,
-                    finishedKeyguardState,
+                    transitionInteractor.isFinishedIn(KeyguardState.LOCKSCREEN),
                     keyguardInteractor.isActiveDreamLockscreenHosted,
                 )
                 .collect {
-                    (
-                        isAbleToDream,
-                        transitionInfo,
-                        finishedKeyguardState,
-                        isActiveDreamLockscreenHosted) ->
-                    val isOnLockscreen = finishedKeyguardState == KeyguardState.LOCKSCREEN
+                    (isAbleToDream, transitionInfo, isOnLockscreen, isActiveDreamLockscreenHosted)
+                    ->
                     val isTransitionInterruptible =
                         transitionInfo.to == KeyguardState.LOCKSCREEN &&
                             !invalidFromStates.contains(transitionInfo.from)
@@ -189,7 +185,7 @@
         scope.launch("$TAG#listenForLockscreenToPrimaryBouncerDragging") {
             shadeRepository.legacyShadeExpansion
                 .sampleCombine(
-                    startedKeyguardTransitionStep,
+                    transitionInteractor.startedKeyguardTransitionStep,
                     internalTransitionInteractor.currentTransitionInfoInternal,
                     keyguardInteractor.statusBarState,
                     keyguardInteractor.isKeyguardDismissible,
@@ -334,7 +330,7 @@
             listenForSleepTransition(
                 modeOnCanceledFromStartedStep = { startedStep ->
                     if (
-                        transitionInteractor.asleepKeyguardState.value == KeyguardState.AOD &&
+                        keyguardInteractor.asleepKeyguardState.value == KeyguardState.AOD &&
                             startedStep.from == KeyguardState.AOD
                     ) {
                         TransitionModeOnCanceled.REVERSE
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/GlanceableHubTransitions.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/GlanceableHubTransitions.kt
index f9ab1bb..bde0f56 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/GlanceableHubTransitions.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/GlanceableHubTransitions.kt
@@ -62,16 +62,16 @@
         communalInteractor
             .transitionProgressToScene(toScene)
             .sample(
-                transitionInteractor.startedKeyguardState,
+                transitionInteractor.startedKeyguardTransitionStep,
                 ::Pair,
             )
-            .collect { (transitionProgress, lastStartedState) ->
+            .collect { (transitionProgress, lastStartedStep) ->
                 val id = transitionId
                 if (id == null) {
                     // No transition started.
                     if (
                         transitionProgress is CommunalTransitionProgressModel.Transition &&
-                            lastStartedState == fromState
+                            lastStartedStep.to == fromState
                     ) {
                         transitionId =
                             transitionRepository.startTransition(
@@ -84,7 +84,7 @@
                             )
                     }
                 } else {
-                    if (lastStartedState != toState) {
+                    if (lastStartedStep.to != toState) {
                         return@collect
                     }
                     // An existing `id` means a transition is started, and calls to
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt
new file mode 100644
index 0000000..c19bbbc
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardDismissTransitionInteractor.kt
@@ -0,0 +1,73 @@
+/*
+ * 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 android.util.Log
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
+import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
+import com.android.systemui.keyguard.shared.model.KeyguardState.DOZING
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
+import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
+import javax.inject.Inject
+
+@SysUISingleton
+class KeyguardDismissTransitionInteractor
+@Inject
+constructor(
+    private val repository: KeyguardTransitionRepository,
+    private val fromLockscreenTransitionInteractor: FromLockscreenTransitionInteractor,
+    private val fromPrimaryBouncerTransitionInteractor: FromPrimaryBouncerTransitionInteractor,
+    private val fromAodTransitionInteractor: FromAodTransitionInteractor,
+    private val fromAlternateBouncerTransitionInteractor: FromAlternateBouncerTransitionInteractor,
+    private val fromDozingTransitionInteractor: FromDozingTransitionInteractor,
+    private val fromOccludedTransitionInteractor: FromOccludedTransitionInteractor,
+) {
+
+    /**
+     * Called to start a transition that will ultimately dismiss the keyguard from the current
+     * state.
+     *
+     * This is called exclusively by sources that can authoritatively say we should be unlocked,
+     * including KeyguardSecurityContainerController and WindowManager.
+     */
+    fun startDismissKeyguardTransition(reason: String = "") {
+        if (SceneContainerFlag.isEnabled) return
+        Log.d(TAG, "#startDismissKeyguardTransition(reason=$reason)")
+        when (val startedState = repository.currentTransitionInfoInternal.value.to) {
+            LOCKSCREEN -> fromLockscreenTransitionInteractor.dismissKeyguard()
+            PRIMARY_BOUNCER -> fromPrimaryBouncerTransitionInteractor.dismissPrimaryBouncer()
+            ALTERNATE_BOUNCER -> fromAlternateBouncerTransitionInteractor.dismissAlternateBouncer()
+            AOD -> fromAodTransitionInteractor.dismissAod()
+            DOZING -> fromDozingTransitionInteractor.dismissFromDozing()
+            KeyguardState.OCCLUDED -> fromOccludedTransitionInteractor.dismissFromOccluded()
+            KeyguardState.GONE ->
+                Log.i(
+                    TAG,
+                    "Already transitioning to GONE; ignoring startDismissKeyguardTransition."
+                )
+            else -> Log.e(TAG, "We don't know how to dismiss keyguard from state $startedState.")
+        }
+    }
+
+    companion object {
+        private val TAG = KeyguardDismissTransitionInteractor::class.simpleName
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
index 4aef808..44aafab 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractor.kt
@@ -48,7 +48,7 @@
     @Application scope: CoroutineScope,
     val repository: KeyguardRepository,
     val biometricSettingsRepository: BiometricSettingsRepository,
-    transitionInteractor: KeyguardTransitionInteractor,
+    keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor,
     internalTransitionInteractor: InternalKeyguardTransitionInteractor,
 ) {
 
@@ -94,7 +94,9 @@
                 showKeyguardWhenReenabled
                     .filter { shouldDismiss -> shouldDismiss }
                     .collect {
-                        transitionInteractor.startDismissKeyguardTransition("keyguard disabled")
+                        keyguardDismissTransitionInteractor.startDismissKeyguardTransition(
+                            "keyguard disabled"
+                        )
                     }
             }
         }
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 a96d7a8..f6f0cc5 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
@@ -36,7 +36,9 @@
 import com.android.systemui.keyguard.shared.model.DozeStateModel.Companion.isDozeOff
 import com.android.systemui.keyguard.shared.model.DozeTransitionModel
 import com.android.systemui.keyguard.shared.model.Edge
+import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
+import com.android.systemui.keyguard.shared.model.KeyguardState.DOZING
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
@@ -409,6 +411,12 @@
         }
     }
 
+    /** Which keyguard state to use when the device goes to sleep. */
+    val asleepKeyguardState: StateFlow<KeyguardState> =
+        repository.isAodAvailable
+            .map { aodAvailable -> if (aodAvailable) AOD else DOZING }
+            .stateIn(applicationScope, SharingStarted.Eagerly, DOZING)
+
     /**
      * Whether the primary authentication is required for the given user due to lockdown or
      * encryption after reboot.
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
index 2af95f2..2c3b481 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardQuickAffordanceInteractor.kt
@@ -119,7 +119,8 @@
                             is ObservableTransitionState.Idle ->
                                 it.currentScene == Scenes.Lockscreen
                             is ObservableTransitionState.Transition ->
-                                it.fromScene == Scenes.Lockscreen || it.toScene == Scenes.Lockscreen
+                                it.fromContent == Scenes.Lockscreen ||
+                                    it.toContent == Scenes.Lockscreen
                         }
                     }
                     .distinctUntilChanged()
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTouchHandlingInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTouchHandlingInteractor.kt
index 4a8ada7..505c749d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTouchHandlingInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTouchHandlingInteractor.kt
@@ -48,7 +48,6 @@
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
@@ -74,9 +73,7 @@
     val isLongPressHandlingEnabled: StateFlow<Boolean> =
         if (isFeatureEnabled()) {
                 combine(
-                    transitionInteractor.finishedKeyguardState.map {
-                        it == KeyguardState.LOCKSCREEN
-                    },
+                    transitionInteractor.isFinishedIn(KeyguardState.LOCKSCREEN),
                     repository.isQuickSettingsVisible,
                 ) { isFullyTransitionedToLockScreen, isQuickSettingsVisible ->
                     isFullyTransitionedToLockScreen && !isQuickSettingsVisible
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
index 6ff369e..e19b72e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractor.kt
@@ -22,21 +22,17 @@
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.keyguard.data.repository.KeyguardRepository
 import com.android.systemui.keyguard.data.repository.KeyguardTransitionRepository
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState
-import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
-import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
-import com.android.systemui.keyguard.shared.model.KeyguardState.DOZING
-import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
-import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
+import com.android.systemui.keyguard.shared.model.KeyguardState.OFF
 import com.android.systemui.keyguard.shared.model.KeyguardState.UNDEFINED
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.util.kotlin.WithPrev
 import com.android.systemui.util.kotlin.pairwise
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -47,7 +43,6 @@
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
@@ -56,7 +51,6 @@
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.shareIn
 import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.flow.transform
 import kotlinx.coroutines.launch
 
 /** Encapsulates business-logic related to the keyguard transitions. */
@@ -66,16 +60,7 @@
 @Inject
 constructor(
     @Application val scope: CoroutineScope,
-    private val keyguardRepository: KeyguardRepository,
     private val repository: KeyguardTransitionRepository,
-    private val fromLockscreenTransitionInteractor: dagger.Lazy<FromLockscreenTransitionInteractor>,
-    private val fromPrimaryBouncerTransitionInteractor:
-        dagger.Lazy<FromPrimaryBouncerTransitionInteractor>,
-    private val fromAodTransitionInteractor: dagger.Lazy<FromAodTransitionInteractor>,
-    private val fromAlternateBouncerTransitionInteractor:
-        dagger.Lazy<FromAlternateBouncerTransitionInteractor>,
-    private val fromDozingTransitionInteractor: dagger.Lazy<FromDozingTransitionInteractor>,
-    private val fromOccludedTransitionInteractor: dagger.Lazy<FromOccludedTransitionInteractor>,
     private val sceneInteractor: SceneInteractor,
 ) {
     private val transitionMap = mutableMapOf<Edge.StateToState, MutableSharedFlow<TransitionStep>>()
@@ -104,6 +89,18 @@
     val transitionState: StateFlow<TransitionStep> =
         transitions.stateIn(scope, SharingStarted.Eagerly, TransitionStep())
 
+    private val sceneTransitionPair =
+        sceneInteractor.transitionState
+            .pairwise()
+            .stateIn(
+                scope,
+                SharingStarted.Eagerly,
+                WithPrev(
+                    sceneInteractor.transitionState.value,
+                    sceneInteractor.transitionState.value
+                )
+            )
+
     /**
      * A pair of the most recent STARTED step, and the transition step immediately preceding it. The
      * transition framework enforces that the previous step is either a CANCELED or FINISHED step,
@@ -126,8 +123,10 @@
             repository.transitions
                 .filter { it.transitionState != TransitionState.CANCELED }
                 .collect { step ->
-                    getTransitionValueFlow(step.from).emit(1f - step.value)
-                    getTransitionValueFlow(step.to).emit(step.value)
+                    val value =
+                        if (step.transitionState == TransitionState.FINISHED) 1f else step.value
+                    getTransitionValueFlow(step.from).emit(1f - value)
+                    getTransitionValueFlow(step.to).emit(value)
                 }
         }
 
@@ -183,8 +182,14 @@
         }
     }
 
-    fun transition(edge: Edge, edgeWithoutSceneContainer: Edge): Flow<TransitionStep> {
-        return transition(if (SceneContainerFlag.isEnabled) edge else edgeWithoutSceneContainer)
+    fun transition(edge: Edge, edgeWithoutSceneContainer: Edge? = null): Flow<TransitionStep> {
+        return transition(
+            if (SceneContainerFlag.isEnabled || edgeWithoutSceneContainer == null) {
+                edge
+            } else {
+                edgeWithoutSceneContainer
+            }
+        )
     }
 
     /** Given an [edge], return a Flow to collect only relevant [TransitionStep]s. */
@@ -202,7 +207,7 @@
             }
 
         return if (SceneContainerFlag.isEnabled) {
-            flow.filter {
+            flow.filter { step ->
                 val fromScene =
                     when (edge) {
                         is Edge.StateToState -> edge.from?.mapToSceneContainerScene()
@@ -219,8 +224,23 @@
 
                 fun SceneKey?.isLockscreenOrNull() = this == Scenes.Lockscreen || this == null
 
-                return@filter (fromScene.isLockscreenOrNull() && toScene.isLockscreenOrNull()) ||
+                val isTransitioningBetweenLockscreenStates =
+                    fromScene.isLockscreenOrNull() && toScene.isLockscreenOrNull()
+                val isTransitioningBetweenDesiredScenes =
                     sceneInteractor.transitionState.value.isTransitioning(fromScene, toScene)
+
+                // We can't compare the terminal step with the current sceneTransition because
+                // a) STL has no guarantee that it will settle in Idle() when finished/canceled
+                // b) Comparing to Idle(toScene) would make any other FINISHED step settling in
+                //    toScene pass as well
+                val terminalStepBelongsToPreviousTransition =
+                    (step.transitionState == TransitionState.FINISHED ||
+                        step.transitionState == TransitionState.CANCELED) &&
+                        sceneTransitionPair.value.previousValue.isTransitioning(fromScene, toScene)
+
+                return@filter isTransitioningBetweenLockscreenStates ||
+                    isTransitioningBetweenDesiredScenes ||
+                    terminalStepBelongsToPreviousTransition
             }
         } else {
             flow
@@ -250,10 +270,10 @@
     }
 
     fun transitionValue(
-        scene: SceneKey,
+        scene: SceneKey? = null,
         stateWithoutSceneContainer: KeyguardState,
     ): Flow<Float> {
-        return if (SceneContainerFlag.isEnabled) {
+        return if (SceneContainerFlag.isEnabled && scene != null) {
             sceneInteractor.transitionProgress(scene)
         } else {
             transitionValue(stateWithoutSceneContainer)
@@ -277,73 +297,10 @@
     }
 
     /** The last [TransitionStep] with a [TransitionState] of STARTED */
-    val startedKeyguardTransitionStep: Flow<TransitionStep> =
-        repository.transitions.filter { step -> step.transitionState == TransitionState.STARTED }
-
-    /** The destination state of the last [TransitionState.STARTED] transition. */
-    @SuppressLint("SharedFlowCreation")
-    val startedKeyguardState: SharedFlow<KeyguardState> =
-        startedKeyguardTransitionStep
-            .map { step -> step.to }
-            .buffer(2, BufferOverflow.DROP_OLDEST)
-            .shareIn(scope, SharingStarted.Eagerly, replay = 1)
-
-    /** The from state of the last [TransitionState.STARTED] transition. */
-    // TODO: is it performant to have several SharedFlows side by side instead of one?
-    @SuppressLint("SharedFlowCreation")
-    val startedKeyguardFromState: SharedFlow<KeyguardState> =
-        startedKeyguardTransitionStep
-            .map { step -> step.from }
-            .buffer(2, BufferOverflow.DROP_OLDEST)
-            .shareIn(scope, SharingStarted.Eagerly, replay = 1)
-
-    /** Which keyguard state to use when the device goes to sleep. */
-    val asleepKeyguardState: StateFlow<KeyguardState> =
-        keyguardRepository.isAodAvailable
-            .map { aodAvailable -> if (aodAvailable) AOD else DOZING }
-            .stateIn(scope, SharingStarted.Eagerly, DOZING)
-
-    /**
-     * The last [KeyguardState] to which we [TransitionState.FINISHED] a transition.
-     *
-     * WARNING: This will NOT emit a value if a transition is CANCELED, and will also not emit a
-     * value when a subsequent transition is STARTED. It will *only* emit once we have finally
-     * FINISHED in a state. This can have unintuitive implications.
-     *
-     * For example, if we're transitioning from GONE -> DOZING, and that transition is CANCELED in
-     * favor of a DOZING -> LOCKSCREEN transition, the FINISHED state is still GONE, and will remain
-     * GONE throughout the DOZING -> LOCKSCREEN transition until the DOZING -> LOCKSCREEN transition
-     * finishes (at which point we'll be FINISHED in LOCKSCREEN).
-     *
-     * Since there's no real limit to how many consecutive transitions can be canceled, it's even
-     * possible for the FINISHED state to be the same as the STARTED state while still
-     * transitioning.
-     *
-     * For example:
-     * 1. We're finished in GONE.
-     * 2. The user presses the power button, starting a GONE -> DOZING transition. We're still
-     *    FINISHED in GONE.
-     * 3. The user changes their mind, pressing the power button to wake up; this starts a DOZING ->
-     *    LOCKSCREEN transition. We're still FINISHED in GONE.
-     * 4. The user quickly swipes away the lockscreen prior to DOZING -> LOCKSCREEN finishing; this
-     *    starts a LOCKSCREEN -> GONE transition. We're still FINISHED in GONE, but we've also
-     *    STARTED a transition *to* GONE.
-     * 5. We'll emit KeyguardState.GONE again once the transition finishes.
-     *
-     * If you just need to know when we eventually settle into a state, this flow is likely
-     * sufficient. However, if you're having issues with state *during* transitions started after
-     * one or more canceled transitions, you probably need to use [currentKeyguardState].
-     */
-    @SuppressLint("SharedFlowCreation")
-    val finishedKeyguardState: SharedFlow<KeyguardState> =
+    val startedKeyguardTransitionStep: StateFlow<TransitionStep> =
         repository.transitions
-            .transform { step ->
-                if (step.transitionState == TransitionState.FINISHED) {
-                    emit(step.to)
-                }
-            }
-            .buffer(2, BufferOverflow.DROP_OLDEST)
-            .shareIn(scope, SharingStarted.Eagerly, replay = 1)
+            .filter { step -> step.transitionState == TransitionState.STARTED }
+            .stateIn(scope, SharingStarted.Eagerly, TransitionStep())
 
     /**
      * The [KeyguardState] we're currently in.
@@ -409,8 +366,7 @@
                     it.from
                 }
             }
-            .distinctUntilChanged()
-            .stateIn(scope, SharingStarted.Eagerly, KeyguardState.OFF)
+            .stateIn(scope, SharingStarted.Eagerly, OFF)
 
     val isInTransition =
         combine(
@@ -422,33 +378,6 @@
         }
 
     /**
-     * Called to start a transition that will ultimately dismiss the keyguard from the current
-     * state.
-     *
-     * This is called exclusively by sources that can authoritatively say we should be unlocked,
-     * including KeyguardSecurityContainerController and WindowManager.
-     */
-    fun startDismissKeyguardTransition(reason: String = "") {
-        if (SceneContainerFlag.isEnabled) return
-        Log.d(TAG, "#startDismissKeyguardTransition(reason=$reason)")
-        when (val startedState = repository.currentTransitionInfoInternal.value.to) {
-            LOCKSCREEN -> fromLockscreenTransitionInteractor.get().dismissKeyguard()
-            PRIMARY_BOUNCER -> fromPrimaryBouncerTransitionInteractor.get().dismissPrimaryBouncer()
-            ALTERNATE_BOUNCER ->
-                fromAlternateBouncerTransitionInteractor.get().dismissAlternateBouncer()
-            AOD -> fromAodTransitionInteractor.get().dismissAod()
-            DOZING -> fromDozingTransitionInteractor.get().dismissFromDozing()
-            KeyguardState.OCCLUDED -> fromOccludedTransitionInteractor.get().dismissFromOccluded()
-            KeyguardState.GONE ->
-                Log.i(
-                    TAG,
-                    "Already transitioning to GONE; ignoring startDismissKeyguardTransition."
-                )
-            else -> Log.e(TAG, "We don't know how to dismiss keyguard from state $startedState.")
-        }
-    }
-
-    /**
      * Whether we're in a transition to and from the given [KeyguardState]s, but haven't yet
      * completed it.
      *
@@ -506,12 +435,13 @@
 
     fun isFinishedIn(scene: SceneKey, stateWithoutSceneContainer: KeyguardState): Flow<Boolean> {
         return if (SceneContainerFlag.isEnabled) {
-            sceneInteractor.transitionState
-                .map { it.isIdle(scene) || it.isTransitioning(from = scene) }
-                .distinctUntilChanged()
-        } else {
-            isFinishedIn(stateWithoutSceneContainer)
-        }
+                sceneInteractor.transitionState.map {
+                    it.isIdle(scene) || it.isTransitioning(from = scene)
+                }
+            } else {
+                isFinishedIn(stateWithoutSceneContainer)
+            }
+            .distinctUntilChanged()
     }
 
     /** Whether we've FINISHED a transition to a state */
@@ -520,17 +450,26 @@
         return finishedKeyguardState.map { it == state }.distinctUntilChanged()
     }
 
+    fun isCurrentlyIn(scene: SceneKey, stateWithoutSceneContainer: KeyguardState): Flow<Boolean> {
+        return if (SceneContainerFlag.isEnabled) {
+                // In STL there is no difference between finished/currentState
+                isFinishedIn(scene, stateWithoutSceneContainer)
+            } else {
+                stateWithoutSceneContainer.checkValidState()
+                currentKeyguardState.map { it == stateWithoutSceneContainer }
+            }
+            .distinctUntilChanged()
+    }
+
     fun getCurrentState(): KeyguardState {
         return currentKeyguardState.replayCache.last()
     }
 
-    fun getStartedFromState(): KeyguardState {
-        return startedKeyguardFromState.replayCache.last()
-    }
-
-    fun getFinishedState(): KeyguardState {
-        return finishedKeyguardState.replayCache.last()
-    }
+    private val finishedKeyguardState: StateFlow<KeyguardState> =
+        repository.transitions
+            .filter { it.transitionState == TransitionState.FINISHED }
+            .map { it.to }
+            .stateIn(scope, SharingStarted.Eagerly, OFF)
 
     companion object {
         private val TAG = KeyguardTransitionInteractor::class.simpleName
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt
index f0bf402..9b8d9ea1 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardWakeDirectlyToGoneInteractor.kt
@@ -37,6 +37,7 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState.Companion.deviceIsAwakeInState
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.power.shared.model.WakeSleepReason
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor
 import com.android.systemui.util.kotlin.sample
 import com.android.systemui.util.settings.SecureSettings
@@ -181,14 +182,20 @@
         scope.launch {
             powerInteractor.detailedWakefulness
                 .distinctUntilChangedBy { it.isAwake() }
-                .sample(transitionInteractor.currentKeyguardState, ::Pair)
-                .collect { (wakefulness, currentState) ->
+                .sample(
+                    transitionInteractor.isCurrentlyIn(
+                        Scenes.Gone,
+                        stateWithoutSceneContainer = KeyguardState.GONE
+                    ),
+                    ::Pair
+                )
+                .collect { (wakefulness, finishedInGone) ->
                     // Save isAwake for use in onDreamingStarted/onDreamingStopped.
                     this@KeyguardWakeDirectlyToGoneInteractor.isAwake = wakefulness.isAwake()
 
                     // If we're sleeping from GONE, check the timeout and lock instantly settings.
                     // These are not relevant if we're coming from non-GONE states.
-                    if (!isAwake && currentState == KeyguardState.GONE) {
+                    if (!isAwake && finishedInGone) {
                         val lockTimeoutDuration = getCanIgnoreAuthAndReturnToGoneDuration()
 
                         // If the screen timed out and went to sleep, and the lock timeout is > 0ms,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/StatusBarDisableFlagsInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/StatusBarDisableFlagsInteractor.kt
index 47818cb..e00e33d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/StatusBarDisableFlagsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/StatusBarDisableFlagsInteractor.kt
@@ -45,6 +45,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 
@@ -76,7 +77,7 @@
     private val disableFlagsForUserId =
         combine(
                 selectedUserInteractor.selectedUser,
-                keyguardTransitionInteractor.startedKeyguardState,
+                keyguardTransitionInteractor.startedKeyguardTransitionStep.map { it.to },
                 deviceConfigInteractor.property(
                     namespace = DeviceConfig.NAMESPACE_SYSTEMUI,
                     name = SystemUiDeviceConfigFlags.NAV_BAR_HANDLE_SHOW_OVER_LOCKSCREEN,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/SwipeToDismissInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/SwipeToDismissInteractor.kt
index 906d586..e404f27 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/SwipeToDismissInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/SwipeToDismissInteractor.kt
@@ -22,12 +22,12 @@
 import com.android.systemui.keyguard.shared.model.StatusBarState
 import com.android.systemui.shade.data.repository.ShadeRepository
 import com.android.systemui.util.kotlin.Utils.Companion.sample
+import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
-import javax.inject.Inject
 
 /**
  * Handles logic around the swipe to dismiss gesture, where the user swipes up on the dismissable
@@ -53,15 +53,15 @@
     val dismissFling =
         shadeRepository.currentFling
             .sample(
-                transitionInteractor.startedKeyguardState,
+                transitionInteractor.startedKeyguardTransitionStep,
                 keyguardInteractor.isKeyguardDismissible,
                 keyguardInteractor.statusBarState,
             )
-            .filter { (flingInfo, startedState, keyguardDismissable, statusBarState) ->
+            .filter { (flingInfo, startedStep, keyguardDismissable, statusBarState) ->
                 flingInfo != null &&
-                        !flingInfo.expand &&
-                        statusBarState != StatusBarState.SHADE_LOCKED &&
-                        startedState == KeyguardState.LOCKSCREEN &&
+                    !flingInfo.expand &&
+                    statusBarState != StatusBarState.SHADE_LOCKED &&
+                    startedStep.to == KeyguardState.LOCKSCREEN &&
                     keyguardDismissable
             }
             .map { (flingInfo, _) -> flingInfo }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
index d06ee64..ba12e93 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/TransitionInteractor.kt
@@ -32,7 +32,6 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.filter
-import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.launch
 
@@ -62,17 +61,6 @@
 
     abstract fun start()
 
-    /* Use background dispatcher for all [KeyguardTransitionInteractor] flows. Necessary because
-     * the [sample] utility internally runs a collect on the Unconfined dispatcher, resulting
-     * in continuations on the main thread. We don't want that for classes that inherit from this.
-     */
-    val startedKeyguardTransitionStep =
-        transitionInteractor.startedKeyguardTransitionStep.flowOn(bgDispatcher)
-    // The following are MutableSharedFlows, and do not require flowOn
-    val startedKeyguardState = transitionInteractor.startedKeyguardState
-    val finishedKeyguardState = transitionInteractor.finishedKeyguardState
-    val currentKeyguardState = transitionInteractor.currentKeyguardState
-
     suspend fun startTransitionTo(
         toState: KeyguardState,
         animator: ValueAnimator? = getDefaultAnimatorForTransitionsToState(toState),
@@ -92,17 +80,6 @@
                     " $fromState. This should never happen - check currentTransitionInfoInternal" +
                     " or use filterRelevantKeyguardState before starting transitions."
             )
-
-            if (fromState == transitionInteractor.finishedKeyguardState.replayCache.last()) {
-                Log.e(
-                    name,
-                    "This transition would not have been ignored prior to ag/26681239, since we " +
-                        "are FINISHED in $fromState (but have since started another transition). " +
-                        "If ignoring this transition has caused a regression, fix it by ensuring " +
-                        "that transitions are exclusively started from the most recently started " +
-                        "state."
-                )
-            }
             return null
         }
 
@@ -207,11 +184,11 @@
         powerInteractor.isAsleep
             .filter { isAsleep -> isAsleep }
             .filterRelevantKeyguardState()
-            .sample(startedKeyguardTransitionStep)
+            .sample(transitionInteractor.startedKeyguardTransitionStep)
             .map(modeOnCanceledFromStartedStep)
             .collect { modeOnCanceled ->
                 startTransitionTo(
-                    toState = transitionInteractor.asleepKeyguardState.value,
+                    toState = keyguardInteractor.asleepKeyguardState.value,
                     modeOnCanceled = modeOnCanceled,
                     ownerReason = "Sleep transition triggered"
                 )
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
index 25b2b7c..a09cd7c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/WindowManagerLockscreenVisibilityInteractor.kt
@@ -63,10 +63,13 @@
 ) {
     private val defaultSurfaceBehindVisibility =
         combine(
-            transitionInteractor.finishedKeyguardState,
+            transitionInteractor.isFinishedIn(
+                scene = Scenes.Gone,
+                stateWithoutSceneContainer = KeyguardState.GONE
+            ),
             wakeToGoneInteractor.canWakeDirectlyToGone,
-        ) { finishedState, canWakeDirectlyToGone ->
-            isSurfaceVisible(finishedState) || canWakeDirectlyToGone
+        ) { isOnGone, canWakeDirectlyToGone ->
+            isOnGone || canWakeDirectlyToGone
         }
 
     /**
@@ -124,8 +127,8 @@
                     when (transitionState) {
                         is ObservableTransitionState.Transition ->
                             when {
-                                transitionState.fromScene == Scenes.Lockscreen &&
-                                    transitionState.toScene == Scenes.Gone ->
+                                transitionState.fromContent == Scenes.Lockscreen &&
+                                    transitionState.toContent == Scenes.Gone ->
                                     sceneInteractor
                                         .get()
                                         .isTransitionUserInputOngoing
@@ -136,8 +139,8 @@
                                                 flowOf(true)
                                             }
                                         }
-                                transitionState.fromScene == Scenes.Bouncer &&
-                                    transitionState.toScene == Scenes.Gone ->
+                                transitionState.fromContent == Scenes.Bouncer &&
+                                    transitionState.toContent == Scenes.Gone ->
                                     transitionState.progress.map { progress ->
                                         progress >
                                             FromPrimaryBouncerTransitionInteractor
@@ -196,18 +199,20 @@
                         edge = Edge.create(to = Scenes.Gone),
                         edgeWithoutSceneContainer = Edge.create(to = KeyguardState.GONE)
                     ),
-                    transitionInteractor.finishedKeyguardState,
+                    transitionInteractor.isFinishedIn(
+                        scene = Scenes.Gone,
+                        stateWithoutSceneContainer = KeyguardState.GONE
+                    ),
                     surfaceBehindInteractor.isAnimatingSurface,
                     notificationLaunchAnimationInteractor.isLaunchAnimationRunning,
-                ) { isInTransitionToGone, finishedState, isAnimatingSurface, notifLaunchRunning ->
+                ) { isInTransitionToGone, isOnGone, isAnimatingSurface, notifLaunchRunning ->
                     // Using the animation if we're animating it directly, or if the
                     // ActivityLaunchAnimator is in the process of animating it.
                     val animationsRunning = isAnimatingSurface || notifLaunchRunning
                     // We may still be animating the surface after the keyguard is fully GONE, since
                     // some animations (like the translation spring) are not tied directly to the
                     // transition step amount.
-                    isInTransitionToGone ||
-                        (finishedState == KeyguardState.GONE && animationsRunning)
+                    isInTransitionToGone || (isOnGone && animationsRunning)
                 }
                 .distinctUntilChanged()
         }
@@ -248,7 +253,7 @@
                         // transition. Same for waking directly to gone, due to the lockscreen being
                         // disabled or because the device was woken back up before the lock timeout
                         // duration elapsed.
-                        KeyguardState.lockscreenVisibleInState(KeyguardState.GONE)
+                        false
                     } else if (canWakeDirectlyToGone) {
                         // Never show the lockscreen if we can wake directly to GONE. This means
                         // that the lock timeout has not yet elapsed, or the keyguard is disabled.
@@ -274,8 +279,7 @@
                         // *not* play the going away animation or related animations.
                         false
                     } else {
-                        // Otherwise, use the visibility of the current state.
-                        KeyguardState.lockscreenVisibleInState(currentState)
+                        currentState != KeyguardState.GONE
                     }
                 }
                 .distinctUntilChanged()
@@ -302,10 +306,4 @@
                     !BiometricUnlockMode.isWakeAndUnlock(biometricUnlockState.mode)
             }
             .distinctUntilChanged()
-
-    companion object {
-        fun isSurfaceVisible(state: KeyguardState): Boolean {
-            return !KeyguardState.lockscreenVisibleInState(state)
-        }
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
index b850095..f3bb829 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/scenetransition/LockscreenSceneTransitionInteractor.kt
@@ -111,12 +111,12 @@
         if (currentTransitionId == null) return
         if (prevTransition !is ObservableTransitionState.Transition) return
 
-        if (idle.currentScene == prevTransition.toScene) {
+        if (idle.currentScene == prevTransition.toContent) {
             finishCurrentTransition()
         } else {
             val targetState =
                 if (idle.currentScene == Scenes.Lockscreen) {
-                    transitionInteractor.getStartedFromState()
+                    transitionInteractor.startedKeyguardTransitionStep.value.from
                 } else {
                     UNDEFINED
                 }
@@ -150,17 +150,17 @@
     }
 
     private suspend fun handleTransition(transition: ObservableTransitionState.Transition) {
-        if (transition.fromScene == Scenes.Lockscreen) {
+        if (transition.fromContent == Scenes.Lockscreen) {
             if (currentTransitionId != null) {
                 val currentToState =
                     internalTransitionInteractor.currentTransitionInfoInternal.value.to
                 if (currentToState == UNDEFINED) {
-                    transitionKtfTo(transitionInteractor.getStartedFromState())
+                    transitionKtfTo(transitionInteractor.startedKeyguardTransitionStep.value.from)
                 }
             }
             startTransitionFromLockscreen()
             collectProgress(transition)
-        } else if (transition.toScene == Scenes.Lockscreen) {
+        } else if (transition.toContent == Scenes.Lockscreen) {
             if (currentTransitionId != null) {
                 transitionKtfTo(UNDEFINED)
             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/KeyguardState.kt b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/KeyguardState.kt
index 24db3c2..080ddfd 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/KeyguardState.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/shared/model/KeyguardState.kt
@@ -156,12 +156,6 @@
 
     companion object {
 
-        /** Whether the lockscreen is visible when we're FINISHED in the given state. */
-        fun lockscreenVisibleInState(state: KeyguardState): Boolean {
-            // TODO(b/349784682): Transform deprecated states for Flexiglass
-            return state != GONE
-        }
-
         /**
          * Whether the device is awake ([PowerInteractor.isAwake]) when we're FINISHED in the given
          * keyguard state.
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardBlueprintViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardBlueprintViewBinder.kt
index bec8f3d..f1b9cba 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardBlueprintViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardBlueprintViewBinder.kt
@@ -93,7 +93,7 @@
                                         blueprint.applyConstraints(this)
                                     }
 
-                                logAlphaVisibilityOfAppliedConstraintSet(cs, clockViewModel)
+                                logAlphaVisibilityScaleOfAppliedConstraintSet(cs, clockViewModel)
                                 cs.applyTo(constraintLayout)
                             }
                         }
@@ -115,7 +115,7 @@
                                     clone(constraintLayout)
                                     blueprint.applyConstraints(this)
                                 }
-                            logAlphaVisibilityOfAppliedConstraintSet(cs, clockViewModel)
+                            logAlphaVisibilityScaleOfAppliedConstraintSet(cs, clockViewModel)
                             cs.applyTo(constraintLayout)
                         }
                     }
@@ -124,7 +124,7 @@
         }
     }
 
-    private fun logAlphaVisibilityOfAppliedConstraintSet(
+    private fun logAlphaVisibilityScaleOfAppliedConstraintSet(
         cs: ConstraintSet,
         viewModel: KeyguardClockViewModel
     ) {
@@ -136,12 +136,15 @@
         Log.i(
             TAG,
             "applyCsToSmallClock: vis=${cs.getVisibility(smallClockViewId)} " +
-                "alpha=${cs.getConstraint(smallClockViewId).propertySet.alpha}"
+                "alpha=${cs.getConstraint(smallClockViewId).propertySet.alpha} " +
+                "scale=${cs.getConstraint(smallClockViewId).transform.scaleX} "
         )
         Log.i(
             TAG,
             "applyCsToLargeClock: vis=${cs.getVisibility(largeClockViewId)} " +
-                "alpha=${cs.getConstraint(largeClockViewId).propertySet.alpha}"
+                "alpha=${cs.getConstraint(largeClockViewId).propertySet.alpha} " +
+                "scale=${cs.getConstraint(largeClockViewId).transform.scaleX} " +
+                "pivotX=${cs.getConstraint(largeClockViewId).transform.transformPivotX} "
         )
         Log.i(
             TAG,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
index a7a8321..7899971 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
@@ -369,8 +369,7 @@
                                 } else {
                                     vibratorHelper.performHapticFeedback(
                                         view,
-                                        HapticFeedbackConstants.CONFIRM,
-                                        HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING,
+                                        HapticFeedbackConstants.BIOMETRIC_CONFIRM,
                                     )
                                 }
                             }
@@ -390,8 +389,7 @@
                                 } else {
                                     vibratorHelper.performHapticFeedback(
                                         view,
-                                        HapticFeedbackConstants.REJECT,
-                                        HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING,
+                                        HapticFeedbackConstants.BIOMETRIC_REJECT,
                                     )
                                 }
                             }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
index 0032c2f..e2ad4635 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/transitions/DeviceEntryIconTransitionModule.kt
@@ -46,6 +46,7 @@
 import com.android.systemui.keyguard.ui.viewmodel.LockscreenToOccludedTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.LockscreenToPrimaryBouncerTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.OccludedToAodTransitionViewModel
+import com.android.systemui.keyguard.ui.viewmodel.OccludedToDozingTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.OccludedToGlanceableHubTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.OccludedToLockscreenTransitionViewModel
 import com.android.systemui.keyguard.ui.viewmodel.OffToLockscreenTransitionViewModel
@@ -205,6 +206,12 @@
 
     @Binds
     @IntoSet
+    abstract fun occludedToDozing(
+        impl: OccludedToDozingTransitionViewModel
+    ): DeviceEntryIconTransition
+
+    @Binds
+    @IntoSet
     abstract fun occludedToLockscreen(
         impl: OccludedToLockscreenTransitionViewModel
     ): DeviceEntryIconTransition
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
index c8fe55d..be6b0eb 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/ClockSection.kt
@@ -59,6 +59,16 @@
     alpha: Float,
 ) = views.forEach { view -> this.setAlpha(view.id, alpha) }
 
+internal fun ConstraintSet.setScaleX(
+    views: Iterable<View>,
+    alpha: Float,
+) = views.forEach { view -> this.setScaleX(view.id, alpha) }
+
+internal fun ConstraintSet.setScaleY(
+    views: Iterable<View>,
+    alpha: Float,
+) = views.forEach { view -> this.setScaleY(view.id, alpha) }
+
 @SysUISingleton
 class ClockSection
 @Inject
@@ -125,6 +135,9 @@
             setAlpha(getNonTargetClockFace(clock).views, 0F)
             if (!keyguardClockViewModel.isLargeClockVisible.value) {
                 connect(sharedR.id.bc_smartspace_view, TOP, sharedR.id.date_smartspace_view, BOTTOM)
+            } else {
+                setScaleX(getTargetClockFace(clock).views, rootViewModel.burnInModel.value.scale)
+                setScaleY(getTargetClockFace(clock).views, rootViewModel.burnInModel.value.scale)
             }
         }
     }
@@ -205,6 +218,9 @@
             create(R.id.small_clock_guideline_top, ConstraintSet.HORIZONTAL_GUIDELINE)
             setGuidelineBegin(R.id.small_clock_guideline_top, smallClockTopMargin)
             connect(R.id.lockscreen_clock_view, TOP, R.id.small_clock_guideline_top, BOTTOM)
+
+            // Explicitly clear pivot to force recalculate pivot instead of using legacy value
+            setTransformPivot(R.id.lockscreen_clock_view_large, Float.NaN, Float.NaN)
         }
 
         constrainWeatherClockDateIconsBarrier(constraints)
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SmartspaceSection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SmartspaceSection.kt
index 55fc718..99160f8 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SmartspaceSection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/SmartspaceSection.kt
@@ -127,6 +127,7 @@
 
             // migrate addSmartspaceView from KeyguardClockSwitchController
             constrainHeight(sharedR.id.bc_smartspace_view, ConstraintSet.WRAP_CONTENT)
+            constrainWidth(sharedR.id.bc_smartspace_view, ConstraintSet.MATCH_CONSTRAINT)
             connect(
                 sharedR.id.bc_smartspace_view,
                 ConstraintSet.START,
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
index 7b0b23f..b5d9e2a 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModel.kt
@@ -19,6 +19,7 @@
 
 import android.graphics.Color
 import com.android.systemui.bouncer.domain.interactor.AlternateBouncerInteractor
+import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor
 import com.android.systemui.keyguard.DismissCallbackRegistry
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState.ALTERNATE_BOUNCER
@@ -41,6 +42,7 @@
     keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val dismissCallbackRegistry: DismissCallbackRegistry,
     alternateBouncerInteractor: Lazy<AlternateBouncerInteractor>,
+    private val primaryBouncerInteractor: PrimaryBouncerInteractor,
 ) {
     // When we're fully transitioned to the AlternateBouncer, the alpha of the scrim should be:
     private val alternateBouncerScrimAlpha = .66f
@@ -73,5 +75,6 @@
     fun onBackRequested() {
         statusBarKeyguardViewManager.hideAlternateBouncer(false)
         dismissCallbackRegistry.notifyDismissCancelled()
+        primaryBouncerInteractor.setDismissAction(null, null)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryBackgroundViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryBackgroundViewModel.kt
index 6f8389f..9f68210 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryBackgroundViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryBackgroundViewModel.kt
@@ -59,6 +59,7 @@
     alternateBouncerToDozingTransitionViewModel: AlternateBouncerToDozingTransitionViewModel,
     dreamingToAodTransitionViewModel: DreamingToAodTransitionViewModel,
     primaryBouncerToLockscreenTransitionViewModel: PrimaryBouncerToLockscreenTransitionViewModel,
+    occludedToDozingTransitionViewModel: OccludedToDozingTransitionViewModel,
 ) {
     val color: Flow<Int> =
         deviceEntryIconViewModel.useBackgroundProtection.flatMapLatest { useBackground ->
@@ -103,6 +104,7 @@
                         dreamingToAodTransitionViewModel.deviceEntryBackgroundViewAlpha,
                         primaryBouncerToLockscreenTransitionViewModel
                             .deviceEntryBackgroundViewAlpha,
+                        occludedToDozingTransitionViewModel.deviceEntryBackgroundViewAlpha,
                     )
                     .merge()
                     .onStart {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryForegroundViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryForegroundViewModel.kt
index 06b76b3..87c32a5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryForegroundViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryForegroundViewModel.kt
@@ -53,10 +53,10 @@
 ) {
     private val isShowingAodOrDozing: Flow<Boolean> =
         combine(
-            transitionInteractor.startedKeyguardState,
+            transitionInteractor.startedKeyguardTransitionStep,
             transitionInteractor.transitionValue(KeyguardState.DOZING),
-        ) { startedKeyguardState, dozingTransitionValue ->
-            startedKeyguardState == KeyguardState.AOD || dozingTransitionValue == 1f
+        ) { startedKeyguardStep, dozingTransitionValue ->
+            startedKeyguardStep.to == KeyguardState.AOD || dozingTransitionValue == 1f
         }
 
     private fun getColor(usingBackgroundProtection: Boolean): Int {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryIconViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryIconViewModel.kt
index 5ce1b5e..d3bb4f5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryIconViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/DeviceEntryIconViewModel.kt
@@ -83,8 +83,8 @@
     private val intEvaluator = IntEvaluator()
     private val floatEvaluator = FloatEvaluator()
     private val showingAlternateBouncer: Flow<Boolean> =
-        transitionInteractor.startedKeyguardState.map { keyguardState ->
-            keyguardState == KeyguardState.ALTERNATE_BOUNCER
+        transitionInteractor.startedKeyguardTransitionStep.map { keyguardStep ->
+            keyguardStep.to == KeyguardState.ALTERNATE_BOUNCER
         }
     private val qsProgress: Flow<Float> = shadeInteractor.qsExpansion.onStart { emit(0f) }
     private val shadeExpansion: Flow<Float> = shadeInteractor.shadeExpansion.onStart { emit(0f) }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModel.kt
index c885c9a..72740d5 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModel.kt
@@ -20,7 +20,6 @@
 import androidx.annotation.VisibleForTesting
 import com.android.app.tracing.FlowTracing.traceEmissionCount
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.keyguard.NewPickerUiKeyguardPreview
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardQuickAffordanceInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
@@ -29,6 +28,7 @@
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shared.Flags
 import com.android.systemui.shared.keyguard.shared.model.KeyguardQuickAffordanceSlots
 import com.android.systemui.utils.coroutines.flow.flatMapLatestConflated
 import javax.inject.Inject
@@ -85,9 +85,7 @@
     private val previewMode = MutableStateFlow(PreviewMode())
 
     private val showingLockscreen: Flow<Boolean> =
-        transitionInteractor.finishedKeyguardState.map { keyguardState ->
-            keyguardState == KeyguardState.LOCKSCREEN
-        }
+        transitionInteractor.isFinishedIn(KeyguardState.LOCKSCREEN)
 
     /** The only time the expansion is important is while lockscreen is actively displayed */
     private val shadeExpansionAlpha =
@@ -171,7 +169,7 @@
 
     /** An observable for the view-model of the "start button" quick affordance. */
     val startButton: Flow<KeyguardQuickAffordanceViewModel> =
-        if (NewPickerUiKeyguardPreview.isEnabled) {
+        if (Flags.newCustomizationPickerUi()) {
             previewAffordances.flatMapLatestConflated {
                 button(
                     position = KeyguardQuickAffordancePosition.BOTTOM_START,
@@ -186,7 +184,7 @@
 
     /** An observable for the view-model of the "end button" quick affordance. */
     val endButton: Flow<KeyguardQuickAffordanceViewModel> =
-        if (NewPickerUiKeyguardPreview.isEnabled) {
+        if (Flags.newCustomizationPickerUi()) {
             previewAffordances.flatMapLatestConflated {
                 button(
                     position = KeyguardQuickAffordancePosition.BOTTOM_END,
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 a96869d..eaa61a1 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
@@ -34,20 +34,18 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
 import com.android.systemui.keyguard.shared.model.KeyguardState.DREAMING
-import com.android.systemui.keyguard.shared.model.KeyguardState.GLANCEABLE_HUB
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
-import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
 import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.keyguard.shared.model.TransitionState.RUNNING
 import com.android.systemui.keyguard.shared.model.TransitionState.STARTED
 import com.android.systemui.keyguard.ui.StateToValue
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.ui.viewmodel.NotificationShadeWindowModel
 import com.android.systemui.statusbar.notification.domain.interactor.NotificationsKeyguardInteractor
 import com.android.systemui.statusbar.phone.DozeParameters
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController
-import com.android.systemui.util.kotlin.BooleanFlowOperators.any
 import com.android.systemui.util.kotlin.BooleanFlowOperators.anyOf
 import com.android.systemui.util.kotlin.pairwise
 import com.android.systemui.util.kotlin.sample
@@ -86,6 +84,7 @@
     private val communalInteractor: CommunalInteractor,
     private val keyguardTransitionInteractor: KeyguardTransitionInteractor,
     private val notificationsKeyguardInteractor: NotificationsKeyguardInteractor,
+    notificationShadeWindowModel: NotificationShadeWindowModel,
     private val alternateBouncerToAodTransitionViewModel: AlternateBouncerToAodTransitionViewModel,
     private val alternateBouncerToGoneTransitionViewModel:
         AlternateBouncerToGoneTransitionViewModel,
@@ -132,8 +131,8 @@
     val burnInModel = _burnInModel.asStateFlow()
 
     val burnInLayerVisibility: Flow<Int> =
-        keyguardTransitionInteractor.startedKeyguardState
-            .filter { it == AOD || it == LOCKSCREEN }
+        keyguardTransitionInteractor.startedKeyguardTransitionStep
+            .filter { it.to == AOD || it.to == LOCKSCREEN }
             .map { VISIBLE }
 
     val goneToAodTransition =
@@ -197,37 +196,18 @@
             .distinctUntilChanged()
 
     /**
-     * Keyguard states which should fully hide the keyguard.
-     *
-     * Note: [GONE] is not included as it is handled separately.
-     */
-    private val hiddenKeyguardStates = listOf(OCCLUDED, DREAMING, GLANCEABLE_HUB)
-
-    /**
      * Keyguard should not show if fully transitioned into a hidden keyguard state or if
      * transitioning between hidden states.
      */
     private val hideKeyguard: Flow<Boolean> =
-        (hiddenKeyguardStates.map { state ->
-                keyguardTransitionInteractor
-                    .transitionValue(state)
-                    .map { it == 1f }
-                    .onStart { emit(false) }
-            } +
-                listOf(
-                    communalInteractor.isIdleOnCommunal,
-                    keyguardTransitionInteractor
-                        .transitionValue(scene = Scenes.Gone, stateWithoutSceneContainer = GONE)
-                        .map { it == 1f }
-                        .onStart { emit(false) },
-                    keyguardTransitionInteractor
-                        .isInTransitionWhere(
-                            fromStatePredicate = { hiddenKeyguardStates.contains(it) },
-                            toStatePredicate = { hiddenKeyguardStates.contains(it) },
-                        )
-                        .onStart { emit(false) },
-                ))
-            .any()
+        anyOf(
+            notificationShadeWindowModel.isKeyguardOccluded,
+            communalInteractor.isIdleOnCommunal,
+            keyguardTransitionInteractor
+                .transitionValue(scene = Scenes.Gone, stateWithoutSceneContainer = GONE)
+                .map { it == 1f }
+                .onStart { emit(false) },
+        )
 
     /** Last point that the root view was tapped */
     val lastRootViewTapPosition: Flow<Point?> = keyguardInteractor.lastRootViewTapPosition
@@ -333,16 +313,17 @@
                     .transitionValue(LOCKSCREEN)
                     .map { it > 0f }
                     .onStart { emit(false) },
-                keyguardTransitionInteractor.finishedKeyguardState.map {
-                    KeyguardState.lockscreenVisibleInState(it)
-                },
+                keyguardTransitionInteractor.isFinishedIn(
+                    scene = Scenes.Gone,
+                    stateWithoutSceneContainer = GONE
+                ),
                 deviceEntryInteractor.isBypassEnabled,
                 areNotifsFullyHiddenAnimated(),
                 isPulseExpandingAnimated(),
             ) { flows ->
                 val goneToAodTransitionRunning = flows[0] as Boolean
                 val isOnLockscreen = flows[1] as Boolean
-                val onKeyguard = flows[2] as Boolean
+                val isOnGone = flows[2] as Boolean
                 val isBypassEnabled = flows[3] as Boolean
                 val notifsFullyHidden = flows[4] as AnimatedValue<Boolean>
                 val pulseExpanding = flows[5] as AnimatedValue<Boolean>
@@ -352,8 +333,7 @@
                     // animation is playing, in which case we want them to be visible if we're
                     // animating in the AOD UI and will be switching to KEYGUARD shortly.
                     goneToAodTransitionRunning ||
-                        (!onKeyguard &&
-                            !screenOffAnimationController.shouldShowAodIconsWhenShade()) ->
+                        (isOnGone && !screenOffAnimationController.shouldShowAodIconsWhenShade()) ->
                         AnimatedValue.NotAnimating(false)
                     else ->
                         zip(notifsFullyHidden, pulseExpanding) {
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenContentViewModel.kt
index 2b6c3c0..adb63b7 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenContentViewModel.kt
@@ -25,7 +25,6 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardClockInteractor
 import com.android.systemui.keyguard.shared.model.ClockSize
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.SceneContainerOcclusionInteractor
 import com.android.systemui.scene.shared.model.Scenes
@@ -41,7 +40,6 @@
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOf
@@ -60,7 +58,7 @@
     private val unfoldTransitionInteractor: UnfoldTransitionInteractor,
     private val occlusionInteractor: SceneContainerOcclusionInteractor,
     private val deviceEntryInteractor: DeviceEntryInteractor,
-) : SysUiViewModel, ExclusiveActivatable() {
+) : ExclusiveActivatable() {
     @VisibleForTesting val clockSize = clockInteractor.clockSize
 
     val isUdfpsVisible: Boolean
@@ -92,13 +90,13 @@
                             end = end,
                         )
                     }
-                    .collectLatest { _unfoldTranslations.value = it }
+                    .collect { _unfoldTranslations.value = it }
             }
 
             launch {
                 occlusionInteractor.isOccludingActivityShown
                     .map { !it }
-                    .collectLatest { _isContentVisible.value = it }
+                    .collect { _isContentVisible.value = it }
             }
 
             awaitCancellation()
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToPrimaryBouncerTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToPrimaryBouncerTransitionViewModel.kt
index e64c614..c0b9efaa 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToPrimaryBouncerTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenToPrimaryBouncerTransitionViewModel.kt
@@ -23,7 +23,9 @@
 import com.android.systemui.keyguard.shared.model.KeyguardState.PRIMARY_BOUNCER
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
 import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.scene.ui.composable.transitions.FROM_LOCK_SCREEN_TO_BOUNCER_FADE_FRACTION
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -51,10 +53,18 @@
                 edge = Edge.create(from = LOCKSCREEN, to = PRIMARY_BOUNCER),
             )
 
+    private val alphaForAnimationStep: (Float) -> Float =
+        when {
+            SceneContainerFlag.isEnabled -> { step ->
+                    1f - Math.min((step / FROM_LOCK_SCREEN_TO_BOUNCER_FADE_FRACTION), 1f)
+                }
+            else -> { step -> 1f - step }
+        }
+
     val shortcutsAlpha: Flow<Float> =
         transitionAnimation.sharedFlow(
             duration = FromLockscreenTransitionInteractor.TO_PRIMARY_BOUNCER_DURATION,
-            onStep = { 1f - it }
+            onStep = alphaForAnimationStep
         )
 
     val lockscreenAlpha: Flow<Float> = shortcutsAlpha
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
similarity index 94%
rename from packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
index 7383f57..dd47678 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/LockscreenUserActionsViewModel.kt
@@ -28,26 +28,25 @@
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.util.kotlin.filterValuesNotNull
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 
 /** Models UI state and handles user input for the lockscreen scene. */
-class LockscreenSceneActionsViewModel
+class LockscreenUserActionsViewModel
 @AssistedInject
 constructor(
     private val deviceEntryInteractor: DeviceEntryInteractor,
     private val communalInteractor: CommunalInteractor,
     private val shadeInteractor: ShadeInteractor,
-) : SceneActionsViewModel() {
+) : UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         shadeInteractor.isShadeTouchable
@@ -100,7 +99,7 @@
                     }
                 }
             }
-            .collectLatest { setActions(it) }
+            .collect { setActions(it) }
     }
 
     private fun swipeDownFromTop(pointerCount: Int): Swipe {
@@ -120,6 +119,6 @@
 
     @AssistedFactory
     interface Factory {
-        fun create(): LockscreenSceneActionsViewModel
+        fun create(): LockscreenUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModel.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModel.kt
index af01930..4fb2b9b 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModel.kt
@@ -17,15 +17,19 @@
 package com.android.systemui.keyguard.ui.viewmodel
 
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryUdfpsInteractor
 import com.android.systemui.keyguard.domain.interactor.FromOccludedTransitionInteractor
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState.DOZING
 import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
 import com.android.systemui.keyguard.ui.KeyguardTransitionAnimationFlow
+import com.android.systemui.keyguard.ui.transitions.DeviceEntryIconTransition
 import javax.inject.Inject
 import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flatMapLatest
 
 /**
  * Breaks down OCCLUDED->DOZING transition into discrete steps for corresponding views to consume.
@@ -35,8 +39,9 @@
 class OccludedToDozingTransitionViewModel
 @Inject
 constructor(
+    deviceEntryUdfpsInteractor: DeviceEntryUdfpsInteractor,
     animationFlow: KeyguardTransitionAnimationFlow,
-) {
+) : DeviceEntryIconTransition {
     private val transitionAnimation =
         animationFlow.setup(
             duration = FromOccludedTransitionInteractor.TO_DOZING_DURATION,
@@ -50,4 +55,17 @@
             duration = 250.milliseconds,
             onStep = { it },
         )
+
+    val deviceEntryBackgroundViewAlpha: Flow<Float> =
+        transitionAnimation.immediatelyTransitionTo(0f)
+
+    override val deviceEntryParentViewAlpha: Flow<Float> =
+        deviceEntryUdfpsInteractor.isUdfpsEnrolledAndEnabled.flatMapLatest { udfpsEnrolledAndEnabled
+            ->
+            if (udfpsEnrolledAndEnabled) {
+                transitionAnimation.immediatelyTransitionTo(1f)
+            } else {
+                emptyFlow()
+            }
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/ShadeDependentFlows.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/ShadeDependentFlows.kt
index e45d537..708b408 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/ShadeDependentFlows.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/viewmodel/ShadeDependentFlows.kt
@@ -34,7 +34,9 @@
 ) {
     /** When the last keyguard state transition started, was the shade fully expanded? */
     private val lastStartedTransitionHadShadeFullyExpanded: Flow<Boolean> =
-        transitionInteractor.startedKeyguardState.sample(shadeInteractor.isAnyFullyExpanded)
+        transitionInteractor.startedKeyguardTransitionStep.sample(
+            shadeInteractor.isAnyFullyExpanded
+        )
 
     /**
      * Decide which flow to use depending on the shade expansion state at the start of the last
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/Activatable.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/Activatable.kt
index bd3d40b..c1768a4 100644
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/Activatable.kt
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/Activatable.kt
@@ -19,6 +19,7 @@
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.remember
+import com.android.app.tracing.coroutines.traceCoroutine
 
 /** Defines interface for classes that can be activated to do coroutine work. */
 interface Activatable {
@@ -66,13 +67,19 @@
  *
  * If the [key] changes, the old [Activatable] is deactivated and a new one will be instantiated,
  * activated, and returned.
+ *
+ * The [traceName] is used for coroutine performance tracing purposes. Please try to use a label
+ * that's unique enough and easy enough to find in code search; this should help correlate
+ * performance findings with actual code. One recommendation: prefer whole string literals instead
+ * of some complex concatenation or templating scheme.
  */
 @Composable
 fun <T : Activatable> rememberActivated(
+    traceName: String,
     key: Any = Unit,
     factory: () -> T,
 ): T {
     val instance = remember(key) { factory() }
-    LaunchedEffect(instance) { instance.activate() }
+    LaunchedEffect(instance) { traceCoroutine(traceName) { instance.activate() } }
     return instance
 }
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/Hydrator.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/Hydrator.kt
index 59ec2af..df1394b 100644
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/Hydrator.kt
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/Hydrator.kt
@@ -19,6 +19,8 @@
 import androidx.compose.runtime.State
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.snapshots.StateFactoryMarker
+import com.android.app.tracing.coroutines.launch
+import com.android.app.tracing.coroutines.traceCoroutine
 import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.Flow
@@ -37,35 +39,55 @@
  * }
  * ```
  */
-class Hydrator : ExclusiveActivatable() {
+class Hydrator(
+    /**
+     * A name for performance tracing purposes.
+     *
+     * Please use a short string literal that's easy to find in code search. Try to avoid
+     * concatenation or templating.
+     */
+    private val traceName: String,
+) : ExclusiveActivatable() {
 
-    private val children = mutableListOf<Activatable>()
+    private val children = mutableListOf<NamedActivatable>()
 
     /**
-     * Returns a snapshot [State] that's kept up-to-date as long as the [SysUiViewModel] is active.
+     * Returns a snapshot [State] that's kept up-to-date as long as its owner is active.
      *
+     * @param traceName Used for coroutine performance tracing purposes. Please try to use a label
+     *   that's unique enough and easy enough to find in code search; this should help correlate
+     *   performance findings with actual code. One recommendation: prefer whole string literals
+     *   instead of some complex concatenation or templating scheme.
      * @param source The upstream [StateFlow] to collect from; values emitted to it will be
      *   automatically set on the returned [State].
      */
     @StateFactoryMarker
     fun <T> hydratedStateOf(
+        traceName: String,
         source: StateFlow<T>,
     ): State<T> {
         return hydratedStateOf(
+            traceName = traceName,
             initialValue = source.value,
             source = source,
         )
     }
 
     /**
-     * Returns a snapshot [State] that's kept up-to-date as long as the [SysUiViewModel] is active.
+     * Returns a snapshot [State] that's kept up-to-date as long as its owner is active.
      *
+     * @param traceName Used for coroutine performance tracing purposes. Please try to use a label
+     *   that's unique enough and easy enough to find in code search; this should help correlate
+     *   performance findings with actual code. One recommendation: prefer whole string literals
+     *   instead of some complex concatenation or templating scheme. Use `null` to disable
+     *   performance tracing for this state.
      * @param initialValue The first value to place on the [State]
      * @param source The upstream [Flow] to collect from; values emitted to it will be automatically
      *   set on the returned [State].
      */
     @StateFactoryMarker
     fun <T> hydratedStateOf(
+        traceName: String?,
         initialValue: T,
         source: Flow<T>,
     ): State<T> {
@@ -73,18 +95,35 @@
 
         val mutableState = mutableStateOf(initialValue)
         children.add(
-            object : ExclusiveActivatable() {
-                override suspend fun onActivated(): Nothing {
-                    source.collect { mutableState.value = it }
-                    awaitCancellation()
-                }
-            }
+            NamedActivatable(
+                traceName = traceName,
+                activatable =
+                    object : ExclusiveActivatable() {
+                        override suspend fun onActivated(): Nothing {
+                            source.collect { mutableState.value = it }
+                            awaitCancellation()
+                        }
+                    },
+            )
         )
         return mutableState
     }
 
     override suspend fun onActivated() = coroutineScope {
-        children.forEach { child -> launch { child.activate() } }
-        awaitCancellation()
+        traceCoroutine(traceName) {
+            children.forEach { child ->
+                if (child.traceName != null) {
+                    launch(spanName = child.traceName) { child.activatable.activate() }
+                } else {
+                    launch { child.activatable.activate() }
+                }
+            }
+            awaitCancellation()
+        }
     }
+
+    private data class NamedActivatable(
+        val traceName: String?,
+        val activatable: Activatable,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt b/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
index 29ffcbd..508b04e 100644
--- a/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/lifecycle/SysUiViewModel.kt
@@ -20,36 +20,46 @@
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.remember
+import com.android.app.tracing.coroutines.traceCoroutine
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
 
-/** Defines interface for all System UI view-models. */
-interface SysUiViewModel
-
 /**
- * Returns a remembered [SysUiViewModel] of the type [T]. If the returned instance is also an
+ * Returns a remembered view-model of the type [T]. If the returned instance is also an
  * [Activatable], it's automatically kept active until this composable leaves the composition; if
- * the [key] changes, the old [SysUiViewModel] is deactivated and a new one will be instantiated,
+ * the [key] changes, the old view-model is deactivated and a new one will be instantiated,
  * activated, and returned.
+ *
+ * The [traceName] is used for coroutine performance tracing purposes. Please try to use a label
+ * that's unique enough and easy enough to find in code search; this should help correlate
+ * performance findings with actual code. One recommendation: prefer whole string literals instead
+ * of some complex concatenation or templating scheme.
  */
 @Composable
-fun <T : SysUiViewModel> rememberViewModel(
+fun <T> rememberViewModel(
+    traceName: String,
     key: Any = Unit,
     factory: () -> T,
 ): T {
     val instance = remember(key) { factory() }
     if (instance is Activatable) {
-        LaunchedEffect(instance) { instance.activate() }
+        LaunchedEffect(instance) { traceCoroutine(traceName) { instance.activate() } }
     }
     return instance
 }
 
 /**
- * Invokes [block] in a new coroutine with a new [SysUiViewModel] that is automatically activated
- * whenever `this` [View]'s Window's [WindowLifecycleState] is at least at
- * [minWindowLifecycleState], and is automatically canceled once that is no longer the case.
+ * Invokes [block] in a new coroutine with a new view-model that is automatically activated whenever
+ * `this` [View]'s Window's [WindowLifecycleState] is at least at [minWindowLifecycleState], and is
+ * automatically canceled once that is no longer the case.
+ *
+ * The [traceName] is used for coroutine performance tracing purposes. Please try to use a label
+ * that's unique enough and easy enough to find in code search; this should help correlate
+ * performance findings with actual code. One recommendation: prefer whole string literals instead
+ * of some complex concatenation or templating scheme.
  */
-suspend fun <T : SysUiViewModel> View.viewModel(
+suspend fun <T> View.viewModel(
+    traceName: String,
     minWindowLifecycleState: WindowLifecycleState,
     factory: () -> T,
     block: suspend CoroutineScope.(T) -> Unit,
@@ -57,7 +67,7 @@
     repeatOnWindowLifecycle(minWindowLifecycleState) {
         val instance = factory()
         if (instance is Activatable) {
-            launch { instance.activate() }
+            launch { traceCoroutine(traceName) { instance.activate() } }
         }
         block(instance)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
index ba3c1d2..ed76646 100644
--- a/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/LogModule.java
@@ -382,6 +382,16 @@
         return factory.create("MediaLog", 20);
     }
 
+    /**
+     * Provides a buffer for media device changes
+     */
+    @Provides
+    @SysUISingleton
+    @MediaDeviceLog
+    public static LogBuffer providesMediaDeviceLogBuffer(LogBufferFactory factory) {
+        return factory.create("MediaDeviceLog", 50);
+    }
+
     /** Allows logging buffers to be tweaked via adb on debug builds but not on prod builds. */
     @Provides
     @SysUISingleton
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt b/packages/SystemUI/src/com/android/systemui/log/dagger/MediaDeviceLog.kt
similarity index 66%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
copy to packages/SystemUI/src/com/android/systemui/log/dagger/MediaDeviceLog.kt
index 60d97d1..06bd269 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
+++ b/packages/SystemUI/src/com/android/systemui/log/dagger/MediaDeviceLog.kt
@@ -14,9 +14,13 @@
  * limitations under the License.
  */
 
-package com.android.systemui.bouncer.shared.flag
+package com.android.systemui.log.dagger
 
-import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.log.LogBuffer
+import javax.inject.Qualifier
 
-var Kosmos.fakeComposeBouncerFlags by Kosmos.Fixture { FakeComposeBouncerFlags() }
-val Kosmos.composeBouncerFlags by Kosmos.Fixture<ComposeBouncerFlags> { fakeComposeBouncerFlags }
+/** A [LogBuffer] for [com.android.systemui.media.controls.domain.pipeline.MediaDeviceLogger] */
+@Qualifier
+@MustBeDocumented
+@Retention(AnnotationRetention.RUNTIME)
+annotation class MediaDeviceLog
diff --git a/packages/SystemUI/src/com/android/systemui/log/table/TableLogBuffer.kt b/packages/SystemUI/src/com/android/systemui/log/table/TableLogBuffer.kt
index 2089cce5..89a599a 100644
--- a/packages/SystemUI/src/com/android/systemui/log/table/TableLogBuffer.kt
+++ b/packages/SystemUI/src/com/android/systemui/log/table/TableLogBuffer.kt
@@ -16,19 +16,17 @@
 
 package com.android.systemui.log.table
 
+import android.annotation.SuppressLint
 import android.icu.text.SimpleDateFormat
 import android.os.Trace
 import com.android.systemui.Dumpable
 import com.android.systemui.common.buffer.RingBuffer
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.log.LogcatEchoTracker
 import com.android.systemui.log.core.LogLevel
 import com.android.systemui.plugins.log.TableLogBufferBase
 import com.android.systemui.util.time.SystemClock
 import java.io.PrintWriter
 import java.util.Locale
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
 
 /**
  * A logger that logs changes in table format.
@@ -75,13 +73,12 @@
  *
  * @param maxSize the maximum size of the buffer. Must be > 0.
  */
+@SuppressLint("DumpableNotRegistered") // Registered as dumpable in [TableLogBufferFactory]
 class TableLogBuffer(
     maxSize: Int,
     private val name: String,
     private val systemClock: SystemClock,
     private val logcatEchoTracker: LogcatEchoTracker,
-    @Background private val bgDispatcher: CoroutineDispatcher,
-    private val coroutineScope: CoroutineScope,
     private val localLogcat: LogProxy = LogProxyDefault(),
 ) : Dumpable, TableLogBufferBase {
     init {
diff --git a/packages/SystemUI/src/com/android/systemui/log/table/TableLogBufferFactory.kt b/packages/SystemUI/src/com/android/systemui/log/table/TableLogBufferFactory.kt
index ff523ae..425e674e 100644
--- a/packages/SystemUI/src/com/android/systemui/log/table/TableLogBufferFactory.kt
+++ b/packages/SystemUI/src/com/android/systemui/log/table/TableLogBufferFactory.kt
@@ -17,15 +17,11 @@
 package com.android.systemui.log.table
 
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.log.LogBufferHelper.Companion.adjustMaxSize
 import com.android.systemui.log.LogcatEchoTracker
 import com.android.systemui.util.time.SystemClock
 import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
 
 @SysUISingleton
 class TableLogBufferFactory
@@ -34,8 +30,6 @@
     private val dumpManager: DumpManager,
     private val systemClock: SystemClock,
     private val logcatEchoTracker: LogcatEchoTracker,
-    @Background private val bgDispatcher: CoroutineDispatcher,
-    @Application private val coroutineScope: CoroutineScope,
 ) {
     private val existingBuffers = mutableMapOf<String, TableLogBuffer>()
 
@@ -58,8 +52,6 @@
                 name,
                 systemClock,
                 logcatEchoTracker,
-                bgDispatcher,
-                coroutineScope,
             )
         dumpManager.registerTableLogBuffer(name, tableBuffer)
         return tableBuffer
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaActions.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaActions.kt
index 70189b7..bcf748e 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaActions.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaActions.kt
@@ -30,6 +30,7 @@
 import androidx.media.utils.MediaConstants
 import com.android.systemui.media.controls.domain.pipeline.LegacyMediaDataManagerImpl.Companion.MAX_COMPACT_ACTIONS
 import com.android.systemui.media.controls.domain.pipeline.LegacyMediaDataManagerImpl.Companion.MAX_NOTIFICATION_ACTIONS
+import com.android.systemui.media.controls.shared.MediaControlDrawables
 import com.android.systemui.media.controls.shared.model.MediaAction
 import com.android.systemui.media.controls.shared.model.MediaButton
 import com.android.systemui.plugins.ActivityStarter
@@ -169,7 +170,7 @@
         }
         PlaybackState.ACTION_SKIP_TO_PREVIOUS -> {
             MediaAction(
-                context.getDrawable(R.drawable.ic_media_prev),
+                MediaControlDrawables.getPrevIcon(context),
                 { controller.transportControls.skipToPrevious() },
                 context.getString(R.string.controls_media_button_prev),
                 null
@@ -177,7 +178,7 @@
         }
         PlaybackState.ACTION_SKIP_TO_NEXT -> {
             MediaAction(
-                context.getDrawable(R.drawable.ic_media_next),
+                MediaControlDrawables.getNextIcon(context),
                 { controller.transportControls.skipToNext() },
                 context.getString(R.string.controls_media_button_next),
                 null
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessor.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessor.kt
index 916f8b2..4555810 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessor.kt
@@ -16,9 +16,8 @@
 
 package com.android.systemui.media.controls.domain.pipeline
 
+import android.annotation.MainThread
 import android.annotation.SuppressLint
-import android.app.ActivityOptions
-import android.app.BroadcastOptions
 import android.app.Notification
 import android.app.Notification.EXTRA_SUBSTITUTE_APP_NAME
 import android.app.PendingIntent
@@ -39,7 +38,6 @@
 import android.content.pm.PackageManager
 import android.graphics.Bitmap
 import android.graphics.ImageDecoder
-import android.graphics.drawable.Animatable
 import android.graphics.drawable.Icon
 import android.media.MediaDescription
 import android.media.MediaMetadata
@@ -47,7 +45,6 @@
 import android.media.session.MediaSession
 import android.media.session.PlaybackState
 import android.net.Uri
-import android.os.Handler
 import android.os.Parcelable
 import android.os.Process
 import android.os.UserHandle
@@ -63,6 +60,7 @@
 import com.android.internal.logging.InstanceId
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.CoreStartable
+import com.android.systemui.Flags
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
@@ -73,7 +71,6 @@
 import com.android.systemui.media.controls.domain.pipeline.MediaDataManager.Companion.isMediaNotification
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.media.controls.domain.resume.ResumeMediaBrowser
-import com.android.systemui.media.controls.shared.MediaControlDrawables
 import com.android.systemui.media.controls.shared.model.EXTRA_KEY_TRIGGER_SOURCE
 import com.android.systemui.media.controls.shared.model.EXTRA_VALUE_TRIGGER_PERIODIC
 import com.android.systemui.media.controls.shared.model.MediaAction
@@ -92,7 +89,6 @@
 import com.android.systemui.plugins.BcSmartspaceDataPlugin
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import com.android.systemui.statusbar.NotificationMediaManager.isConnectingState
 import com.android.systemui.statusbar.NotificationMediaManager.isPlayingState
 import com.android.systemui.statusbar.notification.row.HybridGroupManager
 import com.android.systemui.util.Assert
@@ -137,7 +133,7 @@
     @Background private val backgroundExecutor: Executor,
     @Main private val uiExecutor: Executor,
     @Main private val foregroundExecutor: DelayableExecutor,
-    @Main private val handler: Handler,
+    @Main private val mainDispatcher: CoroutineDispatcher,
     private val mediaControllerFactory: MediaControllerFactory,
     private val broadcastDispatcher: BroadcastDispatcher,
     private val dumpManager: DumpManager,
@@ -152,6 +148,7 @@
     private val smartspaceManager: SmartspaceManager?,
     private val keyguardUpdateMonitor: KeyguardUpdateMonitor,
     private val mediaDataRepository: MediaDataRepository,
+    private val mediaDataLoader: dagger.Lazy<MediaDataLoader>,
 ) : CoreStartable, BcSmartspaceDataPlugin.SmartspaceTargetListener {
 
     companion object {
@@ -217,7 +214,7 @@
         threadFactory: ThreadFactory,
         @Main uiExecutor: Executor,
         @Main foregroundExecutor: DelayableExecutor,
-        @Main handler: Handler,
+        @Main mainDispatcher: CoroutineDispatcher,
         mediaControllerFactory: MediaControllerFactory,
         dumpManager: DumpManager,
         broadcastDispatcher: BroadcastDispatcher,
@@ -230,6 +227,7 @@
         smartspaceManager: SmartspaceManager?,
         keyguardUpdateMonitor: KeyguardUpdateMonitor,
         mediaDataRepository: MediaDataRepository,
+        mediaDataLoader: dagger.Lazy<MediaDataLoader>,
     ) : this(
         context,
         applicationScope,
@@ -239,7 +237,7 @@
         threadFactory.buildExecutorOnNewThread(TAG),
         uiExecutor,
         foregroundExecutor,
-        handler,
+        mainDispatcher,
         mediaControllerFactory,
         broadcastDispatcher,
         dumpManager,
@@ -254,6 +252,7 @@
         smartspaceManager,
         keyguardUpdateMonitor,
         mediaDataRepository,
+        mediaDataLoader,
     )
 
     private val appChangeReceiver =
@@ -436,16 +435,30 @@
             logSingleVsMultipleMediaAdded(appUid, packageName, instanceId)
             logger.logResumeMediaAdded(appUid, packageName, instanceId)
         }
-        backgroundExecutor.execute {
-            loadMediaDataInBgForResumption(
-                userId,
-                desc,
-                action,
-                token,
-                appName,
-                appIntent,
-                packageName
-            )
+        if (Flags.mediaLoadMetadataViaMediaDataLoader()) {
+            applicationScope.launch {
+                loadMediaDataForResumption(
+                    userId,
+                    desc,
+                    action,
+                    token,
+                    appName,
+                    appIntent,
+                    packageName
+                )
+            }
+        } else {
+            backgroundExecutor.execute {
+                loadMediaDataInBgForResumption(
+                    userId,
+                    desc,
+                    action,
+                    token,
+                    appName,
+                    appIntent,
+                    packageName
+                )
+            }
         }
     }
 
@@ -471,7 +484,13 @@
         oldKey: String?,
         isNewlyActiveEntry: Boolean = false,
     ) {
-        backgroundExecutor.execute { loadMediaDataInBg(key, sbn, oldKey, isNewlyActiveEntry) }
+        if (Flags.mediaLoadMetadataViaMediaDataLoader()) {
+            applicationScope.launch {
+                loadMediaDataWithLoader(key, sbn, oldKey, isNewlyActiveEntry)
+            }
+        } else {
+            backgroundExecutor.execute { loadMediaDataInBg(key, sbn, oldKey, isNewlyActiveEntry) }
+        }
     }
 
     /** Add a listener for internal events. */
@@ -646,6 +665,75 @@
         }
     }
 
+    private suspend fun loadMediaDataForResumption(
+        userId: Int,
+        desc: MediaDescription,
+        resumeAction: Runnable,
+        token: MediaSession.Token,
+        appName: String,
+        appIntent: PendingIntent,
+        packageName: String
+    ) =
+        withContext(backgroundDispatcher) {
+            val lastActive = systemClock.elapsedRealtime()
+            val currentEntry = mediaDataRepository.mediaEntries.value[packageName]
+            val createdTimestampMillis = currentEntry?.createdTimestampMillis ?: 0L
+            val result =
+                mediaDataLoader
+                    .get()
+                    .loadMediaDataForResumption(
+                        userId,
+                        desc,
+                        resumeAction,
+                        currentEntry,
+                        token,
+                        appName,
+                        appIntent,
+                        packageName
+                    )
+            if (result == null || desc.title.isNullOrBlank()) {
+                Log.d(TAG, "No MediaData result for resumption")
+                mediaDataRepository.removeMediaEntry(packageName)
+                return@withContext
+            }
+
+            val instanceId = currentEntry?.instanceId ?: logger.getNewInstanceId()
+            withContext(mainDispatcher) {
+                onMediaDataLoaded(
+                    packageName,
+                    null,
+                    MediaData(
+                        userId = userId,
+                        initialized = true,
+                        app = result.appName,
+                        appIcon = null,
+                        artist = result.artist,
+                        song = result.song,
+                        artwork = result.artworkIcon,
+                        actions = result.actionIcons,
+                        actionsToShowInCompact = result.actionsToShowInCompact,
+                        semanticActions = result.semanticActions,
+                        packageName = packageName,
+                        token = result.token,
+                        clickIntent = result.clickIntent,
+                        device = result.device,
+                        active = false,
+                        resumeAction = resumeAction,
+                        resumption = true,
+                        notificationKey = packageName,
+                        hasCheckedForResume = true,
+                        lastActive = lastActive,
+                        createdTimestampMillis = createdTimestampMillis,
+                        instanceId = instanceId,
+                        appUid = result.appUid,
+                        isExplicit = result.isExplicit,
+                        resumeProgress = result.resumeProgress,
+                    )
+                )
+            }
+        }
+
+    @Deprecated("Cleanup when media_load_metadata_via_media_data_loader is cleaned up")
     private fun loadMediaDataInBgForResumption(
         userId: Int,
         desc: MediaDescription,
@@ -730,6 +818,82 @@
         }
     }
 
+    private suspend fun loadMediaDataWithLoader(
+        key: String,
+        sbn: StatusBarNotification,
+        oldKey: String?,
+        isNewlyActiveEntry: Boolean = false,
+    ) =
+        withContext(backgroundDispatcher) {
+            val lastActive = systemClock.elapsedRealtime()
+            val result = mediaDataLoader.get().loadMediaData(key, sbn)
+            if (result == null) {
+                Log.d(TAG, "No result from loadMediaData")
+                return@withContext
+            }
+
+            val currentEntry = mediaDataRepository.mediaEntries.value[key]
+            val instanceId = currentEntry?.instanceId ?: logger.getNewInstanceId()
+            val createdTimestampMillis = currentEntry?.createdTimestampMillis ?: 0L
+            val resumeAction: Runnable? = currentEntry?.resumeAction
+            val hasCheckedForResume = currentEntry?.hasCheckedForResume == true
+            val active = currentEntry?.active ?: true
+
+            // We need to log the correct media added.
+            if (isNewlyActiveEntry) {
+                logSingleVsMultipleMediaAdded(result.appUid, sbn.packageName, instanceId)
+                logger.logActiveMediaAdded(
+                    result.appUid,
+                    sbn.packageName,
+                    instanceId,
+                    result.playbackLocation
+                )
+            } else if (result.playbackLocation != currentEntry?.playbackLocation) {
+                logger.logPlaybackLocationChange(
+                    result.appUid,
+                    sbn.packageName,
+                    instanceId,
+                    result.playbackLocation
+                )
+            }
+
+            withContext(mainDispatcher) {
+                onMediaDataLoaded(
+                    key,
+                    oldKey,
+                    MediaData(
+                        userId = sbn.normalizedUserId,
+                        initialized = true,
+                        app = result.appName,
+                        appIcon = result.appIcon,
+                        artist = result.artist,
+                        song = result.song,
+                        artwork = result.artworkIcon,
+                        actions = result.actionIcons,
+                        actionsToShowInCompact = result.actionsToShowInCompact,
+                        semanticActions = result.semanticActions,
+                        packageName = sbn.packageName,
+                        token = result.token,
+                        clickIntent = result.clickIntent,
+                        device = result.device,
+                        active = active,
+                        resumeAction = resumeAction,
+                        playbackLocation = result.playbackLocation,
+                        notificationKey = key,
+                        hasCheckedForResume = hasCheckedForResume,
+                        isPlaying = result.isPlaying,
+                        isClearable = !sbn.isOngoing,
+                        lastActive = lastActive,
+                        createdTimestampMillis = createdTimestampMillis,
+                        instanceId = instanceId,
+                        appUid = result.appUid,
+                        isExplicit = result.isExplicit,
+                    )
+                )
+            }
+        }
+
+    @Deprecated("Cleanup when media_load_metadata_via_media_data_loader is cleaned up")
     fun loadMediaDataInBg(
         key: String,
         sbn: StatusBarNotification,
@@ -843,7 +1007,7 @@
         var actionsToShowCollapsed: List<Int> = emptyList()
         val semanticActions = createActionsFromState(sbn.packageName, mediaController, sbn.user)
         if (semanticActions == null) {
-            val actions = createActionsFromNotification(sbn)
+            val actions = createActionsFromNotification(context, activityStarter, sbn)
             actionIcons = actions.first
             actionsToShowCollapsed = actions.second
         }
@@ -926,6 +1090,7 @@
         }
     }
 
+    @Deprecated("Cleanup when media_load_metadata_via_media_data_loader is cleaned up")
     private fun getAppInfoFromPackage(packageName: String): ApplicationInfo? {
         try {
             return context.packageManager.getApplicationInfo(packageName, 0)
@@ -935,6 +1100,7 @@
         return null
     }
 
+    @Deprecated("Cleanup when media_load_metadata_via_media_data_loader is cleaned up")
     private fun getAppName(sbn: StatusBarNotification, appInfo: ApplicationInfo?): String {
         val name = sbn.notification.extras.getString(EXTRA_SUBSTITUTE_APP_NAME)
         if (name != null) {
@@ -948,78 +1114,6 @@
         }
     }
 
-    /** Generate action buttons based on notification actions */
-    private fun createActionsFromNotification(
-        sbn: StatusBarNotification
-    ): Pair<List<MediaAction>, List<Int>> {
-        val notif = sbn.notification
-        val actionIcons: MutableList<MediaAction> = ArrayList()
-        val actions = notif.actions
-        var actionsToShowCollapsed =
-            notif.extras.getIntArray(Notification.EXTRA_COMPACT_ACTIONS)?.toMutableList()
-                ?: mutableListOf()
-        if (actionsToShowCollapsed.size > MAX_COMPACT_ACTIONS) {
-            Log.e(
-                TAG,
-                "Too many compact actions for ${sbn.key}," +
-                    "limiting to first $MAX_COMPACT_ACTIONS"
-            )
-            actionsToShowCollapsed = actionsToShowCollapsed.subList(0, MAX_COMPACT_ACTIONS)
-        }
-
-        if (actions != null) {
-            for ((index, action) in actions.withIndex()) {
-                if (index == MAX_NOTIFICATION_ACTIONS) {
-                    Log.w(
-                        TAG,
-                        "Too many notification actions for ${sbn.key}," +
-                            " limiting to first $MAX_NOTIFICATION_ACTIONS"
-                    )
-                    break
-                }
-                if (action.getIcon() == null) {
-                    if (DEBUG) Log.i(TAG, "No icon for action $index ${action.title}")
-                    actionsToShowCollapsed.remove(index)
-                    continue
-                }
-                val runnable =
-                    if (action.actionIntent != null) {
-                        Runnable {
-                            if (action.actionIntent.isActivity) {
-                                activityStarter.startPendingIntentDismissingKeyguard(
-                                    action.actionIntent
-                                )
-                            } else if (action.isAuthenticationRequired()) {
-                                activityStarter.dismissKeyguardThenExecute(
-                                    {
-                                        var result = sendPendingIntent(action.actionIntent)
-                                        result
-                                    },
-                                    {},
-                                    true
-                                )
-                            } else {
-                                sendPendingIntent(action.actionIntent)
-                            }
-                        }
-                    } else {
-                        null
-                    }
-                val mediaActionIcon =
-                    if (action.getIcon()?.getType() == Icon.TYPE_RESOURCE) {
-                            Icon.createWithResource(sbn.packageName, action.getIcon()!!.getResId())
-                        } else {
-                            action.getIcon()
-                        }
-                        .setTint(themeText)
-                        .loadDrawable(context)
-                val mediaAction = MediaAction(mediaActionIcon, runnable, action.title, null)
-                actionIcons.add(mediaAction)
-            }
-        }
-        return Pair(actionIcons, actionsToShowCollapsed)
-    }
-
     /**
      * Generates action button info for this media session based on the PlaybackState
      *
@@ -1036,174 +1130,14 @@
         controller: MediaController,
         user: UserHandle
     ): MediaButton? {
-        val state = controller.playbackState
-        if (state == null || !mediaFlags.areMediaSessionActionsEnabled(packageName, user)) {
+        if (!mediaFlags.areMediaSessionActionsEnabled(packageName, user)) {
             return null
         }
-
-        // First, check for standard actions
-        val playOrPause =
-            if (isConnectingState(state.state)) {
-                // Spinner needs to be animating to render anything. Start it here.
-                val drawable = MediaControlDrawables.getProgress(context)
-                (drawable as Animatable).start()
-                MediaAction(
-                    drawable,
-                    null, // no action to perform when clicked
-                    context.getString(R.string.controls_media_button_connecting),
-                    MediaControlDrawables.getConnecting(context),
-                    // Specify a rebind id to prevent the spinner from restarting on later binds.
-                    com.android.internal.R.drawable.progress_small_material
-                )
-            } else if (isPlayingState(state.state)) {
-                getStandardAction(controller, state.actions, PlaybackState.ACTION_PAUSE)
-            } else {
-                getStandardAction(controller, state.actions, PlaybackState.ACTION_PLAY)
-            }
-        val prevButton =
-            getStandardAction(controller, state.actions, PlaybackState.ACTION_SKIP_TO_PREVIOUS)
-        val nextButton =
-            getStandardAction(controller, state.actions, PlaybackState.ACTION_SKIP_TO_NEXT)
-
-        // Then, create a way to build any custom actions that will be needed
-        val customActions =
-            state.customActions
-                .asSequence()
-                .filterNotNull()
-                .map { getCustomAction(packageName, controller, it) }
-                .iterator()
-        fun nextCustomAction() = if (customActions.hasNext()) customActions.next() else null
-
-        // Finally, assign the remaining button slots: play/pause A B C D
-        // A = previous, else custom action (if not reserved)
-        // B = next, else custom action (if not reserved)
-        // C and D are always custom actions
-        val reservePrev =
-            controller.extras?.getBoolean(
-                MediaConstants.SESSION_EXTRAS_KEY_SLOT_RESERVATION_SKIP_TO_PREV
-            ) == true
-        val reserveNext =
-            controller.extras?.getBoolean(
-                MediaConstants.SESSION_EXTRAS_KEY_SLOT_RESERVATION_SKIP_TO_NEXT
-            ) == true
-
-        val prevOrCustom =
-            if (prevButton != null) {
-                prevButton
-            } else if (!reservePrev) {
-                nextCustomAction()
-            } else {
-                null
-            }
-
-        val nextOrCustom =
-            if (nextButton != null) {
-                nextButton
-            } else if (!reserveNext) {
-                nextCustomAction()
-            } else {
-                null
-            }
-
-        return MediaButton(
-            playOrPause,
-            nextOrCustom,
-            prevOrCustom,
-            nextCustomAction(),
-            nextCustomAction(),
-            reserveNext,
-            reservePrev
-        )
-    }
-
-    /**
-     * Create a [MediaAction] for a given action and media session
-     *
-     * @param controller MediaController for the session
-     * @param stateActions The actions included with the session's [PlaybackState]
-     * @param action A [PlaybackState.Actions] value representing what action to generate. One of:
-     * ```
-     *      [PlaybackState.ACTION_PLAY]
-     *      [PlaybackState.ACTION_PAUSE]
-     *      [PlaybackState.ACTION_SKIP_TO_PREVIOUS]
-     *      [PlaybackState.ACTION_SKIP_TO_NEXT]
-     * @return
-     * ```
-     *
-     * A [MediaAction] with correct values set, or null if the state doesn't support it
-     */
-    private fun getStandardAction(
-        controller: MediaController,
-        stateActions: Long,
-        @PlaybackState.Actions action: Long
-    ): MediaAction? {
-        if (!includesAction(stateActions, action)) {
-            return null
-        }
-
-        return when (action) {
-            PlaybackState.ACTION_PLAY -> {
-                MediaAction(
-                    MediaControlDrawables.getPlayIcon(context),
-                    { controller.transportControls.play() },
-                    context.getString(R.string.controls_media_button_play),
-                    MediaControlDrawables.getPlayBackground(context)
-                )
-            }
-            PlaybackState.ACTION_PAUSE -> {
-                MediaAction(
-                    MediaControlDrawables.getPauseIcon(context),
-                    { controller.transportControls.pause() },
-                    context.getString(R.string.controls_media_button_pause),
-                    MediaControlDrawables.getPauseBackground(context)
-                )
-            }
-            PlaybackState.ACTION_SKIP_TO_PREVIOUS -> {
-                MediaAction(
-                    MediaControlDrawables.getPrevIcon(context),
-                    { controller.transportControls.skipToPrevious() },
-                    context.getString(R.string.controls_media_button_prev),
-                    null
-                )
-            }
-            PlaybackState.ACTION_SKIP_TO_NEXT -> {
-                MediaAction(
-                    MediaControlDrawables.getNextIcon(context),
-                    { controller.transportControls.skipToNext() },
-                    context.getString(R.string.controls_media_button_next),
-                    null
-                )
-            }
-            else -> null
-        }
-    }
-
-    /** Check whether the actions from a [PlaybackState] include a specific action */
-    private fun includesAction(stateActions: Long, @PlaybackState.Actions action: Long): Boolean {
-        if (
-            (action == PlaybackState.ACTION_PLAY || action == PlaybackState.ACTION_PAUSE) &&
-                (stateActions and PlaybackState.ACTION_PLAY_PAUSE > 0L)
-        ) {
-            return true
-        }
-        return (stateActions and action != 0L)
-    }
-
-    /** Get a [MediaAction] representing a [PlaybackState.CustomAction] */
-    private fun getCustomAction(
-        packageName: String,
-        controller: MediaController,
-        customAction: PlaybackState.CustomAction
-    ): MediaAction {
-        return MediaAction(
-            Icon.createWithResource(packageName, customAction.icon).loadDrawable(context),
-            { controller.transportControls.sendCustomAction(customAction, customAction.extras) },
-            customAction.name,
-            null
-        )
+        return createActionsFromState(context, packageName, controller)
     }
 
     /** Load a bitmap from the various Art metadata URIs */
+    @Deprecated("Cleanup when media_load_metadata_via_media_data_loader is cleaned up")
     private fun loadBitmapFromUri(metadata: MediaMetadata): Bitmap? {
         for (uri in ART_URIS) {
             val uriString = metadata.getString(uri)
@@ -1218,21 +1152,6 @@
         return null
     }
 
-    private fun sendPendingIntent(intent: PendingIntent): Boolean {
-        return try {
-            val options = BroadcastOptions.makeBasic()
-            options.setInteractive(true)
-            options.setPendingIntentBackgroundActivityStartMode(
-                ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED
-            )
-            intent.send(options.toBundle())
-            true
-        } catch (e: PendingIntent.CanceledException) {
-            Log.d(TAG, "Intent canceled", e)
-            false
-        }
-    }
-
     /** Returns a bitmap if the user can access the given URI, else null */
     private fun loadBitmapFromUriForUser(
         uri: Uri,
@@ -1309,10 +1228,11 @@
                 .loadDrawable(context),
             action,
             context.getString(R.string.controls_media_resume),
-            MediaControlDrawables.getPlayBackground(context)
+            context.getDrawable(R.drawable.ic_media_play_container)
         )
     }
 
+    @MainThread
     fun onMediaDataLoaded(key: String, oldKey: String?, data: MediaData) =
         traceSection("MediaDataProcessor#onMediaDataLoaded") {
             Assert.isMainThread()
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLogger.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLogger.kt
new file mode 100644
index 0000000..f886166
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLogger.kt
@@ -0,0 +1,116 @@
+/*
+ * 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.domain.pipeline
+
+import android.media.session.MediaController
+import com.android.settingslib.media.MediaDevice
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.LogLevel
+import com.android.systemui.log.dagger.MediaDeviceLog
+import com.android.systemui.media.controls.shared.model.MediaDeviceData
+import javax.inject.Inject
+
+/** A [LogBuffer] for media device changes */
+class MediaDeviceLogger @Inject constructor(@MediaDeviceLog private val buffer: LogBuffer) {
+
+    fun logBroadcastEvent(event: String, reason: Int, broadcastId: Int) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = event
+                int1 = reason
+                int2 = broadcastId
+            },
+            { "$str1, reason = $int1, broadcastId = $int2" }
+        )
+    }
+
+    fun logBroadcastEvent(event: String, reason: Int) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = event
+                int1 = reason
+            },
+            { "$str1, reason = $int1" }
+        )
+    }
+
+    fun logBroadcastMetadataChanged(broadcastId: Int, metadata: String) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                int1 = broadcastId
+                str1 = metadata
+            },
+            { "onBroadcastMetadataChanged, broadcastId = $int1, metadata = $str1" }
+        )
+    }
+
+    fun logNewDeviceName(name: String?) {
+        buffer.log(TAG, LogLevel.DEBUG, { str1 = name }, { "New device name $str1" })
+    }
+
+    fun logLocalDevice(sassDevice: MediaDeviceData?, connectedDevice: MediaDeviceData?) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = sassDevice?.name?.toString()
+                str2 = connectedDevice?.name?.toString()
+            },
+            { "Local device: $str1 or $str2" }
+        )
+    }
+
+    fun logRemoteDevice(routingSessionName: CharSequence?, connectedDevice: MediaDeviceData?) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = routingSessionName?.toString()
+                str2 = connectedDevice?.name?.toString()
+            },
+            { "Remote device: $str1 or $str2 or unknown" }
+        )
+    }
+
+    fun logDeviceName(
+        device: MediaDevice?,
+        controller: MediaController?,
+        routingSessionName: CharSequence?,
+        selectedRouteName: CharSequence?
+    ) {
+        buffer.log(
+            TAG,
+            LogLevel.DEBUG,
+            {
+                str1 = "device $device, controller: $controller"
+                str2 = routingSessionName?.toString()
+                str3 = selectedRouteName?.toString()
+            },
+            { "$str1, routingSession $str2 or selected route $str3" }
+        )
+    }
+
+    companion object {
+        private const val TAG = "MediaDeviceLog"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
index a193f7f..49b53c2 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManager.kt
@@ -71,6 +71,7 @@
     private val localBluetoothManager: Lazy<LocalBluetoothManager?>,
     @Main private val fgExecutor: Executor,
     @Background private val bgExecutor: Executor,
+    private val logger: MediaDeviceLogger,
 ) : MediaDataManager.Listener {
 
     private val listeners: MutableSet<Listener> = mutableSetOf()
@@ -281,59 +282,38 @@
         }
 
         override fun onBroadcastStarted(reason: Int, broadcastId: Int) {
-            if (DEBUG) {
-                Log.d(TAG, "onBroadcastStarted(), reason = $reason , broadcastId = $broadcastId")
-            }
+            logger.logBroadcastEvent("onBroadcastStarted", reason, broadcastId)
             updateCurrent()
         }
 
         override fun onBroadcastStartFailed(reason: Int) {
-            if (DEBUG) {
-                Log.d(TAG, "onBroadcastStartFailed(), reason = $reason")
-            }
+            logger.logBroadcastEvent("onBroadcastStartFailed", reason)
         }
 
         override fun onBroadcastMetadataChanged(
             broadcastId: Int,
             metadata: BluetoothLeBroadcastMetadata
         ) {
-            if (DEBUG) {
-                Log.d(
-                    TAG,
-                    "onBroadcastMetadataChanged(), broadcastId = $broadcastId , " +
-                        "metadata = $metadata"
-                )
-            }
+            logger.logBroadcastMetadataChanged(broadcastId, metadata.toString())
             updateCurrent()
         }
 
         override fun onBroadcastStopped(reason: Int, broadcastId: Int) {
-            if (DEBUG) {
-                Log.d(TAG, "onBroadcastStopped(), reason = $reason , broadcastId = $broadcastId")
-            }
+            logger.logBroadcastEvent("onBroadcastStopped", reason, broadcastId)
             updateCurrent()
         }
 
         override fun onBroadcastStopFailed(reason: Int) {
-            if (DEBUG) {
-                Log.d(TAG, "onBroadcastStopFailed(), reason = $reason")
-            }
+            logger.logBroadcastEvent("onBroadcastStopFailed", reason)
         }
 
         override fun onBroadcastUpdated(reason: Int, broadcastId: Int) {
-            if (DEBUG) {
-                Log.d(TAG, "onBroadcastUpdated(), reason = $reason , broadcastId = $broadcastId")
-            }
+            logger.logBroadcastEvent("onBroadcastUpdated", reason, broadcastId)
             updateCurrent()
         }
 
         override fun onBroadcastUpdateFailed(reason: Int, broadcastId: Int) {
-            if (DEBUG) {
-                Log.d(
-                    TAG,
-                    "onBroadcastUpdateFailed(), reason = $reason , " + "broadcastId = $broadcastId"
-                )
-            }
+            logger.logBroadcastEvent("onBroadcastUpdateFailed", reason, broadcastId)
         }
 
         override fun onPlaybackStarted(reason: Int, broadcastId: Int) {}
@@ -381,12 +361,16 @@
                                 name = context.getString(R.string.media_seamless_other_device),
                                 showBroadcastButton = false
                             )
+                    logger.logRemoteDevice(routingSession?.name, connectedDevice)
                 } else {
                     // Prefer SASS if available when playback is local.
-                    activeDevice = getSassDevice() ?: connectedDevice
+                    val sassDevice = getSassDevice()
+                    activeDevice = sassDevice ?: connectedDevice
+                    logger.logLocalDevice(sassDevice, connectedDevice)
                 }
 
                 current = activeDevice ?: EMPTY_AND_DISABLED_MEDIA_DEVICE_DATA
+                logger.logNewDeviceName(current?.name?.toString())
             } else {
                 val aboutToConnect = aboutToConnectDeviceOverride
                 if (
@@ -407,9 +391,7 @@
                 val enabled = device != null && (controller == null || routingSession != null)
 
                 val name = getDeviceName(device, routingSession)
-                if (DEBUG) {
-                    Log.d(TAG, "new device name $name")
-                }
+                logger.logNewDeviceName(name)
                 current =
                     MediaDeviceData(
                         enabled,
@@ -463,14 +445,12 @@
         ): String? {
             val selectedRoutes = routingSession?.let { mr2manager.get().getSelectedRoutes(it) }
 
-            if (DEBUG) {
-                Log.d(
-                    TAG,
-                    "device is $device, controller $controller," +
-                        " routingSession ${routingSession?.name}" +
-                        " or ${selectedRoutes?.firstOrNull()?.name}"
-                )
-            }
+            logger.logDeviceName(
+                device,
+                controller,
+                routingSession?.name,
+                selectedRoutes?.firstOrNull()?.name
+            )
 
             if (controller == null) {
                 // In resume state, we don't have a controller - just use the device name
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/shared/MediaControlDrawables.kt b/packages/SystemUI/src/com/android/systemui/media/controls/shared/MediaControlDrawables.kt
index c78220e..95ca11c 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/shared/MediaControlDrawables.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/shared/MediaControlDrawables.kt
@@ -17,20 +17,12 @@
 package com.android.systemui.media.controls.shared
 
 import android.content.Context
-import android.graphics.drawable.AnimatedVectorDrawable
 import android.graphics.drawable.Drawable
 import com.android.systemui.Flags.mediaControlsDrawablesReuse
 import com.android.systemui.res.R
 
 object MediaControlDrawables {
 
-    // Play/Pause Button drawables.
-    private var progress: Drawable? = null
-    private var connecting: Drawable? = null
-    private var playIcon: AnimatedVectorDrawable? = null
-    private var playBackground: AnimatedVectorDrawable? = null
-    private var pauseIcon: AnimatedVectorDrawable? = null
-    private var pauseBackground: AnimatedVectorDrawable? = null
     // Prev button.
     private var prevIcon: Drawable? = null
     // Next button.
@@ -40,81 +32,6 @@
     private var antenna: Drawable? = null
     private var groupDevice: Drawable? = null
     private var homeDevices: Drawable? = null
-    // Guts drawables.
-    private var outline: Drawable? = null
-    private var solid: Drawable? = null
-
-    fun getProgress(context: Context): Drawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(com.android.internal.R.drawable.progress_small_material)
-        }
-        return progress?.mutate()
-            ?: context.getDrawable(com.android.internal.R.drawable.progress_small_material).also {
-                progress = it
-            }
-    }
-
-    fun getConnecting(context: Context): Drawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.ic_media_connecting_container)
-        }
-        return connecting?.mutate()
-            ?: context.getDrawable(R.drawable.ic_media_connecting_container).also {
-                connecting = it
-            }
-    }
-
-    fun getPlayIcon(context: Context): AnimatedVectorDrawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.ic_media_play) as AnimatedVectorDrawable?
-        }
-        return playIcon?.let {
-            it.reset()
-            it.mutate() as AnimatedVectorDrawable
-        }
-            ?: (context.getDrawable(R.drawable.ic_media_play) as AnimatedVectorDrawable?).also {
-                playIcon = it
-            }
-    }
-
-    fun getPlayBackground(context: Context): AnimatedVectorDrawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.ic_media_play_container)
-                as AnimatedVectorDrawable?
-        }
-        return playBackground?.let {
-            it.reset()
-            it.mutate() as AnimatedVectorDrawable
-        }
-            ?: (context.getDrawable(R.drawable.ic_media_play_container) as AnimatedVectorDrawable?)
-                .also { playBackground = it }
-    }
-
-    fun getPauseIcon(context: Context): AnimatedVectorDrawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.ic_media_pause) as AnimatedVectorDrawable?
-        }
-        return pauseIcon?.let {
-            it.reset()
-            it.mutate() as AnimatedVectorDrawable
-        }
-            ?: (context.getDrawable(R.drawable.ic_media_pause) as AnimatedVectorDrawable?).also {
-                pauseIcon = it
-            }
-    }
-
-    fun getPauseBackground(context: Context): AnimatedVectorDrawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.ic_media_pause_container)
-                as AnimatedVectorDrawable?
-        }
-        return pauseBackground?.let {
-            it.reset()
-            it.mutate() as AnimatedVectorDrawable
-        }
-            ?: (context.getDrawable(R.drawable.ic_media_pause_container) as AnimatedVectorDrawable?)
-                .also { pauseBackground = it }
-    }
 
     fun getNextIcon(context: Context): Drawable? {
         if (!mediaControlsDrawablesReuse()) {
@@ -165,19 +82,4 @@
         return homeDevices
             ?: context.getDrawable(R.drawable.ic_media_home_devices).also { homeDevices = it }
     }
-
-    fun getOutline(context: Context): Drawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.qs_media_outline_button)
-        }
-        return outline
-            ?: context.getDrawable(R.drawable.qs_media_outline_button).also { outline = it }
-    }
-
-    fun getSolid(context: Context): Drawable? {
-        if (!mediaControlsDrawablesReuse()) {
-            return context.getDrawable(R.drawable.qs_media_solid_button)
-        }
-        return solid ?: context.getDrawable(R.drawable.qs_media_solid_button).also { solid = it }
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaCarouselController.kt b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaCarouselController.kt
index 19cdee7..bf9ef8c 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaCarouselController.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaCarouselController.kt
@@ -45,10 +45,10 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.Edge
-import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
 import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
 import com.android.systemui.keyguard.shared.model.TransitionState
@@ -75,7 +75,6 @@
 import com.android.systemui.plugins.FalsingManager
 import com.android.systemui.qs.PageIndicator
 import com.android.systemui.res.R
-import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shared.system.SysUiStatsLog
@@ -103,13 +102,16 @@
 import javax.inject.Provider
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 
@@ -121,6 +123,7 @@
  * Class that is responsible for keeping the view carousel up to date. This also handles changes in
  * state and applies them to the media carousel like the expansion.
  */
+@ExperimentalCoroutinesApi
 @SysUISingleton
 class MediaCarouselController
 @Inject
@@ -149,7 +152,7 @@
     private val secureSettings: SecureSettings,
     private val mediaCarouselViewModel: MediaCarouselViewModel,
     private val mediaViewControllerFactory: Provider<MediaViewController>,
-    private val sceneInteractor: SceneInteractor,
+    private val deviceEntryInteractor: DeviceEntryInteractor,
 ) : Dumpable {
     /** The current width of the carousel */
     var currentCarouselWidth: Int = 0
@@ -164,6 +167,9 @@
     /** Is the player currently visible (at the end of the transformation */
     private var playersVisible: Boolean = false
 
+    /** Are we currently disabling pagination only allowing one media session to show */
+    private var currentlyDisablePagination: Boolean = false
+
     /**
      * The desired location where we'll be at the end of the transformation. Usually this matches
      * the end location, except when we're still waiting on a state update call.
@@ -327,6 +333,11 @@
     private val controllerById = mutableMapOf<String, MediaViewController>()
     private val commonViewModels = mutableListOf<MediaCommonViewModel>()
 
+    private val isOnGone =
+        keyguardTransitionInteractor
+            .isFinishedIn(Scenes.Gone, GONE)
+            .stateIn(applicationScope, SharingStarted.Eagerly, true)
+
     init {
         dumpManager.registerDumpable(TAG, this)
         mediaFrame = inflateMediaCarousel()
@@ -904,9 +915,13 @@
 
     /** Return true if the carousel should be hidden because lockscreen is currently visible */
     fun isLockedAndHidden(): Boolean {
-        val keyguardState = keyguardTransitionInteractor.getFinishedState()
-        return !allowMediaPlayerOnLockScreen &&
-            KeyguardState.lockscreenVisibleInState(keyguardState)
+        val isOnLockscreen =
+            if (SceneContainerFlag.isEnabled) {
+                !deviceEntryInteractor.isDeviceEntered.value
+            } else {
+                !isOnGone.value
+            }
+        return !allowMediaPlayerOnLockScreen && isOnLockscreen
     }
 
     private fun reorderAllPlayers(
@@ -1347,14 +1362,20 @@
         val endShowsActive = hostStates[currentEndLocation]?.showsOnlyActiveMedia ?: true
         val startShowsActive =
             hostStates[currentStartLocation]?.showsOnlyActiveMedia ?: endShowsActive
+        val startDisablePagination = hostStates[currentStartLocation]?.disablePagination ?: false
+        val endDisablePagination = hostStates[currentEndLocation]?.disablePagination ?: false
+
         if (
             currentlyShowingOnlyActive != endShowsActive ||
+                currentlyDisablePagination != endDisablePagination ||
                 ((currentTransitionProgress != 1.0f && currentTransitionProgress != 0.0f) &&
-                    startShowsActive != endShowsActive)
+                    (startShowsActive != endShowsActive ||
+                        startDisablePagination != endDisablePagination))
         ) {
             // Whenever we're transitioning from between differing states or the endstate differs
             // we reset the translation
             currentlyShowingOnlyActive = endShowsActive
+            currentlyDisablePagination = endDisablePagination
             mediaCarouselScrollHandler.resetTranslation(animate = true)
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/view/MediaHost.kt b/packages/SystemUI/src/com/android/systemui/media/controls/ui/view/MediaHost.kt
index 91050c8..09a6181 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/view/MediaHost.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/view/MediaHost.kt
@@ -44,6 +44,7 @@
     lateinit var hostView: UniqueObjectHostView
     var location: Int = -1
         private set
+
     private var visibleChangedListeners: ArraySet<(Boolean) -> Unit> = ArraySet()
 
     private val tmpLocationOnScreen: IntArray = intArrayOf(0, 0)
@@ -287,6 +288,15 @@
                 changedListener?.invoke()
             }
 
+        override var disablePagination: Boolean = false
+            set(value) {
+                if (field == value) {
+                    return
+                }
+                field = value
+                changedListener?.invoke()
+            }
+
         private var lastDisappearHash = disappearParameters.hashCode()
 
         /** A listener for all changes. This won't be copied over when invoking [copy] */
@@ -303,6 +313,7 @@
             mediaHostState.visible = visible
             mediaHostState.disappearParameters = disappearParameters.deepCopy()
             mediaHostState.falsingProtectionNeeded = falsingProtectionNeeded
+            mediaHostState.disablePagination = disablePagination
             return mediaHostState
         }
 
@@ -331,6 +342,9 @@
             if (!disappearParameters.equals(other.disappearParameters)) {
                 return false
             }
+            if (disablePagination != other.disablePagination) {
+                return false
+            }
             return true
         }
 
@@ -342,6 +356,7 @@
             result = 31 * result + showsOnlyActiveMedia.hashCode()
             result = 31 * result + if (visible) 1 else 2
             result = 31 * result + disappearParameters.hashCode()
+            result = 31 * result + disablePagination.hashCode()
             return result
         }
     }
@@ -400,6 +415,12 @@
      */
     var disappearParameters: DisappearParameters
 
+    /**
+     * Whether pagination should be disabled for this host, meaning that when there are multiple
+     * media sessions, only the first one will appear.
+     */
+    var disablePagination: Boolean
+
     /** Get a copy of this view state, deepcopying all appropriate members */
     fun copy(): MediaHostState
 }
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/viewmodel/MediaControlViewModel.kt b/packages/SystemUI/src/com/android/systemui/media/controls/ui/viewmodel/MediaControlViewModel.kt
index f460134..64820e0 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/viewmodel/MediaControlViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/viewmodel/MediaControlViewModel.kt
@@ -30,7 +30,6 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaControlInteractor
-import com.android.systemui.media.controls.shared.MediaControlDrawables
 import com.android.systemui.media.controls.shared.model.MediaAction
 import com.android.systemui.media.controls.shared.model.MediaButton
 import com.android.systemui.media.controls.shared.model.MediaControlModel
@@ -285,9 +284,9 @@
             },
             cancelTextBackground =
                 if (model.isDismissible) {
-                    MediaControlDrawables.getOutline(applicationContext)
+                    applicationContext.getDrawable(R.drawable.qs_media_outline_button)
                 } else {
-                    MediaControlDrawables.getSolid(applicationContext)
+                    applicationContext.getDrawable(R.drawable.qs_media_solid_button)
                 },
             onSettingsClicked = {
                 logger.logLongPressSettings(model.uid, model.packageName, model.instanceId)
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
index 9b1ca1e..6440205 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/MediaProjectionAppSelectorComponent.kt
@@ -122,8 +122,10 @@
         @Provides
         @MediaProjectionAppSelector
         @MediaProjectionAppSelectorScope
-        fun bindConfigurationController(context: Context): ConfigurationController =
-            ConfigurationControllerImpl(context)
+        fun bindConfigurationController(
+            context: Context,
+            configurationControlleFactory: ConfigurationControllerImpl.Factory
+        ): ConfigurationController = configurationControlleFactory.create(context)
 
         @Provides fun bindIconFactory(context: Context): IconFactory = IconFactory.obtain(context)
 
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavBarHelper.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavBarHelper.java
index 6f82d5d..173a964 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavBarHelper.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavBarHelper.java
@@ -80,6 +80,7 @@
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.settings.DisplayTracker;
 import com.android.systemui.settings.UserTracker;
+import com.android.systemui.shared.Flags;
 import com.android.systemui.shared.rotation.RotationPolicyUtil;
 import com.android.systemui.shared.statusbar.phone.BarTransitions.TransitionMode;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -501,9 +502,11 @@
                 Settings.Secure.ASSIST_TOUCH_GESTURE_ENABLED, gestureDefault ? 1 : 0,
                 mUserTracker.getUserId()) != 0;
 
+        boolean supportsSwipeGesture = QuickStepContract.isGesturalMode(mNavBarMode)
+                || (QuickStepContract.isLegacyMode(mNavBarMode) && Flags.threeButtonCornerSwipe());
         mAssistantAvailable = assistantAvailableForUser
                 && mAssistantTouchGestureEnabled
-                && QuickStepContract.isGesturalMode(mNavBarMode);
+                && supportsSwipeGesture;
         dispatchAssistantEventUpdate(mAssistantAvailable, mLongPressHomeEnabled);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
index 0f82e02..6bd880d 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
@@ -78,6 +78,7 @@
 import com.android.systemui.model.SysUiState;
 import com.android.systemui.navigationbar.NavigationModeController;
 import com.android.systemui.navigationbar.gestural.domain.GestureInteractor;
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.NavigationEdgeBackPlugin;
 import com.android.systemui.plugins.PluginListener;
@@ -474,9 +475,14 @@
                 } else {
                     String[] gestureBlockingActivities = resources.getStringArray(resId);
                     for (String gestureBlockingActivity : gestureBlockingActivities) {
-                        mGestureInteractor.addGestureBlockedActivity(
-                                ComponentName.unflattenFromString(gestureBlockingActivity),
-                                GestureInteractor.Scope.Local);
+                        final ComponentName component =
+                                ComponentName.unflattenFromString(gestureBlockingActivity);
+
+                        if (component != null) {
+                            mGestureInteractor.addGestureBlockedMatcher(
+                                    new TaskMatcher.TopActivityComponent(component),
+                                    GestureInteractor.Scope.Local);
+                        }
                     }
                 }
             } catch (NameNotFoundException e) {
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt
index 8f35343..c1f238a 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/data/respository/GestureRepository.kt
@@ -16,10 +16,9 @@
 
 package com.android.systemui.navigationbar.gestural.data.respository
 
-import android.content.ComponentName
-import android.util.ArraySet
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.navigationbar.gestural.domain.TaskMatcher
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -28,36 +27,43 @@
 
 /** A repository for storing gesture related information */
 interface GestureRepository {
-    /** A {@link StateFlow} tracking activities currently blocked from gestures. */
-    val gestureBlockedActivities: StateFlow<Set<ComponentName>>
+    /** A {@link StateFlow} tracking matchers that can block gestures. */
+    val gestureBlockedMatchers: StateFlow<Set<TaskMatcher>>
 
-    /** Adds an activity to be blocked from gestures. */
-    suspend fun addGestureBlockedActivity(activity: ComponentName)
+    /** Adds a matcher to determine whether a gesture should be blocked. */
+    suspend fun addGestureBlockedMatcher(matcher: TaskMatcher)
 
-    /** Removes an activity from being blocked from gestures. */
-    suspend fun removeGestureBlockedActivity(activity: ComponentName)
+    /** Removes a matcher from blocking from gestures. */
+    suspend fun removeGestureBlockedMatcher(matcher: TaskMatcher)
 }
 
 @SysUISingleton
 class GestureRepositoryImpl
 @Inject
 constructor(@Main private val mainDispatcher: CoroutineDispatcher) : GestureRepository {
-    private val _gestureBlockedActivities = MutableStateFlow<Set<ComponentName>>(ArraySet())
+    private val _gestureBlockedMatchers = MutableStateFlow<Set<TaskMatcher>>(emptySet())
 
-    override val gestureBlockedActivities: StateFlow<Set<ComponentName>>
-        get() = _gestureBlockedActivities
+    override val gestureBlockedMatchers: StateFlow<Set<TaskMatcher>>
+        get() = _gestureBlockedMatchers
 
-    override suspend fun addGestureBlockedActivity(activity: ComponentName) =
+    override suspend fun addGestureBlockedMatcher(matcher: TaskMatcher) =
         withContext(mainDispatcher) {
-            _gestureBlockedActivities.emit(
-                _gestureBlockedActivities.value.toMutableSet().apply { add(activity) }
-            )
+            val existingMatchers = _gestureBlockedMatchers.value
+            if (existingMatchers.contains(matcher)) {
+                return@withContext
+            }
+
+            _gestureBlockedMatchers.value = existingMatchers.toMutableSet().apply { add(matcher) }
         }
 
-    override suspend fun removeGestureBlockedActivity(activity: ComponentName) =
+    override suspend fun removeGestureBlockedMatcher(matcher: TaskMatcher) =
         withContext(mainDispatcher) {
-            _gestureBlockedActivities.emit(
-                _gestureBlockedActivities.value.toMutableSet().apply { remove(activity) }
-            )
+            val existingMatchers = _gestureBlockedMatchers.value
+            if (!existingMatchers.contains(matcher)) {
+                return@withContext
+            }
+
+            _gestureBlockedMatchers.value =
+                existingMatchers.toMutableSet().apply { remove(matcher) }
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt
index 6182878..96386e5 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/GestureInteractor.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.navigationbar.gestural.domain
 
-import android.content.ComponentName
 import com.android.app.tracing.coroutines.flow.flowOn
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
@@ -25,7 +24,6 @@
 import com.android.systemui.shared.system.ActivityManagerWrapper
 import com.android.systemui.shared.system.TaskStackChangeListener
 import com.android.systemui.shared.system.TaskStackChangeListeners
-import com.android.systemui.util.kotlin.combine
 import com.android.systemui.util.kotlin.emitOnStart
 import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import javax.inject.Inject
@@ -60,7 +58,7 @@
         Global
     }
 
-    private val _localGestureBlockedActivities = MutableStateFlow<Set<ComponentName>>(setOf())
+    private val _localGestureBlockedMatchers = MutableStateFlow<Set<TaskMatcher>>(setOf())
 
     private val _topActivity =
         conflatedCallbackFlow {
@@ -79,53 +77,47 @@
             .mapLatest { getTopActivity() }
             .distinctUntilChanged()
 
-    private suspend fun getTopActivity(): ComponentName? =
+    private suspend fun getTopActivity(): TaskInfo? =
         withContext(backgroundCoroutineContext) {
-            val runningTask = activityManagerWrapper.runningTask
-            runningTask?.topActivity
+            activityManagerWrapper.runningTask?.let { TaskInfo(it.topActivity, it.activityType) }
         }
 
     val topActivityBlocked =
         combine(
             _topActivity,
-            gestureRepository.gestureBlockedActivities,
-            _localGestureBlockedActivities.asStateFlow()
-        ) { activity, global, local ->
-            activity != null && (global + local).contains(activity)
+            gestureRepository.gestureBlockedMatchers,
+            _localGestureBlockedMatchers.asStateFlow()
+        ) { runningTask, global, local ->
+            runningTask != null && (global + local).any { it.matches(runningTask) }
         }
 
-    /**
-     * Adds an {@link Activity} to be blocked based on component when the topmost, focused {@link
-     * Activity}.
-     */
-    fun addGestureBlockedActivity(activity: ComponentName, gestureScope: Scope) {
+    /** Adds an [TaskMatcher] to decide whether gestures should be blocked. */
+    fun addGestureBlockedMatcher(matcher: TaskMatcher, gestureScope: Scope) {
         scope.launch {
             when (gestureScope) {
                 Scope.Local -> {
-                    _localGestureBlockedActivities.emit(
-                        _localGestureBlockedActivities.value.toMutableSet().apply { add(activity) }
+                    _localGestureBlockedMatchers.emit(
+                        _localGestureBlockedMatchers.value.toMutableSet().apply { add(matcher) }
                     )
                 }
                 Scope.Global -> {
-                    gestureRepository.addGestureBlockedActivity(activity)
+                    gestureRepository.addGestureBlockedMatcher(matcher)
                 }
             }
         }
     }
 
-    /** Removes an {@link Activity} from being blocked from gestures. */
-    fun removeGestureBlockedActivity(activity: ComponentName, gestureScope: Scope) {
+    /** Removes a gesture from deciding whether gestures should be blocked */
+    fun removeGestureBlockedMatcher(matcher: TaskMatcher, gestureScope: Scope) {
         scope.launch {
             when (gestureScope) {
                 Scope.Local -> {
-                    _localGestureBlockedActivities.emit(
-                        _localGestureBlockedActivities.value.toMutableSet().apply {
-                            remove(activity)
-                        }
+                    _localGestureBlockedMatchers.emit(
+                        _localGestureBlockedMatchers.value.toMutableSet().apply { remove(matcher) }
                     )
                 }
                 Scope.Global -> {
-                    gestureRepository.removeGestureBlockedActivity(activity)
+                    gestureRepository.removeGestureBlockedMatcher(matcher)
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/TaskMatcher.kt b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/TaskMatcher.kt
new file mode 100644
index 0000000..d62b2c0
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/domain/TaskMatcher.kt
@@ -0,0 +1,51 @@
+/*
+ * 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.navigationbar.gestural.domain
+
+import android.content.ComponentName
+
+/**
+ * A simple data class for capturing details around a task. Implements equality to ensure changes
+ * can be identified between emitted values.
+ */
+data class TaskInfo(val topActivity: ComponentName?, val topActivityType: Int) {
+    override fun equals(other: Any?): Boolean {
+        return other is TaskInfo &&
+            other.topActivityType == topActivityType &&
+            other.topActivity == topActivity
+    }
+}
+
+/**
+ * [TaskMatcher] provides a way to identify a task based on particular attributes, such as the top
+ * activity type or component name.
+ */
+sealed interface TaskMatcher {
+    fun matches(info: TaskInfo): Boolean
+
+    class TopActivityType(private val type: Int) : TaskMatcher {
+        override fun matches(info: TaskInfo): Boolean {
+            return info.topActivity != null && info.topActivityType == type
+        }
+    }
+
+    class TopActivityComponent(private val component: ComponentName) : TaskMatcher {
+        override fun matches(info: TaskInfo): Boolean {
+            return component == info.topActivity
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/views/NavigationBar.java b/packages/SystemUI/src/com/android/systemui/navigationbar/views/NavigationBar.java
index c706c3e..e8c90c1 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/views/NavigationBar.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/views/NavigationBar.java
@@ -140,7 +140,6 @@
 import com.android.systemui.settings.UserTracker;
 import com.android.systemui.shade.ShadeViewController;
 import com.android.systemui.shade.domain.interactor.PanelExpansionInteractor;
-import com.android.systemui.shared.navigationbar.RegionSamplingHelper;
 import com.android.systemui.shared.recents.utilities.Utilities;
 import com.android.systemui.shared.rotation.RotationButtonController;
 import com.android.systemui.shared.rotation.RotationPolicyUtil;
@@ -166,6 +165,7 @@
 import com.android.systemui.util.ViewController;
 import com.android.wm.shell.back.BackAnimation;
 import com.android.wm.shell.pip.Pip;
+import com.android.wm.shell.shared.handles.RegionSamplingHelper;
 
 import dagger.Lazy;
 
diff --git a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt
new file mode 100644
index 0000000..b6868c1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayActionsViewModel.kt
@@ -0,0 +1,45 @@
+/*
+ * 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.notifications.ui.viewmodel
+
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/** Models the UI state for the user actions for navigating to other scenes or overlays. */
+class NotificationsShadeOverlayActionsViewModel @AssistedInject constructor() :
+    UserActionsViewModel() {
+
+    override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
+        setActions(
+            mapOf(
+                Swipe.Up to UserActionResult.HideOverlay(Overlays.NotificationsShade),
+                Back to UserActionResult.HideOverlay(Overlays.NotificationsShade),
+            )
+        )
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): NotificationsShadeOverlayActionsViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModel.kt
new file mode 100644
index 0000000..5be225c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeOverlayContentViewModel.kt
@@ -0,0 +1,50 @@
+/*
+ * 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.notifications.ui.viewmodel
+
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
+import com.android.systemui.statusbar.notification.stack.ui.viewmodel.NotificationsPlaceholderViewModel
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/**
+ * Models UI state used to render the content of the notifications shade overlay.
+ *
+ * Different from [NotificationsShadeOverlayActionsViewModel], which only models user actions that
+ * can be performed to navigate to other scenes.
+ */
+class NotificationsShadeOverlayContentViewModel
+@AssistedInject
+constructor(
+    val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
+    val notificationsPlaceholderViewModelFactory: NotificationsPlaceholderViewModel.Factory,
+    private val sceneInteractor: SceneInteractor,
+) {
+    fun onScrimClicked() {
+        sceneInteractor.hideOverlay(
+            overlay = Overlays.NotificationsShade,
+            loggingReason = "Shade scrim clicked",
+        )
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): NotificationsShadeOverlayContentViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModel.kt
similarity index 69%
rename from packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModel.kt
index 9fb09c0..a5c07bc 100644
--- a/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/notifications/ui/viewmodel/NotificationsShadeUserActionsViewModel.kt
@@ -21,29 +21,20 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
-import com.android.systemui.shade.shared.model.ShadeAlignment
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 
 /**
  * Models the UI state for the user actions that the user can perform to navigate to other scenes.
  */
-class NotificationsShadeSceneActionsViewModel
-@AssistedInject
-constructor(
-    private val shadeInteractor: ShadeInteractor,
-) : SceneActionsViewModel() {
+class NotificationsShadeUserActionsViewModel @AssistedInject constructor() :
+    UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         setActions(
             mapOf(
-                if (shadeInteractor.shadeAlignment == ShadeAlignment.Top) {
-                    Swipe.Up
-                } else {
-                    Swipe.Down
-                } to SceneFamilies.Home,
+                Swipe.Up to SceneFamilies.Home,
                 Back to SceneFamilies.Home,
             )
         )
@@ -51,6 +42,6 @@
 
     @AssistedFactory
     interface Factory {
-        fun create(): NotificationsShadeSceneActionsViewModel
+        fun create(): NotificationsShadeUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
index c3274b7..c39ff55 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/QSFragmentCompose.kt
@@ -32,6 +32,7 @@
 import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.WindowInsets
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.foundation.layout.fillMaxWidth
 import androidx.compose.foundation.layout.navigationBars
 import androidx.compose.foundation.layout.windowInsetsPadding
 import androidx.compose.runtime.Composable
@@ -403,30 +404,40 @@
             onDispose { qqsVisible.value = false }
         }
         Column(modifier = Modifier.sysuiResTag("quick_qs_panel")) {
-            QuickQuickSettings(
-                viewModel = viewModel.containerViewModel.quickQuickSettingsViewModel,
-                modifier =
-                    Modifier.onGloballyPositioned { coordinates ->
-                            val (leftFromRoot, topFromRoot) = coordinates.positionInRoot().round()
-                            val (width, height) = coordinates.size
-                            qqsPositionOnRoot.set(
-                                leftFromRoot,
-                                topFromRoot,
-                                leftFromRoot + width,
-                                topFromRoot + height
-                            )
-                        }
-                        .layout { measurable, constraints ->
-                            val placeable = measurable.measure(constraints)
-                            qqsHeight.value = placeable.height
+            Box(modifier = Modifier.fillMaxWidth()) {
+                val qsEnabled by viewModel.qsEnabled.collectAsStateWithLifecycle()
+                if (qsEnabled) {
+                    QuickQuickSettings(
+                        viewModel = viewModel.containerViewModel.quickQuickSettingsViewModel,
+                        modifier =
+                            Modifier.onGloballyPositioned { coordinates ->
+                                    val (leftFromRoot, topFromRoot) =
+                                        coordinates.positionInRoot().round()
+                                    val (width, height) = coordinates.size
+                                    qqsPositionOnRoot.set(
+                                        leftFromRoot,
+                                        topFromRoot,
+                                        leftFromRoot + width,
+                                        topFromRoot + height
+                                    )
+                                }
+                                .layout { measurable, constraints ->
+                                    val placeable = measurable.measure(constraints)
+                                    qqsHeight.value = placeable.height
 
-                            layout(placeable.width, placeable.height) { placeable.place(0, 0) }
-                        }
-                        .padding(top = { qqsPadding })
-                        .collapseExpandSemanticAction(
-                            stringResource(id = R.string.accessibility_quick_settings_expand)
-                        )
-            )
+                                    layout(placeable.width, placeable.height) {
+                                        placeable.place(0, 0)
+                                    }
+                                }
+                                .padding(top = { qqsPadding })
+                                .collapseExpandSemanticAction(
+                                    stringResource(
+                                        id = R.string.accessibility_quick_settings_expand
+                                    )
+                                )
+                    )
+                }
+            }
             Spacer(modifier = Modifier.weight(1f))
         }
     }
@@ -441,18 +452,23 @@
                     stringResource(id = R.string.accessibility_quick_settings_collapse)
                 )
         ) {
-            Box(modifier = Modifier.fillMaxSize().weight(1f)) {
-                Column {
-                    Spacer(modifier = Modifier.height { qqsPadding + qsExtraPadding.roundToPx() })
-                    ShadeBody(viewModel = viewModel.containerViewModel)
+            val qsEnabled by viewModel.qsEnabled.collectAsStateWithLifecycle()
+            if (qsEnabled) {
+                Box(modifier = Modifier.fillMaxSize().weight(1f)) {
+                    Column {
+                        Spacer(
+                            modifier = Modifier.height { qqsPadding + qsExtraPadding.roundToPx() }
+                        )
+                        ShadeBody(viewModel = viewModel.containerViewModel)
+                    }
                 }
-            }
-            QuickSettingsTheme {
-                FooterActions(
-                    viewModel = viewModel.footerActionsViewModel,
-                    qsVisibilityLifecycleOwner = this@QSFragmentCompose,
-                    modifier = Modifier.sysuiResTag("qs_footer_actions")
-                )
+                QuickSettingsTheme {
+                    FooterActions(
+                        viewModel = viewModel.footerActionsViewModel,
+                        qsVisibilityLifecycleOwner = this@QSFragmentCompose,
+                        modifier = Modifier.sysuiResTag("qs_footer_actions")
+                    )
+                }
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
index df77878..16133f4 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/composefragment/viewmodel/QSFragmentComposeViewModel.kt
@@ -132,13 +132,17 @@
             _stackScrollerOverscrolling.value = value
         }
 
-    private val qsDisabled =
+    /**
+     * Whether QS is enabled by policy. This is normally true, except when it's disabled by some
+     * policy. See [DisableFlagsRepository].
+     */
+    val qsEnabled =
         disableFlagsRepository.disableFlags
-            .map { !it.isQuickSettingsEnabled() }
+            .map { it.isQuickSettingsEnabled() }
             .stateIn(
                 lifecycleScope,
                 SharingStarted.WhileSubscribed(),
-                !disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled()
+                disableFlagsRepository.disableFlags.value.isQuickSettingsEnabled()
             )
 
     private val _showCollapsedOnKeyguard = MutableStateFlow(false)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
index d948dfd..c75b601 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
-import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.runtime.Composable
@@ -24,7 +23,6 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.remember
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.res.dimensionResource
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.qs.panels.shared.model.SizedTileImpl
@@ -33,7 +31,6 @@
 import com.android.systemui.qs.panels.ui.viewmodel.IconTilesViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
 import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.android.systemui.res.R
 import javax.inject.Inject
 
 @SysUISingleton
@@ -64,7 +61,7 @@
                 Tile(
                     tile = sizedTiles[index].tile,
                     iconOnly = iconTilesViewModel.isIconTile(sizedTiles[index].tile.spec),
-                    modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
+                    modifier = Modifier
                 )
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
index a9027ff..eeb55ca 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/QuickQuickSettings.kt
@@ -16,18 +16,15 @@
 
 package com.android.systemui.qs.panels.ui.compose
 
-import androidx.compose.foundation.layout.height
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.Modifier
-import androidx.compose.ui.res.dimensionResource
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.systemui.compose.modifiers.sysuiResTag
 import com.android.systemui.qs.panels.ui.viewmodel.QuickQuickSettingsViewModel
-import com.android.systemui.res.R
 
 @Composable
 fun QuickQuickSettings(
@@ -54,11 +51,7 @@
             key = { index -> sizedTiles[index].tile.spec.spec },
             span = { index -> GridItemSpan(sizedTiles[index].width) }
         ) { index ->
-            Tile(
-                tile = tiles[index],
-                iconOnly = sizedTiles[index].isIcon,
-                modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
-            )
+            Tile(tile = tiles[index], iconOnly = sizedTiles[index].isIcon, modifier = Modifier)
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
index 79c2eb9..24af09d 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/Tile.kt
@@ -44,7 +44,6 @@
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
-import androidx.compose.foundation.layout.aspectRatio
 import androidx.compose.foundation.layout.fillMaxHeight
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
@@ -136,23 +135,23 @@
 
     // TODO(b/361789146): Draw the shapes instead of clipping
     val tileShape = TileDefaults.animateTileShape(uiState.state)
-    val iconShape = TileDefaults.animateIconShape(uiState.state)
 
     TileContainer(
         colors = colors,
         showLabels = showLabels,
         label = uiState.label,
         iconOnly = iconOnly,
-        shape = if (iconOnly) iconShape else tileShape,
+        shape = tileShape,
         clickEnabled = true,
         onClick = tile::onClick,
         onLongClick = tile::onLongClick,
-        modifier = modifier,
+        modifier = modifier.height(tileHeight()),
     ) {
         val icon = getTileIcon(icon = uiState.icon)
         if (iconOnly) {
             TileIcon(icon = icon, color = colors.icon, modifier = Modifier.align(Alignment.Center))
         } else {
+            val iconShape = TileDefaults.animateIconShape(uiState.state)
             LargeTileContent(
                 label = uiState.label,
                 secondaryLabel = uiState.secondaryLabel,
@@ -199,7 +198,7 @@
         Expandable(
             color = backgroundColor,
             shape = shape,
-            modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height)).clip(shape)
+            modifier = Modifier.height(tileHeight()).clip(shape)
         ) {
             Box(
                 modifier =
@@ -246,7 +245,7 @@
         // Icon
         Box(
             modifier =
-                Modifier.fillMaxHeight().aspectRatio(1f).thenIf(toggleClickSupported) {
+                Modifier.size(TileDefaults.ToggleTargetSize).thenIf(toggleClickSupported) {
                     Modifier.clip(iconShape)
                         .background(colors.iconBackground, { 1f })
                         .combinedClickable(onClick = onClick, onLongClick = onLongClick)
@@ -673,7 +672,7 @@
     animateToEnd: Boolean = false,
     modifier: Modifier = Modifier,
 ) {
-    val iconModifier = modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
+    val iconModifier = modifier.size(TileDefaults.IconSize)
     val context = LocalContext.current
     val loadedDrawable =
         remember(icon, context) {
@@ -710,17 +709,12 @@
     }
 }
 
-@Composable
 private fun Modifier.tilePadding(): Modifier {
-    return padding(dimensionResource(id = R.dimen.qs_label_container_margin))
+    return padding(TileDefaults.TilePadding)
 }
 
-@Composable
 private fun tileHorizontalArrangement(): Arrangement.Horizontal {
-    return spacedBy(
-        space = dimensionResource(id = R.dimen.qs_label_container_margin),
-        alignment = Alignment.Start
-    )
+    return spacedBy(space = TileDefaults.TileArrangementPadding, alignment = Alignment.Start)
 }
 
 @Composable
@@ -728,7 +722,7 @@
     return if (iconWithLabel) {
         TileDefaults.IconTileWithLabelHeight
     } else {
-        dimensionResource(id = R.dimen.qs_tile_height)
+        TileDefaults.TileHeight
     }
 }
 
@@ -749,6 +743,14 @@
     val InactiveCornerRadius = 50.dp
     val ActiveIconCornerRadius = 16.dp
     val ActiveTileCornerRadius = 24.dp
+
+    val ToggleTargetSize = 56.dp
+    val IconSize = 24.dp
+
+    val TilePadding = 8.dp
+    val TileArrangementPadding = 6.dp
+
+    val TileHeight = 72.dp
     val IconTileWithLabelHeight = 140.dp
 
     /** An active tile without dual target uses the active color as background */
@@ -812,7 +814,7 @@
     fun animateIconShape(state: Int): Shape {
         return animateShape(
             state = state,
-            activeCornerRadius = ActiveTileCornerRadius,
+            activeCornerRadius = ActiveIconCornerRadius,
             label = "QSTileCornerRadius",
         )
     }
@@ -821,7 +823,7 @@
     fun animateTileShape(state: Int): Shape {
         return animateShape(
             state = state,
-            activeCornerRadius = ActiveIconCornerRadius,
+            activeCornerRadius = ActiveTileCornerRadius,
             label = "QSTileIconCornerRadius",
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
index 6d63d26..7d23fbd 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
@@ -21,6 +21,8 @@
 import android.os.Handler
 import android.os.Looper
 import android.service.quicksettings.Tile
+import androidx.annotation.DrawableRes
+import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
@@ -62,7 +64,7 @@
     activityStarter: ActivityStarter,
     qsLogger: QSLogger,
     qsTileConfigProvider: QSTileConfigProvider,
-    dataInteractor: ModesTileDataInteractor,
+    private val dataInteractor: ModesTileDataInteractor,
     private val tileMapper: ModesTileMapper,
     private val userActionInteractor: ModesTileUserActionInteractor,
 ) :
@@ -98,7 +100,7 @@
     override fun newTileState(): QSTile.State {
         return QSTile.State().apply {
             label = mContext.getString(R.string.quick_settings_modes_label)
-            icon = ResourceIcon.get(R.drawable.qs_dnd_icon_off)
+            icon = ResourceIcon.get(ICON_RES_ID)
             state = Tile.STATE_INACTIVE
         }
     }
@@ -109,23 +111,26 @@
 
     override fun getLongClickIntent(): Intent = userActionInteractor.longClickIntent
 
-    override fun handleUpdateState(state: QSTile.State?, arg: Any?) {
-        if (arg is ModesTileModel) {
-            tileState = tileMapper.map(config, arg)
+    @VisibleForTesting
+    public override fun handleUpdateState(state: QSTile.State?, arg: Any?) {
+        // This runBlocking() will block @Background. Due to caches, it's expected to be fast.
+        val model =
+            if (arg is ModesTileModel) arg else runBlocking { dataInteractor.getCurrentTileModel() }
 
-            state?.apply {
-                this.state = tileState.activationState.legacyState
-                val tileStateIcon = tileState.icon()
-                icon = tileStateIcon?.asQSTileIcon() ?: ResourceIcon.get(R.drawable.qs_dnd_icon_off)
-                label = tileLabel
-                secondaryLabel = tileState.secondaryLabel
-                contentDescription = tileState.contentDescription
-                expandedAccessibilityClassName = tileState.expandedAccessibilityClassName
-            }
+        tileState = tileMapper.map(config, model)
+        state?.apply {
+            this.state = tileState.activationState.legacyState
+            val tileStateIcon = tileState.icon()
+            icon = tileStateIcon?.asQSTileIcon() ?: ResourceIcon.get(ICON_RES_ID)
+            label = tileLabel
+            secondaryLabel = tileState.secondaryLabel
+            contentDescription = tileState.contentDescription
+            expandedAccessibilityClassName = tileState.expandedAccessibilityClassName
         }
     }
 
     companion object {
         const val TILE_SPEC = "dnd"
+        @DrawableRes val ICON_RES_ID = com.android.internal.R.drawable.ic_zen_priority_modes
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogDelegate.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogDelegate.java
index 71f8639..89b9eee 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogDelegate.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/dialog/InternetDialogDelegate.java
@@ -306,6 +306,8 @@
                 mInternetDialogController.isAirplaneModeEnabled() ? View.VISIBLE : View.GONE);
         mWifiRecyclerView.setLayoutManager(new LinearLayoutManager(context));
         mWifiRecyclerView.setAdapter(mAdapter);
+
+        updateDialogUI(getWifiNetworkContent());
     }
 
     @Override
@@ -315,6 +317,7 @@
         }
 
         mLifecycleRegistry.setCurrentState(Lifecycle.State.RESUMED);
+
         mInternetDialogController.onStart(this, mCanConfigWifi);
         if (!mCanConfigWifi) {
             hideWifiViews();
@@ -402,10 +405,12 @@
         internetContent.mShouldUpdateMobileNetwork = shouldUpdateMobileNetwork;
         internetContent.mInternetDialogTitleString = getDialogTitleText();
         internetContent.mInternetDialogSubTitle = getSubtitleText();
-        internetContent.mActiveNetworkIsCellular =
-                mInternetDialogController.activeNetworkIsCellular();
-        internetContent.mIsCarrierNetworkActive =
-                mInternetDialogController.isCarrierNetworkActive();
+        if (shouldUpdateMobileNetwork) {
+            internetContent.mActiveNetworkIsCellular =
+                    mInternetDialogController.activeNetworkIsCellular();
+            internetContent.mIsCarrierNetworkActive =
+                    mInternetDialogController.isCarrierNetworkActive();
+        }
         internetContent.mIsAirplaneModeEnabled = mInternetDialogController.isAirplaneModeEnabled();
         internetContent.mHasEthernet = mInternetDialogController.hasEthernet();
         internetContent.mIsWifiEnabled = mInternetDialogController.isWifiEnabled();
@@ -416,6 +421,15 @@
         return internetContent;
     }
 
+    private InternetContent getWifiNetworkContent() {
+        InternetContent internetContent = new InternetContent();
+        internetContent.mInternetDialogTitleString = getDialogTitleText();
+        internetContent.mInternetDialogSubTitle = getSubtitleText();
+        internetContent.mIsWifiEnabled = mInternetDialogController.isWifiEnabled();
+        internetContent.mIsDeviceLocked = mInternetDialogController.isDeviceLocked();
+        return internetContent;
+    }
+
     private void setOnClickListener(SystemUIDialog dialog) {
         mMobileNetworkLayout.setOnClickListener(v -> {
             int autoSwitchNonDdsSubId = mInternetDialogController.getActiveAutoSwitchNonDdsSubId();
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
index 3f18fc2..c2d112e 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
@@ -20,11 +20,13 @@
 import android.content.Context
 import android.os.UserHandle
 import com.android.app.tracing.coroutines.flow.map
+import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.qs.tiles.base.interactor.DataUpdateTrigger
 import com.android.systemui.qs.tiles.base.interactor.QSTileDataInteractor
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor
+import com.android.systemui.statusbar.policy.domain.model.ActiveZenModes
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
@@ -52,18 +54,36 @@
      */
     fun tileData() =
         zenModeInteractor.activeModes
-            .map { modes ->
-                ModesTileModel(
-                    isActivated = modes.isNotEmpty(),
-                    icon =
-                        if (Flags.modesApi() && Flags.modesUi() && Flags.modesUiIcons())
-                            zenModeInteractor.getActiveModeIcon(modes)
-                        else null,
-                    activeModes = modes.map { it.name }
-                )
-            }
+            .map { activeModes -> buildTileData(activeModes) }
             .flowOn(bgDispatcher)
             .distinctUntilChanged()
 
+    suspend fun getCurrentTileModel() = buildTileData(zenModeInteractor.getActiveModes())
+
+    private fun buildTileData(activeModes: ActiveZenModes): ModesTileModel {
+        val modesIconResId = com.android.internal.R.drawable.ic_zen_priority_modes
+
+        if (usesModeIcons()) {
+            val mainModeDrawable = activeModes.mainMode?.icon?.drawable
+            val iconResId = if (mainModeDrawable == null) modesIconResId else null
+
+            return ModesTileModel(
+                isActivated = activeModes.isAnyActive(),
+                icon = (mainModeDrawable ?: context.getDrawable(modesIconResId)!!).asIcon(),
+                iconResId = iconResId,
+                activeModes = activeModes.modeNames
+            )
+        } else {
+            return ModesTileModel(
+                isActivated = activeModes.isAnyActive(),
+                icon = context.getDrawable(modesIconResId)!!.asIcon(),
+                iconResId = modesIconResId,
+                activeModes = activeModes.modeNames
+            )
+        }
+    }
+
     override fun availability(user: UserHandle): Flow<Boolean> = flowOf(Flags.modesUi())
+
+    private fun usesModeIcons() = Flags.modesApi() && Flags.modesUi() && Flags.modesUiIcons()
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
index 904ff3a..db48123 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
@@ -21,5 +21,12 @@
 data class ModesTileModel(
     val isActivated: Boolean,
     val activeModes: List<String>,
-    val icon: Icon? = null
+    val icon: Icon.Loaded,
+
+    /**
+     * Resource id corresponding to [icon]. Will only be present if it's know to correspond to a
+     * resource with a known id in SystemUI (such as resources from `android.R`,
+     * `com.android.internal.R`, or `com.android.systemui.res` itself).
+     */
+    val iconResId: Int? = null
 )
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
index 83c3335..7f571b1 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
@@ -16,11 +16,9 @@
 
 package com.android.systemui.qs.tiles.impl.modes.ui
 
-import android.app.Flags
 import android.content.res.Resources
 import android.icu.text.MessageFormat
 import android.widget.Button
-import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.qs.tiles.base.interactor.QSTileDataToStateMapper
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
@@ -38,15 +36,10 @@
 ) : QSTileDataToStateMapper<ModesTileModel> {
     override fun map(config: QSTileConfig, data: ModesTileModel): QSTileState =
         QSTileState.build(resources, theme, config.uiConfig) {
-            if (Flags.modesApi() && Flags.modesUi() && Flags.modesUiIcons() && data.icon != null) {
-                icon = { data.icon }
-            } else {
-                val iconRes =
-                    if (data.isActivated) R.drawable.qs_dnd_icon_on else R.drawable.qs_dnd_icon_off
-                val icon = resources.getDrawable(iconRes, theme).asIcon()
-                this.iconRes = iconRes
-                this.icon = { icon }
+            if (!android.app.Flags.modesUiIcons()) {
+                iconRes = data.iconResId
             }
+            icon = { data.icon }
             activationState =
                 if (data.isActivated) {
                     QSTileState.ActivationState.ACTIVE
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModel.kt
index 12f3c9c..f77386d 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneContentViewModel.kt
@@ -17,22 +17,29 @@
 package com.android.systemui.qs.ui.viewmodel
 
 import androidx.lifecycle.LifecycleOwner
-import com.android.systemui.lifecycle.SysUiViewModel
+import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.qs.FooterActionsController
 import com.android.systemui.qs.footer.ui.viewmodel.FooterActionsViewModel
 import com.android.systemui.qs.ui.adapter.QSSceneAdapter
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.settings.brightness.ui.viewModel.BrightnessMirrorViewModel
+import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import java.util.concurrent.atomic.AtomicBoolean
+import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.launch
 
 /**
  * Models UI state needed for rendering the content of the quick settings scene.
  *
- * Different from [QuickSettingsSceneActionsViewModel] that models the UI state needed to figure out
+ * Different from [QuickSettingsUserActionsViewModel] that models the UI state needed to figure out
  * which user actions can trigger navigation to other scenes.
  */
 class QuickSettingsSceneContentViewModel
@@ -44,7 +51,9 @@
     private val footerActionsViewModelFactory: FooterActionsViewModel.Factory,
     private val footerActionsController: FooterActionsController,
     val mediaCarouselInteractor: MediaCarouselInteractor,
-) : SysUiViewModel {
+    private val shadeInteractor: ShadeInteractor,
+    private val sceneInteractor: SceneInteractor,
+) : ExclusiveActivatable() {
 
     val isMediaVisible: StateFlow<Boolean> = mediaCarouselInteractor.hasAnyMediaOrRecommendation
 
@@ -57,6 +66,19 @@
         return footerActionsViewModelFactory.create(lifecycleOwner)
     }
 
+    override suspend fun onActivated(): Nothing {
+        coroutineScope {
+            launch {
+                shadeInteractor.shadeMode.collect { shadeMode ->
+                    if (shadeMode == ShadeMode.Split) {
+                        sceneInteractor.snapToScene(Scenes.Shade, "Unfold while on QS")
+                    }
+                }
+            }
+            awaitCancellation()
+        }
+    }
+
     @AssistedFactory
     interface Factory {
         fun create(): QuickSettingsSceneContentViewModel
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt
new file mode 100644
index 0000000..61c4c8c
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModel.kt
@@ -0,0 +1,45 @@
+/*
+ * 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.qs.ui.viewmodel
+
+import com.android.compose.animation.scene.Back
+import com.android.compose.animation.scene.Swipe
+import com.android.compose.animation.scene.UserAction
+import com.android.compose.animation.scene.UserActionResult
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/** Models the UI state for the user actions for navigating to other scenes or overlays. */
+class QuickSettingsShadeOverlayActionsViewModel @AssistedInject constructor() :
+    UserActionsViewModel() {
+
+    override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
+        setActions(
+            buildMap {
+                put(Swipe.Up, UserActionResult.HideOverlay(Overlays.QuickSettingsShade))
+                put(Back, UserActionResult.HideOverlay(Overlays.QuickSettingsShade))
+            }
+        )
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): QuickSettingsShadeOverlayActionsViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModel.kt
new file mode 100644
index 0000000..3b97d82
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModel.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.qs.ui.viewmodel
+
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.model.Overlays
+import com.android.systemui.shade.ui.viewmodel.ShadeHeaderViewModel
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+
+/**
+ * Models UI state used to render the content of the quick settings shade overlay.
+ *
+ * Different from [QuickSettingsShadeOverlayActionsViewModel], which only models user actions that
+ * can be performed to navigate to other scenes.
+ */
+class QuickSettingsShadeOverlayContentViewModel
+@AssistedInject
+constructor(
+    val sceneInteractor: SceneInteractor,
+    val shadeHeaderViewModelFactory: ShadeHeaderViewModel.Factory,
+    val quickSettingsContainerViewModel: QuickSettingsContainerViewModel,
+) {
+    fun onScrimClicked() {
+        sceneInteractor.hideOverlay(
+            overlay = Overlays.QuickSettingsShade,
+            loggingReason = "Shade scrim clicked",
+        )
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): QuickSettingsShadeOverlayContentViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModel.kt
index cb99be4..d01b33b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModel.kt
@@ -14,28 +14,22 @@
  * limitations under the License.
  */
 
-@file:OptIn(ExperimentalCoroutinesApi::class)
-
 package com.android.systemui.qs.ui.viewmodel
 
-import com.android.systemui.lifecycle.SysUiViewModel
-import com.android.systemui.shade.ui.viewmodel.OverlayShadeViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 /**
  * Models UI state used to render the content of the quick settings shade scene.
  *
- * Different from [QuickSettingsShadeSceneActionsViewModel], which only models user actions that can
+ * Different from [QuickSettingsShadeUserActionsViewModel], which only models user actions that can
  * be performed to navigate to other scenes.
  */
 class QuickSettingsShadeSceneContentViewModel
 @AssistedInject
 constructor(
-    val overlayShadeViewModelFactory: OverlayShadeViewModel.Factory,
     val quickSettingsContainerViewModel: QuickSettingsContainerViewModel,
-) : SysUiViewModel {
+) {
 
     @AssistedFactory
     interface Factory {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModel.kt
similarity index 66%
rename from packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModel.kt
index d2967b8..d3dc302 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModel.kt
@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-@file:OptIn(ExperimentalCoroutinesApi::class)
-
 package com.android.systemui.qs.ui.viewmodel
 
 import com.android.compose.animation.scene.Back
@@ -23,13 +21,9 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
-import com.android.systemui.shade.shared.model.ShadeAlignment
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.map
 
 /**
@@ -38,35 +32,27 @@
  * Different from the [QuickSettingsShadeSceneContentViewModel] which models the _content_ of the
  * scene.
  */
-class QuickSettingsShadeSceneActionsViewModel
+class QuickSettingsShadeUserActionsViewModel
 @AssistedInject
 constructor(
-    private val shadeInteractor: ShadeInteractor,
     val quickSettingsContainerViewModel: QuickSettingsContainerViewModel,
-) : SceneActionsViewModel() {
+) : UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         quickSettingsContainerViewModel.editModeViewModel.isEditing
             .map { editing ->
                 buildMap {
-                    put(
-                        if (shadeInteractor.shadeAlignment == ShadeAlignment.Top) {
-                            Swipe.Up
-                        } else {
-                            Swipe.Down
-                        },
-                        UserActionResult(SceneFamilies.Home)
-                    )
+                    put(Swipe.Up, UserActionResult(SceneFamilies.Home))
                     if (!editing) {
                         put(Back, UserActionResult(SceneFamilies.Home))
                     }
                 }
             }
-            .collectLatest { actions -> setActions(actions) }
+            .collect { actions -> setActions(actions) }
     }
 
     @AssistedFactory
     interface Factory {
-        fun create(): QuickSettingsShadeSceneActionsViewModel
+        fun create(): QuickSettingsShadeUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModel.kt
similarity index 91%
rename from packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModel.kt
index af55f5a..54e5cac 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsUserActionsViewModel.kt
@@ -27,11 +27,10 @@
 import com.android.systemui.scene.domain.interactor.SceneBackInteractor
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.map
@@ -43,12 +42,12 @@
  * Different from [QuickSettingsSceneContentViewModel] that models UI state needed for rendering the
  * content of the quick settings scene.
  */
-class QuickSettingsSceneActionsViewModel
+class QuickSettingsUserActionsViewModel
 @AssistedInject
 constructor(
     private val qsSceneAdapter: QSSceneAdapter,
     sceneBackInteractor: SceneBackInteractor,
-) : SceneActionsViewModel() {
+) : UserActionsViewModel() {
 
     private val backScene: Flow<SceneKey> =
         sceneBackInteractor.backScene
@@ -78,11 +77,11 @@
                     }
                 }
             }
-            .collectLatest { actions -> setActions(actions) }
+            .collect { actions -> setActions(actions) }
     }
 
     @AssistedFactory
     interface Factory {
-        fun create(): QuickSettingsSceneActionsViewModel
+        fun create(): QuickSettingsUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index fe5cbb1..000781a 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -894,11 +894,21 @@
             return;
         }
         mHandler.removeCallbacks(mConnectionRunnable);
+
+        // Avoid creating TouchInteractionService because the System user in HSUM mode does not
+        // interact with UI elements
+        UserHandle currentUser = UserHandle.of(mUserTracker.getUserId());
+        if (UserManager.isHeadlessSystemUserMode() && currentUser.isSystem()) {
+            Log.w(TAG_OPS,
+                    "Skipping connection to TouchInteractionService for the System user in HSUM "
+                            + "mode.");
+            return;
+        }
         try {
             mBound = mContext.bindServiceAsUser(mQuickStepIntent,
                     mOverviewServiceConnection,
                     Context.BIND_AUTO_CREATE | Context.BIND_FOREGROUND_SERVICE_WHILE_AWAKE,
-                    UserHandle.of(mUserTracker.getUserId()));
+                    currentUser);
         } catch (SecurityException e) {
             Log.e(TAG_OPS, "Unable to bind because of security error", e);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/recordissue/CustomTraceState.kt b/packages/SystemUI/src/com/android/systemui/recordissue/CustomTraceState.kt
index 14dfcc5..b0eaf9f 100644
--- a/packages/SystemUI/src/com/android/systemui/recordissue/CustomTraceState.kt
+++ b/packages/SystemUI/src/com/android/systemui/recordissue/CustomTraceState.kt
@@ -21,6 +21,10 @@
 import com.android.traceur.PresetTraceConfigs.getDefaultConfig
 import com.android.traceur.TraceConfig
 
+/**
+ * This class encapsulates the values that go into a customized record issue trace config, part of
+ * the RecordIssueTile feature. This class stores the last configuration chosen by power users.
+ */
 class CustomTraceState(private val prefs: SharedPreferences) {
 
     private var enabledTags: Set<String>?
diff --git a/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt b/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
index efb9375..7a57fba 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/EmptySceneModule.kt
@@ -16,7 +16,8 @@
 
 package com.android.systemui.scene
 
-import com.android.systemui.scene.shared.model.Scene
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.scene.ui.composable.Scene
 import dagger.Module
 import dagger.Provides
 import dagger.multibindings.ElementsIntoSet
@@ -29,4 +30,10 @@
     fun emptySceneSet(): Set<Scene> {
         return emptySet()
     }
+
+    @Provides
+    @ElementsIntoSet
+    fun emptyOverlaySet(): Set<Overlay> {
+        return emptySet()
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
index 6e89973..00944b8 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/KeyguardlessSceneContainerFrameworkModule.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.scene.domain.startable.SceneContainerStartable
 import com.android.systemui.scene.domain.startable.ScrimStartable
 import com.android.systemui.scene.domain.startable.StatusBarStartable
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.shared.flag.DualShade
@@ -42,8 +43,10 @@
         [
             EmptySceneModule::class,
             GoneSceneModule::class,
+            NotificationsShadeOverlayModule::class,
             NotificationsShadeSceneModule::class,
             NotificationsShadeSessionModule::class,
+            QuickSettingsShadeOverlayModule::class,
             QuickSettingsSceneModule::class,
             ShadeSceneModule::class,
             SceneDomainModule::class,
@@ -99,6 +102,11 @@
                         Scenes.Shade.takeUnless { DualShade.isEnabled },
                     ),
                 initialSceneKey = Scenes.Gone,
+                overlayKeys =
+                    listOfNotNull(
+                        Overlays.NotificationsShade.takeIf { DualShade.isEnabled },
+                        Overlays.QuickSettingsShade.takeIf { DualShade.isEnabled },
+                    ),
                 navigationDistances =
                     mapOf(
                             Scenes.Gone to 0,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
index 7d63b4c..4061ad8 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/SceneContainerFrameworkModule.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.scene
 
 import com.android.systemui.CoreStartable
-import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlagsModule
 import com.android.systemui.notifications.ui.composable.NotificationsShadeSessionModule
 import com.android.systemui.scene.domain.SceneDomainModule
 import com.android.systemui.scene.domain.interactor.WindowRootViewVisibilityInteractor
@@ -28,6 +27,7 @@
 import com.android.systemui.scene.domain.startable.SceneContainerStartable
 import com.android.systemui.scene.domain.startable.ScrimStartable
 import com.android.systemui.scene.domain.startable.StatusBarStartable
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.shade.shared.flag.DualShade
@@ -43,13 +43,14 @@
         [
             BouncerSceneModule::class,
             CommunalSceneModule::class,
-            ComposeBouncerFlagsModule::class,
             EmptySceneModule::class,
             GoneSceneModule::class,
             LockscreenSceneModule::class,
             QuickSettingsSceneModule::class,
             ShadeSceneModule::class,
+            QuickSettingsShadeOverlayModule::class,
             QuickSettingsShadeSceneModule::class,
+            NotificationsShadeOverlayModule::class,
             NotificationsShadeSceneModule::class,
             NotificationsShadeSessionModule::class,
             SceneDomainModule::class,
@@ -108,6 +109,11 @@
                         Scenes.Shade.takeUnless { DualShade.isEnabled },
                     ),
                 initialSceneKey = Scenes.Lockscreen,
+                overlayKeys =
+                    listOfNotNull(
+                        Overlays.NotificationsShade.takeIf { DualShade.isEnabled },
+                        Overlays.QuickSettingsShade.takeIf { DualShade.isEnabled },
+                    ),
                 navigationDistances =
                     mapOf(
                             Scenes.Gone to 0,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt b/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
index 9a7eef8..16ed59f4 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ShadelessSceneContainerFrameworkModule.kt
@@ -53,6 +53,7 @@
                     Scenes.Bouncer,
                 ),
             initialSceneKey = Scenes.Lockscreen,
+            overlayKeys = emptyList(),
             navigationDistances =
                 mapOf(
                     Scenes.Gone to 0,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
index beb6816..d60f05e 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/data/repository/SceneContainerRepository.kt
@@ -18,7 +18,9 @@
 
 package com.android.systemui.scene.data.repository
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.systemui.dagger.SysUISingleton
@@ -43,11 +45,27 @@
 @Inject
 constructor(
     @Application applicationScope: CoroutineScope,
-    private val config: SceneContainerConfig,
+    config: SceneContainerConfig,
     private val dataSource: SceneDataSource,
 ) {
+    /**
+     * The keys of all scenes and overlays in the container.
+     *
+     * They will be sorted in z-order such that the last one is the one that should be rendered on
+     * top of all previous ones.
+     */
+    val allContentKeys: List<ContentKey> = config.sceneKeys + config.overlayKeys
+
     val currentScene: StateFlow<SceneKey> = dataSource.currentScene
 
+    /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>> = dataSource.currentOverlays
+
     private val _isVisible = MutableStateFlow(true)
     val isVisible: StateFlow<Boolean> = _isVisible.asStateFlow()
 
@@ -72,16 +90,6 @@
                 initialValue = defaultTransitionState,
             )
 
-    /**
-     * Returns the keys to all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    fun allSceneKeys(): List<SceneKey> {
-        return config.sceneKeys
-    }
-
     fun changeScene(
         toScene: SceneKey,
         transitionKey: TransitionKey? = null,
@@ -100,6 +108,48 @@
         )
     }
 
+    /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     */
+    fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     */
+    fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     */
+    fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey? = null) {
+        dataSource.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
     /** Sets whether the container is visible. */
     fun setVisible(isVisible: Boolean) {
         _isVisible.value = isVisible
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneBackInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneBackInteractor.kt
index c176cca..2d40845 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneBackInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneBackInteractor.kt
@@ -58,29 +58,20 @@
 
     fun onSceneChange(from: SceneKey, to: SceneKey) {
         check(from != to) { "from == to, from=${from.debugName}, to=${to.debugName}" }
-        when (stackOperation(from, to)) {
-            Clear -> {
-                _backStack.value = sceneStackOf()
-            }
-            Push -> {
-                _backStack.update { s -> s.push(from) }
-            }
-            Pop -> {
-                _backStack.update { s ->
-                    checkNotNull(s.pop()) { "Cannot pop ${from.debugName} when stack is empty" }
-                        .also {
-                            val popped = s.peek()
-                            check(popped == to) {
-                                "Expected to pop ${to.debugName} but instead popped ${popped?.debugName}"
-                            }
-                        }
-                }
+
+        _backStack.update { stack ->
+            when (stackOperation(from, to, stack)) {
+                null -> stack
+                Clear -> sceneStackOf()
+                Push -> stack.push(from)
+                Pop ->
+                    checkNotNull(stack.pop()) { "Cannot pop ${from.debugName} when stack is empty" }
             }
         }
         logger.logSceneBackStack(backStack.value.asIterable())
     }
 
-    private fun stackOperation(from: SceneKey, to: SceneKey): StackOperation {
+    private fun stackOperation(from: SceneKey, to: SceneKey, stack: SceneStack): StackOperation? {
         val fromDistance =
             checkNotNull(sceneContainerConfig.navigationDistances[from]) {
                 "No distance mapping for scene \"${from.debugName}\"!"
@@ -93,6 +84,7 @@
         return when {
             toDistance == 0 -> Clear
             toDistance > fromDistance -> Push
+            stack.peek() != to -> null
             toDistance < fromDistance -> Pop
             else ->
                 error(
@@ -103,7 +95,10 @@
     }
 
     private sealed interface StackOperation
+
     private data object Clear : StackOperation
+
     private data object Push : StackOperation
+
     private data object Pop : StackOperation
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
index ea61bd3..04620d6 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneContainerOcclusionInteractor.kt
@@ -118,7 +118,7 @@
         get() =
             when (this) {
                 is ObservableTransitionState.Idle -> currentScene.canBeOccluded
-                is ObservableTransitionState.Transition.ChangeCurrentScene ->
+                is ObservableTransitionState.Transition.ChangeScene ->
                     fromScene.canBeOccluded && toScene.canBeOccluded
                 is ObservableTransitionState.Transition.ReplaceOverlay,
                 is ObservableTransitionState.Transition.ShowOrHideOverlay ->
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
index 4c404e2..0d24adc 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/SceneInteractor.kt
@@ -16,7 +16,9 @@
 
 package com.android.systemui.scene.domain.interactor
 
+import com.android.compose.animation.scene.ContentKey
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import com.android.systemui.dagger.SysUISingleton
@@ -51,6 +53,7 @@
  * other feature modules should depend on and call into this class when their parts of the
  * application state change.
  */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SysUISingleton
 class SceneInteractor
 @Inject
@@ -76,6 +79,14 @@
     private val onSceneAboutToChangeListener = mutableSetOf<OnSceneAboutToChangeListener>()
 
     /**
+     * The keys of all scenes and overlays in the container.
+     *
+     * They will be sorted in z-order such that the last one is the one that should be rendered on
+     * top of all previous ones.
+     */
+    val allContentKeys: List<ContentKey> = repository.allContentKeys
+
+    /**
      * The current scene.
      *
      * Note that during a transition between scenes, more than one scene might be rendered but only
@@ -84,6 +95,14 @@
     val currentScene: StateFlow<SceneKey> = repository.currentScene
 
     /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>> = repository.currentOverlays
+
+    /**
      * The current state of the transition.
      *
      * Consumers should use this state to know:
@@ -112,7 +131,7 @@
             .map { state ->
                 when (state) {
                     is ObservableTransitionState.Idle -> null
-                    is ObservableTransitionState.Transition.ChangeCurrentScene -> state.toScene
+                    is ObservableTransitionState.Transition.ChangeScene -> state.toScene
                     is ObservableTransitionState.Transition.ShowOrHideOverlay,
                     is ObservableTransitionState.Transition.ReplaceOverlay ->
                         TODO("b/359173565: Handle overlay transitions")
@@ -183,8 +202,8 @@
                 }
                 is ObservableTransitionState.Transition -> {
                     when {
-                        transition.toScene == scene -> transition.progress
-                        transition.fromScene == scene -> transition.progress.map { 1f - it }
+                        transition.toContent == scene -> transition.progress
+                        transition.fromContent == scene -> transition.progress.map { 1f - it }
                         else -> flowOf(0f)
                     }
                 }
@@ -192,16 +211,6 @@
         }
     }
 
-    /**
-     * Returns the keys of all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    fun allSceneKeys(): List<SceneKey> {
-        return repository.allSceneKeys()
-    }
-
     fun registerSceneStateProcessor(processor: OnSceneAboutToChangeListener) {
         onSceneAboutToChangeListener.add(processor)
     }
@@ -284,6 +293,105 @@
     }
 
     /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     *
+     * @param overlay The overlay to be shown
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun showOverlay(
+        overlay: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(to = overlay, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            to = overlay,
+            reason = loggingReason,
+        )
+
+        repository.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     *
+     * @param overlay The overlay to be hidden
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun hideOverlay(
+        overlay: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(from = overlay, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            from = overlay,
+            reason = loggingReason,
+        )
+
+        repository.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     *
+     * @param from The overlay to be hidden, if any
+     * @param to The overlay to be shown, if any
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @param transitionKey The transition key for this animated transition
+     */
+    @JvmOverloads
+    fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        loggingReason: String,
+        transitionKey: TransitionKey? = null,
+    ) {
+        if (!validateOverlayChange(from = from, to = to, loggingReason = loggingReason)) {
+            return
+        }
+
+        logger.logOverlayChangeRequested(
+            from = from,
+            to = to,
+            reason = loggingReason,
+        )
+
+        repository.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
+    /**
      * Sets the visibility of the container.
      *
      * Please do not call this from outside of the scene framework. If you are trying to force the
@@ -388,12 +496,12 @@
         to: SceneKey,
         loggingReason: String,
     ): Boolean {
-        if (!repository.allSceneKeys().contains(to)) {
+        if (to !in repository.allContentKeys) {
             return false
         }
 
         val inMidTransitionFromGone =
-            (transitionState.value as? ObservableTransitionState.Transition)?.fromScene ==
+            (transitionState.value as? ObservableTransitionState.Transition)?.fromContent ==
                 Scenes.Gone
         val isChangeAllowed =
             to != Scenes.Gone ||
@@ -409,6 +517,34 @@
         return from != to
     }
 
+    /**
+     * Validates that the given overlay change is allowed.
+     *
+     * Will throw a runtime exception for illegal states.
+     *
+     * @param from The overlay to be hidden, if any
+     * @param to The overlay to be shown, if any
+     * @param loggingReason The reason why the transition is requested, for logging purposes
+     * @return `true` if the scene change is valid; `false` if it shouldn't happen
+     */
+    private fun validateOverlayChange(
+        from: OverlayKey? = null,
+        to: OverlayKey? = null,
+        loggingReason: String,
+    ): Boolean {
+        check(from != null || to != null) {
+            "No overlay key provided for requested change." +
+                " Current transition state is ${transitionState.value}." +
+                " Logging reason for overlay change was: $loggingReason"
+        }
+
+        val isFromValid = (from == null) || (from in currentOverlays.value)
+        val isToValid =
+            (to == null) || (to !in currentOverlays.value && to in repository.allContentKeys)
+
+        return isFromValid && isToValid && from != to
+    }
+
     /** Returns a flow indicating if the currently visible scene can be resolved from [family]. */
     fun isCurrentSceneInFamily(family: SceneKey): Flow<Boolean> =
         currentScene.map { currentScene -> isSceneInFamily(currentScene, family) }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/WindowRootViewVisibilityInteractor.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/WindowRootViewVisibilityInteractor.kt
index 9c2b992..e51a8bc 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/WindowRootViewVisibilityInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/interactor/WindowRootViewVisibilityInteractor.kt
@@ -81,14 +81,14 @@
                                 state.currentScene == Scenes.QuickSettingsShade ||
                                 state.currentScene == Scenes.Lockscreen
                         is ObservableTransitionState.Transition ->
-                            state.toScene == Scenes.Shade ||
-                                state.toScene == Scenes.NotificationsShade ||
-                                state.toScene == Scenes.QuickSettingsShade ||
-                                state.toScene == Scenes.Lockscreen ||
-                                state.fromScene == Scenes.Shade ||
-                                state.fromScene == Scenes.NotificationsShade ||
-                                state.fromScene == Scenes.QuickSettingsShade ||
-                                state.fromScene == Scenes.Lockscreen
+                            state.toContent == Scenes.Shade ||
+                                state.toContent == Scenes.NotificationsShade ||
+                                state.toContent == Scenes.QuickSettingsShade ||
+                                state.toContent == Scenes.Lockscreen ||
+                                state.fromContent == Scenes.Shade ||
+                                state.fromContent == Scenes.NotificationsShade ||
+                                state.fromContent == Scenes.QuickSettingsShade ||
+                                state.fromContent == Scenes.Lockscreen
                     }
                 }
                 .distinctUntilChanged()
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 cc46216..e251c9e 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
@@ -36,6 +36,7 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.DisplayId
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor
+import com.android.systemui.deviceentry.domain.interactor.DeviceEntryHapticsInteractor
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.deviceentry.domain.interactor.DeviceUnlockedInteractor
 import com.android.systemui.deviceentry.shared.model.DeviceUnlockSource
@@ -62,6 +63,7 @@
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.statusbar.NotificationShadeWindowController
 import com.android.systemui.statusbar.SysuiStatusBarStateController
+import com.android.systemui.statusbar.VibratorHelper
 import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor
 import com.android.systemui.statusbar.phone.CentralSurfaces
 import com.android.systemui.statusbar.policy.domain.interactor.DeviceProvisioningInteractor
@@ -78,6 +80,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
@@ -107,6 +110,7 @@
     @Application private val applicationScope: CoroutineScope,
     private val sceneInteractor: SceneInteractor,
     private val deviceEntryInteractor: DeviceEntryInteractor,
+    private val deviceEntryHapticsInteractor: DeviceEntryHapticsInteractor,
     private val deviceUnlockedInteractor: DeviceUnlockedInteractor,
     private val bouncerInteractor: BouncerInteractor,
     private val keyguardInteractor: KeyguardInteractor,
@@ -134,6 +138,7 @@
     private val dismissCallbackRegistry: DismissCallbackRegistry,
     private val statusBarStateController: SysuiStatusBarStateController,
     private val alternateBouncerInteractor: AlternateBouncerInteractor,
+    private val vibratorHelper: VibratorHelper,
 ) : CoreStartable {
     private val centralSurfaces: CentralSurfaces?
         get() = centralSurfacesOptLazy.get().getOrNull()
@@ -148,6 +153,7 @@
             respondToFalsingDetections()
             hydrateInteractionState()
             handleBouncerOverscroll()
+            handleDeviceEntryHapticsWhileDeviceLocked()
             hydrateWindowController()
             hydrateBackStack()
             resetShadeSessions()
@@ -189,7 +195,7 @@
                             // current scene
                             when (state) {
                                 is ObservableTransitionState.Idle -> state.currentScene
-                                is ObservableTransitionState.Transition -> state.fromScene
+                                is ObservableTransitionState.Transition -> state.fromContent
                             }.let { it == Scenes.Shade || it == Scenes.QuickSettings }
                         }
                         .distinctUntilChanged()
@@ -220,7 +226,7 @@
                                             }
                                         }
                                         is ObservableTransitionState.Transition -> {
-                                            if (state.fromScene == Scenes.Gone) {
+                                            if (state.fromContent == Scenes.Gone) {
                                                 true to "scene transitioning away from Gone"
                                             } else {
                                                 null
@@ -351,8 +357,8 @@
                             is ObservableTransitionState.Idle -> setOf(transitionState.currentScene)
                             is ObservableTransitionState.Transition ->
                                 setOf(
-                                    transitionState.fromScene,
-                                    transitionState.toScene,
+                                    transitionState.fromContent,
+                                    transitionState.toContent,
                                 )
                         }
                     val isOnLockscreen = renderedScenes.contains(Scenes.Lockscreen)
@@ -461,7 +467,8 @@
                     sceneInteractor.transitionState.value as? ObservableTransitionState.Transition
                         ?: return@collect
                 if (
-                    transition.fromScene == Scenes.Gone && transition.toScene == Scenes.Lockscreen
+                    transition.fromContent == Scenes.Gone &&
+                        transition.toContent == Scenes.Lockscreen
                 ) {
                     switchToScene(
                         targetSceneKey = Scenes.Gone,
@@ -479,7 +486,7 @@
                     switchToScene(
                         targetSceneKey = Scenes.Lockscreen,
                         loggingReason = "device is starting to sleep",
-                        sceneState = keyguardTransitionInteractor.asleepKeyguardState.value,
+                        sceneState = keyguardInteractor.asleepKeyguardState.value,
                     )
                 } else {
                     val canSwipeToEnter = deviceEntryInteractor.canSwipeToEnter.value
@@ -524,6 +531,37 @@
         }
     }
 
+    private fun handleDeviceEntryHapticsWhileDeviceLocked() {
+        applicationScope.launch {
+            deviceEntryInteractor.isDeviceEntered.collectLatest { isDeviceEntered ->
+                // Only check for haptics signals before device is entered
+                if (!isDeviceEntered) {
+                    coroutineScope {
+                        launch {
+                            deviceEntryHapticsInteractor.playSuccessHaptic
+                                .sample(sceneInteractor.currentScene)
+                                .collect { currentScene ->
+                                    vibratorHelper.vibrateAuthSuccess(
+                                        "$TAG, $currentScene device-entry::success"
+                                    )
+                                }
+                        }
+
+                        launch {
+                            deviceEntryHapticsInteractor.playErrorHaptic
+                                .sample(sceneInteractor.currentScene)
+                                .collect { currentScene ->
+                                    vibratorHelper.vibrateAuthError(
+                                        "$TAG, $currentScene device-entry::error"
+                                    )
+                                }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     /** Keeps [SysUiState] up-to-date */
     private fun hydrateSystemUiState() {
         applicationScope.launch {
@@ -570,15 +608,6 @@
         }
 
         applicationScope.launch {
-            sceneInteractor.currentScene
-                .map { it == Scenes.Bouncer }
-                .distinctUntilChanged()
-                .collect { isBouncerShowing ->
-                    windowController.setBouncerShowing(isBouncerShowing)
-                }
-        }
-
-        applicationScope.launch {
             occlusionInteractor.invisibleDueToOcclusion.collect { invisibleDueToOcclusion ->
                 windowController.setKeyguardOccluded(invisibleDueToOcclusion)
             }
@@ -694,8 +723,8 @@
                 .filterIsInstance<ObservableTransitionState.Transition>()
                 // Only consider user-initiated (e.g. drags) that go from bouncer to lockscreen.
                 .filter { transition ->
-                    transition.fromScene == Scenes.Bouncer &&
-                        transition.toScene == Scenes.Lockscreen &&
+                    transition.fromContent == Scenes.Bouncer &&
+                        transition.toContent == Scenes.Lockscreen &&
                         transition.isInitiatedByUserInput
                 }
                 .flatMapLatest { it.progress }
@@ -816,4 +845,8 @@
                 .collectLatest { deviceEntryInteractor.refreshLockscreenEnabled() }
         }
     }
+
+    companion object {
+        private const val TAG = "SceneContainerStartable"
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
index ec743ba..d1629c7 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/domain/startable/ScrimStartable.kt
@@ -112,7 +112,7 @@
                 // It
                 // happens only when unlocking or when dismissing a dismissible lockscreen.
                 val isTransitioningAwayFromKeyguard =
-                    transitionState is ObservableTransitionState.Transition.ChangeCurrentScene &&
+                    transitionState is ObservableTransitionState.Transition.ChangeScene &&
                         transitionState.fromScene.isKeyguard() &&
                         transitionState.toScene == Scenes.Gone
 
@@ -120,7 +120,7 @@
                 val isCurrentSceneShade = currentScene.isShade()
                 // This is true when moving into one of the shade scenes when a non-shade scene.
                 val isTransitioningToShade =
-                    transitionState is ObservableTransitionState.Transition.ChangeCurrentScene &&
+                    transitionState is ObservableTransitionState.Transition.ChangeScene &&
                         !transitionState.fromScene.isShade() &&
                         transitionState.toScene.isShade()
 
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/flag/SceneContainerFlag.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/flag/SceneContainerFlag.kt
index 6c63c97..751448f 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/flag/SceneContainerFlag.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/flag/SceneContainerFlag.kt
@@ -27,7 +27,7 @@
 import com.android.systemui.keyguard.KeyguardWmStateRefactor
 import com.android.systemui.keyguard.MigrateClocksToBlueprint
 import com.android.systemui.keyguard.shared.ComposeLockscreen
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
+import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun
 import com.android.systemui.statusbar.phone.PredictiveBackSysUiFlag
 
 /** Helper for reading or using the scene container flag state. */
@@ -43,7 +43,7 @@
                 KeyguardBottomAreaRefactor.isEnabled &&
                 KeyguardWmStateRefactor.isEnabled &&
                 MigrateClocksToBlueprint.isEnabled &&
-                NotificationsHeadsUpRefactor.isEnabled &&
+                NotificationThrottleHun.isEnabled &&
                 PredictiveBackSysUiFlag.isEnabled &&
                 DeviceEntryUdfpsRefactor.isEnabled
 
@@ -59,7 +59,7 @@
             KeyguardBottomAreaRefactor.token,
             KeyguardWmStateRefactor.token,
             MigrateClocksToBlueprint.token,
-            NotificationsHeadsUpRefactor.token,
+            NotificationThrottleHun.token,
             PredictiveBackSysUiFlag.token,
             DeviceEntryUdfpsRefactor.token,
             // NOTE: Changes should also be made in isEnabled and @EnableSceneContainer
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
index 045a887..fb53ddb 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/logger/SceneLogger.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.scene.shared.logger
 
 import com.android.compose.animation.scene.ObservableTransitionState
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.core.LogLevel
@@ -77,8 +78,8 @@
                     tag = TAG,
                     level = LogLevel.INFO,
                     messageInitializer = {
-                        str1 = transitionState.fromScene.toString()
-                        str2 = transitionState.toScene.toString()
+                        str1 = transitionState.fromContent.toString()
+                        str2 = transitionState.toContent.toString()
                     },
                     messagePrinter = { "Scene transition started: $str1 → $str2" },
                 )
@@ -94,6 +95,34 @@
         }
     }
 
+    fun logOverlayChangeRequested(
+        from: OverlayKey? = null,
+        to: OverlayKey? = null,
+        reason: String,
+    ) {
+        logBuffer.log(
+            tag = TAG,
+            level = LogLevel.INFO,
+            messageInitializer = {
+                str1 = from?.toString()
+                str2 = to?.toString()
+                str3 = reason
+            },
+            messagePrinter = {
+                buildString {
+                    append("Overlay change requested: ")
+                    if (str1 != null) {
+                        append(str1)
+                        append(if (str2 == null) " (hidden)" else " → $str2")
+                    } else {
+                        append("$str2 (shown)")
+                    }
+                    append(", reason: $str3")
+                }
+            },
+        )
+    }
+
     fun logVisibilityChange(
         from: Boolean,
         to: Boolean,
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Overlays.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Overlays.kt
new file mode 100644
index 0000000..c47a850
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Overlays.kt
@@ -0,0 +1,52 @@
+/*
+ * 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.scene.shared.model
+
+import com.android.compose.animation.scene.OverlayKey
+
+/**
+ * Keys of all known overlays.
+ *
+ * PLEASE KEEP THE KEYS SORTED ALPHABETICALLY.
+ */
+object Overlays {
+    /**
+     * The notifications shade overlay primarily shows a scrollable list of notifications.
+     *
+     * It's used only in the dual shade configuration, where there are two separate shades: one for
+     * notifications (this overlay) and another for [QuickSettingsShade].
+     *
+     * It's not used in the single/accordion configuration (swipe down once to reveal the shade,
+     * swipe down again the to expand quick settings) or in the "split" shade configuration (on
+     * large screens or unfolded foldables, where notifications and quick settings are shown
+     * side-by-side in their own columns).
+     */
+    @JvmField val NotificationsShade = OverlayKey("notifications_shade")
+
+    /**
+     * The quick settings shade overlay shows the quick settings tiles UI.
+     *
+     * It's used only in the dual shade configuration, where there are two separate shades: one for
+     * quick settings (this overlay) and another for [NotificationsShade].
+     *
+     * It's not used in the single/accordion configuration (swipe down once to reveal the shade,
+     * swipe down again the to expand quick settings) or in the "split" shade configuration (on
+     * large screens or unfolded foldables, where notifications and quick settings are shown
+     * side-by-side in their own columns).
+     */
+    @JvmField val QuickSettingsShade = OverlayKey("quick_settings_shade")
+}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt
deleted file mode 100644
index 8e2e8a1..0000000
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scene.kt
+++ /dev/null
@@ -1,59 +0,0 @@
-/*
- * Copyright (C) 2023 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.scene.shared.model
-
-import com.android.compose.animation.scene.SceneKey
-import com.android.compose.animation.scene.UserAction
-import com.android.compose.animation.scene.UserActionResult
-import com.android.systemui.lifecycle.Activatable
-import kotlinx.coroutines.flow.Flow
-
-/**
- * Defines interface for classes that can describe a "scene".
- *
- * In the scene framework, there can be multiple scenes in a single scene "container". The container
- * takes care of rendering the current scene and allowing scenes to be switched from one to another
- * based on either user action (for example, swiping down while on the lock screen scene may switch
- * to the shade scene).
- */
-interface Scene : Activatable {
-
-    /** Uniquely-identifying key for this scene. The key must be unique within its container. */
-    val key: SceneKey
-
-    /**
-     * The mapping between [UserAction] and destination [UserActionResult]s.
-     *
-     * When the scene framework detects a user action, if the current scene has a map entry for that
-     * user action, the framework starts a transition to the scene in the map.
-     *
-     * Once the [Scene] becomes the current one, the scene framework will read this property and set
-     * up a collector to watch for new mapping values. If every map entry provided by the scene, the
-     * framework will set up user input handling for its [UserAction] and, if such a user action is
-     * detected, initiate a transition to the specified [UserActionResult].
-     *
-     * Note that reading from this method does _not_ mean that any user action has occurred.
-     * Instead, the property is read before any user action/gesture is detected so that the
-     * framework can decide whether to set up gesture/input detectors/listeners in case user actions
-     * of the given types ever occur.
-     *
-     * Note that a missing value for a specific [UserAction] means that the user action of the given
-     * type is not currently active in the scene and should be ignored by the framework, while the
-     * current scene is this one.
-     */
-    val destinationScenes: Flow<Map<UserAction, UserActionResult>>
-}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
index 0a30c31..2311e47 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneContainerConfig.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 
 /** Models the configuration of the scene container. */
@@ -38,6 +39,13 @@
     val initialSceneKey: SceneKey,
 
     /**
+     * The keys to all overlays in the container, sorted by z-order such that the last one renders
+     * on top of all previous ones. Overlay keys within the same container must not repeat but it's
+     * okay to have the same overlay keys in different containers.
+     */
+    val overlayKeys: List<OverlayKey> = emptyList(),
+
+    /**
      * Navigation distance of each scene.
      *
      * The navigation distance is a measure of how many non-back user action "steps" away from the
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
index 034da25..4538d1c 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSource.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.flow.StateFlow
@@ -33,6 +34,14 @@
     val currentScene: StateFlow<SceneKey>
 
     /**
+     * The current set of overlays to be shown (may be empty).
+     *
+     * Note that during a transition between overlays, a different set of overlays may be rendered -
+     * but only the ones in this set are considered the current overlays.
+     */
+    val currentOverlays: StateFlow<Set<OverlayKey>>
+
+    /**
      * Asks for an asynchronous scene switch to [toScene], which will use the corresponding
      * installed transition or the one specified by [transitionKey], if provided.
      */
@@ -47,4 +56,40 @@
     fun snapToScene(
         toScene: SceneKey,
     )
+
+    /**
+     * Request to show [overlay] so that it animates in from [currentScene] and ends up being
+     * visible on screen.
+     *
+     * After this returns, this overlay will be included in [currentOverlays]. This does nothing if
+     * [overlay] is already shown.
+     */
+    fun showOverlay(
+        overlay: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Request to hide [overlay] so that it animates out to [currentScene] and ends up *not* being
+     * visible on screen.
+     *
+     * After this returns, this overlay will not be included in [currentOverlays]. This does nothing
+     * if [overlay] is already hidden.
+     */
+    fun hideOverlay(
+        overlay: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
+
+    /**
+     * Replace [from] by [to] so that [from] ends up not being visible on screen and [to] ends up
+     * being visible.
+     *
+     * This throws if [from] is not currently shown or if [to] is already shown.
+     */
+    fun replaceOverlay(
+        from: OverlayKey,
+        to: OverlayKey,
+        transitionKey: TransitionKey? = null,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
index 43c3635..eb4c0f2 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/SceneDataSourceDelegator.kt
@@ -18,6 +18,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.CoroutineScope
@@ -49,6 +50,15 @@
                 initialValue = config.initialSceneKey,
             )
 
+    override val currentOverlays: StateFlow<Set<OverlayKey>> =
+        delegateMutable
+            .flatMapLatest { delegate -> delegate.currentOverlays }
+            .stateIn(
+                scope = applicationScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptySet(),
+            )
+
     override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) {
         delegateMutable.value.changeScene(
             toScene = toScene,
@@ -62,6 +72,28 @@
         )
     }
 
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.showOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.hideOverlay(
+            overlay = overlay,
+            transitionKey = transitionKey,
+        )
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        delegateMutable.value.replaceOverlay(
+            from = from,
+            to = to,
+            transitionKey = transitionKey,
+        )
+    }
+
     /**
      * Binds the current, dependency injection provided [SceneDataSource] to the given object.
      *
@@ -82,8 +114,21 @@
         override val currentScene: StateFlow<SceneKey> =
             MutableStateFlow(initialSceneKey).asStateFlow()
 
+        override val currentOverlays: StateFlow<Set<OverlayKey>> =
+            MutableStateFlow(emptySet<OverlayKey>()).asStateFlow()
+
         override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) = Unit
 
         override fun snapToScene(toScene: SceneKey) = Unit
+
+        override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) = Unit
+
+        override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) = Unit
+
+        override fun replaceOverlay(
+            from: OverlayKey,
+            to: OverlayKey,
+            transitionKey: TransitionKey?
+        ) = Unit
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
index ef5290f..fcf6288 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/Scenes.kt
@@ -54,7 +54,9 @@
      * large screens or unfolded foldables, where notifications and quick settings are shown
      * side-by-side in their own columns).
      */
-    @JvmField val NotificationsShade = SceneKey("notifications_shade")
+    @Deprecated("The notifications shade scene has been replaced by an overlay")
+    @JvmField
+    val NotificationsShade = SceneKey("notifications_shade")
 
     /**
      * The quick settings scene shows the quick setting tiles.
@@ -70,7 +72,9 @@
      * and one for quick settings, [NotificationsShade] and [QuickSettingsShade] scenes are used
      * respectively.
      */
-    @JvmField val QuickSettings = SceneKey("quick_settings")
+    @Deprecated("The quick settings shade scene has been replaced by an overlay")
+    @JvmField
+    val QuickSettings = SceneKey("quick_settings")
 
     /**
      * The quick settings shade scene shows the quick setting tiles as an overlay UI.
diff --git a/packages/SystemUI/src/com/android/systemui/scene/shared/model/TransitionKeys.kt b/packages/SystemUI/src/com/android/systemui/scene/shared/model/TransitionKeys.kt
index be95441..b9f57f2 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/shared/model/TransitionKeys.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/shared/model/TransitionKeys.kt
@@ -27,14 +27,6 @@
     /** Reference to the gone/lockscreen to shade transition with split shade enabled. */
     val ToSplitShade = TransitionKey("GoneToSplitShade")
 
-    /** Reference to a scene transition that can collapse the shade scene instantly. */
-    val CollapseShadeInstantly = TransitionKey("CollapseShadeInstantly")
-
-    /**
-     * Reference to a scene transition that brings up the shade from the bottom instead of the top.
-     */
-    val OpenBottomShade = TransitionKey("OpenBottomShade")
-
     /**
      * Reference to a scene transition that can collapse the shade scene slightly faster than a
      * normal collapse would.
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
index 8aa601f..8a2e274 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootView.kt
@@ -6,9 +6,10 @@
 import android.view.View
 import android.view.WindowInsets
 import com.android.systemui.keyguard.ui.viewmodel.AlternateBouncerDependencies
-import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.shade.TouchLogger
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
@@ -35,6 +36,7 @@
         containerConfig: SceneContainerConfig,
         sharedNotificationContainer: SharedNotificationContainer,
         scenes: Set<Scene>,
+        overlays: Set<Overlay>,
         layoutInsetController: LayoutInsetsController,
         sceneDataSourceDelegator: SceneDataSourceDelegator,
         alternateBouncerDependencies: AlternateBouncerDependencies,
@@ -50,6 +52,7 @@
             containerConfig = containerConfig,
             sharedNotificationContainer = sharedNotificationContainer,
             scenes = scenes,
+            overlays = overlays,
             onVisibilityChangedInternal = { isVisible ->
                 super.setVisibility(if (isVisible) View.VISIBLE else View.INVISIBLE)
             },
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
index 0f05af6..075599b 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/view/SceneWindowRootViewBinder.kt
@@ -29,6 +29,7 @@
 import androidx.compose.ui.unit.dp
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.theme.PlatformTheme
 import com.android.internal.policy.ScreenDecorationsUtils
@@ -43,10 +44,10 @@
 import com.android.systemui.lifecycle.viewModel
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
-import com.android.systemui.scene.ui.composable.ComposableScene
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.scene.ui.composable.SceneContainer
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
 import com.android.systemui.statusbar.notification.stack.ui.view.SharedNotificationContainer
@@ -70,6 +71,7 @@
         containerConfig: SceneContainerConfig,
         sharedNotificationContainer: SharedNotificationContainer,
         scenes: Set<Scene>,
+        overlays: Set<Overlay>,
         onVisibilityChangedInternal: (isVisible: Boolean) -> Unit,
         dataSourceDelegator: SceneDataSourceDelegator,
         alternateBouncerDependencies: AlternateBouncerDependencies,
@@ -86,8 +88,22 @@
             }
         }
 
+        val unsortedOverlayByKey: Map<OverlayKey, Overlay> =
+            overlays.associateBy { overlay -> overlay.key }
+        val sortedOverlayByKey: Map<OverlayKey, Overlay> = buildMap {
+            containerConfig.overlayKeys.forEach { overlayKey ->
+                val overlay =
+                    checkNotNull(unsortedOverlayByKey[overlayKey]) {
+                        "Overlay not found for key \"$overlayKey\"!"
+                    }
+
+                put(overlayKey, overlay)
+            }
+        }
+
         view.repeatWhenAttached {
             view.viewModel(
+                traceName = "SceneWindowRootViewBinder",
                 minWindowLifecycleState = WindowLifecycleState.ATTACHED,
                 factory = { viewModelFactory.create(motionEventHandlerReceiver) },
             ) { viewModel ->
@@ -112,6 +128,7 @@
                                 viewModel = viewModel,
                                 windowInsets = windowInsets,
                                 sceneByKey = sortedSceneByKey,
+                                overlayByKey = sortedOverlayByKey,
                                 dataSourceDelegator = dataSourceDelegator,
                                 containerConfig = containerConfig,
                             )
@@ -156,6 +173,7 @@
         viewModel: SceneContainerViewModel,
         windowInsets: StateFlow<WindowInsets?>,
         sceneByKey: Map<SceneKey, Scene>,
+        overlayByKey: Map<OverlayKey, Overlay>,
         dataSourceDelegator: SceneDataSourceDelegator,
         containerConfig: SceneContainerConfig,
     ): View {
@@ -168,8 +186,8 @@
                     ) {
                         SceneContainer(
                             viewModel = viewModel,
-                            sceneByKey =
-                                sceneByKey.mapValues { (_, scene) -> scene as ComposableScene },
+                            sceneByKey = sceneByKey,
+                            overlayByKey = overlayByKey,
                             initialSceneKey = containerConfig.initialSceneKey,
                             dataSourceDelegator = dataSourceDelegator,
                         )
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModel.kt
deleted file mode 100644
index b707a5a..0000000
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneSceneActionsViewModel.kt
+++ /dev/null
@@ -1,92 +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.scene.ui.viewmodel
-
-import androidx.compose.ui.Alignment
-import com.android.compose.animation.scene.Edge
-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.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.shared.model.TransitionKeys.OpenBottomShade
-import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
-import com.android.systemui.shade.shared.model.ShadeMode
-import dagger.assisted.AssistedFactory
-import dagger.assisted.AssistedInject
-import kotlinx.coroutines.flow.collectLatest
-import kotlinx.coroutines.flow.map
-
-class GoneSceneActionsViewModel
-@AssistedInject
-constructor(
-    private val shadeInteractor: ShadeInteractor,
-) : SceneActionsViewModel() {
-
-    override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
-        shadeInteractor.shadeMode
-            .map { shadeMode ->
-                buildMap<UserAction, UserActionResult> {
-                    if (
-                        shadeMode is ShadeMode.Single ||
-                            // TODO(b/338577208): Remove this once we add Dual Shade invocation
-                            // zones.
-                            shadeMode is ShadeMode.Dual
-                    ) {
-                        if (shadeInteractor.shadeAlignment == Alignment.BottomEnd) {
-                            put(
-                                Swipe(
-                                    pointerCount = 2,
-                                    fromSource = Edge.Bottom,
-                                    direction = SwipeDirection.Up,
-                                ),
-                                UserActionResult(SceneFamilies.QuickSettings, OpenBottomShade)
-                            )
-                        } else {
-                            put(
-                                Swipe(
-                                    pointerCount = 2,
-                                    fromSource = Edge.Top,
-                                    direction = SwipeDirection.Down,
-                                ),
-                                UserActionResult(SceneFamilies.QuickSettings)
-                            )
-                        }
-                    }
-
-                    if (shadeInteractor.shadeAlignment == Alignment.BottomEnd) {
-                        put(Swipe.Up, UserActionResult(SceneFamilies.NotifShade, OpenBottomShade))
-                    } else {
-                        put(
-                            Swipe.Down,
-                            UserActionResult(
-                                SceneFamilies.NotifShade,
-                                ToSplitShade.takeIf { shadeMode is ShadeMode.Split }
-                            )
-                        )
-                    }
-                }
-            }
-            .collectLatest { setActions(it) }
-    }
-
-    @AssistedFactory
-    interface Factory {
-        fun create(): GoneSceneActionsViewModel
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModel.kt
new file mode 100644
index 0000000..ea4122a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/GoneUserActionsViewModel.kt
@@ -0,0 +1,74 @@
+/*
+ * 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.scene.ui.viewmodel
+
+import com.android.compose.animation.scene.Edge
+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.scene.shared.model.SceneFamilies
+import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
+import com.android.systemui.shade.domain.interactor.ShadeInteractor
+import com.android.systemui.shade.shared.model.ShadeMode
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
+import kotlinx.coroutines.flow.map
+
+class GoneUserActionsViewModel
+@AssistedInject
+constructor(
+    private val shadeInteractor: ShadeInteractor,
+) : UserActionsViewModel() {
+
+    override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
+        shadeInteractor.shadeMode
+            .map { shadeMode ->
+                buildMap<UserAction, UserActionResult> {
+                    if (
+                        shadeMode is ShadeMode.Single ||
+                            // TODO(b/338577208): Remove this once we add Dual Shade invocation
+                            // zones.
+                            shadeMode is ShadeMode.Dual
+                    ) {
+                        put(
+                            Swipe(
+                                pointerCount = 2,
+                                fromSource = Edge.Top,
+                                direction = SwipeDirection.Down,
+                            ),
+                            UserActionResult(SceneFamilies.QuickSettings)
+                        )
+                    }
+
+                    put(
+                        Swipe.Down,
+                        UserActionResult(
+                            SceneFamilies.NotifShade,
+                            ToSplitShade.takeIf { shadeMode is ShadeMode.Split }
+                        )
+                    )
+                }
+            }
+            .collect { setActions(it) }
+    }
+
+    @AssistedFactory
+    interface Factory {
+        fun create(): GoneUserActionsViewModel
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
index c4be26a..a73c39d 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneContainerViewModel.kt
@@ -26,7 +26,6 @@
 import com.android.systemui.classifier.domain.interactor.FalsingInteractor
 import com.android.systemui.lifecycle.ExclusiveActivatable
 import com.android.systemui.lifecycle.Hydrator
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.logger.SceneLogger
@@ -47,22 +46,15 @@
     private val powerInteractor: PowerInteractor,
     private val logger: SceneLogger,
     @Assisted private val motionEventHandlerReceiver: (MotionEventHandler?) -> Unit,
-) : SysUiViewModel, ExclusiveActivatable() {
-    /**
-     * Keys of all scenes in the container.
-     *
-     * The scenes will be sorted in z-order such that the last one is the one that should be
-     * rendered on top of all previous ones.
-     */
-    val allSceneKeys: List<SceneKey> = sceneInteractor.allSceneKeys()
+) : ExclusiveActivatable() {
 
     /** The scene that should be rendered. */
     val currentScene: StateFlow<SceneKey> = sceneInteractor.currentScene
 
-    private val hydrator = Hydrator()
+    private val hydrator = Hydrator("SceneContainerViewModel.hydrator")
 
     /** Whether the container is visible. */
-    val isVisible: Boolean by hydrator.hydratedStateOf(sceneInteractor.isVisible)
+    val isVisible: Boolean by hydrator.hydratedStateOf("isVisible", sceneInteractor.isVisible)
 
     override suspend fun onActivated(): Nothing {
         try {
diff --git a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModel.kt
similarity index 90%
rename from packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModel.kt
index 9144f16d..57628d0 100644
--- a/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/SceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/scene/ui/viewmodel/UserActionsViewModel.kt
@@ -19,21 +19,19 @@
 import com.android.compose.animation.scene.UserAction
 import com.android.compose.animation.scene.UserActionResult
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 
 /**
- * Base class for view-models that need to keep a map of scene actions (also known as "destination
- * scenes") up-to-date.
+ * Base class for view-models that need to keep a map of user actions up-to-date.
  *
  * Subclasses need only to override [hydrateActions], suspending forever if they need; they don't
  * need to worry about resetting the value of [actions] when the view-model is deactivated/canceled,
  * this base class takes care of it.
  */
-abstract class SceneActionsViewModel : SysUiViewModel, ExclusiveActivatable() {
+abstract class UserActionsViewModel : ExclusiveActivatable() {
 
     private val _actions = MutableStateFlow<Map<UserAction, UserActionResult>>(emptyMap())
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/MessageContainerController.kt b/packages/SystemUI/src/com/android/systemui/screenshot/MessageContainerController.kt
index 474afa8b..56afb79 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/MessageContainerController.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/MessageContainerController.kt
@@ -9,7 +9,6 @@
 import android.view.ViewTreeObserver
 import android.view.animation.AccelerateDecelerateInterpolator
 import androidx.constraintlayout.widget.Guideline
-import com.android.systemui.Flags.screenshotPrivateProfileBehaviorFix
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.res.R
 import com.android.systemui.screenshot.message.ProfileMessageController
@@ -49,44 +48,19 @@
     }
 
     fun onScreenshotTaken(screenshot: ScreenshotData) {
-        if (screenshotPrivateProfileBehaviorFix()) {
-            mainScope.launch {
-                val profileData = profileMessageController.onScreenshotTaken(screenshot.userHandle)
-                var notifiedApps: List<CharSequence> =
-                    screenshotDetectionController.maybeNotifyOfScreenshot(screenshot)
-
-                // If profile first run needs to show, bias towards that, otherwise show screenshot
-                // detection notification if needed.
-                if (profileData != null) {
-                    workProfileFirstRunView.visibility = View.VISIBLE
-                    detectionNoticeView.visibility = View.GONE
-                    profileMessageController.bindView(workProfileFirstRunView, profileData) {
-                        animateOutMessageContainer()
-                    }
-                    animateInMessageContainer()
-                } else if (notifiedApps.isNotEmpty()) {
-                    detectionNoticeView.visibility = View.VISIBLE
-                    workProfileFirstRunView.visibility = View.GONE
-                    screenshotDetectionController.populateView(detectionNoticeView, notifiedApps)
-                    animateInMessageContainer()
-                }
-            }
-        } else {
-            val workProfileData =
-                workProfileMessageController.onScreenshotTaken(screenshot.userHandle)
+        mainScope.launch {
+            val profileData = profileMessageController.onScreenshotTaken(screenshot.userHandle)
             var notifiedApps: List<CharSequence> =
                 screenshotDetectionController.maybeNotifyOfScreenshot(screenshot)
 
-            // If work profile first run needs to show, bias towards that, otherwise show screenshot
+            // If profile first run needs to show, bias towards that, otherwise show screenshot
             // detection notification if needed.
-            if (workProfileData != null) {
+            if (profileData != null) {
                 workProfileFirstRunView.visibility = View.VISIBLE
                 detectionNoticeView.visibility = View.GONE
-                workProfileMessageController.populateView(
-                    workProfileFirstRunView,
-                    workProfileData,
-                    this::animateOutMessageContainer
-                )
+                profileMessageController.bindView(workProfileFirstRunView, profileData) {
+                    animateOutMessageContainer()
+                }
                 animateInMessageContainer()
             } else if (notifiedApps.isNotEmpty()) {
                 detectionNoticeView.visibility = View.VISIBLE
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
deleted file mode 100644
index 922997d..0000000
--- a/packages/SystemUI/src/com/android/systemui/screenshot/RequestProcessor.kt
+++ /dev/null
@@ -1,62 +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.systemui.screenshot
-
-import android.util.Log
-import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
-
-/** Implementation of [ScreenshotRequestProcessor] */
-class RequestProcessor(
-    private val capture: ImageCapture,
-    private val policy: ScreenshotPolicy,
-) : ScreenshotRequestProcessor {
-
-    override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
-        var result = screenshot
-
-        // Apply work profile screenshots policy:
-        //
-        // If the focused app belongs to a work profile, transforms a full screen
-        // (or partial) screenshot request to a task snapshot (provided image) screenshot.
-
-        // Whenever displayContentInfo is fetched, the topComponent is also populated
-        // regardless of the managed profile status.
-
-        if (screenshot.type != TAKE_SCREENSHOT_PROVIDED_IMAGE) {
-            val info = policy.findPrimaryContent(screenshot.displayId)
-            Log.d(TAG, "findPrimaryContent: $info")
-            result.taskId = info.taskId
-            result.topComponent = info.component
-            result.userHandle = info.user
-
-            if (policy.isManagedProfile(info.user.identifier)) {
-                val image =
-                    capture.captureTask(info.taskId)
-                        ?: throw RequestProcessorException("Task snapshot returned a null Bitmap!")
-
-                // Provide the task snapshot as the screenshot
-                result.type = TAKE_SCREENSHOT_PROVIDED_IMAGE
-                result.bitmap = image
-                result.screenBounds = info.bounds
-            }
-        }
-
-        return result
-    }
-}
-
-private const val TAG = "RequestProcessor"
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotRequestProcessor.kt
index 3ad4075a..ee1008d 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotRequestProcessor.kt
@@ -27,5 +27,5 @@
     suspend fun process(original: ScreenshotData): ScreenshotData
 }
 
-/** Exception thrown by [RequestProcessor] if something goes wrong. */
+/** Exception thrown by [ScreenshotRequestProcessor] if something goes wrong. */
 class RequestProcessorException(message: String) : IllegalStateException(message)
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
index 9db1f24..ad5e772 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/AppClipsActivity.java
@@ -71,7 +71,9 @@
 import com.android.systemui.screenshot.scroll.CropView;
 import com.android.systemui.settings.UserTracker;
 
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import javax.inject.Inject;
@@ -344,10 +346,63 @@
 
         // Set up the dropdown when multiple backlinks are available.
         if (backlinksData.size() > 1) {
-            setUpListPopupWindow(backlinksData, mBacklinksDataTextView);
+            setUpListPopupWindow(updateBacklinkLabelsWithDuplicateNames(backlinksData),
+                    mBacklinksDataTextView);
         }
     }
 
+    /**
+     * If there are more than 1 backlinks that have the same app name, then this method appends
+     * a numerical suffix to such backlinks to help users distinguish.
+     */
+    private List<InternalBacklinksData> updateBacklinkLabelsWithDuplicateNames(
+            List<InternalBacklinksData> backlinksData) {
+        // Check if there are multiple backlinks with same name.
+        Map<String, Integer> duplicateNamedBacklinksCountMap = new HashMap<>();
+        for (InternalBacklinksData data : backlinksData) {
+            if (duplicateNamedBacklinksCountMap.containsKey(data.getDisplayLabel())) {
+                int duplicateCount = duplicateNamedBacklinksCountMap.get(data.getDisplayLabel());
+                if (duplicateCount == 0) {
+                    // If this is the first time the loop is coming across a duplicate name, set the
+                    // count to 2. This way the count starts from 1 for all duplicate named
+                    // backlinks.
+                    duplicateNamedBacklinksCountMap.put(data.getDisplayLabel(), 2);
+                } else {
+                    // For all duplicate named backlinks, increase the duplicate count by 1.
+                    duplicateNamedBacklinksCountMap.put(data.getDisplayLabel(), duplicateCount + 1);
+                }
+            } else {
+                // This is the first time the loop is coming across a backlink with this name. Set
+                // its count to 0. The loop will increase its count by 1 when a duplicate is found.
+                duplicateNamedBacklinksCountMap.put(data.getDisplayLabel(), 0);
+            }
+        }
+
+        // Go through the backlinks in reverse order as it is easier to assign the numerical suffix
+        // in descending order of frequency using the duplicate map that was built earlier. For
+        // example, if "App A" is present 3 times, then we assign display label "App A (3)" first
+        // and then "App A (2)", lastly "App A (1)".
+        for (InternalBacklinksData data : backlinksData.reversed()) {
+            String originalBacklinkLabel = data.getDisplayLabel();
+            int duplicateCount = duplicateNamedBacklinksCountMap.get(originalBacklinkLabel);
+
+            // The display label should only be updated if there are multiple backlinks with the
+            // same name.
+            if (duplicateCount > 0) {
+                // Update the display label to: "App name (count)"
+                data.setDisplayLabel(
+                        getString(R.string.backlinks_duplicate_label_format, originalBacklinkLabel,
+                                duplicateCount));
+
+                // Decrease the duplicate count and update the map.
+                duplicateCount--;
+                duplicateNamedBacklinksCountMap.put(originalBacklinkLabel, duplicateCount);
+            }
+        }
+
+        return backlinksData;
+    }
+
     private void setUpListPopupWindow(List<InternalBacklinksData> backlinksData, View anchor) {
         ListPopupWindow listPopupWindow = new ListPopupWindow(this);
         listPopupWindow.setAnchorView(anchor);
@@ -365,7 +420,7 @@
             public View getView(int position, @Nullable View convertView, ViewGroup parent) {
                 TextView itemView = (TextView) super.getView(position, convertView, parent);
                 InternalBacklinksData data = backlinksData.get(position);
-                itemView.setText(data.getClipData().getDescription().getLabel());
+                itemView.setText(data.getDisplayLabel());
 
                 Drawable icon = data.getAppIcon();
                 icon.setBounds(createBacklinksTextViewDrawableBounds());
@@ -387,7 +442,7 @@
      * expected to be already set when this method is called.
      */
     private void updateBacklinksTextView(InternalBacklinksData backlinksData) {
-        mBacklinksDataTextView.setText(backlinksData.getClipData().getDescription().getLabel());
+        mBacklinksDataTextView.setText(backlinksData.getDisplayLabel());
         Drawable appIcon = backlinksData.getAppIcon();
         Rect compoundDrawableBounds = createBacklinksTextViewDrawableBounds();
         appIcon.setBounds(compoundDrawableBounds);
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/InternalBacklinksData.kt b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/InternalBacklinksData.kt
index 0e312f9..30c33c5 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/appclips/InternalBacklinksData.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/appclips/InternalBacklinksData.kt
@@ -20,4 +20,6 @@
 import android.graphics.drawable.Drawable
 
 /** A class to hold the [ClipData] for backlinks and the corresponding app's [Drawable] icon. */
-internal data class InternalBacklinksData(val clipData: ClipData, val appIcon: Drawable)
+internal data class InternalBacklinksData(val clipData: ClipData, val appIcon: Drawable) {
+    var displayLabel: String = clipData.description.label.toString()
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
index 44f767a..2cb9fe7 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
@@ -19,14 +19,11 @@
 import android.content.ComponentName
 import android.content.Context
 import android.os.Process
-import com.android.systemui.Flags.screenshotPrivateProfileBehaviorFix
 import com.android.systemui.SystemUIService
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.screenshot.ImageCapture
-import com.android.systemui.screenshot.RequestProcessor
-import com.android.systemui.screenshot.ScreenshotPolicy
 import com.android.systemui.screenshot.ScreenshotRequestProcessor
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
 import com.android.systemui.screenshot.data.repository.DisplayContentRepositoryImpl
@@ -68,23 +65,18 @@
             @Application context: Context,
             @Background background: CoroutineDispatcher,
             imageCapture: ImageCapture,
-            policyProvider: Provider<ScreenshotPolicy>,
-            displayContentRepoProvider: Provider<DisplayContentRepository>,
+            displayContentRepo: DisplayContentRepository,
             policyListProvider: Provider<List<CapturePolicy>>,
         ): ScreenshotRequestProcessor {
-            return if (screenshotPrivateProfileBehaviorFix()) {
-                PolicyRequestProcessor(
-                    background = background,
-                    capture = imageCapture,
-                    displayTasks = displayContentRepoProvider.get(),
-                    policies = policyListProvider.get(),
-                    defaultOwner = Process.myUserHandle(),
-                    defaultComponent =
-                        ComponentName(context.packageName, SystemUIService::class.java.toString())
-                )
-            } else {
-                RequestProcessor(imageCapture, policyProvider.get())
-            }
+            return PolicyRequestProcessor(
+                background = background,
+                capture = imageCapture,
+                displayTasks = displayContentRepo,
+                policies = policyListProvider.get(),
+                defaultOwner = Process.myUserHandle(),
+                defaultComponent =
+                    ComponentName(context.packageName, SystemUIService::class.java.toString())
+            )
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/settings/DisplayTracker.kt b/packages/SystemUI/src/com/android/systemui/settings/DisplayTracker.kt
index e7ee961..edff4bf 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/DisplayTracker.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/DisplayTracker.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.settings
 
 import android.view.Display
+import com.android.systemui.util.annotations.WeaklyReferencedCallback
 import java.util.concurrent.Executor
 
 /**
@@ -52,6 +53,7 @@
     fun getDisplay(displayId: Int): Display
 
     /** Ćallback for notifying of changes. */
+    @WeaklyReferencedCallback
     interface Callback {
 
         /** Notifies that a display has been added. */
diff --git a/packages/SystemUI/src/com/android/systemui/settings/MultiUserUtilsModule.java b/packages/SystemUI/src/com/android/systemui/settings/MultiUserUtilsModule.java
index 05f19ef..b9f9b92 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/MultiUserUtilsModule.java
+++ b/packages/SystemUI/src/com/android/systemui/settings/MultiUserUtilsModule.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.settings;
 
-import android.app.ActivityManager;
 import android.app.IActivityManager;
 import android.content.Context;
 import android.hardware.display.DisplayManager;
@@ -67,7 +66,7 @@
             @Background CoroutineDispatcher backgroundDispatcher,
             @Background Handler handler
     ) {
-        int startingUser = ActivityManager.getCurrentUser();
+        int startingUser = userManager.getBootUser().getIdentifier();
         UserTrackerImpl tracker = new UserTrackerImpl(context, featureFlagsProvider, userManager,
                 iActivityManager, dumpManager, appScope, backgroundDispatcher, handler);
         tracker.initialize(startingUser);
diff --git a/packages/SystemUI/src/com/android/systemui/settings/brightness/ui/viewModel/BrightnessMirrorViewModel.kt b/packages/SystemUI/src/com/android/systemui/settings/brightness/ui/viewModel/BrightnessMirrorViewModel.kt
index 706797d..52bc25d 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/brightness/ui/viewModel/BrightnessMirrorViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/settings/brightness/ui/viewModel/BrightnessMirrorViewModel.kt
@@ -20,7 +20,6 @@
 import android.util.Log
 import android.view.View
 import com.android.systemui.dagger.qualifiers.Main
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.res.R
 import com.android.systemui.settings.brightness.BrightnessSliderController
 import com.android.systemui.settings.brightness.MirrorController
@@ -37,7 +36,7 @@
     private val brightnessMirrorShowingInteractor: BrightnessMirrorShowingInteractor,
     @Main private val resources: Resources,
     val sliderControllerFactory: BrightnessSliderController.Factory,
-) : SysUiViewModel, MirrorController {
+) : MirrorController {
 
     private val tempPosition = IntArray(2)
 
diff --git a/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt b/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
index 4639e22..3bb494b 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/GlanceableHubContainerController.kt
@@ -203,6 +203,13 @@
      */
     private var shadeConsumingTouches = false
 
+    /**
+     * True if the shade is showing at all.
+     *
+     * Inverse of [ShadeInteractor.isShadeFullyCollapsed]
+     */
+    private var shadeShowing = false
+
     /** True if the keyguard transition state is finished on [KeyguardState.LOCKSCREEN]. */
     private var onLockscreen = false
 
@@ -414,6 +421,7 @@
             ),
             { (isFullyExpanded, isUserInteracting, isShadeFullyCollapsed) ->
                 shadeConsumingTouches = isUserInteracting
+                shadeShowing = !isShadeFullyCollapsed
                 val expandedAndNotInteractive = isFullyExpanded && !isUserInteracting
 
                 // If we ever are fully expanded and not interacting, capture this state as we
@@ -529,7 +537,7 @@
         val isMove = ev.actionMasked == MotionEvent.ACTION_MOVE
         val isCancel = ev.actionMasked == MotionEvent.ACTION_CANCEL
 
-        val hubOccluded = anyBouncerShowing || shadeShowingAndConsumingTouches
+        val hubOccluded = anyBouncerShowing || shadeConsumingTouches || shadeShowing
 
         if ((isDown || isMove) && !hubOccluded) {
             if (isDown) {
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
index c023b83..31813b2 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationPanelViewController.java
@@ -191,12 +191,10 @@
 import com.android.systemui.statusbar.notification.ViewGroupFadeHelper;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.domain.interactor.ActiveNotificationsInteractor;
-import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor;
 import com.android.systemui.statusbar.notification.footer.shared.FooterViewRefactor;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.ExpandableView;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.stack.AmbientState;
 import com.android.systemui.statusbar.notification.stack.AnimationProperties;
 import com.android.systemui.statusbar.notification.stack.NotificationListContainer;
@@ -439,7 +437,6 @@
     private boolean mExpandingFromHeadsUp;
     private boolean mCollapsedOnDown;
     private boolean mClosingWithAlphaFadeOut;
-    private boolean mHeadsUpVisible;
     private boolean mHeadsUpAnimatingAway;
     private final FalsingManager mFalsingManager;
     private final FalsingCollector mFalsingCollector;
@@ -610,7 +607,6 @@
     private final PrimaryBouncerToGoneTransitionViewModel mPrimaryBouncerToGoneTransitionViewModel;
     private final SharedNotificationContainerInteractor mSharedNotificationContainerInteractor;
     private final ActiveNotificationsInteractor mActiveNotificationsInteractor;
-    private final HeadsUpNotificationInteractor mHeadsUpNotificationInteractor;
     private final KeyguardTransitionInteractor mKeyguardTransitionInteractor;
     private final KeyguardInteractor mKeyguardInteractor;
     private final PowerInteractor mPowerInteractor;
@@ -776,7 +772,6 @@
             ActivityStarter activityStarter,
             SharedNotificationContainerInteractor sharedNotificationContainerInteractor,
             ActiveNotificationsInteractor activeNotificationsInteractor,
-            HeadsUpNotificationInteractor headsUpNotificationInteractor,
             ShadeAnimationInteractor shadeAnimationInteractor,
             KeyguardViewConfigurator keyguardViewConfigurator,
             DeviceEntryFaceAuthInteractor deviceEntryFaceAuthInteractor,
@@ -811,7 +806,6 @@
         mKeyguardTransitionInteractor = keyguardTransitionInteractor;
         mSharedNotificationContainerInteractor = sharedNotificationContainerInteractor;
         mActiveNotificationsInteractor = activeNotificationsInteractor;
-        mHeadsUpNotificationInteractor = headsUpNotificationInteractor;
         mKeyguardInteractor = keyguardInteractor;
         mPowerInteractor = powerInteractor;
         mKeyguardViewConfigurator = keyguardViewConfigurator;
@@ -1222,11 +1216,6 @@
                     }
                 },
                 mMainDispatcher);
-
-        if (NotificationsHeadsUpRefactor.isEnabled()) {
-            collectFlow(mView, mHeadsUpNotificationInteractor.isHeadsUpOrAnimatingAway(),
-                    setHeadsUpVisible(), mMainDispatcher);
-        }
     }
 
     @VisibleForTesting
@@ -3077,21 +3066,7 @@
         mPanelAlphaEndAction = r;
     }
 
-    private Consumer<Boolean> setHeadsUpVisible() {
-        return (Boolean isHeadsUpVisible) -> {
-            mHeadsUpVisible = isHeadsUpVisible;
-
-            if (isHeadsUpVisible) {
-                updateNotificationTranslucency();
-            }
-            updateExpansionAndVisibility();
-            updateGestureExclusionRect();
-            mKeyguardStatusBarViewController.updateForHeadsUp();
-        };
-    }
-
     private void setHeadsUpAnimatingAway(boolean headsUpAnimatingAway) {
-        NotificationsHeadsUpRefactor.assertInLegacyMode();
         mHeadsUpAnimatingAway = headsUpAnimatingAway;
         mNotificationStackScrollLayoutController.setHeadsUpAnimatingAway(headsUpAnimatingAway);
         updateVisibility();
@@ -3107,16 +3082,13 @@
     }
 
     private boolean shouldPanelBeVisible() {
-        boolean headsUpVisible = NotificationsHeadsUpRefactor.isEnabled() ? mHeadsUpVisible
-                : (mHeadsUpAnimatingAway || mHeadsUpPinnedMode);
+        boolean headsUpVisible = mHeadsUpAnimatingAway || mHeadsUpPinnedMode;
         return headsUpVisible || isExpanded() || mBouncerShowing;
     }
 
     private void setHeadsUpManager(HeadsUpManager headsUpManager) {
         mHeadsUpManager = headsUpManager;
-        if (!NotificationsHeadsUpRefactor.isEnabled()) {
-            mHeadsUpManager.addListener(mOnHeadsUpChangedListener);
-        }
+        mHeadsUpManager.addListener(mOnHeadsUpChangedListener);
         mHeadsUpTouchHelper = new HeadsUpTouchHelper(
                 headsUpManager,
                 mStatusBarService,
@@ -3204,8 +3176,7 @@
     }
 
     private boolean isPanelVisibleBecauseOfHeadsUp() {
-        boolean headsUpVisible = NotificationsHeadsUpRefactor.isEnabled() ? mHeadsUpVisible
-                : (mHeadsUpManager.hasPinnedHeadsUp() || mHeadsUpAnimatingAway);
+        boolean headsUpVisible = mHeadsUpManager.hasPinnedHeadsUp() || mHeadsUpAnimatingAway;
         return headsUpVisible && mBarState == StatusBarState.SHADE;
     }
 
@@ -3521,7 +3492,6 @@
         ipw.print("mExpandingFromHeadsUp="); ipw.println(mExpandingFromHeadsUp);
         ipw.print("mCollapsedOnDown="); ipw.println(mCollapsedOnDown);
         ipw.print("mClosingWithAlphaFadeOut="); ipw.println(mClosingWithAlphaFadeOut);
-        ipw.print("mHeadsUpVisible="); ipw.println(mHeadsUpVisible);
         ipw.print("mHeadsUpAnimatingAway="); ipw.println(mHeadsUpAnimatingAway);
         ipw.print("mShowIconsWhenExpanded="); ipw.println(mShowIconsWhenExpanded);
         ipw.print("mIndicationBottomPadding="); ipw.println(mIndicationBottomPadding);
@@ -4446,8 +4416,6 @@
     private final class ShadeHeadsUpChangedListener implements OnHeadsUpChangedListener {
         @Override
         public void onHeadsUpPinnedModeChanged(final boolean inPinnedMode) {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
-
             if (inPinnedMode) {
                 mHeadsUpExistenceChangedRunnable.run();
                 updateNotificationTranslucency();
@@ -4464,8 +4432,6 @@
 
         @Override
         public void onHeadsUpPinned(NotificationEntry entry) {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
-
             if (!isKeyguardShowing()) {
                 mNotificationStackScrollLayoutController.generateHeadsUpAnimation(entry, true);
             }
@@ -4473,8 +4439,6 @@
 
         @Override
         public void onHeadsUpUnPinned(NotificationEntry entry) {
-            NotificationsHeadsUpRefactor.assertInLegacyMode();
-
             // When we're unpinning the notification via active edge they remain heads-upped,
             // we need to make sure that an animation happens in this case, otherwise the
             // notification
diff --git a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
index 7e0454c..3f3ad13 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/NotificationShadeWindowControllerImpl.java
@@ -49,6 +49,7 @@
 import com.android.systemui.Dumpable;
 import com.android.systemui.Flags;
 import com.android.systemui.biometrics.AuthController;
+import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlags;
 import com.android.systemui.colorextraction.SysuiColorExtractor;
 import com.android.systemui.communal.domain.interactor.CommunalInteractor;
 import com.android.systemui.dagger.SysUISingleton;
@@ -342,6 +343,12 @@
                     this::setKeyguardOccluded
             );
         }
+        if (ComposeBouncerFlags.INSTANCE.isComposeBouncerOrSceneContainerEnabled()) {
+            collectFlow(mWindowRootView, mNotificationShadeWindowModel.isBouncerShowing(),
+                    this::setBouncerShowing);
+            collectFlow(mWindowRootView, mNotificationShadeWindowModel.getDoesBouncerRequireIme(),
+                    this::setKeyguardNeedsInput);
+        }
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java b/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
index 16aef65..830649b 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/shade/QuickSettingsControllerImpl.java
@@ -1005,7 +1005,7 @@
         // When expanding QS, let's authenticate the user if possible,
         // this will speed up notification actions.
         if (height == 0 && !mKeyguardStateController.canDismissLockScreen()) {
-            mDeviceEntryFaceAuthInteractor.onQsExpansionStared();
+            mDeviceEntryFaceAuthInteractor.onShadeExpansionStarted();
         }
     }
 
@@ -1063,13 +1063,17 @@
         mScrimController.setQsPosition(qsExpansionFraction, qsPanelBottomY);
         setClippingBounds();
 
-        if (mSplitShadeEnabled) {
-            // In split shade we want to pretend that QS are always collapsed so their behaviour and
-            // interactions don't influence notifications as they do in portrait. But we want to set
-            // 0 explicitly in case we're rotating from non-split shade with QS expansion of 1.
-            mNotificationStackScrollLayoutController.setQsExpansionFraction(0);
-        } else {
-            mNotificationStackScrollLayoutController.setQsExpansionFraction(qsExpansionFraction);
+        if (!SceneContainerFlag.isEnabled()) {
+            if (mSplitShadeEnabled) {
+                // In split shade we want to pretend that QS are always collapsed so their
+                // behaviour and interactions don't influence notifications as they do in portrait.
+                // But we want to set 0 explicitly in case we're rotating from non-split shade with
+                // QS expansion of 1.
+                mNotificationStackScrollLayoutController.setQsExpansionFraction(0);
+            } else {
+                mNotificationStackScrollLayoutController.setQsExpansionFraction(
+                        qsExpansionFraction);
+            }
         }
 
         mDepthController.setQsPanelExpansion(qsExpansionFraction);
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeControllerSceneImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeControllerSceneImpl.kt
index 23e2620..5d03a28 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeControllerSceneImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeControllerSceneImpl.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.shade
 
 import android.view.MotionEvent
-import androidx.compose.ui.Alignment
 import com.android.systemui.assist.AssistManager
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
@@ -25,7 +24,6 @@
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.scene.shared.model.TransitionKeys.OpenBottomShade
 import com.android.systemui.scene.shared.model.TransitionKeys.SlightlyFasterShadeCollapse
 import com.android.systemui.shade.ShadeController.ShadeVisibilityListener
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
@@ -177,7 +175,6 @@
         sceneInteractor.changeScene(
             SceneFamilies.NotifShade,
             "ShadeController.animateExpandShade",
-            OpenBottomShade.takeIf { shadeInteractor.shadeAlignment == Alignment.BottomEnd }
         )
     }
 
@@ -185,7 +182,6 @@
         sceneInteractor.changeScene(
             SceneFamilies.QuickSettings,
             "ShadeController.animateExpandQs",
-            OpenBottomShade.takeIf { shadeInteractor.shadeAlignment == Alignment.BottomEnd }
         )
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
index 606fef0..fc8a593 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ShadeViewProviderModule.kt
@@ -14,6 +14,8 @@
  * limitations under the License.
  */
 
+@file:OptIn(ExperimentalCoroutinesApi::class)
+
 package com.android.systemui.shade
 
 import android.annotation.SuppressLint
@@ -35,9 +37,10 @@
 import com.android.systemui.privacy.OngoingPrivacyChip
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
-import com.android.systemui.scene.shared.model.Scene
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.SceneDataSourceDelegator
+import com.android.systemui.scene.ui.composable.Overlay
+import com.android.systemui.scene.ui.composable.Scene
 import com.android.systemui.scene.ui.view.SceneWindowRootView
 import com.android.systemui.scene.ui.view.WindowRootView
 import com.android.systemui.scene.ui.viewmodel.SceneContainerViewModel
@@ -59,6 +62,7 @@
 import dagger.Provides
 import javax.inject.Named
 import javax.inject.Provider
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 /** Module for providing views related to the shade. */
 @Module
@@ -82,6 +86,7 @@
             viewModelFactory: SceneContainerViewModel.Factory,
             containerConfigProvider: Provider<SceneContainerConfig>,
             scenesProvider: Provider<Set<@JvmSuppressWildcards Scene>>,
+            overlaysProvider: Provider<Set<@JvmSuppressWildcards Overlay>>,
             layoutInsetController: NotificationInsetsController,
             sceneDataSourceDelegator: Provider<SceneDataSourceDelegator>,
             alternateBouncerDependencies: Provider<AlternateBouncerDependencies>,
@@ -96,6 +101,7 @@
                     sharedNotificationContainer =
                         sceneWindowRootView.requireViewById(R.id.shared_notification_container),
                     scenes = scenesProvider.get(),
+                    overlays = overlaysProvider.get(),
                     layoutInsetController = layoutInsetController,
                     sceneDataSourceDelegator = sceneDataSourceDelegator.get(),
                     alternateBouncerDependencies = alternateBouncerDependencies.get(),
diff --git a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
index a4fed873..193056c 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/data/repository/ShadeRepository.kt
@@ -18,7 +18,6 @@
 import android.content.Context
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.res.R
 import javax.inject.Inject
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -103,9 +102,6 @@
     @Deprecated("Use ShadeInteractor.isQsBypassingShade instead")
     val legacyExpandImmediate: StateFlow<Boolean>
 
-    /** Whether dual shade should be aligned to the bottom (true) or to the top (false). */
-    val isDualShadeAlignedToBottom: Boolean
-
     /**
      * Whether the shade layout should be wide (true) or narrow (false).
      *
@@ -238,9 +234,6 @@
     private val _isShadeLayoutWide = MutableStateFlow(false)
     override val isShadeLayoutWide: StateFlow<Boolean> = _isShadeLayoutWide.asStateFlow()
 
-    override val isDualShadeAlignedToBottom =
-        applicationContext.resources.getBoolean(R.bool.config_dualShadeAlignedToBottom)
-
     override fun setShadeLayoutWide(isShadeLayoutWide: Boolean) {
         _isShadeLayoutWide.value = isShadeLayoutWide
     }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
index 7d67121..e276f88 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/PanelExpansionInteractorImpl.kt
@@ -64,7 +64,7 @@
                             0f
                         }
                     )
-                is ObservableTransitionState.Transition.ChangeCurrentScene ->
+                is ObservableTransitionState.Transition.ChangeScene ->
                     when {
                         state.fromScene == Scenes.Gone ->
                             if (state.toScene.isExpandable()) {
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeAnimationInteractorSceneContainerImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeAnimationInteractorSceneContainerImpl.kt
index 79a94a5..8467185 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeAnimationInteractorSceneContainerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeAnimationInteractorSceneContainerImpl.kt
@@ -49,10 +49,10 @@
                     is ObservableTransitionState.Idle -> flowOf(false)
                     is ObservableTransitionState.Transition ->
                         if (
-                            (state.fromScene == Scenes.Shade &&
-                                state.toScene != Scenes.QuickSettings) ||
-                                (state.fromScene == Scenes.QuickSettings &&
-                                    state.toScene != Scenes.Shade)
+                            (state.fromContent == Scenes.Shade &&
+                                state.toContent != Scenes.QuickSettings) ||
+                                (state.fromContent == Scenes.QuickSettings &&
+                                    state.toContent != Scenes.Shade)
                         ) {
                             state.isUserInputOngoing.map { !it }
                         } else {
@@ -71,10 +71,10 @@
                     is ObservableTransitionState.Transition ->
                         if (
                             state.isInitiatedByUserInput &&
-                                (state.fromScene == Scenes.Shade ||
-                                    state.toScene == Scenes.Shade ||
-                                    state.fromScene == Scenes.QuickSettings ||
-                                    state.toScene == Scenes.QuickSettings)
+                                (state.fromContent == Scenes.Shade ||
+                                    state.toContent == Scenes.Shade ||
+                                    state.fromContent == Scenes.QuickSettings ||
+                                    state.toContent == Scenes.QuickSettings)
                         ) {
                             state.isUserInputOngoing.map { !it }
                         } else {
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
index 45f359e..73e86a2 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractor.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.shade.domain.interactor
 
-import com.android.systemui.shade.shared.model.ShadeAlignment
 import com.android.systemui.shade.shared.model.ShadeMode
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
@@ -70,9 +69,6 @@
      * wide as the entire screen.
      */
     val isShadeLayoutWide: StateFlow<Boolean>
-
-    /** How to align the shade content. */
-    val shadeAlignment: ShadeAlignment
 }
 
 /** ShadeInteractor methods with implementations that differ between non-empty impls. */
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorEmptyImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorEmptyImpl.kt
index e77aca9..d51fd28 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorEmptyImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorEmptyImpl.kt
@@ -17,7 +17,6 @@
 package com.android.systemui.shade.domain.interactor
 
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.shade.shared.model.ShadeAlignment
 import com.android.systemui.shade.shared.model.ShadeMode
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
@@ -48,5 +47,4 @@
     override val isExpandToQsEnabled: Flow<Boolean> = inactiveFlowBoolean
     override val shadeMode: StateFlow<ShadeMode> = MutableStateFlow(ShadeMode.Single)
     override val isShadeLayoutWide: StateFlow<Boolean> = inactiveFlowBoolean
-    override val shadeAlignment: ShadeAlignment = ShadeAlignment.Top
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
index d64b21f..3552092 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorImpl.kt
@@ -26,7 +26,6 @@
 import com.android.systemui.power.domain.interactor.PowerInteractor
 import com.android.systemui.shade.data.repository.ShadeRepository
 import com.android.systemui.shade.shared.flag.DualShade
-import com.android.systemui.shade.shared.model.ShadeAlignment
 import com.android.systemui.shade.shared.model.ShadeMode
 import com.android.systemui.statusbar.disableflags.data.repository.DisableFlagsRepository
 import com.android.systemui.statusbar.phone.DozeParameters
@@ -114,15 +113,6 @@
                 initialValue = determineShadeMode(isShadeLayoutWide.value)
             )
 
-    override val shadeAlignment: ShadeAlignment
-        get() {
-            return if (shadeRepository.isDualShadeAlignedToBottom) {
-                ShadeAlignment.Bottom
-            } else {
-                ShadeAlignment.Top
-            }
-        }
-
     override val isExpandToQsEnabled: Flow<Boolean> =
         combine(
             disableFlagsRepository.disableFlags,
diff --git a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImpl.kt b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImpl.kt
index 6a21531..e84cfa5 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImpl.kt
@@ -90,8 +90,8 @@
                         when (state) {
                             is ObservableTransitionState.Idle -> false
                             is ObservableTransitionState.Transition ->
-                                state.toScene == quickSettingsScene &&
-                                    state.fromScene != notificationsScene
+                                state.toContent == quickSettingsScene &&
+                                    state.fromContent != notificationsScene
                         }
                     }
                     .distinctUntilChanged()
@@ -99,14 +99,17 @@
             .distinctUntilChanged()
 
     override val isQsFullscreen: Flow<Boolean> =
-        sceneInteractor
-            .resolveSceneFamily(SceneFamilies.QuickSettings)
-            .flatMapLatestConflated { quickSettingsScene ->
+        combine(
+                shadeRepository.isShadeLayoutWide,
+                sceneInteractor.resolveSceneFamily(SceneFamilies.QuickSettings),
+                ::Pair
+            )
+            .flatMapLatestConflated { (isShadeLayoutWide, quickSettingsScene) ->
                 sceneInteractor.transitionState
                     .map { state ->
                         when (state) {
                             is ObservableTransitionState.Idle ->
-                                state.currentScene == quickSettingsScene
+                                !isShadeLayoutWide && state.currentScene == quickSettingsScene
                             is ObservableTransitionState.Transition -> false
                         }
                     }
@@ -147,9 +150,9 @@
                                     flowOf(0f)
                                 }
                             is ObservableTransitionState.Transition ->
-                                if (state.toScene == resolvedSceneKey) {
+                                if (state.toContent == resolvedSceneKey) {
                                     state.progress
-                                } else if (state.fromScene == resolvedSceneKey) {
+                                } else if (state.fromContent == resolvedSceneKey) {
                                     state.progress.map { progress -> 1 - progress }
                                 } else {
                                     flowOf(0f)
@@ -172,8 +175,8 @@
                     is ObservableTransitionState.Transition ->
                         sceneInteractor.resolveSceneFamily(sceneKey).map { resolvedSceneKey ->
                             state.isInitiatedByUserInput &&
-                                (state.toScene == resolvedSceneKey ||
-                                    state.fromScene == resolvedSceneKey)
+                                (state.toContent == resolvedSceneKey ||
+                                    state.fromContent == resolvedSceneKey)
                         }
                 }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/shared/model/ShadeAlignment.kt b/packages/SystemUI/src/com/android/systemui/shade/shared/model/ShadeAlignment.kt
deleted file mode 100644
index 06905379..0000000
--- a/packages/SystemUI/src/com/android/systemui/shade/shared/model/ShadeAlignment.kt
+++ /dev/null
@@ -1,27 +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.shade.shared.model
-
-/** Enumerates all supported alignments of the shade. */
-sealed interface ShadeAlignment {
-
-    /** Aligns the shade to the top. */
-    data object Top : ShadeAlignment
-
-    /** Aligns the shade to the bottom. */
-    data object Bottom : ShadeAlignment
-}
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModel.kt
index f270e82..9655d92 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/NotificationShadeWindowModel.kt
@@ -16,16 +16,26 @@
 
 package com.android.systemui.shade.ui.viewmodel
 
+import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
+import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
+import com.android.systemui.bouncer.domain.interactor.PrimaryBouncerInteractor
+import com.android.systemui.bouncer.shared.flag.ComposeBouncerFlags
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.Edge
 import com.android.systemui.keyguard.shared.model.KeyguardState.DREAMING
 import com.android.systemui.keyguard.shared.model.KeyguardState.GLANCEABLE_HUB
 import com.android.systemui.keyguard.shared.model.KeyguardState.OCCLUDED
+import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.util.kotlin.BooleanFlowOperators.any
+import com.android.systemui.util.kotlin.sample
 import javax.inject.Inject
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.map
 
 /** Models UI state for the shade window. */
 @SysUISingleton
@@ -33,6 +43,9 @@
 @Inject
 constructor(
     keyguardTransitionInteractor: KeyguardTransitionInteractor,
+    sceneInteractor: dagger.Lazy<SceneInteractor>,
+    authenticationInteractor: dagger.Lazy<AuthenticationInteractor>,
+    primaryBouncerInteractor: PrimaryBouncerInteractor,
 ) {
     /**
      * Considered to be occluded if in OCCLUDED, DREAMING, GLANCEABLE_HUB/Communal, or transitioning
@@ -42,9 +55,11 @@
     val isKeyguardOccluded: Flow<Boolean> =
         listOf(
                 // Finished in state...
-                keyguardTransitionInteractor.isFinishedIn(OCCLUDED),
-                keyguardTransitionInteractor.isFinishedIn(DREAMING),
-                keyguardTransitionInteractor.isFinishedIn(Scenes.Communal, GLANCEABLE_HUB),
+                keyguardTransitionInteractor.transitionValue(OCCLUDED).map { it == 1f },
+                keyguardTransitionInteractor.transitionValue(DREAMING).map { it == 1f },
+                keyguardTransitionInteractor.transitionValue(Scenes.Communal, GLANCEABLE_HUB).map {
+                    it == 1f
+                },
 
                 // ... or transitions between those states
                 keyguardTransitionInteractor.isInTransition(Edge.create(OCCLUDED, DREAMING)),
@@ -67,4 +82,53 @@
                 ),
             )
             .any()
+
+    /**
+     * Whether bouncer is currently showing or not.
+     *
+     * Applicable only when either [SceneContainerFlag] or [ComposeBouncerFlags] are enabled,
+     * otherwise it throws an error.
+     */
+    val isBouncerShowing: Flow<Boolean> =
+        when {
+            SceneContainerFlag.isEnabled -> {
+                sceneInteractor.get().transitionState.map { it.isIdle(Scenes.Bouncer) }
+            }
+            ComposeBouncerFlags.isOnlyComposeBouncerEnabled() -> primaryBouncerInteractor.isShowing
+            else ->
+                flow {
+                    error(
+                        "Consume this flow only when SceneContainerFlag " +
+                            "or ComposeBouncerFlags are enabled"
+                    )
+                }
+        }.distinctUntilChanged()
+
+    /**
+     * Whether the bouncer currently require IME for device entry.
+     *
+     * This emits true when the authentication method is set to password and the bouncer is
+     * currently showing. Throws an error when this is used without either [SceneContainerFlag] or
+     * [ComposeBouncerFlags]
+     */
+    val doesBouncerRequireIme: Flow<Boolean> =
+        if (ComposeBouncerFlags.isComposeBouncerOrSceneContainerEnabled()) {
+                // This is required to make the window, where the bouncer resides,
+                // focusable. InputMethodManager allows IME to be shown only for views
+                // in windows that do not have the FLAG_NOT_FOCUSABLE flag.
+
+                isBouncerShowing
+                    .sample(authenticationInteractor.get().authenticationMethod, ::Pair)
+                    .map { (showing, authMethod) ->
+                        showing && authMethod == AuthenticationMethodModel.Password
+                    }
+            } else {
+                flow {
+                    error(
+                        "Consume this flow only when SceneContainerFlag " +
+                            "or ComposeBouncerFlags are enabled"
+                    )
+                }
+            }
+            .distinctUntilChanged()
 }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt
deleted file mode 100644
index 25ae44e..0000000
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModel.kt
+++ /dev/null
@@ -1,68 +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.shade.ui.viewmodel
-
-import com.android.compose.animation.scene.SceneKey
-import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
-import com.android.systemui.scene.domain.interactor.SceneInteractor
-import com.android.systemui.scene.shared.model.SceneFamilies
-import com.android.systemui.scene.shared.model.Scenes
-import com.android.systemui.shade.domain.interactor.ShadeInteractor
-import dagger.assisted.AssistedFactory
-import dagger.assisted.AssistedInject
-import kotlinx.coroutines.awaitCancellation
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
-
-/**
- * Models UI state and handles user input for the overlay shade UI, which shows a shade as an
- * overlay on top of another scene UI.
- */
-class OverlayShadeViewModel
-@AssistedInject
-constructor(private val sceneInteractor: SceneInteractor, shadeInteractor: ShadeInteractor) :
-    SysUiViewModel, ExclusiveActivatable() {
-    private val _backgroundScene = MutableStateFlow(Scenes.Lockscreen)
-    /** The scene to show in the background when the overlay shade is open. */
-    val backgroundScene: StateFlow<SceneKey> = _backgroundScene.asStateFlow()
-
-    /** Dictates the alignment of the overlay shade panel on the screen. */
-    val panelAlignment = shadeInteractor.shadeAlignment
-
-    override suspend fun onActivated(): Nothing {
-        sceneInteractor.resolveSceneFamily(SceneFamilies.Home).collectLatest { sceneKey ->
-            _backgroundScene.value = sceneKey
-        }
-        awaitCancellation()
-    }
-
-    /** Notifies that the user has clicked the semi-transparent background scrim. */
-    fun onScrimClicked() {
-        sceneInteractor.changeScene(
-            toScene = SceneFamilies.Home,
-            loggingReason = "Shade scrim clicked",
-        )
-    }
-
-    @AssistedFactory
-    interface Factory {
-        fun create(): OverlayShadeViewModel
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeHeaderViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeHeaderViewModel.kt
index edfe79a..a154e91 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeHeaderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeHeaderViewModel.kt
@@ -25,7 +25,6 @@
 import android.provider.Settings
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.privacy.OngoingPrivacyChip
 import com.android.systemui.privacy.PrivacyItem
@@ -47,7 +46,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
@@ -66,7 +64,7 @@
     private val privacyChipInteractor: PrivacyChipInteractor,
     private val clockInteractor: ShadeHeaderClockInteractor,
     private val broadcastDispatcher: BroadcastDispatcher,
-) : SysUiViewModel, ExclusiveActivatable() {
+) : ExclusiveActivatable() {
     /** True if there is exactly one mobile connection. */
     val isSingleCarrier: StateFlow<Boolean> = mobileIconsInteractor.isSingleCarrier
 
@@ -133,12 +131,10 @@
             launch {
                 mobileIconsInteractor.filteredSubscriptions
                     .map { list -> list.map { it.subscriptionId } }
-                    .collectLatest { _mobileSubIds.value = it }
+                    .collect { _mobileSubIds.value = it }
             }
 
-            launch {
-                shadeInteractor.isQsEnabled.map { !it }.collectLatest { _isDisabled.value = it }
-            }
+            launch { shadeInteractor.isQsEnabled.map { !it }.collect { _isDisabled.value = it } }
 
             awaitCancellation()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneContentViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneContentViewModel.kt
index f0f2a65..ce4c081 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneContentViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneContentViewModel.kt
@@ -21,7 +21,6 @@
 import androidx.lifecycle.LifecycleOwner
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryInteractor
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.qs.FooterActionsController
 import com.android.systemui.qs.footer.ui.viewmodel.FooterActionsViewModel
@@ -36,17 +35,15 @@
 import dagger.assisted.AssistedInject
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.awaitCancellation
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.collectLatest
 
 /**
  * Models UI state used to render the content of the shade scene.
  *
- * Different from [ShadeSceneActionsViewModel], which only models user actions that can be performed
+ * Different from [ShadeUserActionsViewModel], which only models user actions that can be performed
  * to navigate to other scenes.
  */
 class ShadeSceneContentViewModel
@@ -62,7 +59,7 @@
     private val unfoldTransitionInteractor: UnfoldTransitionInteractor,
     private val deviceEntryInteractor: DeviceEntryInteractor,
     private val sceneInteractor: SceneInteractor,
-) : SysUiViewModel, ExclusiveActivatable() {
+) : ExclusiveActivatable() {
 
     val shadeMode: StateFlow<ShadeMode> = shadeInteractor.shadeMode
 
@@ -76,11 +73,9 @@
     private val footerActionsControllerInitialized = AtomicBoolean(false)
 
     override suspend fun onActivated(): Nothing {
-        deviceEntryInteractor.isDeviceEntered.collectLatest { isDeviceEntered ->
+        deviceEntryInteractor.isDeviceEntered.collect { isDeviceEntered ->
             _isEmptySpaceClickable.value = !isDeviceEntered
         }
-
-        awaitCancellation()
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModel.kt b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModel.kt
similarity index 90%
rename from packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModel.kt
rename to packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModel.kt
index bdc0fdb..f8a850a 100644
--- a/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModel.kt
@@ -24,12 +24,11 @@
 import com.android.systemui.scene.shared.model.SceneFamilies
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.scene.shared.model.TransitionKeys.ToSplitShade
-import com.android.systemui.scene.ui.viewmodel.SceneActionsViewModel
+import com.android.systemui.scene.ui.viewmodel.UserActionsViewModel
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.shade.shared.model.ShadeMode
 import dagger.assisted.AssistedFactory
 import dagger.assisted.AssistedInject
-import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.combine
 
 /**
@@ -37,12 +36,12 @@
  *
  * Different from the [ShadeSceneContentViewModel] which models the _content_ of the scene.
  */
-class ShadeSceneActionsViewModel
+class ShadeUserActionsViewModel
 @AssistedInject
 constructor(
     private val qsSceneAdapter: QSSceneAdapter,
     private val shadeInteractor: ShadeInteractor,
-) : SceneActionsViewModel() {
+) : UserActionsViewModel() {
 
     override suspend fun hydrateActions(setActions: (Map<UserAction, UserActionResult>) -> Unit) {
         combine(
@@ -67,11 +66,11 @@
                     }
                 }
             }
-            .collectLatest { actions -> setActions(actions) }
+            .collect { actions -> setActions(actions) }
     }
 
     @AssistedFactory
     interface Factory {
-        fun create(): ShadeSceneActionsViewModel
+        fun create(): ShadeUserActionsViewModel
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/ImmersiveModeConfirmation.java b/packages/SystemUI/src/com/android/systemui/statusbar/ImmersiveModeConfirmation.java
index 2b9daef..5ef5a7d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/ImmersiveModeConfirmation.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/ImmersiveModeConfirmation.java
@@ -79,6 +79,7 @@
 import com.android.app.viewcapture.ViewCapture;
 import com.android.app.viewcapture.ViewCaptureAwareWindowManager;
 import com.android.systemui.CoreStartable;
+import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.res.R;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
@@ -107,6 +108,7 @@
     private Context mDisplayContext;
     private final Context mSysUiContext;
     private final Handler mHandler = new H(Looper.getMainLooper());
+    private final Handler mBackgroundHandler;
     private long mShowDelayMs = 0L;
     private final IBinder mWindowToken = new Binder();
     private final CommandQueue mCommandQueue;
@@ -139,7 +141,8 @@
     @Inject
     public ImmersiveModeConfirmation(Context context, CommandQueue commandQueue,
                                      SecureSettings secureSettings,
-                                     dagger.Lazy<ViewCapture> daggerLazyViewCapture) {
+                                     dagger.Lazy<ViewCapture> daggerLazyViewCapture,
+                                     @Background Handler backgroundHandler) {
         mSysUiContext = context;
         final Display display = mSysUiContext.getDisplay();
         mDisplayContext = display.getDisplayId() == DEFAULT_DISPLAY
@@ -147,6 +150,7 @@
         mCommandQueue = commandQueue;
         mSecureSettings = secureSettings;
         mLazyViewCapture = toKotlinLazy(daggerLazyViewCapture);
+        mBackgroundHandler = backgroundHandler;
     }
 
     boolean loadSetting(int currentUserId) {
@@ -329,7 +333,7 @@
                 }
             }
             TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
-            mContentObserver = new ContentObserver(mHandler) {
+            mContentObserver = new ContentObserver(mBackgroundHandler) {
                 @Override
                 public void onChange(boolean selfChange) {
                     onSettingChanged(mSysUiContext.getUserId());
@@ -343,6 +347,9 @@
             mSecureSettings.registerContentObserverForUserSync(
                     Settings.Secure.USER_SETUP_COMPLETE, mContentObserver,
                     UserHandle.USER_CURRENT);
+            mBackgroundHandler.post(() -> {
+                loadSetting(UserHandle.USER_CURRENT);
+            });
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
index 696e222..d523bc1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
@@ -255,8 +255,8 @@
         }
 
         final float stackBottom = SceneContainerFlag.isEnabled()
-                ? ambientState.getStackTop() + ambientState.getStackHeight()
-                : ambientState.getStackY() + ambientState.getStackHeight();
+                ? ambientState.getStackTop() + ambientState.getInterpolatedStackHeight()
+                : ambientState.getStackY() + ambientState.getInterpolatedStackHeight();
 
         if (viewState.hidden) {
             // if the shelf is hidden, position it at the end of the stack (plus the clip
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
index 6eadd26..87f360e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
@@ -58,6 +58,7 @@
 import com.android.app.animation.Interpolators;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.internal.util.ContrastColorUtil;
 import com.android.systemui.Flags;
 import com.android.systemui.res.R;
@@ -211,16 +212,19 @@
     /** Should always be preceded by {@link #reloadDimens()} */
     @VisibleForTesting
     public void maybeUpdateIconScaleDimens() {
-        // We do not resize and scale system icons (on the right), only notification icons (on the
-        // left).
-        if (isNotification()) {
-            updateIconScaleForNotifications();
+        // We scale notification icons (on the left) plus icons on the right that explicitly
+        // want FIXED_SPACE.
+        boolean useNonSystemIconScaling = isNotification()
+                || (usesModeIcons() && mIcon != null && mIcon.shape == Shape.FIXED_SPACE);
+
+        if (useNonSystemIconScaling) {
+            updateIconScaleForNonSystemIcons();
         } else {
             updateIconScaleForSystemIcons();
         }
     }
 
-    private void updateIconScaleForNotifications() {
+    private void updateIconScaleForNonSystemIcons() {
         float iconScale;
         // we need to scale the image size to be same as the original size
         // (fit mOriginalStatusBarIconSize), then we can scale it with mScaleToFitNewIconSize
@@ -411,7 +415,9 @@
         if (!levelEquals) {
             setImageLevel(icon.iconLevel);
         }
-
+        if (usesModeIcons() && icon.shape == Shape.FIXED_SPACE) {
+            setScaleType(ScaleType.FIT_CENTER);
+        }
         if (!visibilityEquals) {
             setVisibility(icon.visible && !mBlocked ? VISIBLE : GONE);
         }
@@ -501,7 +507,12 @@
     @Nullable
     private Drawable loadDrawable(Context context, StatusBarIcon statusBarIcon) {
         if (usesModeIcons() && statusBarIcon.preloadedIcon != null) {
-            return statusBarIcon.preloadedIcon.mutate();
+            Drawable.ConstantState cached = statusBarIcon.preloadedIcon.getConstantState();
+            if (cached != null) {
+                return cached.newDrawable(mContext.getResources()).mutate();
+            } else {
+                return statusBarIcon.preloadedIcon.mutate();
+            }
         } else {
             int userId = statusBarIcon.user.getIdentifier();
             if (userId == UserHandle.USER_ALL) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/StatusBarChipsModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/StatusBarChipsModule.kt
index 173ff37..be733d4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/StatusBarChipsModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/StatusBarChipsModule.kt
@@ -16,14 +16,24 @@
 
 package com.android.systemui.statusbar.chips
 
+import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.log.LogBuffer
 import com.android.systemui.log.LogBufferFactory
+import com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel.DemoRonChipViewModel
+import dagger.Binds
 import dagger.Module
 import dagger.Provides
+import dagger.multibindings.ClassKey
+import dagger.multibindings.IntoMap
 
 @Module
 abstract class StatusBarChipsModule {
+    @Binds
+    @IntoMap
+    @ClassKey(DemoRonChipViewModel::class)
+    abstract fun binds(impl: DemoRonChipViewModel): CoreStartable
+
     companion object {
         @Provides
         @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModel.kt
index 18ea0b4..e825258 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModel.kt
@@ -69,7 +69,7 @@
                                     state.notificationIconView
                                 )
                             } else {
-                                OngoingActivityChipModel.ChipIcon.Basic(phoneIcon)
+                                OngoingActivityChipModel.ChipIcon.SingleColorIcon(phoneIcon)
                             }
 
                         // This block mimics OngoingCallController#updateChip.
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
index cf4e707..d4ad6ee 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModel.kt
@@ -190,7 +190,7 @@
     ): OngoingActivityChipModel.Shown {
         return OngoingActivityChipModel.Shown.Timer(
             icon =
-                OngoingActivityChipModel.ChipIcon.Basic(
+                OngoingActivityChipModel.ChipIcon.SingleColorIcon(
                     Icon.Resource(
                         CAST_TO_OTHER_DEVICE_ICON,
                         // This string is "Casting screen"
@@ -215,7 +215,7 @@
     private fun createIconOnlyCastChip(deviceName: String?): OngoingActivityChipModel.Shown {
         return OngoingActivityChipModel.Shown.IconOnly(
             icon =
-                OngoingActivityChipModel.ChipIcon.Basic(
+                OngoingActivityChipModel.ChipIcon.SingleColorIcon(
                     Icon.Resource(
                         CAST_TO_OTHER_DEVICE_ICON,
                         // This string is just "Casting"
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModel.kt
new file mode 100644
index 0000000..cce9a16
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModel.kt
@@ -0,0 +1,174 @@
+/*
+ * 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.chips.ron.demo.ui.viewmodel
+
+import android.content.pm.PackageManager
+import android.content.pm.PackageManager.NameNotFoundException
+import android.graphics.drawable.Drawable
+import com.android.systemui.CoreStartable
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.chips.ron.shared.StatusBarRonChips
+import com.android.systemui.statusbar.chips.ui.model.ColorsModel
+import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
+import com.android.systemui.statusbar.chips.ui.viewmodel.OngoingActivityChipViewModel
+import com.android.systemui.statusbar.commandline.CommandRegistry
+import com.android.systemui.statusbar.commandline.ParseableCommand
+import com.android.systemui.statusbar.commandline.Type
+import com.android.systemui.util.time.SystemClock
+import java.io.PrintWriter
+import javax.inject.Inject
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+/**
+ * A view model that will emit demo RON chips (rich ongoing notification chips) from [chip] based on
+ * adb commands sent by the user.
+ *
+ * Example adb commands:
+ *
+ * To show a chip with the SysUI icon and custom text and color:
+ * ```
+ * adb shell cmd statusbar demo-ron -p com.android.systemui -t 10min -c "\\#434343"
+ * ```
+ *
+ * To hide the chip:
+ * ```
+ * adb shell cmd statusbar demo-ron --hide
+ * ```
+ *
+ * See [DemoRonCommand] for more information on the adb command spec.
+ */
+@SysUISingleton
+class DemoRonChipViewModel
+@Inject
+constructor(
+    private val commandRegistry: CommandRegistry,
+    private val packageManager: PackageManager,
+    private val systemClock: SystemClock,
+) : OngoingActivityChipViewModel, CoreStartable {
+    override fun start() {
+        commandRegistry.registerCommand("demo-ron") { DemoRonCommand() }
+    }
+
+    private val _chip =
+        MutableStateFlow<OngoingActivityChipModel>(OngoingActivityChipModel.Hidden())
+    override val chip: StateFlow<OngoingActivityChipModel> = _chip.asStateFlow()
+
+    private inner class DemoRonCommand : ParseableCommand("demo-ron") {
+        private val packageName: String? by
+            param(
+                longName = "packageName",
+                shortName = "p",
+                description = "The package name for the demo RON app",
+                valueParser = Type.String,
+            )
+
+        private val text: String? by
+            param(
+                longName = "text",
+                shortName = "t",
+                description = "Text to display in the chip",
+                valueParser = Type.String,
+            )
+
+        private val backgroundColor: Int? by
+            param(
+                longName = "color",
+                shortName = "c",
+                description =
+                    "The color to show as the chip background color. " +
+                        "You can either just write a basic color like 'red' or 'green', " +
+                        "or you can include a #RRGGBB string in this format: \"\\\\#434343\".",
+                valueParser = Type.Color,
+            )
+
+        private val hide by
+            flag(
+                longName = "hide",
+                description = "Hides any existing demo RON chip",
+            )
+
+        override fun execute(pw: PrintWriter) {
+            if (!StatusBarRonChips.isEnabled) {
+                pw.println(
+                    "Error: com.android.systemui.status_bar_ron_chips must be enabled " +
+                        "before using this demo feature"
+                )
+                return
+            }
+
+            if (hide) {
+                _chip.value = OngoingActivityChipModel.Hidden()
+                return
+            }
+
+            val currentPackageName = packageName
+            if (currentPackageName == null) {
+                pw.println("--packageName (or -p) must be included")
+                return
+            }
+
+            val appIcon = getAppIcon(currentPackageName)
+            if (appIcon == null) {
+                pw.println("Package $currentPackageName could not be found")
+                return
+            }
+
+            val colors =
+                if (backgroundColor != null) {
+                    ColorsModel.Custom(backgroundColorInt = backgroundColor!!)
+                } else {
+                    ColorsModel.Themed
+                }
+
+            val currentText = text
+            if (currentText != null) {
+                _chip.value =
+                    OngoingActivityChipModel.Shown.Text(
+                        icon = appIcon,
+                        colors = colors,
+                        text = currentText,
+                    )
+            } else {
+                _chip.value =
+                    OngoingActivityChipModel.Shown.Timer(
+                        icon = appIcon,
+                        colors = colors,
+                        startTimeMs = systemClock.elapsedRealtime(),
+                        onClickListener = null,
+                    )
+            }
+        }
+
+        private fun getAppIcon(packageName: String): OngoingActivityChipModel.ChipIcon? {
+            lateinit var iconDrawable: Drawable
+            try {
+                // Note: For the real implementation, we should check if applicationInfo exists
+                // before fetching the icon, so that we either don't show the chip or show a good
+                // backup icon in case the app info can't be found for some reason.
+                iconDrawable = packageManager.getApplicationIcon(packageName)
+            } catch (e: NameNotFoundException) {
+                return null
+            }
+            return OngoingActivityChipModel.ChipIcon.FullColorAppIcon(
+                Icon.Loaded(drawable = iconDrawable, contentDescription = null),
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/shared/StatusBarRonChips.kt
similarity index 72%
copy from packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt
copy to packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/shared/StatusBarRonChips.kt
index 62641fe..4ef1909 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ron/shared/StatusBarRonChips.kt
@@ -14,17 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.notification.shared
+package com.android.systemui.statusbar.chips.ron.shared
 
 import com.android.systemui.Flags
 import com.android.systemui.flags.FlagToken
 import com.android.systemui.flags.RefactorFlagUtils
 
-/** Helper for reading or using the notifications heads up refactor flag state. */
+/** Helper for reading or using the status bar RON chips flag state. */
 @Suppress("NOTHING_TO_INLINE")
-object NotificationsHeadsUpRefactor {
+object StatusBarRonChips {
     /** The aconfig flag name */
-    const val FLAG_NAME = Flags.FLAG_NOTIFICATIONS_HEADS_UP_REFACTOR
+    const val FLAG_NAME = Flags.FLAG_STATUS_BAR_RON_CHIPS
 
     /** A token used for dependency declaration */
     val token: FlagToken
@@ -33,7 +33,7 @@
     /** Is the refactor enabled */
     @JvmStatic
     inline val isEnabled
-        get() = Flags.notificationsHeadsUpRefactor()
+        get() = Flags.statusBarRonChips()
 
     /**
      * Called to ensure code is only run when the flag is enabled. This protects users from the
@@ -46,6 +46,14 @@
 
     /**
      * Called to ensure code is only run when the flag is disabled. This will throw an exception if
+     * the flag is not enabled to ensure that the refactor author catches issues in testing.
+     * Caution!! Using this check incorrectly will cause crashes in nextfood builds!
+     */
+    @JvmStatic
+    inline fun assertInNewMode() = RefactorFlagUtils.assertInNewMode(isEnabled, FLAG_NAME)
+
+    /**
+     * Called to ensure code is only run when the flag is disabled. This will throw an exception if
      * the flag is enabled to ensure that the refactor author catches issues in testing.
      */
     @JvmStatic
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
index 9e6cacb..eb73521 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModel.kt
@@ -80,7 +80,7 @@
                     is ScreenRecordChipModel.Recording -> {
                         OngoingActivityChipModel.Shown.Timer(
                             icon =
-                                OngoingActivityChipModel.ChipIcon.Basic(
+                                OngoingActivityChipModel.ChipIcon.SingleColorIcon(
                                     Icon.Resource(
                                         ICON,
                                         ContentDescription.Resource(
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
index 7897f93..d99a916 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModel.kt
@@ -110,7 +110,7 @@
     ): OngoingActivityChipModel.Shown {
         return OngoingActivityChipModel.Shown.Timer(
             icon =
-                OngoingActivityChipModel.ChipIcon.Basic(
+                OngoingActivityChipModel.ChipIcon.SingleColorIcon(
                     Icon.Resource(
                         SHARE_TO_APP_ICON,
                         ContentDescription.Resource(R.string.share_to_app_chip_accessibility_label),
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/ColorsModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/ColorsModel.kt
index 8a5165d8..4b0fc5a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/ColorsModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/ColorsModel.kt
@@ -39,6 +39,24 @@
             Utils.getColorAttrDefaultColor(context, com.android.internal.R.attr.colorPrimary)
     }
 
+    /**
+     * The chip should have the given background color, and text color that matches dark/light
+     * theme.
+     */
+    data class Custom(val backgroundColorInt: Int) : ColorsModel {
+        override fun background(context: Context): ColorStateList =
+            ColorStateList.valueOf(backgroundColorInt)
+
+        // TODO(b/361346412): When dark theme changes, the chip should automatically re-render with
+        // the right text color. Right now, it has the right text color when the chip is first
+        // created but the color doesn't update if dark theme changes.
+        override fun text(context: Context) =
+            Utils.getColorAttrDefaultColor(
+                context,
+                com.android.internal.R.attr.materialColorOnSurface,
+            )
+    }
+
     /** The chip should have a red background with white text. */
     data object Red : ColorsModel {
         override fun background(context: Context): ColorStateList {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/OngoingActivityChipModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/OngoingActivityChipModel.kt
index 26a2f91..62622a8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/OngoingActivityChipModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/model/OngoingActivityChipModel.kt
@@ -89,6 +89,16 @@
         ) : Shown(icon = null, colors, onClickListener = null) {
             override val logName = "Shown.Countdown"
         }
+
+        /** This chip shows the specified [text] in the chip. */
+        data class Text(
+            override val icon: ChipIcon,
+            override val colors: ColorsModel,
+            // TODO(b/361346412): Enforce a max length requirement?
+            val text: String,
+        ) : Shown(icon, colors, onClickListener = null) {
+            override val logName = "Shown.Text"
+        }
     }
 
     /** Represents an icon to show on the chip. */
@@ -106,7 +116,13 @@
             }
         }
 
-        /** The icon is a basic resource or drawable icon that System UI created internally. */
-        data class Basic(val impl: Icon) : ChipIcon
+        /**
+         * This icon is a single color and it came from basic resource or drawable icon that System
+         * UI created internally.
+         */
+        data class SingleColorIcon(val impl: Icon) : ChipIcon
+
+        /** This icon is an app icon in full color (so it should not get tinted in any way). */
+        data class FullColorAppIcon(val impl: Icon) : ChipIcon
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModel.kt
index b0d897d..04c4516 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModel.kt
@@ -23,6 +23,8 @@
 import com.android.systemui.statusbar.chips.StatusBarChipsLog
 import com.android.systemui.statusbar.chips.call.ui.viewmodel.CallChipViewModel
 import com.android.systemui.statusbar.chips.casttootherdevice.ui.viewmodel.CastToOtherDeviceChipViewModel
+import com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel.DemoRonChipViewModel
+import com.android.systemui.statusbar.chips.ron.shared.StatusBarRonChips
 import com.android.systemui.statusbar.chips.screenrecord.ui.viewmodel.ScreenRecordChipViewModel
 import com.android.systemui.statusbar.chips.sharetoapp.ui.viewmodel.ShareToAppChipViewModel
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
@@ -51,6 +53,7 @@
     shareToAppChipViewModel: ShareToAppChipViewModel,
     castToOtherDeviceChipViewModel: CastToOtherDeviceChipViewModel,
     callChipViewModel: CallChipViewModel,
+    demoRonChipViewModel: DemoRonChipViewModel,
     @StatusBarChipsLog private val logger: LogBuffer,
 ) {
     private enum class ChipType {
@@ -58,6 +61,8 @@
         ShareToApp,
         CastToOtherDevice,
         Call,
+        /** A demo of a RON chip (rich ongoing notification chip), used just for testing. */
+        DemoRon,
     }
 
     /** Model that helps us internally track the various chip states from each of the types. */
@@ -78,6 +83,7 @@
             val shareToApp: OngoingActivityChipModel.Hidden,
             val castToOtherDevice: OngoingActivityChipModel.Hidden,
             val call: OngoingActivityChipModel.Hidden,
+            val demoRon: OngoingActivityChipModel.Hidden,
         ) : InternalChipModel
     }
 
@@ -87,7 +93,8 @@
             shareToAppChipViewModel.chip,
             castToOtherDeviceChipViewModel.chip,
             callChipViewModel.chip,
-        ) { screenRecord, shareToApp, castToOtherDevice, call ->
+            demoRonChipViewModel.chip,
+        ) { screenRecord, shareToApp, castToOtherDevice, call, demoRon ->
             logger.log(
                 TAG,
                 LogLevel.INFO,
@@ -98,7 +105,15 @@
                 },
                 { "Chips: ScreenRecord=$str1 > ShareToApp=$str2 > CastToOther=$str3..." },
             )
-            logger.log(TAG, LogLevel.INFO, { str1 = call.logName }, { "... > Call=$str1" })
+            logger.log(
+                TAG,
+                LogLevel.INFO,
+                {
+                    str1 = call.logName
+                    str2 = demoRon.logName
+                },
+                { "... > Call=$str1 > DemoRon=$str2" }
+            )
             // This `when` statement shows the priority order of the chips.
             when {
                 // Screen recording also activates the media projection APIs, so whenever the
@@ -113,17 +128,23 @@
                     InternalChipModel.Shown(ChipType.CastToOtherDevice, castToOtherDevice)
                 call is OngoingActivityChipModel.Shown ->
                     InternalChipModel.Shown(ChipType.Call, call)
+                demoRon is OngoingActivityChipModel.Shown -> {
+                    StatusBarRonChips.assertInNewMode()
+                    InternalChipModel.Shown(ChipType.DemoRon, demoRon)
+                }
                 else -> {
                     // We should only get here if all chip types are hidden
                     check(screenRecord is OngoingActivityChipModel.Hidden)
                     check(shareToApp is OngoingActivityChipModel.Hidden)
                     check(castToOtherDevice is OngoingActivityChipModel.Hidden)
                     check(call is OngoingActivityChipModel.Hidden)
+                    check(demoRon is OngoingActivityChipModel.Hidden)
                     InternalChipModel.Hidden(
                         screenRecord = screenRecord,
                         shareToApp = shareToApp,
                         castToOtherDevice = castToOtherDevice,
                         call = call,
+                        demoRon = demoRon,
                     )
                 }
             }
@@ -154,6 +175,7 @@
                         ChipType.ShareToApp -> new.shareToApp
                         ChipType.CastToOtherDevice -> new.castToOtherDevice
                         ChipType.Call -> new.call
+                        ChipType.DemoRon -> new.demoRon
                     }
                 } else if (new is InternalChipModel.Shown) {
                     // If we have a chip to show, always show it.
@@ -179,6 +201,7 @@
                 shareToApp = OngoingActivityChipModel.Hidden(),
                 castToOtherDevice = OngoingActivityChipModel.Hidden(),
                 call = OngoingActivityChipModel.Hidden(),
+                demoRon = OngoingActivityChipModel.Hidden(),
             )
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/commandline/ValueParser.kt b/packages/SystemUI/src/com/android/systemui/statusbar/commandline/ValueParser.kt
index 01083d9..412c8c6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/commandline/ValueParser.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/commandline/ValueParser.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.commandline
 
+import androidx.core.graphics.toColorInt
 import kotlin.contracts.ExperimentalContracts
 import kotlin.contracts.InvocationKind
 import kotlin.contracts.contract
@@ -164,10 +165,23 @@
         ?: Result.failure(ArgParseError("Failed to parse $value as a float"))
 }
 
+// See https://developer.android.com/reference/android/graphics/Color#parseColor(java.lang.String)
+// for the supported formats of the color string. tl;dr: #RRGGBB, #AARRGGBB, or a basic color name
+// like "red" or "green". For the RRGGBB values, the `#` needs to be escaped. Use `"\\#RRGGBB"` in
+// the command to escape the `#` correctly.
+private val parseColor: ValueParser<Int> = ValueParser { value ->
+    try {
+        Result.success(value.toColorInt())
+    } catch (e: IllegalArgumentException) {
+        Result.failure(ArgParseError("Failed to parse $value as a color: $e"))
+    }
+}
+
 /** Default parsers that can be use as-is, or [map]ped to another type */
 object Type {
     val Boolean = parseBoolean
     val Int = parseInt
     val Float = parseFloat
     val String = parseString
+    val Color = parseColor
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
index ecb6d7f..406a664 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/StatusBarModule.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.log.LogBufferFactory
 import com.android.systemui.statusbar.data.StatusBarDataLayerModule
 import com.android.systemui.statusbar.phone.LightBarController
+import com.android.systemui.statusbar.phone.StatusBarSignalPolicy
 import com.android.systemui.statusbar.phone.ongoingcall.OngoingCallController
 import com.android.systemui.statusbar.phone.ongoingcall.OngoingCallLog
 import com.android.systemui.statusbar.ui.SystemBarUtilsProxyImpl
@@ -51,6 +52,11 @@
     @ClassKey(LightBarController::class)
     abstract fun bindLightBarController(impl: LightBarController): CoreStartable
 
+    @Binds
+    @IntoMap
+    @ClassKey(StatusBarSignalPolicy::class)
+    abstract fun bindStatusBarSignalPolicy(impl: StatusBarSignalPolicy): CoreStartable
+
     companion object {
         @Provides
         @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractor.kt
index 74ec7ed..aa203d7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/domain/interactor/HeadsUpNotificationInteractor.kt
@@ -21,11 +21,11 @@
 import com.android.systemui.deviceentry.domain.interactor.DeviceEntryFaceAuthInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.statusbar.notification.data.repository.HeadsUpRepository
 import com.android.systemui.statusbar.notification.data.repository.HeadsUpRowRepository
 import com.android.systemui.statusbar.notification.shared.HeadsUpRowKey
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import javax.inject.Inject
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
@@ -58,7 +58,7 @@
 
     /** Set of currently pinned top-level heads up rows to be displayed. */
     val pinnedHeadsUpRows: Flow<Set<HeadsUpRowKey>> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(emptySet())
         } else {
             headsUpRepository.activeHeadsUpRows.flatMapLatest { repositories ->
@@ -80,7 +80,7 @@
 
     /** Are there any pinned heads up rows to display? */
     val hasPinnedRows: Flow<Boolean> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(false)
         } else {
             headsUpRepository.activeHeadsUpRows.flatMapLatest { rows ->
@@ -95,7 +95,7 @@
     }
 
     val isHeadsUpOrAnimatingAway: Flow<Boolean> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(false)
         } else {
             combine(hasPinnedRows, headsUpRepository.isHeadsUpAnimatingAway) {
@@ -123,7 +123,7 @@
         }
 
     val showHeadsUpStatusBar: Flow<Boolean> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(false)
         } else {
             combine(hasPinnedRows, canShowHeadsUp) { hasPinnedRows, canShowHeadsUp ->
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
index ce4356a..18d4a04 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
@@ -18,7 +18,7 @@
 
 /**
  * NotificationList proto from atoms.proto, duplicated here so that it's accessible in the build.
- * Must be kept in sync with the version in atoms.proto.
+ * Must be kept in sync with the version in stats/atoms/sysui/sysui_atoms.proto.
  */
 
 message Notification {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
index 4be638f..1431b28 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/AmbientState.java
@@ -173,7 +173,8 @@
     }
 
     /**
-     * @return Height of the notifications panel without top padding when expansion completes.
+     * @return Height of the available space for the notification content, when the shade
+     * expansion completes.
      */
     public float getStackEndHeight() {
         return mStackEndHeight;
@@ -276,17 +277,18 @@
     }
 
     /**
-     * @see #getStackHeight()
+     * @return Height of the notification content returned by {@link #getStackEndHeight()}, but
+     * interpolated by the shade expansion fraction.
      */
-    public void setStackHeight(float stackHeight) {
-        mStackHeight = stackHeight;
+    public float getInterpolatedStackHeight() {
+        return mStackHeight;
     }
 
     /**
-     * @return Height of notifications panel interpolated by the expansion fraction.
+     * @see #getInterpolatedStackHeight()
      */
-    public float getStackHeight() {
-        return mStackHeight;
+    public void setInterpolatedStackHeight(float stackHeight) {
+        mStackHeight = stackHeight;
     }
 
     @Inject
@@ -531,8 +533,15 @@
         if (mDozeAmount == 1.0f && !isPulseExpanding()) {
             return mShelf.getHeight();
         }
-        int height = (int) Math.max(mLayoutMinHeight,
-                Math.min(mLayoutHeight, mContentHeight) - mTopPadding);
+        int height;
+        if (SceneContainerFlag.isEnabled()) {
+            // TODO(b/192348384): This is probably incorrect as mContentHeight is not up to date.
+            //  Consider removing usages of getInnerHeight in flexiglass if possible.
+            height = (int) Math.min(mLayoutHeight, mContentHeight) - mTopPadding;
+        } else {
+            height = (int) Math.max(mLayoutMinHeight,
+                    Math.min(mLayoutHeight, mContentHeight) - mTopPadding);
+        }
         if (ignorePulseHeight) {
             return height;
         }
@@ -569,6 +578,7 @@
     }
 
     public void setLayoutMinHeight(int layoutMinHeight) {
+        SceneContainerFlag.assertInLegacyMode();
         mLayoutMinHeight = layoutMinHeight;
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index 3af1f3c..b9628e9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -114,7 +114,6 @@
 import com.android.systemui.statusbar.notification.row.StackScrollerDecorView;
 import com.android.systemui.statusbar.notification.shared.NotificationHeadsUpCycling;
 import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.shared.NotificationsImprovedHunAnimation;
 import com.android.systemui.statusbar.notification.shared.NotificationsLiveDataStoreRefactor;
 import com.android.systemui.statusbar.notification.stack.shared.model.ShadeScrimBounds;
@@ -790,7 +789,6 @@
     private void onJustBeforeDraw() {
         if (SceneContainerFlag.isEnabled()) {
             if (mChildrenUpdateRequested) {
-                updateForcedScroll();
                 updateChildren();
                 mChildrenUpdateRequested = false;
             }
@@ -875,7 +873,7 @@
         y = (int) (mAmbientState.getStackY());
         drawDebugInfo(canvas, y, Color.CYAN, /* label= */ "mAmbientState.getStackY() = " + y);
 
-        y = (int) (mAmbientState.getStackY() + mAmbientState.getStackHeight());
+        y = (int) (mAmbientState.getStackY() + mAmbientState.getInterpolatedStackHeight());
         drawDebugInfo(canvas, y, Color.LTGRAY,
                 /* label= */ "mAmbientState.getStackY() + mAmbientState.getStackHeight() = " + y);
 
@@ -1124,11 +1122,13 @@
 
     @Override
     public void addStackHeightChangedListener(@NonNull Runnable runnable) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return;
         mStackHeightChangedListeners.addIfAbsent(runnable);
     }
 
     @Override
     public void removeStackHeightChangedListener(@NonNull Runnable runnable) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return;
         mStackHeightChangedListeners.remove(runnable);
     }
 
@@ -1234,7 +1234,6 @@
         if (mAmbientState.getStackTop() != stackTop) {
             mAmbientState.setStackTop(stackTop);
             onTopPaddingChanged(/* animate = */ isAddOrRemoveAnimationPending());
-            setExpandedHeight(mExpandedHeight);
         }
     }
 
@@ -1309,8 +1308,10 @@
     }
 
     private void updateAlgorithmLayoutMinHeight() {
-        mAmbientState.setLayoutMinHeight(mQsFullScreen || isHeadsUpTransition()
-                ? getLayoutMinHeightInternal() : 0);
+        if (!SceneContainerFlag.isEnabled()) {
+            mAmbientState.setLayoutMinHeight(mQsFullScreen || isHeadsUpTransition()
+                    ? getLayoutMinHeightInternal() : 0);
+        }
     }
 
     /**
@@ -1477,7 +1478,7 @@
 
     @VisibleForTesting
     public void updateStackEndHeightAndStackHeight(float fraction) {
-        final float oldStackHeight = mAmbientState.getStackHeight();
+        final float oldStackHeight = mAmbientState.getInterpolatedStackHeight();
         if (SceneContainerFlag.isEnabled()) {
             final float endHeight;
             if (!shouldSkipHeightUpdate()) {
@@ -1485,20 +1486,20 @@
             } else {
                 endHeight = mAmbientState.getStackEndHeight();
             }
-            updateStackHeight(endHeight, fraction);
+            updateInterpolatedStackHeight(endHeight, fraction);
         } else {
             if (mQsExpansionFraction <= 0 && !shouldSkipHeightUpdate()) {
                 final float endHeight = updateStackEndHeight(
                         getHeight(), getEmptyBottomMarginInternal(), getTopPadding());
-                updateStackHeight(endHeight, fraction);
+                updateInterpolatedStackHeight(endHeight, fraction);
             } else {
                 // Always updateStackHeight to prevent jumps in the stack height when this fraction
                 // suddenly reapplies after a freeze.
                 final float endHeight = mAmbientState.getStackEndHeight();
-                updateStackHeight(endHeight, fraction);
+                updateInterpolatedStackHeight(endHeight, fraction);
             }
         }
-        if (oldStackHeight != mAmbientState.getStackHeight()) {
+        if (oldStackHeight != mAmbientState.getInterpolatedStackHeight()) {
             requestChildrenUpdate();
         }
     }
@@ -1532,7 +1533,7 @@
     }
 
     @VisibleForTesting
-    public void updateStackHeight(float endHeight, float fraction) {
+    public void updateInterpolatedStackHeight(float endHeight, float fraction) {
         if (!newAodTransition()) {
             // During the (AOD<=>LS) transition where dozeAmount is changing,
             // apply dozeAmount to stack height instead of expansionFraction
@@ -1542,7 +1543,7 @@
                 fraction = 1f - dozeAmount;
             }
         }
-        mAmbientState.setStackHeight(
+        mAmbientState.setInterpolatedStackHeight(
                 MathUtils.lerp(endHeight * StackScrollAlgorithm.START_FRACTION,
                         endHeight, fraction));
     }
@@ -1571,8 +1572,11 @@
 
         // Update the expand progress between started/stopped events
         mAmbientState.setExpansionFraction(expandFraction);
-        // TODO(b/332577544): don't convert to height which then converts to the fraction again
-        setExpandedHeight(expandFraction * getHeight());
+
+        if (!shouldSkipHeightUpdate()) {
+            updateStackEndHeightAndStackHeight(expandFraction);
+            updateExpandedHeight(expandFraction);
+        }
 
         // expansion stopped event requires that the expandFraction has already been updated
         if (!nowExpanding && wasExpanding) {
@@ -1581,6 +1585,19 @@
         }
     }
 
+    private void updateExpandedHeight(float expandFraction) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return;
+        float expandedHeight = expandFraction * getHeight();
+        setIsExpanded(expandedHeight > 0);
+
+        if (mExpandedHeight != expandedHeight) {
+            mExpandedHeight = expandedHeight;
+            updateAlgorithmHeightAndPadding();
+            requestChildrenUpdate();
+            notifyAppearChangedListeners();
+        }
+    }
+
     @Override
     public void setQsExpandFraction(float expandFraction) {
         if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return;
@@ -1593,6 +1610,7 @@
      * @param height the expanded height of the panel
      */
     public void setExpandedHeight(float height) {
+        SceneContainerFlag.assertInLegacyMode();
         final boolean skipHeightUpdate = shouldSkipHeightUpdate();
 
         updateStackPosition();
@@ -1724,6 +1742,7 @@
      * Measured relative to the resting position.
      */
     private float getExpandTranslationStart() {
+        SceneContainerFlag.assertInLegacyMode();
         return -getTopPadding() + getMinExpansionHeight() - mShelf.getIntrinsicHeight();
     }
 
@@ -1732,6 +1751,7 @@
      * Measured in absolute height.
      */
     private float getAppearStartPosition() {
+        SceneContainerFlag.assertInLegacyMode();
         if (isHeadsUpTransition()) {
             final NotificationSection firstVisibleSection = getFirstVisibleSection();
             final int pinnedHeight = firstVisibleSection != null
@@ -1787,6 +1807,7 @@
      *    have the shelf on its own)
      */
     private float getAppearEndPosition() {
+        SceneContainerFlag.assertInLegacyMode();
         if (FooterViewRefactor.isUnexpectedlyInLegacyMode()) {
             return getAppearEndPositionLegacy();
         }
@@ -1847,7 +1868,7 @@
      */
     @FloatRange(from = -1.0, to = 1.0)
     public float calculateAppearFraction(float height) {
-        if (isHeadsUpTransition()) {
+        if (isHeadsUpTransition() && !SceneContainerFlag.isEnabled()) {
             // HUN is a special case because fraction can go negative if swiping up. And for now
             // it must go negative as other pieces responsible for proper translation up assume
             // negative value for HUN going up.
@@ -1978,7 +1999,8 @@
     }
 
     public void lockScrollTo(View v) {
-        if (mForcedScroll == v) {
+        // NSSL shouldn't handle scrolling with SceneContainer enabled.
+        if (mForcedScroll == v || SceneContainerFlag.isEnabled()) {
             return;
         }
         mForcedScroll = v;
@@ -1986,6 +2008,10 @@
     }
 
     public boolean scrollTo(View v) {
+        // NSSL shouldn't handle scrolling with SceneContainer enabled.
+        if (SceneContainerFlag.isEnabled()) {
+            return false;
+        }
         ExpandableView expandableView = (ExpandableView) v;
         int positionInLinearLayout = getPositionInLinearLayout(v);
         int targetScroll = targetScrollForView(expandableView, positionInLinearLayout);
@@ -2007,6 +2033,7 @@
      * the IME.
      */
     private int targetScrollForView(ExpandableView v, int positionInLinearLayout) {
+        SceneContainerFlag.assertInLegacyMode();
         return positionInLinearLayout + v.getIntrinsicHeight() +
                 getImeInset() - getHeight()
                 + ((!isExpanded() && isPinnedHeadsUp(v)) ? mHeadsUpInset : getTopPadding());
@@ -2316,6 +2343,7 @@
 
     private void setOverScrollAmountInternal(float amount, boolean onTop, boolean animate,
                                              boolean isRubberbanded) {
+        SceneContainerFlag.assertInLegacyMode();
         amount = Math.max(0, amount);
         if (animate) {
             mStateAnimator.animateOverScrollToAmount(amount, onTop, isRubberbanded);
@@ -2522,10 +2550,33 @@
     }
 
     @VisibleForTesting
-    void updateContentHeight() {
+    void updateStackHeight() {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return;
+
+        final int shelfIntrinsicHeight = mShelf != null ? mShelf.getIntrinsicHeight() : 0;
+        final int footerIntrinsicHeight =
+                mFooterView != null ? mFooterView.getIntrinsicHeight() : 0;
+        final int notificationsHeight = (int) mNotificationStackSizeCalculator.computeHeight(
+                /* notificationStackScrollLayout= */ this,
+                mMaxDisplayedNotifications,
+                shelfIntrinsicHeight
+        );
+        mIntrinsicContentHeight = notificationsHeight;
+        final int fullStackHeight = notificationsHeight + footerIntrinsicHeight + mBottomPadding;
+        if (mScrollViewFields.getIntrinsicStackHeight() != fullStackHeight) {
+            mScrollViewFields.setIntrinsicStackHeight(fullStackHeight);
+            notifyStackHeightChangedListeners();
+        }
+    }
+
+    private void updateContentHeight() {
+        if (SceneContainerFlag.isEnabled()) {
+            updateStackHeight();
+            return;
+        }
+
         final float scrimTopPadding = getScrimTopPaddingOrZero();
         final int shelfIntrinsicHeight = mShelf != null ? mShelf.getIntrinsicHeight() : 0;
-        final int footerIntrinsicHeight = mFooterView != null ? mFooterView.getIntrinsicHeight() : 0;
         final float height =
                 (int) scrimTopPadding + (int) mNotificationStackSizeCalculator.computeHeight(
                         /* notificationStackScrollLayout= */ this, mMaxDisplayedNotifications,
@@ -2536,19 +2587,15 @@
         // state the maxPanelHeight and the contentHeight should be bigger
         mContentHeight =
                 (int) (height + Math.max(getIntrinsicPadding(), getTopPadding()) + mBottomPadding);
-        mScrollViewFields.setIntrinsicStackHeight(
-                (int) (getIntrinsicPadding() + mIntrinsicContentHeight + footerIntrinsicHeight
-                        + mBottomPadding));
         updateScrollability();
         clampScrollPosition();
         updateStackPosition();
         mAmbientState.setContentHeight(mContentHeight);
-
-        notifyStackHeightChangedListeners();
     }
 
     @Override
     public int getIntrinsicStackHeight() {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) return 0;
         return mScrollViewFields.getIntrinsicStackHeight();
     }
 
@@ -2584,6 +2631,9 @@
     }
 
     private void updateScrollability() {
+        if (SceneContainerFlag.isEnabled()) {
+            return;
+        }
         boolean scrollable = !mQsFullScreen && getScrollRange() > 0;
         if (scrollable != mScrollable) {
             mScrollable = scrollable;
@@ -2593,6 +2643,7 @@
     }
 
     private void updateForwardAndBackwardScrollability() {
+        SceneContainerFlag.assertInLegacyMode();
         boolean forwardScrollable = mScrollable && !mScrollAdapter.isScrolledToBottom();
         boolean backwardsScrollable = mScrollable && !mScrollAdapter.isScrolledToTop();
         boolean changed = forwardScrollable != mForwardScrollable
@@ -2750,6 +2801,7 @@
     }
 
     private int getLayoutMinHeightInternal() {
+        SceneContainerFlag.assertInLegacyMode();
         if (isHeadsUpTransition()) {
             ExpandableNotificationRow trackedHeadsUpRow = mAmbientState.getTrackedHeadsUpRow();
             if (trackedHeadsUpRow.isAboveShelf()) {
@@ -4132,6 +4184,11 @@
      */
     @Override
     public boolean performAccessibilityActionInternal(int action, Bundle arguments) {
+        // Don't handle scroll accessibility events from the NSSL, when SceneContainer enabled.
+        if (SceneContainerFlag.isEnabled()) {
+            return super.performAccessibilityActionInternal(action, arguments);
+        }
+
         if (super.performAccessibilityActionInternal(action, arguments)) {
             return true;
         }
@@ -4303,7 +4360,7 @@
                 // Resetting headsUpAnimatingAway on Shade expansion avoids delays caused by
                 // waiting for all child animations to finish.
                 // TODO(b/328390331) Do we need to reset this on QS expanded as well?
-                if (NotificationsHeadsUpRefactor.isEnabled()) {
+                if (SceneContainerFlag.isEnabled()) {
                     setHeadsUpAnimatingAway(false);
                 }
             } else {
@@ -4414,7 +4471,7 @@
 
     void onChildAnimationFinished() {
         setAnimationRunning(false);
-        if (NotificationsHeadsUpRefactor.isEnabled()) {
+        if (SceneContainerFlag.isEnabled()) {
             setHeadsUpAnimatingAway(false);
         }
         requestChildrenUpdate();
@@ -4893,6 +4950,11 @@
     @Override
     public void onInitializeAccessibilityEventInternal(AccessibilityEvent event) {
         super.onInitializeAccessibilityEventInternal(event);
+        // Don't handle scroll accessibility events from the NSSL, when SceneContainer enabled.
+        if (SceneContainerFlag.isEnabled()) {
+            return;
+        }
+
         event.setScrollable(mScrollable);
         event.setMaxScrollX(mScrollX);
         event.setScrollY(mOwnScrollY);
@@ -4902,6 +4964,11 @@
     @Override
     public void onInitializeAccessibilityNodeInfoInternal(AccessibilityNodeInfo info) {
         super.onInitializeAccessibilityNodeInfoInternal(info);
+        // Don't handle scroll accessibility events from the NSSL, when SceneContainer enabled.
+        if (SceneContainerFlag.isEnabled()) {
+            return;
+        }
+
         if (mScrollable) {
             info.setScrollable(true);
             if (mBackwardScrollable) {
@@ -4960,7 +5027,7 @@
     }
 
     public void generateHeadsUpAnimation(NotificationEntry entry, boolean isHeadsUp) {
-        NotificationsHeadsUpRefactor.assertInLegacyMode();
+        SceneContainerFlag.assertInLegacyMode();
         ExpandableNotificationRow row = entry.getHeadsUpAnimationView();
         generateHeadsUpAnimation(row, isHeadsUp);
     }
@@ -5003,7 +5070,7 @@
             mNeedsAnimation = true;
             if (!mIsExpanded && !mWillExpand && !isHeadsUp) {
                 row.setHeadsUpAnimatingAway(true);
-                if (NotificationsHeadsUpRefactor.isEnabled()) {
+                if (SceneContainerFlag.isEnabled()) {
                     setHeadsUpAnimatingAway(true);
                 }
             }
@@ -5087,10 +5154,12 @@
     }
 
     boolean isQsFullScreen() {
+        SceneContainerFlag.assertInLegacyMode();
         return mQsFullScreen;
     }
 
     public void setQsExpansionFraction(float qsExpansionFraction) {
+        SceneContainerFlag.assertInLegacyMode();
         boolean footerAffected = mQsExpansionFraction != qsExpansionFraction
                 && (mQsExpansionFraction == 1 || qsExpansionFraction == 1);
         mQsExpansionFraction = qsExpansionFraction;
@@ -5134,6 +5203,7 @@
     }
 
     private void updateOnScrollChange() {
+        SceneContainerFlag.assertInLegacyMode();
         if (mScrollListener != null) {
             mScrollListener.accept(mOwnScrollY);
         }
@@ -5197,7 +5267,7 @@
         updateClipping();
     }
 
-    /** TODO(b/328390331) make this private, when {@link NotificationsHeadsUpRefactor} is removed */
+    /** TODO(b/328390331) make this private, when {@link SceneContainerFlag} is removed */
     public void setHeadsUpAnimatingAway(boolean headsUpAnimatingAway) {
         if (mHeadsUpAnimatingAway != headsUpAnimatingAway) {
             mHeadsUpAnimatingAway = headsUpAnimatingAway;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 4e73529..bcdc3bc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -127,7 +127,6 @@
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
 import com.android.systemui.statusbar.notification.row.NotificationSnooze;
 import com.android.systemui.statusbar.notification.shared.GroupHunAnimationFix;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.stack.ui.viewbinder.NotificationListViewBinder;
 import com.android.systemui.statusbar.phone.HeadsUpAppearanceController;
 import com.android.systemui.statusbar.phone.HeadsUpNotificationViewControllerEmptyImpl;
@@ -686,13 +685,13 @@
             new OnHeadsUpChangedListener() {
                 @Override
                 public void onHeadsUpPinnedModeChanged(boolean inPinnedMode) {
-                    NotificationsHeadsUpRefactor.assertInLegacyMode();
+                    SceneContainerFlag.assertInLegacyMode();
                     mView.setInHeadsUpPinnedMode(inPinnedMode);
                 }
 
                 @Override
                 public void onHeadsUpStateChanged(NotificationEntry entry, boolean isHeadsUp) {
-                    NotificationsHeadsUpRefactor.assertInLegacyMode();
+                    SceneContainerFlag.assertInLegacyMode();
                     NotificationEntry topEntry = mHeadsUpManager.getTopEntry();
                     mView.setTopHeadsUpRow(topEntry != null ? topEntry.getRow() : null);
                     generateHeadsUpAnimation(entry, isHeadsUp);
@@ -880,7 +879,7 @@
             });
         }
 
-        if (!NotificationsHeadsUpRefactor.isEnabled()) {
+        if (!SceneContainerFlag.isEnabled()) {
             mHeadsUpManager.addListener(mOnHeadsUpChangedListener);
         }
         mHeadsUpManager.setAnimationStateHandler(mView::setHeadsUpGoingAwayAnimationsAllowed);
@@ -1276,6 +1275,7 @@
     }
 
     public void setQsExpansionFraction(float expansionFraction) {
+        SceneContainerFlag.assertInLegacyMode();
         mView.setQsExpansionFraction(expansionFraction);
     }
 
@@ -1409,6 +1409,7 @@
     }
 
     public float calculateAppearFraction(float height) {
+        SceneContainerFlag.assertInLegacyMode();
         return mView.calculateAppearFraction(height);
     }
 
@@ -1508,7 +1509,7 @@
     }
 
     public void setHeadsUpAnimatingAway(boolean headsUpAnimatingAway) {
-        NotificationsHeadsUpRefactor.assertInLegacyMode();
+        SceneContainerFlag.assertInLegacyMode();
         mView.setHeadsUpAnimatingAway(headsUpAnimatingAway);
     }
 
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 0c2b5ae..ef1bcfc 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
@@ -575,7 +575,8 @@
         final float shelfHeight = showingShelf ? ambientState.getShelf().getIntrinsicHeight() : 0f;
         final float scrimPadding = getScrimTopPaddingOrZero(ambientState);
 
-        final float stackHeight = ambientState.getStackHeight() - shelfHeight - scrimPadding;
+        final float stackHeight =
+                ambientState.getInterpolatedStackHeight() - shelfHeight - scrimPadding;
         final float stackEndHeight = ambientState.getStackEndHeight() - shelfHeight - scrimPadding;
         if (stackEndHeight == 0f) {
             // This should not happen, since even when the shade is empty we show EmptyShadeView
@@ -734,7 +735,7 @@
                             || ambientState.getDozeAmount() == 1f
                             || bypassPulseNotExpanding
                             ? ambientState.getInnerHeight()
-                            : ambientState.getStackHeight();
+                            : ambientState.getInterpolatedStackHeight();
                     final float shelfStart = stackBottom
                             - ambientState.getShelf().getIntrinsicHeight()
                             - mPaddingBetweenElements;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationListViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationListViewBinder.kt
index 5572f8e..d770b20 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationListViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationListViewBinder.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.lifecycle.repeatWhenAttached
 import com.android.systemui.plugins.FalsingManager
 import com.android.systemui.res.R
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.statusbar.NotificationShelf
 import com.android.systemui.statusbar.notification.NotificationActivityStarter
 import com.android.systemui.statusbar.notification.collection.render.SectionHeaderController
@@ -36,7 +37,6 @@
 import com.android.systemui.statusbar.notification.footer.ui.viewbinder.FooterViewBinder
 import com.android.systemui.statusbar.notification.footer.ui.viewmodel.FooterViewModel
 import com.android.systemui.statusbar.notification.icon.ui.viewbinder.NotificationIconContainerShelfViewBinder
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.shared.NotificationsLiveDataStoreRefactor
 import com.android.systemui.statusbar.notification.shelf.ui.viewbinder.NotificationShelfViewBinder
 import com.android.systemui.statusbar.notification.stack.DisplaySwitchNotificationsHiderTracker
@@ -93,7 +93,7 @@
 
         view.repeatWhenAttached {
             lifecycleScope.launch {
-                if (NotificationsHeadsUpRefactor.isEnabled) {
+                if (SceneContainerFlag.isEnabled) {
                     launch { hunBinder.bindHeadsUpNotifications(view) }
                 }
                 launch { bindShelf(shelf) }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
index 3cc6e81..6d5553f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewbinder/NotificationScrollViewBinder.kt
@@ -67,6 +67,7 @@
 
     suspend fun bind(): Nothing =
         view.asView().viewModel(
+            traceName = "NotificationScrollViewBinder",
             minWindowLifecycleState = WindowLifecycleState.ATTACHED,
             factory = viewModelFactory::create,
         ) { viewModel ->
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModel.kt
index 5fba615..e55492e6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationListViewModel.kt
@@ -18,6 +18,7 @@
 
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
 import com.android.systemui.statusbar.domain.interactor.RemoteInputInteractor
 import com.android.systemui.statusbar.notification.domain.interactor.ActiveNotificationsInteractor
@@ -26,7 +27,6 @@
 import com.android.systemui.statusbar.notification.footer.shared.FooterViewRefactor
 import com.android.systemui.statusbar.notification.footer.ui.viewmodel.FooterViewModel
 import com.android.systemui.statusbar.notification.shared.HeadsUpRowKey
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.shelf.ui.viewmodel.NotificationShelfViewModel
 import com.android.systemui.statusbar.notification.stack.domain.interactor.NotificationStackInteractor
 import com.android.systemui.statusbar.policy.domain.interactor.UserSetupInteractor
@@ -256,7 +256,7 @@
     }
 
     val topHeadsUpRow: Flow<HeadsUpRowKey?> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(null)
         } else {
             headsUpNotificationInteractor.topHeadsUpRow.dumpWhileCollecting("topHeadsUpRow")
@@ -264,7 +264,7 @@
     }
 
     val pinnedHeadsUpRows: Flow<Set<HeadsUpRowKey>> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(emptySet())
         } else {
             headsUpNotificationInteractor.pinnedHeadsUpRows.dumpWhileCollecting("pinnedHeadsUpRows")
@@ -272,7 +272,7 @@
     }
 
     val headsUpAnimationsEnabled: Flow<Boolean> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(false)
         } else {
             flowOf(true).dumpWhileCollecting("headsUpAnimationsEnabled")
@@ -280,7 +280,7 @@
     }
 
     val hasPinnedHeadsUpRow: Flow<Boolean> by lazy {
-        if (NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode()) {
+        if (SceneContainerFlag.isUnexpectedlyInLegacyMode()) {
             flowOf(false)
         } else {
             headsUpNotificationInteractor.hasPinnedRows.dumpWhileCollecting("hasPinnedHeadsUpRow")
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
index 3999578..3e42413 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationScrollViewModel.kt
@@ -19,12 +19,11 @@
 
 import com.android.compose.animation.scene.ObservableTransitionState.Idle
 import com.android.compose.animation.scene.ObservableTransitionState.Transition
-import com.android.compose.animation.scene.ObservableTransitionState.Transition.ChangeCurrentScene
+import com.android.compose.animation.scene.ObservableTransitionState.Transition.ChangeScene
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.SceneFamilies
@@ -63,7 +62,6 @@
     keyguardInteractor: Lazy<KeyguardInteractor>,
 ) :
     ActivatableFlowDumper by ActivatableFlowDumperImpl(dumpManager, "NotificationScrollViewModel"),
-    SysUiViewModel,
     ExclusiveActivatable() {
 
     override suspend fun onActivated(): Nothing {
@@ -79,7 +77,7 @@
         }
     }
 
-    private fun fullyExpandedDuringSceneChange(change: ChangeCurrentScene): Boolean {
+    private fun fullyExpandedDuringSceneChange(change: ChangeScene): Boolean {
         // The lockscreen stack is visible during all transitions away from the lockscreen, so keep
         // the stack expanded until those transitions finish.
         return (expandedInScene(change.fromScene) && expandedInScene(change.toScene)) ||
@@ -87,7 +85,7 @@
     }
 
     private fun expandFractionDuringSceneChange(
-        change: ChangeCurrentScene,
+        change: ChangeScene,
         shadeExpansion: Float,
         qsExpansion: Float,
     ): Float {
@@ -120,7 +118,7 @@
             ) { shadeExpansion, _, qsExpansion, transitionState, _ ->
                 when (transitionState) {
                     is Idle -> if (expandedInScene(transitionState.currentScene)) 1f else 0f
-                    is ChangeCurrentScene ->
+                    is ChangeScene ->
                         expandFractionDuringSceneChange(
                             transitionState,
                             shadeExpansion,
@@ -250,7 +248,7 @@
     }
 }
 
-private fun ChangeCurrentScene.isBetween(
+private fun ChangeScene.isBetween(
     a: (SceneKey) -> Boolean,
-    b: (SceneKey) -> Boolean
+    b: (SceneKey) -> Boolean,
 ): Boolean = (a(fromScene) && b(toScene)) || (b(fromScene) && a(toScene))
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
index d891f62..69c1bf3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModel.kt
@@ -21,7 +21,6 @@
 import com.android.systemui.flags.FeatureFlagsClassic
 import com.android.systemui.flags.Flags
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.scene.domain.interactor.SceneInteractor
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
@@ -53,7 +52,6 @@
     featureFlags: FeatureFlagsClassic,
     dumpManager: DumpManager,
 ) :
-    SysUiViewModel,
     ExclusiveActivatable(),
     ActivatableFlowDumper by ActivatableFlowDumperImpl(
         dumpManager = dumpManager,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
index f63ee7b..aed00d8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/SharedNotificationContainerViewModel.kt
@@ -20,6 +20,7 @@
 package com.android.systemui.statusbar.notification.stack.ui.viewmodel
 
 import androidx.annotation.VisibleForTesting
+import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.common.shared.model.NotificationContainerBounds
 import com.android.systemui.communal.domain.interactor.CommunalSceneInteractor
 import com.android.systemui.dagger.SysUISingleton
@@ -141,10 +142,6 @@
     private val communalSceneInteractor: CommunalSceneInteractor,
     unfoldTransitionInteractor: UnfoldTransitionInteractor,
 ) : FlowDumperImpl(dumpManager) {
-    // TODO(b/349784682): Transform deprecated states for Flexiglass
-    private val statesForConstrainedNotifications: Set<KeyguardState> =
-        setOf(AOD, LOCKSCREEN, DOZING, ALTERNATE_BOUNCER, PRIMARY_BOUNCER)
-    private val statesForHiddenKeyguard: Set<KeyguardState> = setOf(GONE, OCCLUDED)
 
     /**
      * Is either shade/qs expanded? This intentionally does not use the [ShadeInteractor] version,
@@ -217,14 +214,16 @@
 
     /** If the user is visually on one of the unoccluded lockscreen states. */
     val isOnLockscreen: Flow<Boolean> =
-        combine(
-                keyguardTransitionInteractor.finishedKeyguardState.map {
-                    statesForConstrainedNotifications.contains(it)
-                },
+        anyOf(
+                keyguardTransitionInteractor.isFinishedIn(AOD),
+                keyguardTransitionInteractor.isFinishedIn(DOZING),
+                keyguardTransitionInteractor.isFinishedIn(ALTERNATE_BOUNCER),
+                keyguardTransitionInteractor.isFinishedIn(
+                    scene = Scenes.Bouncer,
+                    stateWithoutSceneContainer = PRIMARY_BOUNCER
+                ),
                 keyguardTransitionInteractor.transitionValue(LOCKSCREEN).map { it > 0f },
-            ) { constrainedNotificationState, transitioningToOrFromLockscreen ->
-                constrainedNotificationState || transitioningToOrFromLockscreen
-            }
+            )
             .stateIn(
                 scope = applicationScope,
                 started = SharingStarted.Eagerly,
@@ -250,9 +249,10 @@
     /** If the user is visually on the glanceable hub or transitioning to/from it */
     private val isOnGlanceableHub: Flow<Boolean> =
         combine(
-                keyguardTransitionInteractor.finishedKeyguardState.map { state ->
-                    state == GLANCEABLE_HUB
-                },
+                keyguardTransitionInteractor.isFinishedIn(
+                    scene = Scenes.Communal,
+                    stateWithoutSceneContainer = GLANCEABLE_HUB
+                ),
                 anyOf(
                     keyguardTransitionInteractor.isInTransition(
                         edge = Edge.create(to = Scenes.Communal),
@@ -424,32 +424,19 @@
             .onStart { emit(1f) }
             .dumpWhileCollecting("alphaForShadeAndQsExpansion")
 
-    private fun toFlowArray(
-        states: Set<KeyguardState>,
-        flow: (KeyguardState) -> Flow<Boolean>
-    ): Array<Flow<Boolean>> {
-        return states.map { flow(it) }.toTypedArray()
-    }
-
     private val isTransitioningToHiddenKeyguard: Flow<Boolean> =
         flow {
                 while (currentCoroutineContext().isActive) {
                     emit(false)
                     // Ensure states are inactive to start
-                    allOf(
-                            *toFlowArray(statesForHiddenKeyguard) { state ->
-                                keyguardTransitionInteractor.transitionValue(state).map { it == 0f }
-                            }
-                        )
-                        .first { it }
+                    allOf(isNotOnState(OCCLUDED), isNotOnState(GONE, Scenes.Gone)).first { it }
                     // Wait for a qualifying transition to begin
                     anyOf(
-                            *toFlowArray(statesForHiddenKeyguard) { state ->
-                                keyguardTransitionInteractor
-                                    .transition(Edge.create(to = state))
-                                    .map { it.value > 0f && it.transitionState == RUNNING }
-                                    .onStart { emit(false) }
-                            }
+                            transitionToIsRunning(Edge.create(to = OCCLUDED)),
+                            transitionToIsRunning(
+                                edge = Edge.create(to = Scenes.Gone),
+                                edgeWithoutSceneContainer = Edge.create(to = GONE)
+                            )
                         )
                         .first { it }
                     emit(true)
@@ -458,13 +445,7 @@
                     // it is considered safe to reset alpha to 1f for HUNs.
                     combine(
                             keyguardInteractor.statusBarState,
-                            allOf(
-                                *toFlowArray(statesForHiddenKeyguard) { state ->
-                                    keyguardTransitionInteractor.transitionValue(state).map {
-                                        it == 0f
-                                    }
-                                }
-                            )
+                            allOf(isNotOnState(OCCLUDED), isNotOnState(GONE, Scenes.Gone))
                         ) { statusBarState, stateIsReversed ->
                             statusBarState == SHADE || stateIsReversed
                         }
@@ -473,6 +454,17 @@
             }
             .dumpWhileCollecting("isTransitioningToHiddenKeyguard")
 
+    private fun isNotOnState(stateWithoutSceneContainer: KeyguardState, scene: SceneKey? = null) =
+        keyguardTransitionInteractor
+            .transitionValue(scene = scene, stateWithoutSceneContainer = stateWithoutSceneContainer)
+            .map { it == 0f }
+
+    private fun transitionToIsRunning(edge: Edge, edgeWithoutSceneContainer: Edge? = null) =
+        keyguardTransitionInteractor
+            .transition(edge = edge, edgeWithoutSceneContainer = edgeWithoutSceneContainer)
+            .map { it.value > 0f && it.transitionState == RUNNING }
+            .onStart { emit(false) }
+
     val panelAlpha = keyguardInteractor.panelAlpha
 
     private fun bouncerToGoneNotificationAlpha(viewState: ViewStateAccessor): Flow<Float> =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
index e3242d1..7227b93 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/CentralSurfacesImpl.java
@@ -31,6 +31,7 @@
 import static com.android.systemui.Flags.lightRevealMigration;
 import static com.android.systemui.Flags.newAodTransition;
 import static com.android.systemui.Flags.relockWithPowerButtonImmediately;
+import static com.android.systemui.Flags.statusBarSignalPolicyRefactor;
 import static com.android.systemui.charging.WirelessChargingAnimation.UNKNOWN_BATTERY_LEVEL;
 import static com.android.systemui.flags.Flags.SHORTCUT_LIST_SEARCH_LAYOUT;
 import static com.android.systemui.statusbar.NotificationLockscreenUserManager.PERMISSION_SELF;
@@ -870,7 +871,10 @@
         mBubblesOptional.ifPresent(this::initBubbles);
         mKeyguardBypassController.listenForQsExpandedChange();
 
-        mStatusBarSignalPolicy.init();
+        if (!statusBarSignalPolicyRefactor()) {
+            mStatusBarSignalPolicy.init();
+        }
+
         mKeyguardIndicationController.init();
 
         mColorExtractor.addOnColorsChangedListener(mOnColorsChangedListener);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerImpl.kt
index 2e1ab38..bb5aa23 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerImpl.kt
@@ -20,16 +20,17 @@
 import android.graphics.Rect
 import android.os.LocaleList
 import android.view.View.LAYOUT_DIRECTION_RTL
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener
-import javax.inject.Inject
+import dagger.assisted.Assisted
+import dagger.assisted.AssistedFactory
+import dagger.assisted.AssistedInject
 
-@SysUISingleton
-class ConfigurationControllerImpl @Inject constructor(
-        @Application context: Context,
-        ) : ConfigurationController {
+class ConfigurationControllerImpl
+@AssistedInject
+constructor(
+    @Assisted private val context: Context,
+) : ConfigurationController {
 
     private val listeners: MutableList<ConfigurationListener> = ArrayList()
     private val lastConfig = Configuration()
@@ -40,18 +41,17 @@
     private val inCarMode: Boolean
     private var uiMode: Int = 0
     private var localeList: LocaleList? = null
-    private val context: Context
     private var layoutDirection: Int
     private var orientation = Configuration.ORIENTATION_UNDEFINED
 
     init {
         val currentConfig = context.resources.configuration
-        this.context = context
         fontScale = currentConfig.fontScale
         density = currentConfig.densityDpi
         smallestScreenWidth = currentConfig.smallestScreenWidthDp
         maxBounds.set(currentConfig.windowConfiguration.maxBounds)
-        inCarMode = currentConfig.uiMode and Configuration.UI_MODE_TYPE_MASK ==
+        inCarMode =
+            currentConfig.uiMode and Configuration.UI_MODE_TYPE_MASK ==
                 Configuration.UI_MODE_TYPE_CAR
         uiMode = currentConfig.uiMode and Configuration.UI_MODE_NIGHT_MASK
         localeList = currentConfig.locales
@@ -60,29 +60,20 @@
 
     override fun notifyThemeChanged() {
         // Avoid concurrent modification exception
-        val listeners = synchronized(this.listeners) {
-           ArrayList(this.listeners)
-        }
+        val listeners = synchronized(this.listeners) { ArrayList(this.listeners) }
 
-        listeners.filterForEach({ this.listeners.contains(it) }) {
-            it.onThemeChanged()
-        }
+        listeners.filterForEach({ this.listeners.contains(it) }) { it.onThemeChanged() }
     }
 
     override fun onConfigurationChanged(newConfig: Configuration) {
         // Avoid concurrent modification exception
-        val listeners = synchronized(this.listeners) {
-           ArrayList(this.listeners)
-        }
-        listeners.filterForEach({ this.listeners.contains(it) }) {
-            it.onConfigChanged(newConfig)
-        }
+        val listeners = synchronized(this.listeners) { ArrayList(this.listeners) }
+        listeners.filterForEach({ this.listeners.contains(it) }) { it.onConfigChanged(newConfig) }
         val fontScale = newConfig.fontScale
         val density = newConfig.densityDpi
         val uiMode = newConfig.uiMode and Configuration.UI_MODE_NIGHT_MASK
         val uiModeChanged = uiMode != this.uiMode
-        if (density != this.density || fontScale != this.fontScale ||
-                inCarMode && uiModeChanged) {
+        if (density != this.density || fontScale != this.fontScale || inCarMode && uiModeChanged) {
             listeners.filterForEach({ this.listeners.contains(it) }) {
                 it.onDensityOrFontScaleChanged()
             }
@@ -105,17 +96,13 @@
             // would be a direct reference to windowConfiguration.maxBounds, so the if statement
             // above would always fail. See b/245799099 for more information.
             this.maxBounds.set(maxBounds)
-            listeners.filterForEach({ this.listeners.contains(it) }) {
-                it.onMaxBoundsChanged()
-            }
+            listeners.filterForEach({ this.listeners.contains(it) }) { it.onMaxBoundsChanged() }
         }
 
         val localeList = newConfig.locales
         if (localeList != this.localeList) {
             this.localeList = localeList
-            listeners.filterForEach({ this.listeners.contains(it) }) {
-                it.onLocaleListChanged()
-            }
+            listeners.filterForEach({ this.listeners.contains(it) }) { it.onLocaleListChanged() }
         }
 
         if (uiModeChanged) {
@@ -124,9 +111,7 @@
             context.theme.applyStyle(context.themeResId, true)
 
             this.uiMode = uiMode
-            listeners.filterForEach({ this.listeners.contains(it) }) {
-                it.onUiModeChanged()
-            }
+            listeners.filterForEach({ this.listeners.contains(it) }) { it.onUiModeChanged() }
         }
 
         if (layoutDirection != newConfig.layoutDirection) {
@@ -137,9 +122,7 @@
         }
 
         if (lastConfig.updateFrom(newConfig) and ActivityInfo.CONFIG_ASSETS_PATHS != 0) {
-            listeners.filterForEach({ this.listeners.contains(it) }) {
-                it.onThemeChanged()
-            }
+            listeners.filterForEach({ this.listeners.contains(it) }) { it.onThemeChanged() }
         }
 
         val newOrientation = newConfig.orientation
@@ -152,16 +135,12 @@
     }
 
     override fun addCallback(listener: ConfigurationListener) {
-        synchronized(listeners) {
-            listeners.add(listener)
-        }
+        synchronized(listeners) { listeners.add(listener) }
         listener.onDensityOrFontScaleChanged()
     }
 
     override fun removeCallback(listener: ConfigurationListener) {
-        synchronized(listeners) {
-            listeners.remove(listener)
-        }
+        synchronized(listeners) { listeners.remove(listener) }
     }
 
     override fun isLayoutRtl(): Boolean {
@@ -176,6 +155,15 @@
             else -> "err"
         }
     }
+
+    @AssistedFactory
+    interface Factory {
+        /**
+         * Creates a [ConfigurationController] that uses [context] to resolve the current
+         * configuration and resources.
+         */
+        fun create(context: Context): ConfigurationControllerImpl
+    }
 }
 
 // This could be done with a Collection.filter and Collection.forEach, but Collection.filter
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerStartable.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerStartable.kt
index 90ebaf2..8f4279e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerStartable.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ConfigurationControllerStartable.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.statusbar.phone
 
 import com.android.systemui.CoreStartable
+import com.android.systemui.common.ui.GlobalConfig
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.statusbar.policy.ConfigurationController
 import com.android.systemui.statusbar.policy.ConfigurationController.ConfigurationListener
@@ -26,7 +27,7 @@
 class ConfigurationControllerStartable
 @Inject
 constructor(
-    private val configurationController: ConfigurationController,
+    @GlobalConfig private val configurationController: ConfigurationController,
     private val listeners: Set<@JvmSuppressWildcards ConfigurationListener>
 ) : CoreStartable {
     override fun start() {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
index 0ea28a7..1efad3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpManagerPhone.java
@@ -47,7 +47,6 @@
 import com.android.systemui.statusbar.notification.data.repository.HeadsUpRowRepository;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.policy.AccessibilityManagerWrapper;
 import com.android.systemui.statusbar.policy.AnimationStateHandler;
 import com.android.systemui.statusbar.policy.AvalancheController;
@@ -284,7 +283,7 @@
     private void onShadeOrQsExpanded(Boolean isExpanded) {
         if (isExpanded != mIsExpanded) {
             mIsExpanded = isExpanded;
-            if (!NotificationsHeadsUpRefactor.isEnabled() && isExpanded) {
+            if (!SceneContainerFlag.isEnabled() && isExpanded) {
                 mHeadsUpAnimatingAway.setValue(false);
             }
         }
@@ -517,7 +516,7 @@
 
     @Nullable
     private HeadsUpEntryPhone getTopHeadsUpEntryPhone() {
-        if (NotificationsHeadsUpRefactor.isEnabled()) {
+        if (SceneContainerFlag.isEnabled()) {
             return (HeadsUpEntryPhone) mTopHeadsUpRow.getValue();
         } else {
             return (HeadsUpEntryPhone) getTopHeadsUpEntry();
@@ -711,7 +710,7 @@
         }
 
         private NotificationEntry requireEntry() {
-            /* check if */ NotificationsHeadsUpRefactor.isUnexpectedlyInLegacyMode();
+            /* check if */ SceneContainerFlag.isUnexpectedlyInLegacyMode();
             return Objects.requireNonNull(mEntry);
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
index 56ea00c..7ef1e41 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
@@ -22,6 +22,7 @@
 
 import androidx.annotation.NonNull;
 
+import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.settings.UserTracker;
@@ -43,17 +44,20 @@
     private final UserManager mUserManager;
     private final UserTracker mUserTracker;
     private final LinkedList<UserInfo> mProfiles;
+    private final KeyguardUpdateMonitor mKeyguardUpdateMonitor;
 
     private boolean mListening;
     private int mCurrentUser;
 
     @Inject
     public ManagedProfileControllerImpl(Context context, @Main Executor mainExecutor,
-            UserTracker userTracker, UserManager userManager) {
+            UserTracker userTracker, UserManager userManager,
+            KeyguardUpdateMonitor keyguardUpdateMonitor) {
         mContext = context;
         mMainExecutor = mainExecutor;
         mUserManager = userManager;
         mUserTracker = userTracker;
+        mKeyguardUpdateMonitor = keyguardUpdateMonitor;
         mProfiles = new LinkedList<>();
     }
 
@@ -80,6 +84,7 @@
                     StatusBarManager statusBarManager = (StatusBarManager) mContext
                             .getSystemService(android.app.Service.STATUS_BAR_SERVICE);
                     statusBarManager.collapsePanels();
+                    mKeyguardUpdateMonitor.awakenFromDream();
                 }
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
index 05bd1a7..ba39c3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
@@ -44,7 +44,7 @@
 
 import androidx.lifecycle.Observer;
 
-import com.android.settingslib.notification.modes.ZenMode;
+import com.android.internal.statusbar.StatusBarIcon;
 import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.qualifiers.DisplayId;
@@ -80,6 +80,7 @@
 import com.android.systemui.statusbar.policy.UserInfoController;
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor;
+import com.android.systemui.statusbar.policy.domain.model.ZenModeInfo;
 import com.android.systemui.util.RingerModeTracker;
 import com.android.systemui.util.kotlin.JavaAdapter;
 import com.android.systemui.util.time.DateFormatUtil;
@@ -363,7 +364,7 @@
             // Note that we're not fully replacing ZenModeController with ZenModeInteractor, so
             // we listen for the extra event here but still add the ZMC callback.
             mJavaAdapter.alwaysCollectFlow(mZenModeInteractor.getMainActiveMode(),
-                    this::onActiveModeChanged);
+                    this::onMainActiveModeChanged);
         }
         mZenController.addCallback(mZenControllerCallback);
         if (!Flags.statusBarScreenSharingChips()) {
@@ -395,20 +396,23 @@
                 () -> mResources.getString(R.string.accessibility_managed_profile));
     }
 
-    private void onActiveModeChanged(@Nullable ZenMode mode) {
+    private void onMainActiveModeChanged(@Nullable ZenModeInfo mainActiveMode) {
         if (!usesModeIcons()) {
-            Log.wtf(TAG, "onActiveModeChanged shouldn't be called if MODES_UI_ICONS is disabled");
+            Log.wtf(TAG, "onMainActiveModeChanged shouldn't run if MODES_UI_ICONS is disabled");
             return;
         }
-        boolean visible = mode != null;
-        if (visible) {
-            // TODO: b/360399800 - Get the resource id, package, and cached drawable from the mode;
-            //  this is a shortcut for testing.
-            String resPackage = mode.getIconKey().resPackage();
-            int iconResId = mode.getIconKey().resId();
 
-            mIconController.setResourceIcon(mSlotZen, resPackage, iconResId,
-                    /* preloadedIcon= */ null, mode.getName());
+        boolean visible = mainActiveMode != null;
+        if (visible) {
+            // Shape=FIXED_SPACE because mode icons can be from 3P packages and may not be square;
+            // we don't want to allow apps to set incredibly wide icons and take up too much space
+            // in the status bar.
+            mIconController.setResourceIcon(mSlotZen,
+                    mainActiveMode.getIcon().key().resPackage(),
+                    mainActiveMode.getIcon().key().resId(),
+                    mainActiveMode.getIcon().drawable(),
+                    mainActiveMode.getName(),
+                    StatusBarIcon.Shape.FIXED_SPACE);
         }
         if (visible != mZenVisible) {
             mIconController.setIconVisibility(mSlotZen, visible);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarHeadsUpChangeListener.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarHeadsUpChangeListener.java
index da5877b..8f2d4f9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarHeadsUpChangeListener.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarHeadsUpChangeListener.java
@@ -19,12 +19,12 @@
 import com.android.systemui.CoreStartable;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
+import com.android.systemui.scene.shared.flag.SceneContainerFlag;
 import com.android.systemui.shade.ShadeViewController;
 import com.android.systemui.shade.domain.interactor.PanelExpansionInteractor;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.NotificationShadeWindowController;
 import com.android.systemui.statusbar.StatusBarState;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController;
 import com.android.systemui.statusbar.policy.HeadsUpManager;
 import com.android.systemui.statusbar.policy.OnHeadsUpChangedListener;
@@ -112,7 +112,7 @@
     }
 
     private void setHeadsAnimatingAway(boolean headsUpAnimatingAway) {
-        if (!NotificationsHeadsUpRefactor.isEnabled()) {
+        if (!SceneContainerFlag.isEnabled()) {
             mHeadsUpManager.setHeadsUpAnimatingAway(headsUpAnimatingAway);
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
index f11fd7b..43f9af6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
@@ -68,11 +68,11 @@
 import com.android.systemui.deviceentry.shared.DeviceEntryUdfpsRefactor;
 import com.android.systemui.dock.DockManager;
 import com.android.systemui.dreams.DreamOverlayStateController;
+import com.android.systemui.keyguard.DismissCallbackRegistry;
 import com.android.systemui.keyguard.KeyguardWmStateRefactor;
+import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissActionInteractor;
-import com.android.systemui.keyguard.domain.interactor.KeyguardSurfaceBehindInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
-import com.android.systemui.keyguard.domain.interactor.WindowManagerLockscreenVisibilityInteractor;
 import com.android.systemui.keyguard.shared.model.DismissAction;
 import com.android.systemui.keyguard.shared.model.Edge;
 import com.android.systemui.keyguard.shared.model.KeyguardDone;
@@ -170,6 +170,7 @@
     private final Lazy<ShadeController> mShadeController;
     private final Lazy<SceneInteractor> mSceneInteractorLazy;
     private final Lazy<DeviceEntryInteractor> mDeviceEntryInteractorLazy;
+    private final DismissCallbackRegistry mDismissCallbackRegistry;
 
     private Job mListenForAlternateBouncerTransitionSteps = null;
     private Job mListenForKeyguardAuthenticatedBiometricsHandled = null;
@@ -360,8 +361,6 @@
             }
         }
     };
-    private Lazy<WindowManagerLockscreenVisibilityInteractor> mWmLockscreenVisibilityInteractor;
-    private Lazy<KeyguardSurfaceBehindInteractor> mSurfaceBehindInteractor;
     private Lazy<KeyguardDismissActionInteractor> mKeyguardDismissActionInteractor;
     private final JavaAdapter mJavaAdapter;
     private StatusBarKeyguardViewManagerInteractor mStatusBarKeyguardViewManagerInteractor;
@@ -391,16 +390,16 @@
             UdfpsOverlayInteractor udfpsOverlayInteractor,
             ActivityStarter activityStarter,
             KeyguardTransitionInteractor keyguardTransitionInteractor,
+            KeyguardDismissTransitionInteractor keyguardDismissTransitionInteractor,
             @Main CoroutineDispatcher mainDispatcher,
-            Lazy<WindowManagerLockscreenVisibilityInteractor> wmLockscreenVisibilityInteractor,
             Lazy<KeyguardDismissActionInteractor> keyguardDismissActionInteractorLazy,
             SelectedUserInteractor selectedUserInteractor,
-            Lazy<KeyguardSurfaceBehindInteractor> surfaceBehindInteractor,
             JavaAdapter javaAdapter,
             Lazy<SceneInteractor> sceneInteractorLazy,
             StatusBarKeyguardViewManagerInteractor statusBarKeyguardViewManagerInteractor,
             @Main DelayableExecutor executor,
-            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy
+            Lazy<DeviceEntryInteractor> deviceEntryInteractorLazy,
+            DismissCallbackRegistry dismissCallbackRegistry
     ) {
         mContext = context;
         mExecutor = executor;
@@ -428,18 +427,19 @@
         mUdfpsOverlayInteractor = udfpsOverlayInteractor;
         mActivityStarter = activityStarter;
         mKeyguardTransitionInteractor = keyguardTransitionInteractor;
+        mKeyguardDismissTransitionInteractor = keyguardDismissTransitionInteractor;
         mMainDispatcher = mainDispatcher;
-        mWmLockscreenVisibilityInteractor = wmLockscreenVisibilityInteractor;
         mKeyguardDismissActionInteractor = keyguardDismissActionInteractorLazy;
         mSelectedUserInteractor = selectedUserInteractor;
-        mSurfaceBehindInteractor = surfaceBehindInteractor;
         mJavaAdapter = javaAdapter;
         mSceneInteractorLazy = sceneInteractorLazy;
         mStatusBarKeyguardViewManagerInteractor = statusBarKeyguardViewManagerInteractor;
         mDeviceEntryInteractorLazy = deviceEntryInteractorLazy;
+        mDismissCallbackRegistry = dismissCallbackRegistry;
     }
 
     KeyguardTransitionInteractor mKeyguardTransitionInteractor;
+    KeyguardDismissTransitionInteractor mKeyguardDismissTransitionInteractor;
     CoroutineDispatcher mMainDispatcher;
 
     @Override
@@ -994,6 +994,8 @@
             }
             if (!SceneContainerFlag.isEnabled() && hideBouncerWhenShowing) {
                 hideAlternateBouncer(true);
+                mDismissCallbackRegistry.notifyDismissCancelled();
+                mPrimaryBouncerInteractor.setDismissAction(null, null);
             }
             mKeyguardUpdateManager.sendKeyguardReset();
             updateStates();
@@ -1609,7 +1611,7 @@
         }
 
         if (KeyguardWmStateRefactor.isEnabled()) {
-            mKeyguardTransitionInteractor.startDismissKeyguardTransition(
+            mKeyguardDismissTransitionInteractor.startDismissKeyguardTransition(
                     "SBKVM#keyguardAuthenticated");
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicy.java
index ba59398..d5fafe2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicy.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.statusbar.phone;
 
+import static com.android.systemui.Flags.statusBarSignalPolicyRefactor;
+
 import android.annotation.NonNull;
 import android.content.Context;
 import android.os.Handler;
@@ -23,16 +25,19 @@
 import android.util.Log;
 
 import com.android.settingslib.mobile.TelephonyIcons;
+import com.android.systemui.CoreStartable;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.res.R;
 import com.android.systemui.statusbar.connectivity.IconState;
 import com.android.systemui.statusbar.connectivity.NetworkController;
 import com.android.systemui.statusbar.connectivity.SignalCallback;
 import com.android.systemui.statusbar.phone.ui.StatusBarIconController;
+import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor;
 import com.android.systemui.statusbar.policy.SecurityController;
 import com.android.systemui.tuner.TunerService;
 import com.android.systemui.tuner.TunerService.Tunable;
 import com.android.systemui.util.CarrierConfigTracker;
+import com.android.systemui.util.kotlin.JavaAdapter;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -40,10 +45,13 @@
 
 import javax.inject.Inject;
 
-/** Controls the signal policies for icons shown in the statusbar. **/
+/** Controls the signal policies for icons shown in the statusbar. */
 @SysUISingleton
-public class StatusBarSignalPolicy implements SignalCallback,
-        SecurityController.SecurityControllerCallback, Tunable {
+public class StatusBarSignalPolicy
+        implements SignalCallback,
+                SecurityController.SecurityControllerCallback,
+                Tunable,
+                CoreStartable {
     private static final String TAG = "StatusBarSignalPolicy";
     private static final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
@@ -61,16 +69,15 @@
     private final Handler mHandler = Handler.getMain();
     private final CarrierConfigTracker mCarrierConfigTracker;
     private final TunerService mTunerService;
+    private final JavaAdapter mJavaAdapter;
+    private final AirplaneModeInteractor mAirplaneModeInteractor;
 
     private boolean mHideAirplane;
     private boolean mHideMobile;
     private boolean mHideEthernet;
-    private boolean mActivityEnabled;
+    private final boolean mActivityEnabled;
 
-    // Track as little state as possible, and only for padding purposes
-    private boolean mIsAirplaneMode = false;
-
-    private ArrayList<CallIndicatorIconState> mCallIndicatorStates = new ArrayList<>();
+    private final ArrayList<CallIndicatorIconState> mCallIndicatorStates = new ArrayList<>();
     private boolean mInitialized;
 
     @Inject
@@ -80,15 +87,19 @@
             CarrierConfigTracker carrierConfigTracker,
             NetworkController networkController,
             SecurityController securityController,
-            TunerService tunerService
+            TunerService tunerService,
+            JavaAdapter javaAdapter,
+            AirplaneModeInteractor airplaneModeInteractor
     ) {
         mContext = context;
 
         mIconController = iconController;
         mCarrierConfigTracker = carrierConfigTracker;
+        mJavaAdapter = javaAdapter;
         mNetworkController = networkController;
         mSecurityController = securityController;
         mTunerService = tunerService;
+        mAirplaneModeInteractor = airplaneModeInteractor;
 
         mSlotAirplane = mContext.getString(com.android.internal.R.string.status_bar_airplane);
         mSlotMobile   = mContext.getString(com.android.internal.R.string.status_bar_mobile);
@@ -100,15 +111,35 @@
         mActivityEnabled = mContext.getResources().getBoolean(R.bool.config_showActivity);
     }
 
+    @Override
+    public void start() {
+        if (!statusBarSignalPolicyRefactor()) {
+            return;
+        }
+
+        mTunerService.addTunable(this, StatusBarIconController.ICON_HIDE_LIST);
+        mNetworkController.addCallback(this);
+        mSecurityController.addCallback(this);
+
+        mJavaAdapter.alwaysCollectFlow(
+                mAirplaneModeInteractor.isAirplaneMode(), this::updateAirplaneModeIcon);
+    }
+
     /** Call to initialize and register this class with the system. */
     public void init() {
-        if (mInitialized) {
+        if (mInitialized || statusBarSignalPolicyRefactor()) {
             return;
         }
         mInitialized = true;
         mTunerService.addTunable(this, StatusBarIconController.ICON_HIDE_LIST);
         mNetworkController.addCallback(this);
         mSecurityController.addCallback(this);
+
+        if (statusBarSignalPolicyRefactor()) {
+            mJavaAdapter.alwaysCollectFlow(
+                    mAirplaneModeInteractor.isAirplaneMode(),
+                    this::updateAirplaneModeIcon);
+        }
     }
 
     public void destroy() {
@@ -222,15 +253,19 @@
 
     @Override
     public void setIsAirplaneMode(IconState icon) {
+        if (statusBarSignalPolicyRefactor()) {
+            return;
+        }
+
         if (DEBUG) {
             Log.d(TAG, "setIsAirplaneMode: "
                     + "icon = " + (icon == null ? "" : icon.toString()));
         }
-        mIsAirplaneMode = icon.visible && !mHideAirplane;
+        boolean isAirplaneMode = icon.visible && !mHideAirplane;
         int resId = icon.icon;
         String description = icon.contentDescription;
 
-        if (mIsAirplaneMode && resId > 0) {
+        if (isAirplaneMode && resId > 0) {
             mIconController.setIcon(mSlotAirplane, resId, description);
             mIconController.setIconVisibility(mSlotAirplane, true);
         } else {
@@ -238,6 +273,21 @@
         }
     }
 
+    public void updateAirplaneModeIcon(boolean isAirplaneModeOn) {
+        if (StatusBarSignalPolicyRefactor.isUnexpectedlyInLegacyMode()) {
+            return;
+        }
+
+        boolean isAirplaneMode = isAirplaneModeOn && !mHideAirplane;
+        mIconController.setIconVisibility(mSlotAirplane, isAirplaneMode);
+        if (isAirplaneMode) {
+            mIconController.setIcon(
+                    mSlotAirplane,
+                    TelephonyIcons.FLIGHT_MODE_ICON,
+                    mContext.getString(R.string.accessibility_airplane_mode));
+        }
+    }
+
     /**
      * Stores the statusbar state for no Calling & SMS.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicyRefactor.kt
similarity index 84%
rename from packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt
rename to packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicyRefactor.kt
index 62641fe..0577f49 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/shared/NotificationsHeadsUpRefactor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarSignalPolicyRefactor.kt
@@ -14,17 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.notification.shared
+package com.android.systemui.statusbar.phone
 
 import com.android.systemui.Flags
 import com.android.systemui.flags.FlagToken
 import com.android.systemui.flags.RefactorFlagUtils
 
-/** Helper for reading or using the notifications heads up refactor flag state. */
+/** Helper for reading or using the status_bar_signal_policy_refactor flag state. */
 @Suppress("NOTHING_TO_INLINE")
-object NotificationsHeadsUpRefactor {
+object StatusBarSignalPolicyRefactor {
     /** The aconfig flag name */
-    const val FLAG_NAME = Flags.FLAG_NOTIFICATIONS_HEADS_UP_REFACTOR
+    const val FLAG_NAME = Flags.FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR
 
     /** A token used for dependency declaration */
     val token: FlagToken
@@ -33,7 +33,7 @@
     /** Is the refactor enabled */
     @JvmStatic
     inline val isEnabled
-        get() = Flags.notificationsHeadsUpRefactor()
+        get() = Flags.statusBarSignalPolicyRefactor()
 
     /**
      * Called to ensure code is only run when the flag is enabled. This protects users from the
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
index 8871dae..6c30330 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.phone.ui;
 
-import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
 import com.android.internal.statusbar.StatusBarIcon;
@@ -64,9 +63,8 @@
     }
 
     @Override
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(
-                ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(StatusBarIcon.Shape shape) {
+        LinearLayout.LayoutParams lp = super.onCreateLayoutParams(shape);
         lp.setMargins(mIconHorizontalMargin, 0, mIconHorizontalMargin, 0);
         return lp;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
index 5ad7376..91ead61 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
@@ -20,6 +20,7 @@
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_ICON;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_MOBILE_NEW;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_WIFI_NEW;
+import static com.android.systemui.statusbar.phone.ui.StatusBarIconControllerImpl.usesModeIcons;
 
 import android.annotation.Nullable;
 import android.content.Context;
@@ -27,9 +28,8 @@
 import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
-import androidx.annotation.VisibleForTesting;
-
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.systemui.demomode.DemoModeCommandReceiver;
 import com.android.systemui.statusbar.BaseStatusBarFrameLayout;
 import com.android.systemui.statusbar.StatusBarIconView;
@@ -155,12 +155,11 @@
         };
     }
 
-    @VisibleForTesting
     protected StatusBarIconView addIcon(int index, String slot, boolean blocked,
             StatusBarIcon icon) {
         StatusBarIconView view = onCreateStatusBarIconView(slot, blocked);
         view.set(icon);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(icon.shape));
         return view;
     }
 
@@ -174,7 +173,7 @@
             int index) {
         mBindableIcons.put(holder.getSlot(), holder);
         ModernStatusBarView view = holder.getInitializer().createAndBind(mContext);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
         if (mIsInDemoMode) {
             mDemoStatusIcons.addBindableIcon(holder);
         }
@@ -183,7 +182,7 @@
 
     protected StatusIconDisplayable addNewWifiIcon(int index, String slot) {
         ModernStatusBarWifiView view = onCreateModernStatusBarWifiView(slot);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             mDemoStatusIcons.addModernWifiView(mWifiViewModel);
@@ -199,7 +198,7 @@
             int subId
     ) {
         BaseStatusBarFrameLayout view = onCreateModernStatusBarMobileView(slot, subId);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             Context mobileContext = mMobileContextProvider
@@ -233,8 +232,12 @@
                 );
     }
 
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        return new LinearLayout.LayoutParams(ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(Shape shape) {
+        int width = usesModeIcons() && shape == StatusBarIcon.Shape.FIXED_SPACE
+                ? mIconSize
+                : ViewGroup.LayoutParams.WRAP_CONTENT;
+
+        return new LinearLayout.LayoutParams(width, mIconSize);
     }
 
     protected void destroy() {
@@ -256,6 +259,13 @@
     /** Called once an icon has been set. */
     public void onSetIcon(int viewIndex, StatusBarIcon icon) {
         StatusBarIconView view = (StatusBarIconView) mGroup.getChildAt(viewIndex);
+        if (usesModeIcons()) {
+            ViewGroup.LayoutParams current = view.getLayoutParams();
+            ViewGroup.LayoutParams desired = onCreateLayoutParams(icon.shape);
+            if (desired.width != current.width || desired.height != current.height) {
+                view.setLayoutParams(desired);
+            }
+        }
         view.set(icon);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
index ee528e9..0459b97 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
@@ -70,7 +70,8 @@
      * @param preloadedIcon optional drawable corresponding to {@code iconResId}, if known
      */
     void setResourceIcon(String slot, @Nullable String resPackage, @DrawableRes int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription);
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape);
 
     /**
      * Sets up a wifi icon using the new data pipeline. No effect if the wifi icon has already been
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
index ad3a9e3..9b6d32b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
@@ -234,13 +234,14 @@
                 Icon.createWithResource(mContext, resourceId),
                 /* preloadedIcon= */ null,
                 contentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.WRAP_CONTENT);
     }
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage,
             @DrawableRes int iconResId, @Nullable Drawable preloadedIcon,
-            CharSequence contentDescription) {
+            CharSequence contentDescription, StatusBarIcon.Shape shape) {
         if (!usesModeIcons()) {
             Log.wtf("TAG",
                     "StatusBarIconController.setResourceIcon() should not be called without "
@@ -260,12 +261,13 @@
                 icon,
                 preloadedIcon,
                 contentDescription,
-                StatusBarIcon.Type.ResourceIcon);
+                StatusBarIcon.Type.ResourceIcon,
+                shape);
     }
 
     private void setResourceIconInternal(String slot, Icon resourceIcon,
             @Nullable Drawable preloadedIcon, CharSequence contentDescription,
-            StatusBarIcon.Type type) {
+            StatusBarIcon.Type type, StatusBarIcon.Shape shape) {
         checkArgument(resourceIcon.getType() == Icon.TYPE_RESOURCE,
                 "Expected Icon of TYPE_RESOURCE, but got " + resourceIcon.getType());
         String resPackage = resourceIcon.getResPackage();
@@ -277,7 +279,7 @@
         if (holder == null) {
             StatusBarIcon icon = new StatusBarIcon(UserHandle.SYSTEM, resPackage,
                     resourceIcon, /* iconLevel= */ 0, /* number=*/ 0,
-                    contentDescription, type);
+                    contentDescription, type, shape);
             icon.preloadedIcon = preloadedIcon;
             holder = StatusBarIconHolder.fromIcon(icon);
             setIcon(slot, holder);
@@ -286,6 +288,7 @@
             holder.getIcon().icon = resourceIcon;
             holder.getIcon().contentDescription = contentDescription;
             holder.getIcon().type = type;
+            holder.getIcon().shape = shape;
             holder.getIcon().preloadedIcon = preloadedIcon;
             handleSet(slot, holder);
         }
@@ -578,7 +581,7 @@
         }
     }
 
-    private static boolean usesModeIcons() {
+    static boolean usesModeIcons() {
         return android.app.Flags.modesApi() && android.app.Flags.modesUi()
                 && android.app.Flags.modesUiIcons();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/ServiceStateModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/ServiceStateModel.kt
deleted file mode 100644
index cce3eb0..0000000
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/ServiceStateModel.kt
+++ /dev/null
@@ -1,31 +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.statusbar.pipeline.mobile.data.model
-
-import android.telephony.ServiceState
-
-/**
- * Simplified representation of a [ServiceState] for use in SystemUI. Add any fields that we need to
- * extract from service state here for consumption downstream
- */
-data class ServiceStateModel(val isEmergencyOnly: Boolean) {
-    companion object {
-        fun fromServiceState(serviceState: ServiceState): ServiceStateModel {
-            return ServiceStateModel(isEmergencyOnly = serviceState.isEmergencyOnly)
-        }
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
index 5ad8bf1..32e9c85 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
@@ -21,7 +21,6 @@
 import com.android.settingslib.SignalIcon.MobileIconGroup
 import com.android.settingslib.mobile.MobileMappings
 import com.android.settingslib.mobile.MobileMappings.Config
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
@@ -93,17 +92,15 @@
     val defaultMobileIconGroup: Flow<MobileIconGroup>
 
     /**
-     * [deviceServiceState] is equivalent to the last [Intent.ACTION_SERVICE_STATE] broadcast with a
-     * subscriptionId of -1 (aka [SubscriptionManager.INVALID_SUBSCRIPTION_ID]).
+     * Can the device make emergency calls using the device-based service state? This field is only
+     * useful when all known active subscriptions are OOS and not emergency call capable.
      *
-     * While each [MobileConnectionsRepository] listens for the service state of each subscription,
-     * there is potentially a service state associated with the device itself. This value can be
-     * used to calculate e.g., the emergency calling capability of the device (as opposed to the
-     * emergency calling capability of an individual mobile connection)
+     * Specifically, this checks every [ServiceState] of the device, and looks for any that report
+     * [ServiceState.isEmergencyOnly].
      *
-     * Note: this is a [StateFlow] using an eager sharing strategy.
+     * This is an eager flow, and re-evaluates whenever ACTION_SERVICE_STATE is sent for subId = -1.
      */
-    val deviceServiceState: StateFlow<ServiceStateModel?>
+    val isDeviceEmergencyCallCapable: StateFlow<Boolean>
 
     /**
      * If any active SIM on the device is in
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
index b068152..b247da4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
@@ -25,7 +25,6 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.demomode.DemoMode
 import com.android.systemui.demomode.DemoModeController
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileConnectionsRepositoryImpl
@@ -152,16 +151,17 @@
     override val defaultMobileIconGroup: Flow<SignalIcon.MobileIconGroup> =
         activeRepo.flatMapLatest { it.defaultMobileIconGroup }
 
-    override val deviceServiceState: StateFlow<ServiceStateModel?> =
+    override val isDeviceEmergencyCallCapable: StateFlow<Boolean> =
         activeRepo
-            .flatMapLatest { it.deviceServiceState }
+            .flatMapLatest { it.isDeviceEmergencyCallCapable }
             .stateIn(
                 scope,
                 SharingStarted.WhileSubscribed(),
-                realRepository.deviceServiceState.value
+                realRepository.isDeviceEmergencyCallCapable.value
             )
 
     override val isAnySimSecure: Flow<Boolean> = activeRepo.flatMapLatest { it.isAnySimSecure }
+
     override fun getIsAnySimSecure(): Boolean = activeRepo.value.getIsAnySimSecure()
 
     override val defaultDataSubId: StateFlow<Int> =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
index a944e91..3a79f3f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
@@ -27,7 +27,6 @@
 import com.android.systemui.log.table.TableLogBufferFactory
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType.DefaultNetworkType
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
@@ -137,10 +136,11 @@
 
     override val defaultMobileIconGroup = flowOf(TelephonyIcons.THREE_G)
 
-    // TODO(b/339023069): demo command for device-based connectivity state
-    override val deviceServiceState: StateFlow<ServiceStateModel?> = MutableStateFlow(null)
+    // TODO(b/339023069): demo command for device-based emergency calls state
+    override val isDeviceEmergencyCallCapable: StateFlow<Boolean> = MutableStateFlow(false)
 
     override val isAnySimSecure: Flow<Boolean> = flowOf(getIsAnySimSecure())
+
     override fun getIsAnySimSecure(): Boolean = false
 
     override val defaultMobileIconMapping = MutableStateFlow(TelephonyIcons.ICON_NAME_TO_ICON)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
index 261258a..b756a05 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
@@ -21,7 +21,6 @@
 import android.content.Intent
 import android.content.IntentFilter
 import android.telephony.CarrierConfigManager
-import android.telephony.ServiceState
 import android.telephony.SubscriptionInfo
 import android.telephony.SubscriptionManager
 import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
@@ -49,7 +48,6 @@
 import com.android.systemui.statusbar.pipeline.dagger.MobileSummaryLog
 import com.android.systemui.statusbar.pipeline.mobile.data.MobileInputLogger
 import com.android.systemui.statusbar.pipeline.mobile.data.model.NetworkNameModel
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.util.MobileMappingsProxy
@@ -72,7 +70,6 @@
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -175,8 +172,8 @@
             }
             .flowOn(bgDispatcher)
 
-    /** Note that this flow is eager, so we don't miss any state */
-    override val deviceServiceState: StateFlow<ServiceStateModel?> =
+    /** Turn ACTION_SERVICE_STATE (for subId = -1) into an event */
+    private val serviceStateChangedEvent: Flow<Unit> =
         broadcastDispatcher
             .broadcastFlow(IntentFilter(Intent.ACTION_SERVICE_STATE)) { intent, _ ->
                 val subId =
@@ -185,24 +182,34 @@
                         INVALID_SUBSCRIPTION_ID
                     )
 
-                val extras = intent.extras
-                if (extras == null) {
-                    logger.logTopLevelServiceStateBroadcastMissingExtras(subId)
-                    return@broadcastFlow null
-                }
-
-                val serviceState = ServiceState.newFromBundle(extras)
-                logger.logTopLevelServiceStateBroadcastEmergencyOnly(subId, serviceState)
+                // Only emit if the subId is not associated with an active subscription
                 if (subId == INVALID_SUBSCRIPTION_ID) {
-                    // Assume that -1 here is the device's service state. We don't care about
-                    // other ones.
-                    ServiceStateModel.fromServiceState(serviceState)
-                } else {
-                    null
+                    Unit
                 }
             }
-            .filterNotNull()
-            .stateIn(scope, SharingStarted.Eagerly, null)
+            // Emit on start so that we always check the state at least once
+            .onStart { emit(Unit) }
+
+    /** Eager flow to determine the device-based emergency calls only state */
+    override val isDeviceEmergencyCallCapable: StateFlow<Boolean> =
+        serviceStateChangedEvent
+            .mapLatest {
+                val modems = telephonyManager.activeModemCount
+                // Check the service state for every modem. If any state reports emergency calling
+                // capable, then consider the device to have emergency call capabilities
+                (0..<modems)
+                    .map { telephonyManager.getServiceStateForSlot(it) }
+                    .any { it?.isEmergencyOnly == true }
+            }
+            .flowOn(bgDispatcher)
+            .distinctUntilChanged()
+            .logDiffsForTable(
+                tableLogger,
+                columnPrefix = LOGGING_PREFIX,
+                columnName = "deviceEmergencyOnly",
+                initialValue = false,
+            )
+            .stateIn(scope, SharingStarted.Eagerly, false)
 
     /**
      * State flow that emits the set of mobile data subscriptions, each represented by its own
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractor.kt
index 26553e6..28fff4e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractor.kt
@@ -385,15 +385,7 @@
             .stateIn(scope, SharingStarted.WhileSubscribed(), false)
 
     override val isDeviceInEmergencyCallsOnlyMode: Flow<Boolean> =
-        mobileConnectionsRepo.deviceServiceState
-            .map { it?.isEmergencyOnly ?: false }
-            .distinctUntilChanged()
-            .logDiffsForTable(
-                tableLogger,
-                columnPrefix = LOGGING_PREFIX,
-                columnName = "deviceEmergencyOnly",
-                initialValue = false,
-            )
+        mobileConnectionsRepo.isDeviceEmergencyCallCapable
 
     /** Vends out new [MobileIconInteractor] for a particular subId */
     override fun getMobileConnectionInteractorForSubId(subId: Int): MobileIconInteractor =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepository.kt
index 1a55f7d..f5cfc8c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepository.kt
@@ -31,7 +31,7 @@
 import androidx.annotation.ArrayRes
 import androidx.annotation.VisibleForTesting
 import com.android.systemui.Dumpable
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.Flags
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dump.DumpManager
@@ -47,6 +47,7 @@
 import com.android.systemui.statusbar.pipeline.shared.data.model.DefaultConnectionModel.Wifi
 import com.android.systemui.statusbar.pipeline.shared.data.repository.ConnectivityRepositoryImpl.Companion.getMainOrUnderlyingWifiInfo
 import com.android.systemui.tuner.TunerService
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import java.io.PrintWriter
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -252,7 +253,10 @@
             }
             // Only CELLULAR networks may have underlying wifi information that's relevant to SysUI,
             // so skip the underlying network check if it's not CELLULAR.
-            if (!this.hasTransport(TRANSPORT_CELLULAR)) {
+            if (
+                !this.hasTransport(TRANSPORT_CELLULAR) &&
+                    !Flags.statusBarAlwaysCheckUnderlyingNetworks()
+            ) {
                 return mainWifiInfo
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/binder/CollapsedStatusBarViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/binder/CollapsedStatusBarViewBinder.kt
index d46aaf4..c24d694 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/binder/CollapsedStatusBarViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/shared/ui/binder/CollapsedStatusBarViewBinder.kt
@@ -35,6 +35,7 @@
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.chips.ron.shared.StatusBarRonChips
 import com.android.systemui.statusbar.chips.ui.binder.ChipChronometerBinder
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
 import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
@@ -124,10 +125,6 @@
 
                                     // Colors
                                     val textColor = chipModel.colors.text(chipContext)
-                                    chipDefaultIconView.imageTintList =
-                                        ColorStateList.valueOf(textColor)
-                                    chipBackgroundView.getCustomIconView()?.imageTintList =
-                                        ColorStateList.valueOf(textColor)
                                     chipTimeView.setTextColor(textColor)
                                     chipTextView.setTextColor(textColor)
                                     (chipBackgroundView.background as GradientDrawable).color =
@@ -173,13 +170,22 @@
         // it.
         backgroundView.removeView(backgroundView.getCustomIconView())
 
+        val iconTint = chipModel.colors.text(defaultIconView.context)
+
         when (val icon = chipModel.icon) {
             null -> {
                 defaultIconView.visibility = View.GONE
             }
-            is OngoingActivityChipModel.ChipIcon.Basic -> {
+            is OngoingActivityChipModel.ChipIcon.SingleColorIcon -> {
                 IconViewBinder.bind(icon.impl, defaultIconView)
                 defaultIconView.visibility = View.VISIBLE
+                defaultIconView.tintView(iconTint)
+            }
+            is OngoingActivityChipModel.ChipIcon.FullColorAppIcon -> {
+                StatusBarRonChips.assertInNewMode()
+                IconViewBinder.bind(icon.impl, defaultIconView)
+                defaultIconView.visibility = View.VISIBLE
+                defaultIconView.untintView()
             }
             is OngoingActivityChipModel.ChipIcon.StatusBarView -> {
                 // Hide the default icon since we'll show this custom icon instead.
@@ -194,6 +200,7 @@
                     // maybe include the app name.
                     contentDescription =
                         context.resources.getString(R.string.ongoing_phone_call_content_description)
+                    tintView(iconTint)
                 }
 
                 // 2. If we just reinflated the view, we may need to detach the icon view from the
@@ -219,6 +226,14 @@
         return this.findViewById(CUSTOM_ICON_VIEW_ID)
     }
 
+    private fun ImageView.tintView(color: Int) {
+        this.imageTintList = ColorStateList.valueOf(color)
+    }
+
+    private fun ImageView.untintView() {
+        this.imageTintList = null
+    }
+
     private fun generateCustomIconLayoutParams(iconView: ImageView): FrameLayout.LayoutParams {
         val customIconSize =
             iconView.context.resources.getDimensionPixelSize(
@@ -237,10 +252,13 @@
                 chipTextView.text = chipModel.secondsUntilStarted.toString()
                 chipTextView.visibility = View.VISIBLE
 
-                // The Chronometer should be stopped to prevent leaks -- see b/192243808 and
-                // [Chronometer.start].
-                chipTimeView.stop()
-                chipTimeView.visibility = View.GONE
+                chipTimeView.hide()
+            }
+            is OngoingActivityChipModel.Shown.Text -> {
+                chipTextView.text = chipModel.text
+                chipTextView.visibility = View.VISIBLE
+
+                chipTimeView.hide()
             }
             is OngoingActivityChipModel.Shown.Timer -> {
                 ChipChronometerBinder.bind(chipModel.startTimeMs, chipTimeView)
@@ -250,14 +268,18 @@
             }
             is OngoingActivityChipModel.Shown.IconOnly -> {
                 chipTextView.visibility = View.GONE
-                // The Chronometer should be stopped to prevent leaks -- see b/192243808 and
-                // [Chronometer.start].
-                chipTimeView.stop()
-                chipTimeView.visibility = View.GONE
+                chipTimeView.hide()
             }
         }
     }
 
+    private fun ChipChronometer.hide() {
+        // The Chronometer should be stopped to prevent leaks -- see b/192243808 and
+        // [Chronometer.start].
+        this.stop()
+        this.visibility = View.GONE
+    }
+
     private fun updateChipPadding(
         chipModel: OngoingActivityChipModel.Shown,
         backgroundView: View,
@@ -356,6 +378,7 @@
                 chipView.accessibilityLiveRegion = View.ACCESSIBILITY_LIVE_REGION_ASSERTIVE
             }
             is OngoingActivityChipModel.Shown.Timer,
+            is OngoingActivityChipModel.Shown.Text,
             is OngoingActivityChipModel.Shown.IconOnly -> {
                 chipView.accessibilityLiveRegion = View.ACCESSIBILITY_LIVE_REGION_NONE
             }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
index 21ec14f..591d7af 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/PolicyModule.kt
@@ -403,7 +403,7 @@
                     tileSpec = TileSpec.create(DND_TILE_SPEC),
                     uiConfig =
                         QSTileUIConfig.Resource(
-                            iconRes = R.drawable.qs_dnd_icon_off,
+                            iconRes = com.android.internal.R.drawable.ic_zen_priority_modes,
                             labelRes = R.string.quick_settings_modes_label,
                         ),
                     instanceId = uiEventLogger.getNewInstanceId(),
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
index 71bcdfcb..b81af86 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/StatusBarPolicyModule.java
@@ -22,8 +22,12 @@
 
 import com.android.internal.R;
 import com.android.settingslib.devicestate.DeviceStateRotationLockSettingsManager;
+import com.android.settingslib.notification.modes.ZenIconLoader;
+import com.android.systemui.common.ui.GlobalConfig;
 import com.android.systemui.dagger.SysUISingleton;
+import com.android.systemui.dagger.qualifiers.Application;
 import com.android.systemui.dagger.qualifiers.Main;
+import com.android.systemui.dagger.qualifiers.UiBackground;
 import com.android.systemui.log.LogBuffer;
 import com.android.systemui.log.LogBufferFactory;
 import com.android.systemui.settings.UserTracker;
@@ -79,6 +83,7 @@
 import dagger.Provides;
 
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 
 import javax.inject.Named;
 
@@ -100,9 +105,12 @@
     @Binds
     CastController provideCastController(CastControllerImpl controllerImpl);
 
-    /** */
+    /**
+     * @deprecated: unscoped configuration controller shouldn't be injected as it might lead to
+     * wrong updates in case of secondary displays.
+     */
     @Binds
-    ConfigurationController bindConfigurationController(ConfigurationControllerImpl impl);
+    ConfigurationController bindConfigurationController(@GlobalConfig ConfigurationController impl);
 
     /** */
     @Binds
@@ -178,6 +186,15 @@
             DevicePostureControllerImpl devicePostureControllerImpl);
 
     /** */
+    @Provides
+    @SysUISingleton
+    @GlobalConfig
+    static ConfigurationController provideGlobalConfigurationController(
+            @Application Context context, ConfigurationControllerImpl.Factory factory) {
+        return factory.create(context);
+    }
+
+    /** */
     @SysUISingleton
     @Provides
     static AccessPointControllerImpl  provideAccessPointControllerImpl(
@@ -236,4 +253,12 @@
     static LogBuffer provideCastControllerLog(LogBufferFactory factory) {
         return factory.create("CastControllerLog", 50);
     }
+
+    /** Provides a {@link ZenIconLoader} that fetches icons in a background thread. */
+    @Provides
+    @SysUISingleton
+    static ZenIconLoader provideZenIconLoader(
+            @UiBackground ExecutorService backgroundExecutorService) {
+        return new ZenIconLoader(backgroundExecutorService);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
index a67b47a..dbeaa59 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
@@ -23,16 +23,20 @@
 import android.util.Log
 import androidx.concurrent.futures.await
 import com.android.settingslib.notification.data.repository.ZenModeRepository
+import com.android.settingslib.notification.modes.ZenIcon
 import com.android.settingslib.notification.modes.ZenIconLoader
 import com.android.settingslib.notification.modes.ZenMode
-import com.android.systemui.common.shared.model.Icon
-import com.android.systemui.common.shared.model.asIcon
+import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
+import com.android.systemui.statusbar.policy.domain.model.ActiveZenModes
+import com.android.systemui.statusbar.policy.domain.model.ZenModeInfo
 import java.time.Duration
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 
 /**
@@ -45,9 +49,9 @@
     private val context: Context,
     private val zenModeRepository: ZenModeRepository,
     private val notificationSettingsRepository: NotificationSettingsRepository,
+    @Background private val bgDispatcher: CoroutineDispatcher,
+    private val iconLoader: ZenIconLoader,
 ) {
-    private val iconLoader: ZenIconLoader = ZenIconLoader.getInstance()
-
     val isZenModeEnabled: Flow<Boolean> =
         zenModeRepository.globalZenMode
             .map {
@@ -76,34 +80,29 @@
 
     val modes: Flow<List<ZenMode>> = zenModeRepository.modes
 
-    val activeModes: Flow<List<ZenMode>> =
-        modes.map { modes -> modes.filter { mode -> mode.isActive } }.distinctUntilChanged()
+    /** Flow returning the currently active mode(s), if any. */
+    val activeModes: Flow<ActiveZenModes> =
+        modes
+            .map { modes -> buildActiveZenModes(modes) }
+            .flowOn(bgDispatcher)
+            .distinctUntilChanged()
 
-    /** Flow returning the most prioritized of the active modes, if any. */
-    val mainActiveMode: Flow<ZenMode?> =
-        activeModes.map { modes -> getMainActiveMode(modes) }.distinctUntilChanged()
+    suspend fun getActiveModes() = buildActiveZenModes(zenModeRepository.getModes())
 
-    /**
-     * Given the list of modes (which may include zero or more currently active modes), returns the
-     * most prioritized of the active modes, if any.
-     */
-    private fun getMainActiveMode(modes: List<ZenMode>): ZenMode? {
-        return modes.sortedWith(ZenMode.PRIORITIZING_COMPARATOR).firstOrNull { it.isActive }
+    private suspend fun buildActiveZenModes(modes: List<ZenMode>): ActiveZenModes {
+        val activeModesList =
+            modes.filter { mode -> mode.isActive }.sortedWith(ZenMode.PRIORITIZING_COMPARATOR)
+        val mainActiveMode =
+            activeModesList.firstOrNull()?.let { ZenModeInfo(it.name, getModeIcon(it)) }
+
+        return ActiveZenModes(activeModesList.map { m -> m.name }, mainActiveMode)
     }
 
-    suspend fun getModeIcon(mode: ZenMode): Icon {
-        return iconLoader.getIcon(context, mode).await().drawable().asIcon()
-    }
+    val mainActiveMode: Flow<ZenModeInfo?> =
+        activeModes.map { a -> a.mainMode }.distinctUntilChanged()
 
-    /**
-     * Given the list of modes (which may include zero or more currently active modes), returns an
-     * icon representing the active mode, if any (or, if multiple modes are active, to the most
-     * prioritized one). This icon is suitable for use in the status bar or lockscreen (uses the
-     * standard DND icon for implicit modes, instead of the launcher icon of the associated
-     * package).
-     */
-    suspend fun getActiveModeIcon(modes: List<ZenMode>): Icon? {
-        return getMainActiveMode(modes)?.let { m -> getModeIcon(m) }
+    suspend fun getModeIcon(mode: ZenMode): ZenIcon {
+        return iconLoader.getIcon(context, mode).await()
     }
 
     fun activateMode(zenMode: ZenMode) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ActiveZenModes.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ActiveZenModes.kt
new file mode 100644
index 0000000..569e517
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ActiveZenModes.kt
@@ -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 com.android.systemui.statusbar.policy.domain.model
+
+import com.android.settingslib.notification.modes.ZenMode
+
+/**
+ * Represents the list of [ZenMode] instances that are currently active.
+ *
+ * @property modeNames Names of all the active modes, sorted by their priority.
+ * @property mainMode The most prioritized active mode, if any modes active. Guaranteed to be
+ *   non-null if [modeNames] is not empty.
+ */
+data class ActiveZenModes(val modeNames: List<String>, val mainMode: ZenModeInfo?) {
+    fun isAnyActive(): Boolean = modeNames.isNotEmpty()
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ZenModeInfo.kt
similarity index 69%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
copy to packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ZenModeInfo.kt
index 60d97d1..5004f4c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/model/ZenModeInfo.kt
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.bouncer.shared.flag
+package com.android.systemui.statusbar.policy.domain.model
 
-import com.android.systemui.kosmos.Kosmos
+import com.android.settingslib.notification.modes.ZenIcon
+import com.android.settingslib.notification.modes.ZenMode
 
-var Kosmos.fakeComposeBouncerFlags by Kosmos.Fixture { FakeComposeBouncerFlags() }
-val Kosmos.composeBouncerFlags by Kosmos.Fixture<ComposeBouncerFlags> { fakeComposeBouncerFlags }
+/** Name and icon of a [ZenMode] */
+data class ZenModeInfo(val name: String, val icon: ZenIcon)
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/composable/ModeTile.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/composable/ModeTile.kt
index 3fffd9f..0e88f44 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/composable/ModeTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/composable/ModeTile.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.policy.ui.dialog.composable
 
+import androidx.compose.animation.animateColorAsState
 import androidx.compose.foundation.basicMarquee
 import androidx.compose.foundation.combinedClickable
 import androidx.compose.foundation.layout.Arrangement
@@ -30,9 +31,15 @@
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.getValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.platform.testTag
+import androidx.compose.ui.semantics.clearAndSetSemantics
+import androidx.compose.ui.semantics.contentDescription
+import androidx.compose.ui.semantics.semantics
+import androidx.compose.ui.semantics.stateDescription
 import androidx.compose.ui.text.font.FontWeight
 import androidx.compose.ui.unit.dp
 import com.android.systemui.common.ui.compose.Icon
@@ -40,12 +47,16 @@
 
 @Composable
 fun ModeTile(viewModel: ModeTileViewModel) {
-    val tileColor =
-        if (viewModel.enabled) MaterialTheme.colorScheme.primary
-        else MaterialTheme.colorScheme.surfaceVariant
-    val contentColor =
-        if (viewModel.enabled) MaterialTheme.colorScheme.onPrimary
-        else MaterialTheme.colorScheme.onSurfaceVariant
+    val tileColor: Color by
+        animateColorAsState(
+            if (viewModel.enabled) MaterialTheme.colorScheme.primary
+            else MaterialTheme.colorScheme.surfaceVariant
+        )
+    val contentColor: Color by
+        animateColorAsState(
+            if (viewModel.enabled) MaterialTheme.colorScheme.onPrimary
+            else MaterialTheme.colorScheme.onSurfaceVariant
+        )
 
     CompositionLocalProvider(LocalContentColor provides contentColor) {
         Surface(
@@ -56,9 +67,11 @@
                 modifier =
                     Modifier.combinedClickable(
                             onClick = viewModel.onClick,
-                            onLongClick = viewModel.onLongClick
+                            onLongClick = viewModel.onLongClick,
+                            onLongClickLabel = viewModel.onLongClickLabel
                         )
-                        .padding(20.dp),
+                        .padding(20.dp)
+                        .semantics { stateDescription = viewModel.stateDescription },
                 verticalAlignment = Alignment.CenterVertically,
                 horizontalArrangement =
                     Arrangement.spacedBy(
@@ -76,7 +89,10 @@
                     Text(
                         viewModel.subtext,
                         fontWeight = FontWeight.W400,
-                        modifier = Modifier.tileMarquee().testTag("state")
+                        modifier =
+                            Modifier.tileMarquee().testTag("state").clearAndSetSemantics {
+                                contentDescription = viewModel.subtextDescription
+                            }
                     )
                 }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
index 7c1cb6a..abd2453 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModeTileViewModel.kt
@@ -28,7 +28,10 @@
     val icon: Icon,
     val text: String,
     val subtext: String,
+    val subtextDescription: String, // version of subtext without "on"/"off" for screen readers
     val enabled: Boolean,
+    val stateDescription: String, // "on"/"off" state of the tile, for screen readers
     val onClick: () -> Unit,
     val onLongClick: () -> Unit,
+    val onLongClickLabel: String, // for screen readers
 )
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
index be90bec..6764839c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/ui/dialog/viewmodel/ModesDialogViewModel.kt
@@ -23,6 +23,7 @@
 import android.provider.Settings.EXTRA_AUTOMATIC_ZEN_RULE_ID
 import com.android.settingslib.notification.modes.EnableZenModeDialog
 import com.android.settingslib.notification.modes.ZenMode
+import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.qs.tiles.dialog.QSZenModeDialogMetricsLogger
@@ -88,10 +89,15 @@
                 modesList.map { mode ->
                     ModeTileViewModel(
                         id = mode.id,
-                        icon = zenModeInteractor.getModeIcon(mode),
+                        icon = zenModeInteractor.getModeIcon(mode).drawable().asIcon(),
                         text = mode.name,
                         subtext = getTileSubtext(mode),
+                        subtextDescription = getModeDescription(mode) ?: "",
                         enabled = mode.isActive,
+                        stateDescription =
+                            context.getString(
+                                if (mode.isActive) R.string.zen_mode_on else R.string.zen_mode_off
+                            ),
                         onClick = {
                             if (!mode.rule.isEnabled) {
                                 openSettings(mode)
@@ -112,7 +118,9 @@
                                 }
                             }
                         },
-                        onLongClick = { openSettings(mode) }
+                        onLongClick = { openSettings(mode) },
+                        onLongClickLabel =
+                            context.resources.getString(R.string.accessibility_long_click_tile)
                     )
                 }
             }
@@ -127,23 +135,36 @@
         dialogDelegate.launchFromDialog(intent)
     }
 
-    private fun getTileSubtext(mode: ZenMode): String {
+    /**
+     * Returns a description of the mode, which is:
+     *   * a prompt to set up the mode if it is not enabled
+     *   * if it cannot be manually activated, text that says so
+     *   * otherwise, the trigger description of the mode if it exists...
+     *   * ...or null if it doesn't
+     *
+     * This description is used directly for the content description of a mode tile for screen
+     * readers, and for the tile subtext will be augmented with the current status of the mode.
+     */
+    private fun getModeDescription(mode: ZenMode): String? {
         if (!mode.rule.isEnabled) {
             return context.resources.getString(R.string.zen_mode_set_up)
         }
         if (!mode.rule.isManualInvocationAllowed && !mode.isActive) {
             return context.resources.getString(R.string.zen_mode_no_manual_invocation)
         }
+        return mode.getDynamicDescription(context)
+    }
 
-        val modeSubtext = mode.getDynamicDescription(context)
+    private fun getTileSubtext(mode: ZenMode): String {
+        val modeDescription = getModeDescription(mode)
         return if (mode.isActive) {
-            if (modeSubtext != null) {
-                context.getString(R.string.zen_mode_on_with_details, modeSubtext)
+            if (modeDescription != null) {
+                context.getString(R.string.zen_mode_on_with_details, modeDescription)
             } else {
                 context.getString(R.string.zen_mode_on)
             }
         } else {
-            modeSubtext ?: context.getString(R.string.zen_mode_off)
+            modeDescription ?: context.getString(R.string.zen_mode_off)
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
index 89227cf..fe1d647 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModel.kt
@@ -21,10 +21,10 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.scene.domain.interactor.SceneInteractor
+import com.android.systemui.scene.shared.flag.SceneContainerFlag
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.domain.interactor.KeyguardStatusBarInteractor
 import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.policy.BatteryController
 import com.android.systemui.statusbar.policy.BatteryController.BatteryStateChangeCallback
 import javax.inject.Inject
@@ -58,7 +58,7 @@
 ) {
 
     private val showingHeadsUpStatusBar: Flow<Boolean> =
-        if (NotificationsHeadsUpRefactor.isEnabled) {
+        if (SceneContainerFlag.isEnabled) {
             headsUpNotificationInteractor.showHeadsUpStatusBar
         } else {
             flowOf(false)
diff --git a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
index c7fc445..9c8ef04 100644
--- a/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
+++ b/packages/SystemUI/src/com/android/systemui/theme/ThemeOverlayController.java
@@ -89,6 +89,9 @@
 import com.google.ux.material.libmonet.dynamiccolor.DynamicColor;
 import com.google.ux.material.libmonet.dynamiccolor.MaterialDynamicColors;
 
+import kotlinx.coroutines.flow.Flow;
+import kotlinx.coroutines.flow.StateFlow;
+
 import org.json.JSONException;
 import org.json.JSONObject;
 
@@ -162,6 +165,7 @@
     private final WakefulnessLifecycle mWakefulnessLifecycle;
     private final JavaAdapter mJavaAdapter;
     private final KeyguardTransitionInteractor mKeyguardTransitionInteractor;
+    private final StateFlow<Boolean> mIsKeyguardOnAsleepState;
     private final UiModeManager mUiModeManager;
     private ColorScheme mDarkColorScheme;
     private ColorScheme mLightColorScheme;
@@ -202,8 +206,7 @@
             }
             boolean currentUser = userId == mUserTracker.getUserId();
             boolean isAsleep = themeOverlayControllerWakefulnessDeprecation()
-                    ? KeyguardState.Companion.deviceIsAsleepInState(
-                            mKeyguardTransitionInteractor.getFinishedState())
+                    ? ThemeOverlayController.this.mIsKeyguardOnAsleepState.getValue()
                     : mWakefulnessLifecycle.getWakefulness() != WAKEFULNESS_ASLEEP;
 
             if (currentUser && !mAcceptColorEvents && isAsleep) {
@@ -434,6 +437,10 @@
         mUiModeManager = uiModeManager;
         mActivityManager = activityManager;
         dumpManager.registerDumpable(TAG, this);
+
+        Flow<Boolean> isFinishedInAsleepStateFlow = mKeyguardTransitionInteractor
+                .isFinishedInStateWhere(KeyguardState.Companion::deviceIsAsleepInState);
+        mIsKeyguardOnAsleepState = mJavaAdapter.stateInApp(isFinishedInAsleepStateFlow, false);
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/util/concurrency/GlobalConcurrencyModule.java b/packages/SystemUI/src/com/android/systemui/util/concurrency/GlobalConcurrencyModule.java
index ecf1165..70774f13 100644
--- a/packages/SystemUI/src/com/android/systemui/util/concurrency/GlobalConcurrencyModule.java
+++ b/packages/SystemUI/src/com/android/systemui/util/concurrency/GlobalConcurrencyModule.java
@@ -28,6 +28,7 @@
 import dagger.Provides;
 
 import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 
 import javax.inject.Singleton;
@@ -81,6 +82,18 @@
     @Singleton
     @UiBackground
     public static Executor provideUiBackgroundExecutor() {
+        return provideUiBackgroundExecutorService();
+    }
+
+    /**
+     * Provide an ExecutorService specifically for running UI operations on a separate thread.
+     *
+     * <p>Keep submitted runnables short and to the point, just as with any other UI code.
+     */
+    @Provides
+    @Singleton
+    @UiBackground
+    public static ExecutorService provideUiBackgroundExecutorService() {
         return Executors.newSingleThreadExecutor();
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/util/kotlin/FlowDumper.kt b/packages/SystemUI/src/com/android/systemui/util/kotlin/FlowDumper.kt
index 727e51f..315a89b 100644
--- a/packages/SystemUI/src/com/android/systemui/util/kotlin/FlowDumper.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/kotlin/FlowDumper.kt
@@ -20,7 +20,6 @@
 import com.android.systemui.Dumpable
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.lifecycle.ExclusiveActivatable
-import com.android.systemui.lifecycle.SysUiViewModel
 import com.android.systemui.util.asIndenting
 import com.android.systemui.util.printCollection
 import java.io.PrintWriter
diff --git a/packages/SystemUI/src/com/android/systemui/util/kotlin/JavaAdapter.kt b/packages/SystemUI/src/com/android/systemui/util/kotlin/JavaAdapter.kt
index 055671c..64e056d 100644
--- a/packages/SystemUI/src/com/android/systemui/util/kotlin/JavaAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/util/kotlin/JavaAdapter.kt
@@ -31,7 +31,10 @@
 import kotlinx.coroutines.DisposableHandle
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 
 /** A class allowing Java classes to collect on Kotlin flows. */
@@ -58,6 +61,15 @@
     ): Job {
         return scope.launch { flow.collect { consumer.accept(it) } }
     }
+
+    @JvmOverloads
+    fun <T> stateInApp(
+        flow: Flow<T>,
+        initialValue: T,
+        started: SharingStarted = SharingStarted.Eagerly
+    ): StateFlow<T> {
+        return flow.stateIn(scope, started, initialValue)
+    }
 }
 
 /**
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt b/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerAdapter.kt
similarity index 78%
rename from packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt
rename to packages/SystemUI/src/com/android/systemui/volume/VolumeControllerAdapter.kt
index 6859191..e836731 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerCollector.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeControllerAdapter.kt
@@ -17,7 +17,8 @@
 package com.android.systemui.volume
 
 import android.media.IVolumeController
-import com.android.settingslib.media.data.repository.VolumeControllerEvent
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
+import com.android.settingslib.volume.data.repository.AudioRepository
 import com.android.systemui.dagger.qualifiers.Application
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
@@ -29,17 +30,17 @@
  * [com.android.settingslib.volume.data.repository.AudioRepository.volumeControllerEvents] and the
  * old code that uses [IVolumeController] interface directly.
  */
-class VolumeControllerCollector
+class VolumeControllerAdapter
 @Inject
-constructor(@Application private val coroutineScope: CoroutineScope) {
+constructor(
+    @Application private val coroutineScope: CoroutineScope,
+    private val audioRepository: AudioRepository,
+) {
 
     /** Collects [Flow] of [VolumeControllerEvent] into [IVolumeController]. */
-    fun collectToController(
-        eventsFlow: Flow<VolumeControllerEvent>,
-        controller: IVolumeController
-    ) =
+    fun collectToController(controller: IVolumeController) {
         coroutineScope.launch {
-            eventsFlow.collect { event ->
+            audioRepository.volumeControllerEvents.collect { event ->
                 when (event) {
                     is VolumeControllerEvent.VolumeChanged ->
                         controller.volumeChanged(event.streamType, event.flags)
@@ -56,4 +57,9 @@
                 }
             }
         }
+    }
+
+    fun notifyVolumeControllerVisible(isVisible: Boolean) {
+        coroutineScope.launch { audioRepository.notifyVolumeControllerVisible(isVisible) }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogControllerImpl.java b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogControllerImpl.java
index 1522cc4..8934d8f 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogControllerImpl.java
@@ -68,6 +68,7 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.settingslib.volume.MediaSessions;
 import com.android.systemui.Dumpable;
+import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dump.DumpManager;
@@ -124,7 +125,6 @@
     static final ArrayMap<Integer, Integer> STREAMS = new ArrayMap<>();
     static {
         STREAMS.put(AudioSystem.STREAM_ALARM, R.string.stream_alarm);
-        STREAMS.put(AudioSystem.STREAM_BLUETOOTH_SCO, R.string.stream_bluetooth_sco);
         STREAMS.put(AudioSystem.STREAM_DTMF, R.string.stream_dtmf);
         STREAMS.put(AudioSystem.STREAM_MUSIC, R.string.stream_music);
         STREAMS.put(AudioSystem.STREAM_ACCESSIBILITY, R.string.stream_accessibility);
@@ -153,6 +153,7 @@
     private final KeyguardManager mKeyguardManager;
     private final ActivityManager mActivityManager;
     private final UserTracker mUserTracker;
+    private final VolumeControllerAdapter mVolumeControllerAdapter;
     protected C mCallbacks = new C();
     private final State mState = new State();
     protected final MediaSessionsCallbacks mMediaSessionsCallbacksW;
@@ -197,6 +198,7 @@
             NotificationManager notificationManager,
             VibratorHelper vibrator,
             IAudioService iAudioService,
+            VolumeControllerAdapter volumeControllerAdapter,
             AccessibilityManager accessibilityManager,
             PackageManager packageManager,
             WakefulnessLifecycle wakefulnessLifecycle,
@@ -233,6 +235,7 @@
         mVibrator = vibrator;
         mHasVibrator = mVibrator.hasVibrator();
         mAudioService = iAudioService;
+        mVolumeControllerAdapter = volumeControllerAdapter;
         mKeyguardManager = keyguardManager;
         mActivityManager = activityManager;
         mUserTracker = userTracker;
@@ -259,10 +262,14 @@
     }
 
     protected void setVolumeController() {
-        try {
-            mAudio.setVolumeController(mVolumeController);
-        } catch (SecurityException e) {
-            Log.w(TAG, "Unable to set the volume controller", e);
+        if (Flags.useVolumeController()) {
+            mVolumeControllerAdapter.collectToController(mVolumeController);
+        } else {
+            try {
+                mAudio.setVolumeController(mVolumeController);
+            } catch (SecurityException e) {
+                Log.w(TAG, "Unable to set the volume controller", e);
+            }
         }
     }
 
@@ -384,7 +391,11 @@
     }
 
     public void notifyVisible(boolean visible) {
-        mWorker.obtainMessage(W.NOTIFY_VISIBLE, visible ? 1 : 0, 0).sendToTarget();
+        if (Flags.useVolumeController()) {
+            mVolumeControllerAdapter.notifyVolumeControllerVisible(visible);
+        } else {
+            mWorker.obtainMessage(W.NOTIFY_VISIBLE, visible ? 1 : 0, 0).sendToTarget();
+        }
     }
 
     public void userActivity() {
@@ -642,7 +653,6 @@
     private static boolean isLogWorthy(int stream) {
         switch (stream) {
             case AudioSystem.STREAM_ALARM:
-            case AudioSystem.STREAM_BLUETOOTH_SCO:
             case AudioSystem.STREAM_MUSIC:
             case AudioSystem.STREAM_RING:
             case AudioSystem.STREAM_SYSTEM:
@@ -1270,6 +1280,8 @@
 
     private final class Receiver extends BroadcastReceiver {
 
+        private static final int STREAM_UNKNOWN = -1;
+
         public void init() {
             final IntentFilter filter = new IntentFilter();
             filter.addAction(AudioManager.VOLUME_CHANGED_ACTION);
@@ -1291,30 +1303,39 @@
             final String action = intent.getAction();
             boolean changed = false;
             if (action.equals(AudioManager.VOLUME_CHANGED_ACTION)) {
-                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE, -1);
+                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE,
+                        STREAM_UNKNOWN);
                 final int level = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_VALUE, -1);
                 final int oldLevel = intent
                         .getIntExtra(AudioManager.EXTRA_PREV_VOLUME_STREAM_VALUE, -1);
                 if (D.BUG) Log.d(TAG, "onReceive VOLUME_CHANGED_ACTION stream=" + stream
                         + " level=" + level + " oldLevel=" + oldLevel);
-                changed = updateStreamLevelW(stream, level);
+                if (stream != STREAM_UNKNOWN) {
+                    changed = updateStreamLevelW(stream, level);
+                }
             } else if (action.equals(AudioManager.STREAM_DEVICES_CHANGED_ACTION)) {
-                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE, -1);
+                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE,
+                        STREAM_UNKNOWN);
                 final int devices = intent
                         .getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_DEVICES, -1);
                 final int oldDevices = intent
                         .getIntExtra(AudioManager.EXTRA_PREV_VOLUME_STREAM_DEVICES, -1);
                 if (D.BUG) Log.d(TAG, "onReceive STREAM_DEVICES_CHANGED_ACTION stream="
                         + stream + " devices=" + devices + " oldDevices=" + oldDevices);
-                changed = checkRoutedToBluetoothW(stream);
-                changed |= onVolumeChangedW(stream, 0);
+                if (stream != STREAM_UNKNOWN) {
+                    changed |= checkRoutedToBluetoothW(stream);
+                    changed |= onVolumeChangedW(stream, 0);
+                }
             } else if (action.equals(AudioManager.STREAM_MUTE_CHANGED_ACTION)) {
-                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE, -1);
+                final int stream = intent.getIntExtra(AudioManager.EXTRA_VOLUME_STREAM_TYPE,
+                        STREAM_UNKNOWN);
                 final boolean muted = intent
                         .getBooleanExtra(AudioManager.EXTRA_STREAM_VOLUME_MUTED, false);
                 if (D.BUG) Log.d(TAG, "onReceive STREAM_MUTE_CHANGED_ACTION stream=" + stream
                         + " muted=" + muted);
-                changed = updateStreamMuteW(stream, muted);
+                if (stream != STREAM_UNKNOWN) {
+                    changed = updateStreamMuteW(stream, muted);
+                }
             } else if (action.equals(NotificationManager.ACTION_EFFECTS_SUPPRESSOR_CHANGED)) {
                 if (D.BUG) Log.d(TAG, "onReceive ACTION_EFFECTS_SUPPRESSOR_CHANGED");
                 changed = updateEffectsSuppressorW(mNoMan.getEffectsSuppressor());
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
index eb91518..7786453 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeDialogImpl.java
@@ -704,8 +704,6 @@
                 addRow(AudioManager.STREAM_VOICE_CALL,
                         com.android.internal.R.drawable.ic_phone,
                         com.android.internal.R.drawable.ic_phone, false, false);
-                addRow(AudioManager.STREAM_BLUETOOTH_SCO,
-                        R.drawable.ic_volume_bt_sco, R.drawable.ic_volume_bt_sco, false, false);
                 addRow(AudioManager.STREAM_SYSTEM, R.drawable.ic_volume_system,
                         R.drawable.ic_volume_system_mute, false, false);
             }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java b/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
index 68d12f6..536403c 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
+++ b/packages/SystemUI/src/com/android/systemui/volume/VolumeUI.java
@@ -20,9 +20,9 @@
 
 import android.content.Context;
 import android.content.res.Configuration;
-import android.os.Handler;
 import android.util.Log;
 
+import com.android.settingslib.volume.data.repository.AudioRepository;
 import com.android.systemui.CoreStartable;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.qs.tiles.DndTile;
@@ -39,23 +39,26 @@
     private static final String TAG = "VolumeUI";
     private static boolean LOGD = Log.isLoggable(TAG, Log.DEBUG);
 
-    private final Handler mHandler = new Handler();
-
     private boolean mEnabled;
     private final Context mContext;
     private VolumeDialogComponent mVolumeComponent;
     private AudioSharingInteractor mAudioSharingInteractor;
+    private AudioRepository mAudioRepository;
 
     @Inject
-    public VolumeUI(Context context, VolumeDialogComponent volumeDialogComponent,
+    public VolumeUI(Context context,
+            VolumeDialogComponent volumeDialogComponent,
+            AudioRepository audioRepository,
             AudioSharingInteractor audioSharingInteractor) {
         mContext = context;
         mVolumeComponent = volumeDialogComponent;
+        mAudioRepository = audioRepository;
         mAudioSharingInteractor = audioSharingInteractor;
     }
 
     @Override
     public void start() {
+        mAudioRepository.init();
         boolean enableVolumeUi = mContext.getResources().getBoolean(R.bool.enable_volume_ui);
         boolean enableSafetyWarning =
                 mContext.getResources().getBoolean(R.bool.enable_safety_warning);
@@ -77,7 +80,8 @@
 
     @Override
     public void dump(PrintWriter pw, String[] args) {
-        pw.print("mEnabled="); pw.println(mEnabled);
+        pw.print("mEnabled=");
+        pw.println(mEnabled);
         if (!mEnabled) return;
         mVolumeComponent.dump(pw, args);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt b/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
index d39daaf..20d598a 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/dagger/AudioModule.kt
@@ -70,6 +70,7 @@
                 coroutineContext,
                 coroutineScope,
                 volumeLogger,
+                com.android.systemui.Flags.useVolumeController(),
             )
 
         @Provides
diff --git a/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioOutputInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioOutputInteractor.kt
index 4f77cd0..73728e6 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioOutputInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/domain/interactor/AudioOutputInteractor.kt
@@ -75,7 +75,7 @@
             }
             .map { it ?: AudioOutputDevice.Unknown }
             .flowOn(backgroundCoroutineContext)
-            .stateIn(scope, SharingStarted.Eagerly, AudioOutputDevice.Unknown)
+            .stateIn(scope, SharingStarted.Eagerly, AudioOutputDevice.Unavailable)
 
     private fun AudioDeviceInfo.toAudioOutputDevice(): AudioOutputDevice {
         if (
@@ -120,6 +120,11 @@
                     name = name,
                     icon = icon,
                 )
+            deviceType == MediaDeviceType.TYPE_CAST_DEVICE ->
+                AudioOutputDevice.Remote(
+                    name = name,
+                    icon = icon,
+                )
             else ->
                 AudioOutputDevice.BuiltIn(
                     name = name,
diff --git a/packages/SystemUI/src/com/android/systemui/volume/domain/model/AudioOutputDevice.kt b/packages/SystemUI/src/com/android/systemui/volume/domain/model/AudioOutputDevice.kt
index ba0b082..0e4cac0b 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/domain/model/AudioOutputDevice.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/domain/model/AudioOutputDevice.kt
@@ -31,6 +31,12 @@
         override val icon: Drawable?,
     ) : AudioOutputDevice
 
+    /** Models a cast audio output device. */
+    data class Remote(
+        override val name: String,
+        override val icon: Drawable?,
+    ) : AudioOutputDevice
+
     /** Models a wired audio output device. */
     data class Wired(
         override val name: String,
@@ -52,4 +58,16 @@
         override val icon: Drawable
             get() = error("Unsupported for unknown device")
     }
+
+    /**
+     * Models a state when current audio output device is not loaded yet or the system failed to
+     * load it.
+     */
+    data object Unavailable : AudioOutputDevice {
+        override val name: String
+            get() = error("Unsupported for unavailable device")
+
+        override val icon: Drawable
+            get() = error("Unsupported for unavailable device")
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractor.kt
index a270d5f..f94cbda 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputComponentInteractor.kt
@@ -74,34 +74,51 @@
             )
 
     private val currentAudioDevice: Flow<AudioOutputDevice> =
-        audioOutputInteractor.currentAudioDevice.filter { it !is AudioOutputDevice.Unknown }
+        audioOutputInteractor.currentAudioDevice.filter { it !is AudioOutputDevice.Unavailable }
 
+    /**
+     * Model for the Media Output component in the Volume Panel. It's guaranteed to have an
+     * available device if it's loaded.
+     */
     val mediaOutputModel: StateFlow<Result<MediaOutputComponentModel>> =
-        audioModeInteractor.isOngoingCall
-            .flatMapLatest { isOngoingCall ->
-                audioSharingInteractor.isInAudioSharing.flatMapLatest { isInAudioSharing ->
-                    if (isOngoingCall) {
-                        currentAudioDevice.map {
-                            MediaOutputComponentModel.Calling(it, isInAudioSharing)
-                        }
-                    } else {
-                        combine(sessionWithPlaybackState.filterData(), currentAudioDevice) {
-                            sessionWithPlaybackState,
-                            currentAudioDevice ->
-                            if (sessionWithPlaybackState == null) {
-                                MediaOutputComponentModel.Idle(currentAudioDevice, isInAudioSharing)
-                            } else {
-                                MediaOutputComponentModel.MediaSession(
-                                    sessionWithPlaybackState.session,
-                                    sessionWithPlaybackState.isPlaybackActive,
-                                    currentAudioDevice,
-                                    isInAudioSharing,
-                                )
-                            }
+        combine(
+                audioSharingInteractor.isInAudioSharing,
+                audioModeInteractor.isOngoingCall,
+                currentAudioDevice,
+            ) { isInAudioSharing, isOngoingCall, currentAudioDevice ->
+                if (isOngoingCall) {
+                    flowOf(
+                        MediaOutputComponentModel.Calling(
+                            device = currentAudioDevice,
+                            isInAudioSharing = isInAudioSharing,
+                            canOpenAudioSwitcher = false,
+                        )
+                    )
+                } else {
+                    sessionWithPlaybackState.filterData().map { sessionWithPlaybackState ->
+                        if (sessionWithPlaybackState == null) {
+                            MediaOutputComponentModel.Idle(
+                                device = currentAudioDevice,
+                                isInAudioSharing = isInAudioSharing,
+                                canOpenAudioSwitcher =
+                                    !isInAudioSharing &&
+                                        currentAudioDevice !is AudioOutputDevice.Unknown,
+                            )
+                        } else {
+                            MediaOutputComponentModel.MediaSession(
+                                session = sessionWithPlaybackState.session,
+                                isPlaybackActive = sessionWithPlaybackState.isPlaybackActive,
+                                device = currentAudioDevice,
+                                isInAudioSharing = isInAudioSharing,
+                                canOpenAudioSwitcher =
+                                    !isInAudioSharing &&
+                                        currentAudioDevice !is AudioOutputDevice.Unknown,
+                            )
                         }
                     }
                 }
             }
+            .flatMapLatest { it }
             .wrapInResult()
             .stateIn(coroutineScope, SharingStarted.Eagerly, Result.Loading())
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
index 31a8977..aa07cfd 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
@@ -179,9 +179,7 @@
         return MediaDeviceSession(
             packageName = packageName,
             sessionToken = sessionToken,
-            canAdjustVolume =
-                playbackInfo != null &&
-                    playbackInfo?.volumeControl != VolumeProvider.VOLUME_CONTROL_FIXED,
+            canAdjustVolume = playbackInfo.volumeControl != VolumeProvider.VOLUME_CONTROL_FIXED,
             appLabel = getApplicationLabel(packageName) ?: return null
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/model/MediaOutputComponentModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/model/MediaOutputComponentModel.kt
index 220fb2b..6588b44 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/model/MediaOutputComponentModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/model/MediaOutputComponentModel.kt
@@ -24,11 +24,13 @@
 
     val device: AudioOutputDevice
     val isInAudioSharing: Boolean
+    val canOpenAudioSwitcher: Boolean
 
     /** There is an ongoing call on the device. */
     data class Calling(
         override val device: AudioOutputDevice,
         override val isInAudioSharing: Boolean,
+        override val canOpenAudioSwitcher: Boolean,
     ) : MediaOutputComponentModel
 
     /** There is media playing on the device. */
@@ -37,11 +39,13 @@
         val isPlaybackActive: Boolean,
         override val device: AudioOutputDevice,
         override val isInAudioSharing: Boolean,
+        override val canOpenAudioSwitcher: Boolean,
     ) : MediaOutputComponentModel
 
     /** There is nothing playing on the device. */
     data class Idle(
         override val device: AudioOutputDevice,
         override val isInAudioSharing: Boolean,
+        override val canOpenAudioSwitcher: Boolean,
     ) : MediaOutputComponentModel
 }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/ConnectedDeviceViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/ConnectedDeviceViewModel.kt
index 8ba672d..42f88b4 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/ConnectedDeviceViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/ConnectedDeviceViewModel.kt
@@ -16,11 +16,15 @@
 
 package com.android.systemui.volume.panel.component.mediaoutput.ui.viewmodel
 
+import com.android.systemui.common.shared.model.Color
+
 /**
  * Models part of the Media Session Volume Panel component that displays connected device
  * information.
  */
 data class ConnectedDeviceViewModel(
     val label: CharSequence,
+    val labelColor: Color,
     val deviceName: CharSequence?,
+    val deviceNameColor: Color,
 )
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 36b42f2..e565de5 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
@@ -75,12 +75,25 @@
                         }
                     }
                 ConnectedDeviceViewModel(
-                    label,
-                    if (mediaOutputModel.isInAudioSharing) {
-                        context.getString(R.string.audio_sharing_description)
-                    } else {
-                        mediaOutputModel.device.name
-                    },
+                    label = label,
+                    labelColor =
+                        Color.Attribute(com.android.internal.R.attr.materialColorOnSurfaceVariant),
+                    deviceName =
+                        if (mediaOutputModel.isInAudioSharing) {
+                            context.getString(R.string.audio_sharing_description)
+                        } else {
+                            mediaOutputModel.device
+                                .takeIf { it !is AudioOutputDevice.Unknown }
+                                ?.name ?: context.getString(R.string.media_seamless_other_device)
+                        },
+                    deviceNameColor =
+                        if (mediaOutputModel.canOpenAudioSwitcher) {
+                            Color.Attribute(com.android.internal.R.attr.materialColorOnSurface)
+                        } else {
+                            Color.Attribute(
+                                com.android.internal.R.attr.materialColorOnSurfaceVariant
+                            )
+                        },
                 )
             }
             .stateIn(
@@ -107,19 +120,39 @@
                     DeviceIconViewModel.IsPlaying(
                         icon = icon,
                         iconColor =
-                            Color.Attribute(com.android.internal.R.attr.materialColorSurface),
+                            if (mediaOutputModel.canOpenAudioSwitcher) {
+                                Color.Attribute(com.android.internal.R.attr.materialColorSurface)
+                            } else {
+                                Color.Attribute(
+                                    com.android.internal.R.attr.materialColorSurfaceContainerHighest
+                                )
+                            },
                         backgroundColor =
-                            Color.Attribute(com.android.internal.R.attr.materialColorSecondary),
+                            if (mediaOutputModel.canOpenAudioSwitcher) {
+                                Color.Attribute(com.android.internal.R.attr.materialColorSecondary)
+                            } else {
+                                Color.Attribute(com.android.internal.R.attr.materialColorOutline)
+                            },
                     )
                 } else {
                     DeviceIconViewModel.IsNotPlaying(
                         icon = icon,
                         iconColor =
-                            Color.Attribute(
-                                com.android.internal.R.attr.materialColorOnSurfaceVariant
-                            ),
+                            if (mediaOutputModel.canOpenAudioSwitcher) {
+                                Color.Attribute(
+                                    com.android.internal.R.attr.materialColorOnSurfaceVariant
+                                )
+                            } else {
+                                Color.Attribute(com.android.internal.R.attr.materialColorOutline)
+                            },
                         backgroundColor =
-                            Color.Attribute(com.android.internal.R.attr.materialColorSurface),
+                            if (mediaOutputModel.canOpenAudioSwitcher) {
+                                Color.Attribute(com.android.internal.R.attr.materialColorSurface)
+                            } else {
+                                Color.Attribute(
+                                    com.android.internal.R.attr.materialColorSurfaceContainerHighest
+                                )
+                            },
                     )
                 }
             }
@@ -132,7 +165,7 @@
     val enabled: StateFlow<Boolean> =
         mediaOutputComponentInteractor.mediaOutputModel
             .filterData()
-            .map { !it.isInAudioSharing }
+            .map { it.canOpenAudioSwitcher }
             .stateIn(
                 coroutineScope,
                 SharingStarted.Eagerly,
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/domain/interactor/SpatialAudioComponentInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/domain/interactor/SpatialAudioComponentInteractor.kt
index cfcd6b1..56d0bce 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/domain/interactor/SpatialAudioComponentInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/spatial/domain/interactor/SpatialAudioComponentInteractor.kt
@@ -63,13 +63,7 @@
     private val changes = MutableSharedFlow<Unit>()
     private val currentAudioDeviceAttributes: StateFlow<AudioDeviceAttributes?> =
         audioOutputInteractor.currentAudioDevice
-            .map { audioDevice ->
-                if (audioDevice is AudioOutputDevice.Unknown) {
-                    builtinSpeaker
-                } else {
-                    audioDevice.getAudioDeviceAttributes()
-                }
-            }
+            .map { audioDevice -> audioDevice.getAudioDeviceAttributes() }
             .stateIn(coroutineScope, SharingStarted.WhileSubscribed(), builtinSpeaker)
 
     /**
@@ -185,7 +179,10 @@
                         .firstOrNull { spatializerInteractor.isSpatialAudioAvailable(it) }
                 }
             }
-            else -> null
+            is AudioOutputDevice.Wired -> null
+            is AudioOutputDevice.Remote -> null
+            is AudioOutputDevice.Unknown -> builtinSpeaker
+            is AudioOutputDevice.Unavailable -> builtinSpeaker
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractor.kt
index 0451ce6..4be680e 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractor.kt
@@ -16,9 +16,7 @@
 
 package com.android.systemui.volume.panel.component.volume.domain.interactor
 
-import android.media.AudioDeviceInfo
 import android.media.AudioManager
-import com.android.settingslib.volume.data.repository.AudioRepository
 import com.android.settingslib.volume.domain.interactor.AudioModeInteractor
 import com.android.settingslib.volume.shared.model.AudioStream
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaOutputInteractor
@@ -42,7 +40,6 @@
 constructor(
     @VolumePanelScope scope: CoroutineScope,
     mediaOutputInteractor: MediaOutputInteractor,
-    audioRepository: AudioRepository,
     audioModeInteractor: AudioModeInteractor,
 ) {
 
@@ -50,13 +47,12 @@
         combineTransform(
                 mediaOutputInteractor.activeMediaDeviceSessions,
                 mediaOutputInteractor.defaultActiveMediaSession.filterData(),
-                audioRepository.communicationDevice,
                 audioModeInteractor.isOngoingCall,
-            ) { activeSessions, defaultSession, communicationDevice, isOngoingCall ->
+            ) { activeSessions, defaultSession, isOngoingCall ->
                 coroutineScope {
                     val viewModels = buildList {
                         if (isOngoingCall) {
-                            addCall(communicationDevice?.type)
+                            addStream(AudioManager.STREAM_VOICE_CALL)
                         }
 
                         if (defaultSession?.isTheSameSession(activeSessions.remote) == true) {
@@ -68,7 +64,7 @@
                         }
 
                         if (!isOngoingCall) {
-                            addCall(communicationDevice?.type)
+                            addStream(AudioManager.STREAM_VOICE_CALL)
                         }
 
                         addStream(AudioManager.STREAM_RING)
@@ -80,14 +76,6 @@
             }
             .stateIn(scope, SharingStarted.Eagerly, emptyList())
 
-    private fun MutableList<SliderType>.addCall(communicationDeviceType: Int?) {
-        if (communicationDeviceType == AudioDeviceInfo.TYPE_BLUETOOTH_SCO) {
-            addStream(AudioManager.STREAM_BLUETOOTH_SCO)
-        } else {
-            addStream(AudioManager.STREAM_VOICE_CALL)
-        }
-    }
-
     private fun MutableList<SliderType>.addSession(remoteMediaDeviceSession: MediaDeviceSession?) {
         if (remoteMediaDeviceSession?.canAdjustVolume == true) {
             add(SliderType.MediaDeviceCast(remoteMediaDeviceSession))
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 521f608..ffb1f11 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
@@ -66,7 +66,6 @@
         mapOf(
             AudioStream(AudioManager.STREAM_MUSIC) to R.drawable.ic_music_note,
             AudioStream(AudioManager.STREAM_VOICE_CALL) to R.drawable.ic_call,
-            AudioStream(AudioManager.STREAM_BLUETOOTH_SCO) to R.drawable.ic_call,
             AudioStream(AudioManager.STREAM_RING) to R.drawable.ic_ring_volume,
             AudioStream(AudioManager.STREAM_NOTIFICATION) to R.drawable.ic_volume_ringer,
             AudioStream(AudioManager.STREAM_ALARM) to R.drawable.ic_volume_alarm,
@@ -75,7 +74,6 @@
         mapOf(
             AudioStream(AudioManager.STREAM_MUSIC) to R.string.stream_music,
             AudioStream(AudioManager.STREAM_VOICE_CALL) to R.string.stream_voice_call,
-            AudioStream(AudioManager.STREAM_BLUETOOTH_SCO) to R.string.stream_voice_call,
             AudioStream(AudioManager.STREAM_RING) to R.string.stream_ring,
             AudioStream(AudioManager.STREAM_NOTIFICATION) to R.string.stream_notification,
             AudioStream(AudioManager.STREAM_ALARM) to R.string.stream_alarm,
@@ -91,8 +89,6 @@
                 VolumePanelUiEvent.VOLUME_PANEL_MUSIC_SLIDER_TOUCHED,
             AudioStream(AudioManager.STREAM_VOICE_CALL) to
                 VolumePanelUiEvent.VOLUME_PANEL_VOICE_CALL_SLIDER_TOUCHED,
-            AudioStream(AudioManager.STREAM_BLUETOOTH_SCO) to
-                VolumePanelUiEvent.VOLUME_PANEL_VOICE_CALL_SLIDER_TOUCHED,
             AudioStream(AudioManager.STREAM_RING) to
                 VolumePanelUiEvent.VOLUME_PANEL_RING_SLIDER_TOUCHED,
             AudioStream(AudioManager.STREAM_NOTIFICATION) to
diff --git a/packages/SystemUI/src/com/android/systemui/wallet/dagger/WalletModule.java b/packages/SystemUI/src/com/android/systemui/wallet/dagger/WalletModule.java
index 4841c78..ea213cb 100644
--- a/packages/SystemUI/src/com/android/systemui/wallet/dagger/WalletModule.java
+++ b/packages/SystemUI/src/com/android/systemui/wallet/dagger/WalletModule.java
@@ -23,11 +23,19 @@
 
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Background;
+import com.android.systemui.qs.QsEventLogger;
+import com.android.systemui.qs.pipeline.shared.TileSpec;
 import com.android.systemui.qs.tileimpl.QSTileImpl;
 import com.android.systemui.qs.tiles.QuickAccessWalletTile;
+import com.android.systemui.qs.tiles.viewmodel.QSTileConfig;
+import com.android.systemui.qs.tiles.viewmodel.QSTilePolicy;
+import com.android.systemui.qs.tiles.viewmodel.QSTileUIConfig;
+import com.android.systemui.res.R;
 import com.android.systemui.wallet.controller.WalletContextualLocationsService;
 import com.android.systemui.wallet.ui.WalletActivity;
 
+import java.util.concurrent.Executor;
+
 import dagger.Binds;
 import dagger.Module;
 import dagger.Provides;
@@ -35,14 +43,14 @@
 import dagger.multibindings.IntoMap;
 import dagger.multibindings.StringKey;
 
-import java.util.concurrent.Executor;
-
 /**
  * Module for injecting classes in Wallet.
  */
 @Module
 public abstract class WalletModule {
 
+    public static final String WALLET_TILE_SPEC = "wallet";
+
     @Binds
     @IntoMap
     @ClassKey(WalletContextualLocationsService.class)
@@ -69,4 +77,21 @@
     @StringKey(QuickAccessWalletTile.TILE_SPEC)
     public abstract QSTileImpl<?> bindQuickAccessWalletTile(
             QuickAccessWalletTile quickAccessWalletTile);
+
+    @Provides
+    @IntoMap
+    @StringKey(WALLET_TILE_SPEC)
+    public static QSTileConfig provideQuickAccessWalletTileConfig(QsEventLogger uiEventLogger) {
+        TileSpec tileSpec = TileSpec.create(WALLET_TILE_SPEC);
+        return new QSTileConfig(
+                tileSpec,
+                new QSTileUIConfig.Resource(
+                        R.drawable.ic_wallet_lockscreen,
+                        R.string.wallet_title
+                ),
+                uiEventLogger.getNewInstanceId(),
+                tileSpec.getSpec(),
+                QSTilePolicy.NoRestrictions.INSTANCE
+        );
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardAbsKeyInputViewControllerTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardAbsKeyInputViewControllerTest.java
index e724c60..c0d8be3 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardAbsKeyInputViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardAbsKeyInputViewControllerTest.java
@@ -16,6 +16,8 @@
 
 package com.android.keyguard;
 
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
@@ -27,6 +29,7 @@
 import static org.mockito.Mockito.when;
 
 import android.os.SystemClock;
+import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
 import android.testing.TestableLooper.RunWithLooper;
 import android.view.KeyEvent;
@@ -43,9 +46,13 @@
 import com.android.systemui.classifier.FalsingCollector;
 import com.android.systemui.classifier.FalsingCollectorFake;
 import com.android.systemui.flags.FakeFeatureFlags;
+import com.android.systemui.haptics.msdl.FakeMSDLPlayer;
+import com.android.systemui.kosmos.KosmosJavaAdapter;
 import com.android.systemui.res.R;
 import com.android.systemui.user.domain.interactor.SelectedUserInteractor;
 
+import com.google.android.msdl.data.model.MSDLToken;
+
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -85,7 +92,11 @@
     private FakeFeatureFlags mFeatureFlags;
     @Mock
     private SelectedUserInteractor mSelectedUserInteractor;
+    @Mock
+    private UserActivityNotifier mUserActivityNotifier;
     private KeyguardAbsKeyInputViewController mKeyguardAbsKeyInputViewController;
+    private KosmosJavaAdapter mKosmosJavaAdapter = new KosmosJavaAdapter(this);
+    private final FakeMSDLPlayer mMSDLPlayer = mKosmosJavaAdapter.getMsdlPlayer();
 
     @Before
     public void setup() {
@@ -108,7 +119,8 @@
         return new KeyguardAbsKeyInputViewController(mAbsKeyInputView,
                 mKeyguardUpdateMonitor, mSecurityMode, mLockPatternUtils, mKeyguardSecurityCallback,
                 mKeyguardMessageAreaControllerFactory, mLatencyTracker, mFalsingCollector,
-                mEmergencyButtonController, mFeatureFlags, mSelectedUserInteractor) {
+                mEmergencyButtonController, mFeatureFlags, mSelectedUserInteractor, mMSDLPlayer,
+                mUserActivityNotifier) {
             @Override
             void resetState() {
             }
@@ -197,4 +209,32 @@
         verify(mAbsKeyInputView, never()).setPasswordEntryInputEnabled(true);
         verify(mAbsKeyInputView, never()).setPasswordEntryEnabled(true);
     }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MSDL_FEEDBACK)
+    public void onPasswordChecked_withMSDLFeedback_withMatch_playsUnlockToken() {
+        mKeyguardAbsKeyInputViewController.onPasswordChecked(0, true, 100, true);
+        assertThat(mMSDLPlayer.getLatestTokenPlayed()).isEqualTo(MSDLToken.UNLOCK);
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MSDL_FEEDBACK)
+    public void onPasswordChecked_withoutMSDLFeedback_withMatch_doesNotPlayToken() {
+        mKeyguardAbsKeyInputViewController.onPasswordChecked(0, true, 100, true);
+        assertThat(mMSDLPlayer.getLatestTokenPlayed()).isNull();
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MSDL_FEEDBACK)
+    public void onPasswordChecked_withMSDLFeedback_withoutMatch_playsFailureToken() {
+        mKeyguardAbsKeyInputViewController.onPasswordChecked(0, false, 100, true);
+        assertThat(mMSDLPlayer.getLatestTokenPlayed()).isEqualTo(MSDLToken.FAILURE);
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MSDL_FEEDBACK)
+    public void onPasswordChecked_withoutMSDLFeedback_withoutMatch_doesNotPlayToken() {
+        mKeyguardAbsKeyInputViewController.onPasswordChecked(0, false, 100, true);
+        assertThat(mMSDLPlayer.getLatestTokenPlayed()).isNull();
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinViewControllerTest.kt
index 36d4d12..873bc2c 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardPinViewControllerTest.kt
@@ -103,6 +103,7 @@
     @Mock lateinit var deleteButton: NumPadButton
     @Mock lateinit var enterButton: View
     @Mock lateinit var uiEventLogger: UiEventLogger
+    @Mock lateinit var mUserActivityNotifier: UserActivityNotifier
 
     @Captor lateinit var postureCallbackCaptor: ArgumentCaptor<DevicePostureController.Callback>
 
@@ -149,7 +150,9 @@
             featureFlags,
             mSelectedUserInteractor,
             uiEventLogger,
-            keyguardKeyboardInteractor
+            keyguardKeyboardInteractor,
+            null,
+            mUserActivityNotifier
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
index 7151c42..f141a49 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPinViewControllerTest.kt
@@ -69,6 +69,7 @@
     @Mock
     private lateinit var keyguardMessageAreaController:
         KeyguardMessageAreaController<BouncerKeyguardMessageArea>
+    @Mock private lateinit var mUserActivityNotifier: UserActivityNotifier
     private val updateMonitorCallbackArgumentCaptor =
         ArgumentCaptor.forClass(KeyguardUpdateMonitorCallback::class.java)
 
@@ -101,7 +102,9 @@
                 emergencyButtonController,
                 fakeFeatureFlags,
                 mSelectedUserInteractor,
-                keyguardKeyboardInteractor
+                keyguardKeyboardInteractor,
+                null,
+                mUserActivityNotifier
             )
         underTest.init()
         underTest.onViewAttached()
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
index acae913..a03c839 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardSimPukViewControllerTest.kt
@@ -63,6 +63,7 @@
     @Mock
     private lateinit var keyguardMessageAreaController:
         KeyguardMessageAreaController<BouncerKeyguardMessageArea>
+    @Mock private lateinit var mUserActivityNotifier: UserActivityNotifier
 
     @Before
     fun setup() {
@@ -96,7 +97,9 @@
                 emergencyButtonController,
                 fakeFeatureFlags,
                 mSelectedUserInteractor,
-                keyguardKeyboardInteractor
+                keyguardKeyboardInteractor,
+                null,
+                mUserActivityNotifier
             )
         underTest.init()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuControllerTest.java
index 113a8c0..5e37d4c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/accessibility/floatingmenu/AccessibilityFloatingMenuControllerTest.java
@@ -28,6 +28,7 @@
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.hardware.display.DisplayManager;
+import android.os.Handler;
 import android.os.UserHandle;
 import android.provider.Settings;
 import android.testing.TestableLooper;
@@ -80,6 +81,7 @@
     private AccessibilityManager mAccessibilityManager;
     private KeyguardUpdateMonitor mKeyguardUpdateMonitor;
     private AccessibilityFloatingMenuController mController;
+    private TestableLooper mTestableLooper;
     @Mock
     private AccessibilityButtonTargetsObserver mTargetsObserver;
     @Mock
@@ -108,6 +110,7 @@
         mViewCaptureAwareWindowManager = new ViewCaptureAwareWindowManager(mWindowManager,
                 mLazyViewCapture, /* isViewCaptureEnabled= */ false);
         mAccessibilityManager = mContext.getSystemService(AccessibilityManager.class);
+        mTestableLooper = TestableLooper.get(this);
 
         when(mTargetsObserver.getCurrentAccessibilityButtonTargets())
                 .thenReturn(Settings.Secure.getStringForUser(mContextWrapper.getContentResolver(),
@@ -231,7 +234,8 @@
         mKeyguardCallback.onKeyguardVisibilityChanged(false);
 
         mKeyguardCallback.onUserSwitching(fakeUserId);
-        mKeyguardCallback.onUserSwitchComplete(fakeUserId);
+        mController.mUserInitializationCompleteCallback.onUserInitializationComplete(1);
+        mTestableLooper.processAllMessages();
 
         assertThat(mController.mFloatingMenu).isNotNull();
     }
@@ -346,7 +350,8 @@
                 new AccessibilityFloatingMenuController(mContextWrapper, windowManager,
                         viewCaptureAwareWindowManager, displayManager, mAccessibilityManager,
                         mTargetsObserver, mModeObserver, mKeyguardUpdateMonitor, mSecureSettings,
-                        displayTracker, mNavigationModeController);
+                        displayTracker, mNavigationModeController, new Handler(
+                                mTestableLooper.getLooper()));
         controller.init();
 
         return controller;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
index 6047e7d..4fc4166 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/biometrics/ui/viewmodel/PromptViewModelTest.kt
@@ -717,13 +717,9 @@
             assertThat(confirmHaptics?.hapticFeedbackConstant)
                 .isEqualTo(
                     if (expectConfirmation) HapticFeedbackConstants.NO_HAPTICS
-                    else HapticFeedbackConstants.CONFIRM
+                    else HapticFeedbackConstants.BIOMETRIC_CONFIRM
                 )
-            assertThat(confirmHaptics?.flag)
-                .isEqualTo(
-                    if (expectConfirmation) null
-                    else HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING
-                )
+            assertThat(confirmHaptics?.flag).isNull()
 
             if (expectConfirmation) {
                 kosmos.promptViewModel.confirmAuthenticated()
@@ -731,9 +727,8 @@
 
             val confirmedHaptics by collectLastValue(kosmos.promptViewModel.hapticsToPlay)
             assertThat(confirmedHaptics?.hapticFeedbackConstant)
-                .isEqualTo(HapticFeedbackConstants.CONFIRM)
-            assertThat(confirmedHaptics?.flag)
-                .isEqualTo(HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING)
+                .isEqualTo(HapticFeedbackConstants.BIOMETRIC_CONFIRM)
+            assertThat(confirmedHaptics?.flag).isNull()
         }
 
     @Test
@@ -747,9 +742,8 @@
 
         val currentHaptics by collectLastValue(kosmos.promptViewModel.hapticsToPlay)
         assertThat(currentHaptics?.hapticFeedbackConstant)
-            .isEqualTo(HapticFeedbackConstants.CONFIRM)
-        assertThat(currentHaptics?.flag)
-            .isEqualTo(HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING)
+            .isEqualTo(HapticFeedbackConstants.BIOMETRIC_CONFIRM)
+        assertThat(currentHaptics?.flag).isNull()
     }
 
     @Test
@@ -757,9 +751,9 @@
         kosmos.promptViewModel.showTemporaryError("test", "messageAfterError", false)
 
         val currentHaptics by collectLastValue(kosmos.promptViewModel.hapticsToPlay)
-        assertThat(currentHaptics?.hapticFeedbackConstant).isEqualTo(HapticFeedbackConstants.REJECT)
-        assertThat(currentHaptics?.flag)
-            .isEqualTo(HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING)
+        assertThat(currentHaptics?.hapticFeedbackConstant)
+            .isEqualTo(HapticFeedbackConstants.BIOMETRIC_REJECT)
+        assertThat(currentHaptics?.flag).isNull()
     }
 
     // biometricprompt_sfps_fingerprint_authenticating reused across rotations
@@ -870,8 +864,9 @@
         )
 
         val haptics by collectLastValue(kosmos.promptViewModel.hapticsToPlay)
-        assertThat(haptics?.hapticFeedbackConstant).isEqualTo(HapticFeedbackConstants.REJECT)
-        assertThat(haptics?.flag).isEqualTo(HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING)
+        assertThat(haptics?.hapticFeedbackConstant)
+            .isEqualTo(HapticFeedbackConstants.BIOMETRIC_REJECT)
+        assertThat(haptics?.flag).isNull()
     }
 
     @Test
@@ -901,10 +896,12 @@
 
         val haptics by collectLastValue(kosmos.promptViewModel.hapticsToPlay)
         if (expectConfirmation) {
-            assertThat(haptics?.hapticFeedbackConstant).isEqualTo(HapticFeedbackConstants.REJECT)
-            assertThat(haptics?.flag).isEqualTo(HapticFeedbackConstants.FLAG_IGNORE_GLOBAL_SETTING)
+            assertThat(haptics?.hapticFeedbackConstant)
+                .isEqualTo(HapticFeedbackConstants.BIOMETRIC_REJECT)
+            assertThat(haptics?.flag).isNull()
         } else {
-            assertThat(haptics?.hapticFeedbackConstant).isEqualTo(HapticFeedbackConstants.CONFIRM)
+            assertThat(haptics?.hapticFeedbackConstant)
+                .isEqualTo(HapticFeedbackConstants.BIOMETRIC_CONFIRM)
         }
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractorTest.kt
index 74bc928..681ea75 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemActionInteractorTest.kt
@@ -59,6 +59,7 @@
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var activeMediaDeviceItem: DeviceItem
     private lateinit var notConnectedDeviceItem: DeviceItem
+    private lateinit var connectedAudioSharingMediaDeviceItem: DeviceItem
     private lateinit var connectedMediaDeviceItem: DeviceItem
     private lateinit var connectedOtherDeviceItem: DeviceItem
     @Mock private lateinit var dialog: SystemUIDialog
@@ -100,6 +101,15 @@
                 iconWithDescription = null,
                 background = null
             )
+        connectedAudioSharingMediaDeviceItem =
+            DeviceItem(
+                type = DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE,
+                cachedBluetoothDevice = cachedBluetoothDevice,
+                deviceName = DEVICE_NAME,
+                connectionSummary = DEVICE_CONNECTION_SUMMARY,
+                iconWithDescription = null,
+                background = null
+            )
         connectedOtherDeviceItem =
             DeviceItem(
                 type = DeviceItemType.CONNECTED_BLUETOOTH_DEVICE,
@@ -186,6 +196,21 @@
     }
 
     @Test
+    fun testOnClick_connectedAudioSharingMediaDevice_logClick() {
+        with(kosmos) {
+            testScope.runTest {
+                whenever(cachedBluetoothDevice.address).thenReturn(DEVICE_ADDRESS)
+                actionInteractorImpl.onClick(connectedAudioSharingMediaDeviceItem, dialog)
+                verify(bluetoothTileDialogLogger)
+                    .logDeviceClick(
+                        cachedBluetoothDevice.address,
+                        DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE
+                    )
+            }
+        }
+    }
+
+    @Test
     fun testOnClick_audioSharingDisabled_shouldNotLaunchSettings() {
         with(kosmos) {
             testScope.runTest {
@@ -415,7 +440,7 @@
     }
 
     @Test
-    fun testOnClick_hasTwoConnectedLeDevice_clickedConnectedLe_shouldLaunchSettings() {
+    fun testOnClick_hasTwoConnectedLeDevice_clickedActiveLe_shouldLaunchSettings() {
         with(kosmos) {
             testScope.runTest {
                 whenever(cachedBluetoothDevice.device).thenReturn(bluetoothDevice)
@@ -438,7 +463,7 @@
                     if (device == bluetoothDevice) GROUP_ID_1 else GROUP_ID_2
                 }
 
-                actionInteractorImpl.onClick(connectedMediaDeviceItem, dialog)
+                actionInteractorImpl.onClick(activeMediaDeviceItem, dialog)
                 verify(activityStarter)
                     .postStartActivityDismissingKeyguard(
                         ArgumentMatchers.any(),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactoryTest.kt
index a27ccc6..ef441c1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bluetooth/qsdialog/DeviceItemFactoryTest.kt
@@ -17,18 +17,24 @@
 package com.android.systemui.bluetooth.qsdialog
 
 import android.bluetooth.BluetoothDevice
-import android.content.pm.ApplicationInfo
-import android.content.pm.PackageManager
+import android.graphics.drawable.Drawable
 import android.media.AudioManager
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.testing.TestableLooper
+import android.util.Pair
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.settingslib.bluetooth.BluetoothUtils
 import com.android.settingslib.bluetooth.CachedBluetoothDevice
+import com.android.settingslib.bluetooth.LocalBluetoothManager
 import com.android.settingslib.flags.Flags
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.res.R
 import com.google.common.truth.Truth.assertThat
+import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -37,6 +43,7 @@
 import org.mockito.Mockito.`when`
 import org.mockito.junit.MockitoJUnit
 import org.mockito.junit.MockitoRule
+import org.mockito.kotlin.any
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -44,9 +51,11 @@
 class DeviceItemFactoryTest : SysuiTestCase() {
     @get:Rule val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
+    private lateinit var mockitoSession: StaticMockitoSession
     @Mock private lateinit var cachedDevice: CachedBluetoothDevice
     @Mock private lateinit var bluetoothDevice: BluetoothDevice
-    @Mock private lateinit var packageManager: PackageManager
+    @Mock private lateinit var localBluetoothManager: LocalBluetoothManager
+    @Mock private lateinit var drawable: Drawable
 
     private val availableMediaDeviceItemFactory = AvailableMediaDeviceItemFactory()
     private val connectedDeviceItemFactory = ConnectedDeviceItemFactory()
@@ -56,16 +65,21 @@
 
     @Before
     fun setup() {
-        `when`(cachedDevice.name).thenReturn(DEVICE_NAME)
-        `when`(cachedDevice.address).thenReturn(DEVICE_ADDRESS)
-        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
-        `when`(cachedDevice.connectionSummary).thenReturn(CONNECTION_SUMMARY)
+        mockitoSession =
+            mockitoSession().initMocks(this).mockStatic(BluetoothUtils::class.java).startMocking()
+    }
 
-        context.setMockPackageManager(packageManager)
+    @After
+    fun tearDown() {
+        mockitoSession.finishMocking()
     }
 
     @Test
     fun testAvailableMediaDeviceItemFactory_createFromCachedDevice() {
+        `when`(cachedDevice.name).thenReturn(DEVICE_NAME)
+        `when`(cachedDevice.connectionSummary).thenReturn(CONNECTION_SUMMARY)
+        `when`(BluetoothUtils.getBtClassDrawableWithDescription(any(), any()))
+            .thenReturn(Pair.create(drawable, ""))
         val deviceItem = availableMediaDeviceItemFactory.create(context, cachedDevice)
 
         assertDeviceItem(deviceItem, DeviceItemType.AVAILABLE_MEDIA_BLUETOOTH_DEVICE)
@@ -73,6 +87,10 @@
 
     @Test
     fun testConnectedDeviceItemFactory_createFromCachedDevice() {
+        `when`(cachedDevice.name).thenReturn(DEVICE_NAME)
+        `when`(cachedDevice.connectionSummary).thenReturn(CONNECTION_SUMMARY)
+        `when`(BluetoothUtils.getBtClassDrawableWithDescription(any(), any()))
+            .thenReturn(Pair.create(drawable, ""))
         val deviceItem = connectedDeviceItemFactory.create(context, cachedDevice)
 
         assertDeviceItem(deviceItem, DeviceItemType.CONNECTED_BLUETOOTH_DEVICE)
@@ -80,6 +98,10 @@
 
     @Test
     fun testSavedDeviceItemFactory_createFromCachedDevice() {
+        `when`(cachedDevice.name).thenReturn(DEVICE_NAME)
+        `when`(cachedDevice.connectionSummary).thenReturn(CONNECTION_SUMMARY)
+        `when`(BluetoothUtils.getBtClassDrawableWithDescription(any(), any()))
+            .thenReturn(Pair.create(drawable, ""))
         val deviceItem = savedDeviceItemFactory.create(context, cachedDevice)
 
         assertDeviceItem(deviceItem, DeviceItemType.SAVED_BLUETOOTH_DEVICE)
@@ -87,6 +109,90 @@
     }
 
     @Test
+    fun testAvailableAudioSharingMediaDeviceItemFactory_createFromCachedDevice() {
+        `when`(cachedDevice.name).thenReturn(DEVICE_NAME)
+        `when`(BluetoothUtils.getBtClassDrawableWithDescription(any(), any()))
+            .thenReturn(Pair.create(drawable, ""))
+        val deviceItem =
+            AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager)
+                .create(context, cachedDevice)
+
+        assertThat(deviceItem).isNotNull()
+        assertThat(deviceItem.type)
+            .isEqualTo(DeviceItemType.AVAILABLE_AUDIO_SHARING_MEDIA_BLUETOOTH_DEVICE)
+        assertThat(deviceItem.cachedBluetoothDevice).isEqualTo(cachedDevice)
+        assertThat(deviceItem.deviceName).isEqualTo(DEVICE_NAME)
+        assertThat(deviceItem.isActive).isFalse()
+        assertThat(deviceItem.connectionSummary)
+            .isEqualTo(
+                context.getString(
+                    R.string.quick_settings_bluetooth_device_audio_sharing_or_switch_active
+                )
+            )
+    }
+
+    @Test
+    fun testAvailableAudioSharingMediaDeviceItemFactory_isFilterMatched_flagOff_returnsFalse() {
+         // Flags.FLAG_ENABLE_LE_AUDIO_SHARING off or the device doesn't support broadcast
+         // source or assistant.
+        `when`(BluetoothUtils.isAudioSharingEnabled()).thenReturn(false)
+
+        assertThat(
+                AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager)
+                    .isFilterMatched(context, cachedDevice, audioManager)
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testAvailableAudioSharingMediaDeviceItemFactory_isFilterMatched_isActiveDevice_returnsFalse() {
+         // Flags.FLAG_ENABLE_LE_AUDIO_SHARING on and the device support broadcast source and
+         // assistant.
+        `when`(BluetoothUtils.isAudioSharingEnabled()).thenReturn(true)
+        `when`(BluetoothUtils.isActiveMediaDevice(any())).thenReturn(true)
+
+        assertThat(
+                AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager)
+                    .isFilterMatched(context, cachedDevice, audioManager)
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testAvailableAudioSharingMediaDeviceItemFactory_isFilterMatched_isNotAvailable_returnsFalse() {
+         // Flags.FLAG_ENABLE_LE_AUDIO_SHARING on and the device support broadcast source and
+         // assistant.
+        `when`(BluetoothUtils.isAudioSharingEnabled()).thenReturn(true)
+        `when`(BluetoothUtils.isActiveMediaDevice(any())).thenReturn(false)
+        `when`(BluetoothUtils.isAvailableMediaBluetoothDevice(any(), any())).thenReturn(true)
+        `when`(BluetoothUtils.isAvailableAudioSharingMediaBluetoothDevice(any(), any()))
+            .thenReturn(false)
+
+        assertThat(
+                AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager)
+                    .isFilterMatched(context, cachedDevice, audioManager)
+            )
+            .isFalse()
+    }
+
+    @Test
+    fun testAvailableAudioSharingMediaDeviceItemFactory_isFilterMatched_returnsTrue() {
+         // Flags.FLAG_ENABLE_LE_AUDIO_SHARING on and the device support broadcast source and
+         // assistant.
+        `when`(BluetoothUtils.isAudioSharingEnabled()).thenReturn(true)
+        `when`(BluetoothUtils.isActiveMediaDevice(any())).thenReturn(false)
+        `when`(BluetoothUtils.isAvailableMediaBluetoothDevice(any(), any())).thenReturn(true)
+        `when`(BluetoothUtils.isAvailableAudioSharingMediaBluetoothDevice(any(), any()))
+            .thenReturn(true)
+
+        assertThat(
+                AvailableAudioSharingMediaDeviceItemFactory(localBluetoothManager)
+                    .isFilterMatched(context, cachedDevice, audioManager)
+            )
+            .isTrue()
+    }
+
+    @Test
     @DisableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testSavedFactory_isFilterMatched_bondedAndNotConnected_returnsTrue() {
         `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
@@ -110,7 +216,6 @@
     @DisableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testSavedFactory_isFilterMatched_notBonded_returnsFalse() {
         `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_NONE)
-        `when`(cachedDevice.isConnected).thenReturn(false)
 
         assertThat(savedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isFalse()
@@ -119,12 +224,8 @@
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testSavedFactory_isFilterMatched_exclusivelyManaged_returnsFalse() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenReturn(ApplicationInfo())
-        `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(cachedDevice.isConnected).thenReturn(false)
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(true)
 
         assertThat(savedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isFalse()
@@ -132,7 +233,9 @@
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testSavedFactory_isFilterMatched_noExclusiveManager_returnsTrue() {
+    fun testSavedFactory_isFilterMatched_notExclusiveManaged_returnsTrue() {
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(false)
         `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
         `when`(cachedDevice.isConnected).thenReturn(false)
 
@@ -142,45 +245,9 @@
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testSavedFactory_isFilterMatched_exclusiveManagerNotEnabled_returnsTrue() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenReturn(ApplicationInfo().also { it.enabled = false })
-        `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(cachedDevice.isConnected).thenReturn(false)
-
-        assertThat(savedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isTrue()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testSavedFactory_isFilterMatched_exclusiveManagerNotInstalled_returnsTrue() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenThrow(PackageManager.NameNotFoundException("Test!"))
-        `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(cachedDevice.isConnected).thenReturn(false)
-
-        assertThat(savedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isTrue()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testSavedFactory_isFilterMatched_notExclusivelyManaged_notBonded_returnsFalse() {
-        `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_NONE)
-        `when`(cachedDevice.isConnected).thenReturn(false)
-
-        assertThat(savedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isFalse()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testSavedFactory_isFilterMatched_notExclusivelyManaged_connected_returnsFalse() {
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(false)
         `when`(cachedDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
         `when`(cachedDevice.isConnected).thenReturn(true)
 
@@ -191,9 +258,7 @@
     @Test
     @DisableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testConnectedFactory_isFilterMatched_bondedAndConnected_returnsTrue() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
+        `when`(BluetoothUtils.isConnectedBluetoothDevice(any(), any())).thenReturn(true)
 
         assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isTrue()
@@ -202,21 +267,6 @@
     @Test
     @DisableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testConnectedFactory_isFilterMatched_notConnected_returnsFalse() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(false)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
-
-        assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isFalse()
-    }
-
-    @Test
-    @DisableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testConnectedFactory_isFilterMatched_notBonded_returnsFalse() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_NONE)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
-
         assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isFalse()
     }
@@ -224,13 +274,8 @@
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testConnectedFactory_isFilterMatched_exclusivelyManaged_returnsFalse() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenReturn(ApplicationInfo())
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(true)
 
         assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isFalse()
@@ -239,9 +284,9 @@
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testConnectedFactory_isFilterMatched_noExclusiveManager_returnsTrue() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(false)
+        `when`(BluetoothUtils.isConnectedBluetoothDevice(any(), any())).thenReturn(true)
 
         assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isTrue()
@@ -249,51 +294,10 @@
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testConnectedFactory_isFilterMatched_exclusiveManagerNotEnabled_returnsTrue() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenReturn(ApplicationInfo().also { it.enabled = false })
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
-
-        assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isTrue()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testConnectedFactory_isFilterMatched_exclusiveManagerNotInstalled_returnsTrue() {
-        `when`(bluetoothDevice.getMetadata(BluetoothDevice.METADATA_EXCLUSIVE_MANAGER))
-            .thenReturn(TEST_EXCLUSIVE_MANAGER.toByteArray())
-        `when`(packageManager.getApplicationInfo(TEST_EXCLUSIVE_MANAGER, 0))
-            .thenThrow(PackageManager.NameNotFoundException("Test!"))
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
-
-        assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isTrue()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
-    fun testConnectedFactory_isFilterMatched_notExclusivelyManaged_notBonded_returnsFalse() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_NONE)
-        `when`(bluetoothDevice.isConnected).thenReturn(true)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
-
-        assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
-            .isFalse()
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_HIDE_EXCLUSIVELY_MANAGED_BLUETOOTH_DEVICE)
     fun testConnectedFactory_isFilterMatched_notExclusivelyManaged_notConnected_returnsFalse() {
-        `when`(bluetoothDevice.bondState).thenReturn(BluetoothDevice.BOND_BONDED)
-        `when`(bluetoothDevice.isConnected).thenReturn(false)
-        audioManager.setMode(AudioManager.MODE_NORMAL)
+        `when`(cachedDevice.device).thenReturn(bluetoothDevice)
+        `when`(BluetoothUtils.isExclusivelyManagedBluetoothDevice(any(), any())).thenReturn(false)
+        `when`(BluetoothUtils.isConnectedBluetoothDevice(any(), any())).thenReturn(false)
 
         assertThat(connectedDeviceItemFactory.isFilterMatched(context, cachedDevice, audioManager))
             .isFalse()
@@ -310,7 +314,5 @@
     companion object {
         const val DEVICE_NAME = "DeviceName"
         const val CONNECTION_SUMMARY = "ConnectionSummary"
-        private const val TEST_EXCLUSIVE_MANAGER = "com.test.manager"
-        private const val DEVICE_ADDRESS = "04:52:C7:0B:D8:3C"
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt b/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
index 8e215f9..fd550b0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/bouncer/ui/composable/BouncerContentTest.kt
@@ -83,7 +83,9 @@
         PlatformTheme {
             BouncerContent(
                 viewModel =
-                    rememberViewModel { kosmos.bouncerSceneContentViewModelFactory.create() },
+                    rememberViewModel("test") {
+                        kosmos.bouncerSceneContentViewModelFactory.create()
+                    },
                 layout = BouncerSceneLayout.BESIDE_USER_SWITCHER,
                 modifier = Modifier.fillMaxSize().testTag("BouncerContent"),
                 dialogFactory = bouncerDialogFactory
diff --git a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractorTest.kt
index 6aecc0e..40b2a08 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryFaceAuthInteractorTest.kt
@@ -63,6 +63,7 @@
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
@@ -361,15 +362,88 @@
         }
 
     @Test
-    fun faceAuthIsRequestedWhenQsExpansionStared() =
+    fun faceAuthIsRequestedWhenShadeExpansionStarted() =
         testScope.runTest {
             underTest.start()
 
-            underTest.onQsExpansionStared()
+            underTest.onShadeExpansionStarted()
 
             runCurrent()
             assertThat(faceAuthRepository.runningAuthRequest.value)
-                .isEqualTo(Pair(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, true))
+                .isEqualTo(Pair(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, false))
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun faceAuthIsRequestedWhenShadeExpansionIsStarted() =
+        testScope.runTest {
+            underTest.start()
+            faceAuthRepository.canRunFaceAuth.value = true
+            kosmos.sceneInteractor.snapToScene(toScene = Scenes.Lockscreen, "for-test")
+            runCurrent()
+
+            kosmos.sceneInteractor.changeScene(toScene = Scenes.Shade, loggingReason = "for-test")
+            kosmos.sceneInteractor.setTransitionState(
+                MutableStateFlow(
+                    ObservableTransitionState.Transition(
+                        fromScene = Scenes.Lockscreen,
+                        toScene = Scenes.Shade,
+                        currentScene = flowOf(Scenes.Lockscreen),
+                        progress = MutableStateFlow(0.2f),
+                        isInitiatedByUserInput = true,
+                        isUserInputOngoing = flowOf(false),
+                    )
+                )
+            )
+
+            runCurrent()
+            assertThat(faceAuthRepository.runningAuthRequest.value)
+                .isEqualTo(Pair(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, false))
+        }
+
+    @Test
+    @EnableSceneContainer
+    fun faceAuthIsRequestedOnlyOnceWhenShadeExpansionStarts() =
+        testScope.runTest {
+            underTest.start()
+            faceAuthRepository.canRunFaceAuth.value = true
+            kosmos.sceneInteractor.snapToScene(toScene = Scenes.Lockscreen, "for-test")
+            runCurrent()
+
+            kosmos.sceneInteractor.changeScene(toScene = Scenes.Shade, loggingReason = "for-test")
+            kosmos.sceneInteractor.setTransitionState(
+                MutableStateFlow(
+                    ObservableTransitionState.Transition(
+                        fromScene = Scenes.Lockscreen,
+                        toScene = Scenes.Shade,
+                        currentScene = flowOf(Scenes.Lockscreen),
+                        progress = MutableStateFlow(0.2f),
+                        isInitiatedByUserInput = true,
+                        isUserInputOngoing = flowOf(false),
+                    )
+                )
+            )
+
+            runCurrent()
+            assertThat(faceAuthRepository.runningAuthRequest.value)
+                .isEqualTo(Pair(FaceAuthUiEvent.FACE_AUTH_TRIGGERED_QS_EXPANDED, false))
+            faceAuthRepository.runningAuthRequest.value = null
+
+            // expansion progress shouldn't trigger face auth again
+            kosmos.sceneInteractor.setTransitionState(
+                MutableStateFlow(
+                    ObservableTransitionState.Transition(
+                        fromScene = Scenes.Lockscreen,
+                        toScene = Scenes.Shade,
+                        currentScene = flowOf(Scenes.Lockscreen),
+                        progress = MutableStateFlow(0.5f),
+                        isInitiatedByUserInput = true,
+                        isUserInputOngoing = flowOf(false),
+                    )
+                )
+            )
+
+            assertThat(faceAuthRepository.runningAuthRequest.value).isNull()
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
index 0ac04b6..76539d7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/display/data/repository/DisplayRepositoryTest.kt
@@ -467,6 +467,16 @@
             assertThat(values.toIdSets()).containsExactly(setOf(0, 1, 2))
         }
 
+    @Test
+    fun displayFlow_onlyDefaultDisplayAvailable_neverEmitsEmptySet() =
+        testScope.runTest {
+            setDisplays(0)
+
+            val values: List<Set<Display>> by collectValues(displayRepository.displays)
+
+            assertThat(values.toIdSets()).containsExactly(setOf(0))
+        }
+
     private fun Iterable<Display>.ids(): List<Int> = map { it.displayId }
 
     private fun Iterable<Set<Display>>.toIdSets(): List<Set<Int>> = map { it.ids().toSet() }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/haptics/slider/SliderHapticFeedbackProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/haptics/slider/SliderHapticFeedbackProviderTest.kt
index 933ddb5..4a80d72 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/haptics/slider/SliderHapticFeedbackProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/haptics/slider/SliderHapticFeedbackProviderTest.kt
@@ -21,7 +21,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.haptics.vibratorHelper
+import com.android.systemui.haptics.fakeVibratorHelper
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.fakeSystemClock
@@ -47,6 +47,7 @@
     private val lowTickDuration = 12 // Mocked duration of a low tick
     private val dragTextureThresholdMillis =
         lowTickDuration * config.numberOfLowTicks + config.deltaMillisForDragInterval
+    private val vibratorHelper = kosmos.fakeVibratorHelper
     private lateinit var sliderHapticFeedbackProvider: SliderHapticFeedbackProvider
 
     @Before
@@ -56,11 +57,11 @@
         whenever(velocityTracker.getAxisVelocity(config.velocityAxis))
             .thenReturn(config.maxVelocityToScale)
 
-        kosmos.vibratorHelper.primitiveDurations[VibrationEffect.Composition.PRIMITIVE_LOW_TICK] =
+        vibratorHelper.primitiveDurations[VibrationEffect.Composition.PRIMITIVE_LOW_TICK] =
             lowTickDuration
         sliderHapticFeedbackProvider =
             SliderHapticFeedbackProvider(
-                kosmos.vibratorHelper,
+                vibratorHelper,
                 velocityTracker,
                 config,
                 kosmos.fakeSystemClock,
@@ -136,7 +137,7 @@
             sliderHapticFeedbackProvider.onUpperBookend()
             sliderHapticFeedbackProvider.onUpperBookend()
 
-            assertEquals(/* expected=*/ 1, vibratorHelper.timesVibratedWithEffect(vibration))
+            assertEquals(/* expected= */ 1, vibratorHelper.timesVibratedWithEffect(vibration))
         }
 
     @Test
@@ -162,7 +163,7 @@
             sliderHapticFeedbackProvider.onProgress(progress)
 
             // THEN the correct composition only plays once
-            assertEquals(/* expected=*/ 1, vibratorHelper.timesVibratedWithEffect(ticks.compose()))
+            assertEquals(/* expected= */ 1, vibratorHelper.timesVibratedWithEffect(ticks.compose()))
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyboard/stickykeys/ui/viewmodel/StickyKeysIndicatorViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyboard/stickykeys/ui/viewmodel/StickyKeysIndicatorViewModelTest.kt
index 9d9e5be62..3ccb989 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyboard/stickykeys/ui/viewmodel/StickyKeysIndicatorViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyboard/stickykeys/ui/viewmodel/StickyKeysIndicatorViewModelTest.kt
@@ -34,14 +34,16 @@
 import com.android.systemui.keyboard.stickykeys.shared.model.ModifierKey.META
 import com.android.systemui.keyboard.stickykeys.shared.model.ModifierKey.SHIFT
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.fakeUserRepository
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.android.systemui.util.settings.repository.UserAwareSecureSettingsRepositoryImpl
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
@@ -57,29 +59,28 @@
 @RunWith(AndroidJUnit4::class)
 class StickyKeysIndicatorViewModelTest : SysuiTestCase() {
 
-    private val dispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(dispatcher)
+    private val kosmos = testKosmos()
+    private val dispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
     private lateinit var viewModel: StickyKeysIndicatorViewModel
     private val inputManager = mock<InputManager>()
     private val keyboardRepository = FakeKeyboardRepository()
-    private val secureSettings = FakeSettings()
+    private val secureSettings = kosmos.fakeSettings
     private val userRepository = Kosmos().fakeUserRepository
     private val captor =
         ArgumentCaptor.forClass(InputManager.StickyModifierStateListener::class.java)
 
     @Before
     fun setup() {
-        val settingsRepository = UserAwareSecureSettingsRepositoryImpl(
-            secureSettings,
-            userRepository,
-            dispatcher
-        )
-        val stickyKeysRepository = StickyKeysRepositoryImpl(
-            inputManager,
-            dispatcher,
-            settingsRepository,
-            mock<StickyKeysLogger>()
-        )
+        val settingsRepository =
+            UserAwareSecureSettingsRepositoryImpl(secureSettings, userRepository, dispatcher)
+        val stickyKeysRepository =
+            StickyKeysRepositoryImpl(
+                inputManager,
+                dispatcher,
+                settingsRepository,
+                mock<StickyKeysLogger>()
+            )
         setStickyKeySetting(enabled = false)
         viewModel =
             StickyKeysIndicatorViewModel(
@@ -182,16 +183,16 @@
             val stickyKeys by collectLastValue(viewModel.indicatorContent)
             setStickyKeysActive()
 
-            setStickyKeys(mapOf(
-                ALT to false,
-                META to false,
-                SHIFT to false))
+            setStickyKeys(mapOf(ALT to false, META to false, SHIFT to false))
 
-            assertThat(stickyKeys).isEqualTo(mapOf(
-                ALT to Locked(false),
-                META to Locked(false),
-                SHIFT to Locked(false),
-            ))
+            assertThat(stickyKeys)
+                .isEqualTo(
+                    mapOf(
+                        ALT to Locked(false),
+                        META to Locked(false),
+                        SHIFT to Locked(false),
+                    )
+                )
         }
     }
 
@@ -201,9 +202,7 @@
             val stickyKeys by collectLastValue(viewModel.indicatorContent)
             setStickyKeysActive()
 
-            setStickyKeys(mapOf(
-                ALT to false,
-                ALT to true))
+            setStickyKeys(mapOf(ALT to false, ALT to true))
 
             assertThat(stickyKeys).isEqualTo(mapOf(ALT to Locked(true)))
         }
@@ -215,17 +214,23 @@
             val stickyKeys by collectLastValue(viewModel.indicatorContent)
             setStickyKeysActive()
 
-            setStickyKeys(mapOf(
-                META to false,
-                SHIFT to false, // shift is sticky but not locked
-                CTRL to false))
+            setStickyKeys(
+                mapOf(
+                    META to false,
+                    SHIFT to false, // shift is sticky but not locked
+                    CTRL to false
+                )
+            )
             val previousShiftIndex = stickyKeys?.toList()?.indexOf(SHIFT to Locked(false))
 
-            setStickyKeys(mapOf(
-                SHIFT to false,
-                SHIFT to true, // shift is now locked
-                META to false,
-                CTRL to false))
+            setStickyKeys(
+                mapOf(
+                    SHIFT to false,
+                    SHIFT to true, // shift is now locked
+                    META to false,
+                    CTRL to false
+                )
+            )
             assertThat(stickyKeys?.toList()?.indexOf(SHIFT to Locked(true)))
                 .isEqualTo(previousShiftIndex)
         }
@@ -247,17 +252,27 @@
         StickyModifierState() {
 
         private fun isOn(key: ModifierKey) = keys.any { it.key == key && !it.value }
+
         private fun isLocked(key: ModifierKey) = keys.any { it.key == key && it.value }
 
         override fun isAltGrModifierLocked() = isLocked(ALT_GR)
+
         override fun isAltGrModifierOn() = isOn(ALT_GR)
+
         override fun isAltModifierLocked() = isLocked(ALT)
+
         override fun isAltModifierOn() = isOn(ALT)
+
         override fun isCtrlModifierLocked() = isLocked(CTRL)
+
         override fun isCtrlModifierOn() = isOn(CTRL)
+
         override fun isMetaModifierLocked() = isLocked(META)
+
         override fun isMetaModifierOn() = isOn(META)
+
         override fun isShiftModifierLocked() = isLocked(SHIFT)
+
         override fun isShiftModifierOn() = isOn(SHIFT)
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
index 29cd9a2..fa69fdd 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/CustomizationProviderTest.kt
@@ -50,6 +50,8 @@
 import com.android.systemui.keyguard.ui.preview.KeyguardPreviewRenderer
 import com.android.systemui.keyguard.ui.preview.KeyguardPreviewRendererFactory
 import com.android.systemui.keyguard.ui.preview.KeyguardRemotePreviewManager
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.kosmos.unconfinedTestScope
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.sceneInteractor
@@ -64,12 +66,12 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
+import com.android.systemui.util.settings.unconfinedDispatcherFakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.After
 import org.junit.Before
@@ -87,6 +89,11 @@
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 class CustomizationProviderTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.unconfinedTestDispatcher
+    private val testScope = kosmos.unconfinedTestScope
+    private val fakeSettings = kosmos.unconfinedDispatcherFakeSettings
+
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
     @Mock private lateinit var keyguardStateController: KeyguardStateController
     @Mock private lateinit var userTracker: UserTracker
@@ -104,9 +111,6 @@
     private lateinit var biometricSettingsRepository: FakeBiometricSettingsRepository
 
     private lateinit var underTest: CustomizationProvider
-    private lateinit var testScope: TestScope
-
-    private val kosmos = testKosmos()
 
     @Before
     fun setUp() {
@@ -120,8 +124,6 @@
         biometricSettingsRepository = FakeBiometricSettingsRepository()
 
         underTest = CustomizationProvider()
-        val testDispatcher = UnconfinedTestDispatcher()
-        testScope = TestScope(testDispatcher)
         val localUserSelectionManager =
             KeyguardQuickAffordanceLocalUserSelectionManager(
                 context = context,
@@ -170,7 +172,7 @@
                     KeyguardQuickAffordanceLegacySettingSyncer(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = testDispatcher,
-                        secureSettings = FakeSettings(),
+                        secureSettings = fakeSettings,
                         selectionsManager = localUserSelectionManager,
                     ),
                 dumpManager = mock(),
@@ -216,7 +218,7 @@
                 mainDispatcher = testDispatcher,
                 backgroundHandler = backgroundHandler,
             )
-        underTest.mainDispatcher = UnconfinedTestDispatcher()
+        underTest.mainDispatcher = testDispatcher
 
         underTest.attachInfoForTesting(
             context,
@@ -319,6 +321,7 @@
                         ),
                     )
                 )
+            runCurrent()
         }
 
     @Test
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 af5187d..1e9db64 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
@@ -25,15 +25,14 @@
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
 import com.android.systemui.keyguard.shared.model.ClockSizeSetting
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.res.R
 import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth
 import kotlin.test.Test
-import kotlinx.coroutines.CoroutineDispatcher
-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
@@ -44,12 +43,12 @@
 @SmallTest
 class KeyguardClockRepositoryTest : SysuiTestCase() {
 
-    private lateinit var scheduler: TestCoroutineScheduler
-    private lateinit var dispatcher: CoroutineDispatcher
-    private lateinit var scope: TestScope
+    private val kosmos = testKosmos()
+    private val dispatcher = kosmos.testDispatcher
+    private val scope = kosmos.testScope
+    private val fakeSettings = kosmos.fakeSettings
 
     private lateinit var underTest: KeyguardClockRepository
-    private lateinit var fakeSettings: FakeSettings
     @Mock private lateinit var clockRegistry: ClockRegistry
     @Mock private lateinit var clockEventController: ClockEventController
     private val fakeFeatureFlagsClassic = FakeFeatureFlagsClassic()
@@ -57,10 +56,6 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        fakeSettings = FakeSettings()
-        scheduler = TestCoroutineScheduler()
-        dispatcher = StandardTestDispatcher(scheduler)
-        scope = TestScope(dispatcher)
         underTest =
             KeyguardClockRepositoryImpl(
                 fakeSettings,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepositoryImplTest.kt
index 8b8a6cb..5a597fe 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardSmartspaceRepositoryImplTest.kt
@@ -21,14 +21,12 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.settings.FakeUserTracker
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.testKosmos
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth
 import kotlin.test.Test
-import kotlinx.coroutines.CoroutineDispatcher
-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
@@ -38,23 +36,18 @@
 @SmallTest
 class KeyguardSmartspaceRepositoryImplTest : SysuiTestCase() {
 
-    private lateinit var scheduler: TestCoroutineScheduler
-    private lateinit var dispatcher: CoroutineDispatcher
-    private lateinit var scope: TestScope
+    private val kosmos = testKosmos()
+    private val scope = kosmos.testScope
+    private val fakeSettings = kosmos.fakeSettings
 
     private lateinit var underTest: KeyguardSmartspaceRepository
-    private lateinit var fakeSettings: FakeSettings
     private lateinit var fakeUserTracker: FakeUserTracker
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        fakeSettings = FakeSettings()
         fakeUserTracker = FakeUserTracker()
         fakeSettings.userId = fakeUserTracker.userId
-        scheduler = TestCoroutineScheduler()
-        dispatcher = StandardTestDispatcher(scheduler)
-        scope = TestScope(dispatcher)
         underTest =
             KeyguardSmartspaceRepositoryImpl(
                 context = context,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
index 4e1b12f..43c7ed6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
@@ -21,7 +21,7 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.keyguard.WindowManagerLockscreenVisibilityManager
-import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
+import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
@@ -44,7 +44,8 @@
     @Mock private lateinit var activityTaskManagerService: IActivityTaskManager
     @Mock private lateinit var keyguardStateController: KeyguardStateController
     @Mock private lateinit var keyguardSurfaceBehindAnimator: KeyguardSurfaceBehindParamsApplier
-    @Mock private lateinit var keyguardTransitionInteractor: KeyguardTransitionInteractor
+    @Mock
+    private lateinit var keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor
 
     @Before
     fun setUp() {
@@ -57,7 +58,7 @@
                 activityTaskManagerService = activityTaskManagerService,
                 keyguardStateController = keyguardStateController,
                 keyguardSurfaceBehindAnimator = keyguardSurfaceBehindAnimator,
-                keyguardTransitionInteractor = keyguardTransitionInteractor,
+                keyguardDismissTransitionInteractor = keyguardDismissTransitionInteractor,
             )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
index 664a0bd..844a166 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelTest.kt
@@ -21,6 +21,7 @@
 import androidx.test.filters.SmallTest
 import com.android.internal.policy.IKeyguardDismissCallback
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.bouncer.domain.interactor.primaryBouncerInteractor
 import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -30,6 +31,7 @@
 import com.android.systemui.keyguard.shared.model.TransitionState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.statusbar.phone.statusBarKeyguardViewManager
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
@@ -69,6 +71,12 @@
     @Test
     fun onBackRequested() =
         testScope.runTest {
+            kosmos.primaryBouncerInteractor.setDismissAction(
+                mock(ActivityStarter.OnDismissAction::class.java),
+                {},
+            )
+            assertThat(kosmos.primaryBouncerInteractor.bouncerDismissAction).isNotNull()
+
             val dismissCallback = mock(IKeyguardDismissCallback::class.java)
             kosmos.dismissCallbackRegistry.addCallback(dismissCallback)
 
@@ -76,6 +84,7 @@
             kosmos.fakeExecutor.runAllReady()
             verify(statusBarKeyguardViewManager).hideAlternateBouncer(any())
             verify(dismissCallback).onDismissCancelled()
+            assertThat(kosmos.primaryBouncerInteractor.bouncerDismissAction).isNull()
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
index d13419e..1929cd1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardBottomAreaViewModelTest.kt
@@ -54,6 +54,8 @@
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancesMetricsLogger
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.res.R
 import com.android.systemui.scene.domain.interactor.sceneInteractor
@@ -69,13 +71,11 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlin.math.max
 import kotlin.math.min
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -93,6 +93,11 @@
 @RunWith(ParameterizedAndroidJunit4::class)
 class KeyguardBottomAreaViewModelTest(flags: FlagsParameterization) : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
+
     @Mock private lateinit var expandable: Expandable
     @Mock private lateinit var burnInHelperWrapper: BurnInHelperWrapper
     @Mock private lateinit var lockPatternUtils: LockPatternUtils
@@ -108,7 +113,6 @@
 
     private lateinit var underTest: KeyguardBottomAreaViewModel
 
-    private lateinit var testScope: TestScope
     private lateinit var repository: FakeKeyguardRepository
     private lateinit var homeControlsQuickAffordanceConfig: FakeKeyguardQuickAffordanceConfig
     private lateinit var quickAccessWalletAffordanceConfig: FakeKeyguardQuickAffordanceConfig
@@ -116,8 +120,6 @@
     private lateinit var dockManager: DockManagerFake
     private lateinit var biometricSettingsRepository: FakeBiometricSettingsRepository
 
-    private val kosmos = testKosmos()
-
     init {
         mSetFlagsRule.setFlagsParameterization(flags)
     }
@@ -162,8 +164,6 @@
         whenever(userTracker.userHandle).thenReturn(mock())
         whenever(lockPatternUtils.getStrongAuthForUser(anyInt()))
             .thenReturn(LockPatternUtils.StrongAuthTracker.STRONG_AUTH_NOT_REQUIRED)
-        val testDispatcher = StandardTestDispatcher()
-        testScope = TestScope(testDispatcher)
         val localUserSelectionManager =
             KeyguardQuickAffordanceLocalUserSelectionManager(
                 context = context,
@@ -199,7 +199,7 @@
                     KeyguardQuickAffordanceLegacySettingSyncer(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = testDispatcher,
-                        secureSettings = FakeSettings(),
+                        secureSettings = settings,
                         selectionsManager = localUserSelectionManager,
                     ),
                 configs =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
index 73b9f57..720f2e1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardQuickAffordancesCombinedViewModelTest.kt
@@ -47,8 +47,11 @@
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractor
 import com.android.systemui.keyguard.domain.interactor.KeyguardInteractorFactory
 import com.android.systemui.keyguard.domain.interactor.KeyguardQuickAffordanceInteractor
-import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor
-import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
+import com.android.systemui.keyguard.shared.model.KeyguardState.AOD
+import com.android.systemui.keyguard.shared.model.KeyguardState.GONE
+import com.android.systemui.keyguard.shared.model.KeyguardState.LOCKSCREEN
+import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.keyguard.shared.quickaffordance.ActivationState
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancePosition
 import com.android.systemui.keyguard.shared.quickaffordance.KeyguardQuickAffordancesMetricsLogger
@@ -56,7 +59,10 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.res.R
+import com.android.systemui.scene.data.repository.Idle
+import com.android.systemui.scene.data.repository.setTransition
 import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.settings.UserFileManager
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.shade.domain.interactor.ShadeInteractor
@@ -66,7 +72,7 @@
 import com.android.systemui.util.FakeSharedPreferences
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth
 import kotlin.math.min
 import kotlin.test.assertEquals
@@ -88,6 +94,10 @@
 @RunWith(AndroidJUnit4::class)
 class KeyguardQuickAffordancesCombinedViewModelTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
+
     @Mock private lateinit var activityStarter: ActivityStarter
     @Mock private lateinit var devicePolicyManager: DevicePolicyManager
     @Mock private lateinit var expandable: Expandable
@@ -144,13 +154,9 @@
     @Mock
     private lateinit var glanceableHubToLockscreenTransitionViewModel:
         GlanceableHubToLockscreenTransitionViewModel
-    @Mock private lateinit var transitionInteractor: KeyguardTransitionInteractor
-
-    private val kosmos = testKosmos()
 
     private lateinit var underTest: KeyguardQuickAffordancesCombinedViewModel
 
-    private val testScope = kosmos.testScope
     private lateinit var repository: FakeKeyguardRepository
     private lateinit var homeControlsQuickAffordanceConfig: FakeKeyguardQuickAffordanceConfig
     private lateinit var quickAccessWalletAffordanceConfig: FakeKeyguardQuickAffordanceConfig
@@ -163,8 +169,6 @@
     // the viewModel does a `map { 1 - it }` on this value, which is why it's different
     private val intendedShadeAlphaMutableStateFlow: MutableStateFlow<Float> = MutableStateFlow(0f)
 
-    private val intendedFinishedKeyguardStateFlow = MutableStateFlow(KeyguardState.LOCKSCREEN)
-
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
@@ -241,7 +245,7 @@
                     KeyguardQuickAffordanceLegacySettingSyncer(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = kosmos.testDispatcher,
-                        secureSettings = FakeSettings(),
+                        secureSettings = settings,
                         selectionsManager = localUserSelectionManager,
                     ),
                 configs =
@@ -256,7 +260,6 @@
 
         intendedAlphaMutableStateFlow.value = 1f
         intendedShadeAlphaMutableStateFlow.value = 0f
-        intendedFinishedKeyguardStateFlow.value = KeyguardState.LOCKSCREEN
         whenever(aodToLockscreenTransitionViewModel.shortcutsAlpha)
             .thenReturn(intendedAlphaMutableStateFlow)
         whenever(dozingToLockscreenTransitionViewModel.shortcutsAlpha).thenReturn(emptyFlow())
@@ -283,8 +286,6 @@
         whenever(glanceableHubToLockscreenTransitionViewModel.shortcutsAlpha)
             .thenReturn(emptyFlow())
         whenever(shadeInteractor.anyExpansion).thenReturn(intendedShadeAlphaMutableStateFlow)
-        whenever(transitionInteractor.finishedKeyguardState)
-            .thenReturn(intendedFinishedKeyguardStateFlow)
 
         underTest =
             KeyguardQuickAffordancesCombinedViewModel(
@@ -334,7 +335,7 @@
                     lockscreenToPrimaryBouncerTransitionViewModel,
                 lockscreenToGlanceableHubTransitionViewModel =
                     lockscreenToGlanceableHubTransitionViewModel,
-                transitionInteractor = transitionInteractor,
+                transitionInteractor = kosmos.keyguardTransitionInteractor,
             )
     }
 
@@ -403,7 +404,7 @@
         }
 
     @Test
-    @EnableFlags(com.android.systemui.Flags.FLAG_NEW_PICKER_UI)
+    @EnableFlags(com.android.systemui.shared.Flags.FLAG_NEW_CUSTOMIZATION_PICKER_UI)
     fun startButton_inPreviewMode_onPreviewQuickAffordanceSelected() =
         testScope.runTest {
             underTest.onPreviewSlotSelected(KeyguardQuickAffordanceSlots.SLOT_ID_BOTTOM_START)
@@ -776,7 +777,10 @@
     @Test
     fun shadeExpansionAlpha_changes_whenOnLockscreen() =
         testScope.runTest {
-            intendedFinishedKeyguardStateFlow.value = KeyguardState.LOCKSCREEN
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Lockscreen),
+                stateTransition = TransitionStep(from = AOD, to = LOCKSCREEN)
+            )
             intendedShadeAlphaMutableStateFlow.value = 0.25f
             val underTest = collectLastValue(underTest.transitionAlpha)
             assertEquals(0.75f, underTest())
@@ -788,7 +792,10 @@
     @Test
     fun shadeExpansionAlpha_alwaysZero_whenNotOnLockscreen() =
         testScope.runTest {
-            intendedFinishedKeyguardStateFlow.value = KeyguardState.GONE
+            kosmos.setTransition(
+                sceneTransition = Idle(Scenes.Gone),
+                stateTransition = TransitionStep(from = AOD, to = GONE)
+            )
             intendedShadeAlphaMutableStateFlow.value = 0.5f
             val underTest = collectLastValue(underTest.transitionAlpha)
             assertEquals(0f, underTest())
diff --git a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/ActivatableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/ActivatableTest.kt
index 67517a2..2ba670c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/ActivatableTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/ActivatableTest.kt
@@ -40,7 +40,7 @@
         composeRule.setContent {
             val keepAlive by keepAliveMutable
             if (keepAlive) {
-                rememberActivated {
+                rememberActivated("test") {
                     FakeActivatable(
                         onActivation = { isActive = true },
                         onDeactivation = { isActive = false },
@@ -58,7 +58,7 @@
         composeRule.setContent {
             val keepAlive by keepAliveMutable
             if (keepAlive) {
-                rememberActivated {
+                rememberActivated("name") {
                     FakeActivatable(
                         onActivation = { isActive = true },
                         onDeactivation = { isActive = false },
diff --git a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
index c7acd78..73f724e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/lifecycle/SysUiViewModelTest.kt
@@ -51,7 +51,7 @@
         composeRule.setContent {
             val keepAlive by keepAliveMutable
             if (keepAlive) {
-                rememberViewModel {
+                rememberViewModel("test") {
                     FakeSysUiViewModel(
                         onActivation = { isActive = true },
                         onDeactivation = { isActive = false },
@@ -69,21 +69,25 @@
         var isActive2 = false
         composeRule.setContent {
             val key by keyMutable
-            rememberViewModel(key) {
-                when (key) {
-                    1 ->
-                        FakeSysUiViewModel(
-                            onActivation = { isActive1 = true },
-                            onDeactivation = { isActive1 = false },
-                        )
-                    2 ->
-                        FakeSysUiViewModel(
-                            onActivation = { isActive2 = true },
-                            onDeactivation = { isActive2 = false },
-                        )
-                    else -> error("unsupported key $key")
+            // Need to explicitly state the type to avoid a weird issue where the factory seems to
+            // return Unit instead of FakeSysUiViewModel. It might be an issue with the compose
+            // compiler.
+            val unused: FakeSysUiViewModel =
+                rememberViewModel("test", key) {
+                    when (key) {
+                        1 ->
+                            FakeSysUiViewModel(
+                                onActivation = { isActive1 = true },
+                                onDeactivation = { isActive1 = false },
+                            )
+                        2 ->
+                            FakeSysUiViewModel(
+                                onActivation = { isActive2 = true },
+                                onDeactivation = { isActive2 = false },
+                            )
+                        else -> error("unsupported key $key")
+                    }
                 }
-            }
         }
         assertThat(isActive1).isTrue()
         assertThat(isActive2).isFalse()
@@ -106,7 +110,7 @@
         composeRule.setContent {
             val keepAlive by keepAliveMutable
             if (keepAlive) {
-                rememberViewModel {
+                rememberViewModel("test") {
                     FakeSysUiViewModel(
                         onActivation = { isActive = true },
                         onDeactivation = { isActive = false },
@@ -130,6 +134,7 @@
         val viewModel = FakeViewModel()
         backgroundScope.launch {
             view.viewModel(
+                traceName = "test",
                 minWindowLifecycleState = WindowLifecycleState.ATTACHED,
                 factory = { viewModel },
             ) {
@@ -151,7 +156,7 @@
     }
 }
 
-private class FakeViewModel : SysUiViewModel, ExclusiveActivatable() {
+private class FakeViewModel : ExclusiveActivatable() {
     var isActivated = false
 
     override suspend fun onActivated(): Nothing {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/log/table/LogDiffsForTableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/log/table/LogDiffsForTableTest.kt
index e55cb12..030b172 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/log/table/LogDiffsForTableTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/log/table/LogDiffsForTableTest.kt
@@ -19,8 +19,8 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.log.LogcatEchoTrackerAlways
 import com.android.systemui.log.table.TableChange.Companion.IS_INITIAL_PREFIX
-import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.io.PrintWriter
@@ -57,9 +57,7 @@
                 MAX_SIZE,
                 BUFFER_NAME,
                 systemClock,
-                mock(),
-                testDispatcher,
-                testScope.backgroundScope,
+                LogcatEchoTrackerAlways(),
             )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferFactoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferFactoryTest.kt
index 8c62bc2..dfd964f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferFactoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferFactoryTest.kt
@@ -20,25 +20,20 @@
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.log.LogcatEchoTrackerAlways
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class TableLogBufferFactoryTest : SysuiTestCase() {
     private val dumpManager: DumpManager = mock()
     private val systemClock = FakeSystemClock()
-    private val testDispatcher = UnconfinedTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
     private val underTest =
-        TableLogBufferFactory(dumpManager, systemClock, mock(), testDispatcher, testScope)
+        TableLogBufferFactory(dumpManager, systemClock, LogcatEchoTrackerAlways())
 
     @Test
     fun create_alwaysCreatesNewInstance() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferTest.kt b/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferTest.kt
index ace562b..9c4c862 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/log/table/TableLogBufferTest.kt
@@ -23,22 +23,18 @@
 import com.android.systemui.log.core.LogLevel
 import com.android.systemui.log.table.TableChange.Companion.IS_INITIAL_PREFIX
 import com.android.systemui.log.table.TableChange.Companion.MAX_STRING_LENGTH
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.eq
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.io.PrintWriter
 import java.io.StringWriter
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class TableLogBufferTest : SysuiTestCase() {
@@ -49,9 +45,6 @@
     private lateinit var logcatEchoTracker: LogcatEchoTracker
     private lateinit var localLogcat: FakeLogProxy
 
-    private val testDispatcher = UnconfinedTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
-
     @Before
     fun setup() {
         localLogcat = FakeLogProxy()
@@ -65,8 +58,6 @@
                 NAME,
                 systemClock,
                 logcatEchoTracker,
-                testDispatcher,
-                testScope.backgroundScope,
                 localLogcat = localLogcat,
             )
     }
@@ -78,8 +69,6 @@
             "name",
             systemClock,
             logcatEchoTracker,
-            testDispatcher,
-            testScope.backgroundScope,
             localLogcat = localLogcat,
         )
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorTest.kt
index 99c5b7c..f4c2b47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorTest.kt
@@ -41,12 +41,11 @@
 import android.os.UserHandle
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.FlagsParameterization
 import android.provider.Settings
 import android.service.notification.StatusBarNotification
-import android.testing.TestableLooper
 import android.testing.TestableLooper.RunWithLooper
 import androidx.media.utils.MediaConstants
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.dx.mockito.inline.extended.ExtendedMockito
 import com.android.internal.logging.InstanceId
@@ -58,9 +57,15 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.flags.EnableSceneContainer
+import com.android.systemui.flags.Flags.MEDIA_REMOTE_RESUME
+import com.android.systemui.flags.Flags.MEDIA_RESUME_PROGRESS
+import com.android.systemui.flags.Flags.MEDIA_RETAIN_RECOMMENDATIONS
+import com.android.systemui.flags.Flags.MEDIA_RETAIN_SESSIONS
+import com.android.systemui.flags.Flags.MEDIA_SESSION_ACTIONS
+import com.android.systemui.flags.fakeFeatureFlagsClassic
+import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
-import com.android.systemui.media.controls.data.repository.MediaDataRepository
-import com.android.systemui.media.controls.data.repository.MediaFilterRepository
+import com.android.systemui.media.controls.data.repository.mediaDataRepository
 import com.android.systemui.media.controls.data.repository.mediaFilterRepository
 import com.android.systemui.media.controls.domain.pipeline.interactor.MediaCarouselInteractor
 import com.android.systemui.media.controls.domain.resume.MediaResumeListener
@@ -70,24 +75,21 @@
 import com.android.systemui.media.controls.shared.model.MediaData
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaData
 import com.android.systemui.media.controls.shared.model.SmartspaceMediaDataProvider
-import com.android.systemui.media.controls.util.MediaControllerFactory
-import com.android.systemui.media.controls.util.MediaFlags
 import com.android.systemui.media.controls.util.MediaUiEventLogger
-import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.media.controls.util.fakeMediaControllerFactory
+import com.android.systemui.media.controls.util.mediaFlags
+import com.android.systemui.plugins.activityStarter
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.SbnBuilder
 import com.android.systemui.statusbar.notificationLockscreenUserManager
 import com.android.systemui.testKosmos
 import com.android.systemui.util.concurrency.FakeExecutor
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.android.systemui.util.time.FakeSystemClock
-import com.android.systemui.utils.os.FakeHandler
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestDispatcher
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
-import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.runCurrent
 import org.junit.After
 import org.junit.Before
 import org.junit.Rule
@@ -111,6 +113,8 @@
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 private const val KEY = "KEY"
 private const val KEY_2 = "KEY_2"
@@ -134,13 +138,15 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWithLooper(setAsMainLooper = true)
-@RunWith(AndroidJUnit4::class)
+@RunWith(ParameterizedAndroidJunit4::class)
 @EnableSceneContainer
-class MediaDataProcessorTest : SysuiTestCase() {
-    val kosmos = testKosmos()
+class MediaDataProcessorTest(flags: FlagsParameterization) : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val settings = kosmos.fakeSettings
 
     @JvmField @Rule val mockito = MockitoJUnit.rule()
-    @Mock lateinit var mediaControllerFactory: MediaControllerFactory
     @Mock lateinit var controller: MediaController
     @Mock lateinit var transportControls: MediaController.TransportControls
     @Mock lateinit var playbackInfo: MediaController.PlaybackInfo
@@ -158,7 +164,6 @@
     @Mock lateinit var mediaDataCombineLatest: MediaDataCombineLatest
     @Mock lateinit var listener: MediaDataManager.Listener
     @Mock lateinit var pendingIntent: PendingIntent
-    @Mock lateinit var activityStarter: ActivityStarter
     @Mock lateinit var smartspaceManager: SmartspaceManager
     @Mock lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
     private lateinit var smartspaceMediaDataProvider: SmartspaceMediaDataProvider
@@ -166,7 +171,6 @@
     @Mock private lateinit var mediaRecommendationItem: SmartspaceAction
     private lateinit var validRecommendationList: List<SmartspaceAction>
     @Mock private lateinit var mediaSmartspaceBaseAction: SmartspaceAction
-    @Mock private lateinit var mediaFlags: MediaFlags
     @Mock private lateinit var logger: MediaUiEventLogger
     private lateinit var mediaCarouselInteractor: MediaCarouselInteractor
     private lateinit var mediaDataProcessor: MediaDataProcessor
@@ -179,13 +183,28 @@
     @Captor lateinit var smartSpaceConfigBuilderCaptor: ArgumentCaptor<SmartspaceConfig>
     @Mock private lateinit var ugm: IUriGrantsManager
     @Mock private lateinit var imageSource: ImageDecoder.Source
-    private lateinit var mediaDataRepository: MediaDataRepository
-    private lateinit var testScope: TestScope
-    private lateinit var testDispatcher: TestDispatcher
-    private lateinit var testableLooper: TestableLooper
-    private lateinit var fakeHandler: FakeHandler
 
-    private val settings = FakeSettings()
+    companion object {
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams(): List<FlagsParameterization> {
+            return FlagsParameterization.progressionOf(
+                Flags.FLAG_MEDIA_LOAD_METADATA_VIA_MEDIA_DATA_LOADER
+            )
+        }
+    }
+
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags)
+    }
+
+    private val fakeFeatureFlags = kosmos.fakeFeatureFlagsClassic
+    private val activityStarter = kosmos.activityStarter
+    private val mediaControllerFactory = kosmos.fakeMediaControllerFactory
+    private val notificationLockscreenUserManager = kosmos.notificationLockscreenUserManager
+    private val mediaFilterRepository = kosmos.mediaFilterRepository
+    private val mediaDataFilter = kosmos.mediaDataFilter
+
     private val instanceIdSequence = InstanceIdSequenceFake(1 shl 20)
 
     private val originalSmartspaceSetting =
@@ -194,9 +213,6 @@
             Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION,
             1
         )
-    private val notificationLockscreenUserManager = kosmos.notificationLockscreenUserManager
-    private val mediaFilterRepository: MediaFilterRepository = kosmos.mediaFilterRepository
-    private val mediaDataFilter: MediaDataFilterImpl = kosmos.mediaDataFilter
 
     private lateinit var staticMockSession: MockitoSession
 
@@ -212,17 +228,12 @@
         foregroundExecutor = FakeExecutor(clock)
         backgroundExecutor = FakeExecutor(clock)
         uiExecutor = FakeExecutor(clock)
-        testableLooper = TestableLooper.get(this)
-        fakeHandler = FakeHandler(testableLooper.looper)
         smartspaceMediaDataProvider = SmartspaceMediaDataProvider()
         Settings.Secure.putInt(
             context.contentResolver,
             Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION,
             1
         )
-        testDispatcher = UnconfinedTestDispatcher()
-        testScope = TestScope(testDispatcher)
-        mediaDataRepository = MediaDataRepository(mediaFlags, dumpManager)
         mediaDataProcessor =
             MediaDataProcessor(
                 context = context,
@@ -231,7 +242,7 @@
                 backgroundExecutor = backgroundExecutor,
                 uiExecutor = uiExecutor,
                 foregroundExecutor = foregroundExecutor,
-                handler = fakeHandler,
+                mainDispatcher = testDispatcher,
                 mediaControllerFactory = mediaControllerFactory,
                 broadcastDispatcher = broadcastDispatcher,
                 dumpManager = dumpManager,
@@ -241,13 +252,15 @@
                 useQsMediaPlayer = true,
                 systemClock = clock,
                 secureSettings = settings,
-                mediaFlags = mediaFlags,
+                mediaFlags = kosmos.mediaFlags,
                 logger = logger,
                 smartspaceManager = smartspaceManager,
                 keyguardUpdateMonitor = keyguardUpdateMonitor,
-                mediaDataRepository = mediaDataRepository,
+                mediaDataRepository = kosmos.mediaDataRepository,
+                mediaDataLoader = { kosmos.mediaDataLoader },
             )
         mediaDataProcessor.start()
+        testScope.runCurrent()
         mediaCarouselInteractor =
             MediaCarouselInteractor(
                 applicationScope = testScope.backgroundScope,
@@ -259,7 +272,7 @@
                 mediaDataCombineLatest = mediaDataCombineLatest,
                 mediaDataFilter = mediaDataFilter,
                 mediaFilterRepository = mediaFilterRepository,
-                mediaFlags = mediaFlags
+                mediaFlags = kosmos.mediaFlags
             )
         mediaCarouselInteractor.start()
         verify(mediaTimeoutListener).stateCallback = capture(stateCallbackCaptor)
@@ -295,7 +308,7 @@
                 putString(MediaMetadata.METADATA_KEY_TITLE, SESSION_TITLE)
             }
         verify(smartspaceManager).createSmartspaceSession(capture(smartSpaceConfigBuilderCaptor))
-        whenever(mediaControllerFactory.create(eq(session.sessionToken))).thenReturn(controller)
+        mediaControllerFactory.setControllerForToken(session.sessionToken, controller)
         whenever(controller.transportControls).thenReturn(transportControls)
         whenever(controller.playbackInfo).thenReturn(playbackInfo)
         whenever(controller.metadata).thenReturn(metadataBuilder.build())
@@ -325,10 +338,11 @@
         whenever(mediaSmartspaceTarget.iconGrid).thenReturn(validRecommendationList)
         whenever(mediaSmartspaceTarget.creationTimeMillis).thenReturn(SMARTSPACE_CREATION_TIME)
         whenever(mediaSmartspaceTarget.expiryTimeMillis).thenReturn(SMARTSPACE_EXPIRY_TIME)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(false)
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(false)
-        whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(false)
-        whenever(mediaFlags.isRemoteResumeAllowed()).thenReturn(false)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, false)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, false)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, false)
+        fakeFeatureFlags.set(MEDIA_REMOTE_RESUME, false)
+        fakeFeatureFlags.set(MEDIA_RETAIN_RECOMMENDATIONS, false)
         whenever(logger.getNewInstanceId()).thenReturn(instanceIdSequence.newInstanceId())
         whenever(keyguardUpdateMonitor.isUserInLockdown(any())).thenReturn(false)
     }
@@ -374,6 +388,7 @@
             PACKAGE_NAME
         )
 
+        testScope.runCurrent()
         backgroundExecutor.runAllReady()
         foregroundExecutor.runAllReady()
         verify(listener)
@@ -399,7 +414,7 @@
     @Test
     fun testLoadsMetadataOnBackground() {
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
-        assertThat(backgroundExecutor.numPending()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 0, background = 1)
     }
 
     @Test
@@ -416,8 +431,7 @@
 
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
 
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -434,8 +448,7 @@
     fun testOnMetaDataLoaded_withoutExplicitIndicator() {
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
 
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -462,10 +475,8 @@
 
     @Test
     fun testOnMetaDataLoaded_conservesActiveFlag() {
-        whenever(mediaControllerFactory.create(anyObject())).thenReturn(controller)
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -509,8 +520,7 @@
             }
 
         mediaDataProcessor.onNotificationAdded(KEY, notif)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -596,8 +606,7 @@
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
 
         // Then a media control is created with a placeholder title string
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -627,8 +636,7 @@
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
 
         // Then a media control is created with a placeholder title string
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -669,8 +677,7 @@
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
 
         // Then the media control is added using the notification's title
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -778,8 +785,7 @@
         // GIVEN that the manager has two notifications with resume actions
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
         mediaDataProcessor.onNotificationAdded(KEY_2, mediaNotification)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(2)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(2)
+        testScope.assertRunAllReady(foreground = 2, background = 2)
 
         verify(listener)
             .onMediaDataLoaded(
@@ -866,7 +872,7 @@
     @Test
     fun testOnNotificationRemoved_withResumption_isRemoteAndRemoteAllowed() {
         // With the flag enabled to allow remote media to resume
-        whenever(mediaFlags.isRemoteResumeAllowed()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_REMOTE_RESUME, true)
 
         // GIVEN that the manager has a notification with a resume action, but is not local
         whenever(controller.metadata).thenReturn(metadataBuilder.build())
@@ -897,7 +903,7 @@
     @Test
     fun testOnNotificationRemoved_withResumption_isRcnAndRemoteAllowed() {
         // With the flag enabled to allow remote media to resume
-        whenever(mediaFlags.isRemoteResumeAllowed()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_REMOTE_RESUME, true)
 
         // GIVEN that the manager has a remote cast notification
         addNotificationAndLoad(remoteCastNotification)
@@ -1016,7 +1022,7 @@
 
     @Test
     fun testAddResumptionControls_hasPartialProgress() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added with partial progress
         val progress = 0.5
@@ -1043,7 +1049,7 @@
 
     @Test
     fun testAddResumptionControls_hasNotPlayedProgress() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added that have not been played
         val extras =
@@ -1068,7 +1074,7 @@
 
     @Test
     fun testAddResumptionControls_hasFullProgress() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added with progress info
         val extras =
@@ -1094,7 +1100,7 @@
 
     @Test
     fun testAddResumptionControls_hasNoExtras() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added that do not have any extras
         val desc =
@@ -1112,7 +1118,7 @@
 
     @Test
     fun testAddResumptionControls_hasEmptyTitle() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added that have empty title
         val desc =
@@ -1131,8 +1137,7 @@
         )
 
         // Resumption controls are not added.
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(0)
+        testScope.assertRunAllReady(foreground = 0, background = 1)
         verify(listener, never())
             .onMediaDataLoaded(
                 eq(PACKAGE_NAME),
@@ -1146,7 +1151,7 @@
 
     @Test
     fun testAddResumptionControls_hasBlankTitle() {
-        whenever(mediaFlags.isResumeProgressEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RESUME_PROGRESS, true)
 
         // WHEN resumption controls are added that have a blank title
         val desc =
@@ -1165,8 +1170,7 @@
         )
 
         // Resumption controls are not added.
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(0)
+        testScope.assertRunAllReady(foreground = 0, background = 1)
         verify(listener, never())
             .onMediaDataLoaded(
                 eq(PACKAGE_NAME),
@@ -1233,8 +1237,7 @@
         mediaDataProcessor.onNotificationAdded(KEY, notif)
 
         // THEN it still loads
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -1351,7 +1354,7 @@
 
     @Test
     fun testOnSmartspaceMediaDataLoaded_persistentEnabled_headphoneTrigger_isActive() {
-        whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_RECOMMENDATIONS, true)
         smartspaceMediaDataProvider.onTargetsAvailable(listOf(mediaSmartspaceTarget))
         val instanceId = instanceIdSequence.lastInstanceId
 
@@ -1377,7 +1380,7 @@
 
     @Test
     fun testOnSmartspaceMediaDataLoaded_persistentEnabled_periodicTrigger_notActive() {
-        whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_RECOMMENDATIONS, true)
         val extras =
             Bundle().apply {
                 putString("package_name", PACKAGE_NAME)
@@ -1411,7 +1414,7 @@
 
     @Test
     fun testOnSmartspaceMediaDataLoaded_persistentEnabled_noTargets_inactive() {
-        whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_RECOMMENDATIONS, true)
 
         smartspaceMediaDataProvider.onTargetsAvailable(listOf(mediaSmartspaceTarget))
         val instanceId = instanceIdSequence.lastInstanceId
@@ -1443,7 +1446,7 @@
 
     @Test
     fun testSetRecommendationInactive_notifiesListeners() {
-        whenever(mediaFlags.isPersistentSsCardEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_RECOMMENDATIONS, true)
 
         smartspaceMediaDataProvider.onTargetsAvailable(listOf(mediaSmartspaceTarget))
         val instanceId = instanceIdSequence.lastInstanceId
@@ -1476,6 +1479,7 @@
     fun testOnSmartspaceMediaDataLoaded_settingDisabled_doesNothing() {
         // WHEN media recommendation setting is off
         settings.putInt(Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION, 0)
+        testScope.runCurrent()
 
         smartspaceMediaDataProvider.onTargetsAvailable(listOf(mediaSmartspaceTarget))
 
@@ -1493,6 +1497,7 @@
 
         // WHEN the media recommendation setting is turned off
         settings.putInt(Settings.Secure.MEDIA_CONTROLS_RECOMMENDATION, 0)
+        testScope.runCurrent()
 
         // THEN listeners are notified
         uiExecutor.advanceClockToLast()
@@ -1513,8 +1518,7 @@
     fun testOnMediaDataTimedOut_updatesLastActiveTime() {
         // GIVEN that the manager has a notification
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
 
         // WHEN the notification times out
         clock.advanceTime(100)
@@ -1622,8 +1626,7 @@
 
         // WHEN the notification is loaded
         mediaDataProcessor.onNotificationAdded(KEY, notif)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
 
         // THEN only the first MAX_COMPACT_ACTIONS are actually set
         verify(listener)
@@ -1658,8 +1661,7 @@
 
         // WHEN the notification is loaded
         mediaDataProcessor.onNotificationAdded(KEY, notif)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
 
         // THEN only the first MAX_NOTIFICATION_ACTIONS are actually included
         verify(listener)
@@ -1678,7 +1680,7 @@
     @Test
     fun testPlaybackActions_noState_usesNotification() {
         val desc = "Notification Action"
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         whenever(controller.playbackState).thenReturn(null)
 
         val notifWithAction =
@@ -1693,8 +1695,7 @@
             }
         mediaDataProcessor.onNotificationAdded(KEY, notifWithAction)
 
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -1713,7 +1714,7 @@
     @Test
     fun testPlaybackActions_hasPrevNext() {
         val customDesc = arrayOf("custom 1", "custom 2", "custom 3", "custom 4")
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val stateActions =
             PlaybackState.ACTION_PLAY or
                 PlaybackState.ACTION_SKIP_TO_PREVIOUS or
@@ -1757,7 +1758,7 @@
     @Test
     fun testPlaybackActions_noPrevNext_usesCustom() {
         val customDesc = arrayOf("custom 1", "custom 2", "custom 3", "custom 4", "custom 5")
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val stateActions = PlaybackState.ACTION_PLAY
         val stateBuilder = PlaybackState.Builder().setActions(stateActions)
         customDesc.forEach {
@@ -1789,7 +1790,7 @@
 
     @Test
     fun testPlaybackActions_connecting() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val stateActions = PlaybackState.ACTION_PLAY
         val stateBuilder =
             PlaybackState.Builder()
@@ -1809,87 +1810,76 @@
 
     @Test
     @EnableFlags(Flags.FLAG_MEDIA_CONTROLS_DRAWABLES_REUSE)
-    fun postWithPlaybackActions_drawablesReused() =
-        kosmos.testScope.runTest {
-            whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
-            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
-            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
-            val stateActions =
-                PlaybackState.ACTION_PAUSE or
-                    PlaybackState.ACTION_SKIP_TO_PREVIOUS or
-                    PlaybackState.ACTION_SKIP_TO_NEXT
-            val stateBuilder =
-                PlaybackState.Builder()
-                    .setState(PlaybackState.STATE_PLAYING, 0, 10f)
-                    .setActions(stateActions)
-            whenever(controller.playbackState).thenReturn(stateBuilder.build())
-            val userEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+    fun postWithPlaybackActions_drawablesReused() {
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
+        whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+        whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+        val stateActions =
+            PlaybackState.ACTION_PAUSE or
+                PlaybackState.ACTION_SKIP_TO_PREVIOUS or
+                PlaybackState.ACTION_SKIP_TO_NEXT
+        val stateBuilder =
+            PlaybackState.Builder()
+                .setState(PlaybackState.STATE_PLAYING, 0, 10f)
+                .setActions(stateActions)
+        whenever(controller.playbackState).thenReturn(stateBuilder.build())
+        val userEntries by testScope.collectLastValue(mediaFilterRepository.selectedUserEntries)
 
-            mediaDataProcessor.addInternalListener(mediaDataFilter)
-            mediaDataFilter.mediaDataProcessor = mediaDataProcessor
-            addNotificationAndLoad()
+        mediaDataProcessor.addInternalListener(mediaDataFilter)
+        mediaDataFilter.mediaDataProcessor = mediaDataProcessor
+        addNotificationAndLoad()
 
-            assertThat(userEntries).hasSize(1)
-            val firstSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
+        assertThat(userEntries).hasSize(1)
+        val firstSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
 
-            addNotificationAndLoad()
+        addNotificationAndLoad()
 
-            assertThat(userEntries).hasSize(1)
-            val secondSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
-            assertThat(secondSemanticActions.playOrPause?.icon)
-                .isEqualTo(firstSemanticActions.playOrPause?.icon)
-            assertThat(secondSemanticActions.playOrPause?.background)
-                .isEqualTo(firstSemanticActions.playOrPause?.background)
-            assertThat(secondSemanticActions.nextOrCustom?.icon)
-                .isEqualTo(firstSemanticActions.nextOrCustom?.icon)
-            assertThat(secondSemanticActions.prevOrCustom?.icon)
-                .isEqualTo(firstSemanticActions.prevOrCustom?.icon)
-        }
+        assertThat(userEntries).hasSize(1)
+        val secondSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
+        assertThat(secondSemanticActions.nextOrCustom?.icon)
+            .isEqualTo(firstSemanticActions.nextOrCustom?.icon)
+        assertThat(secondSemanticActions.prevOrCustom?.icon)
+            .isEqualTo(firstSemanticActions.prevOrCustom?.icon)
+    }
 
     @Test
     @DisableFlags(Flags.FLAG_MEDIA_CONTROLS_DRAWABLES_REUSE)
-    fun postWithPlaybackActions_drawablesNotReused() =
-        kosmos.testScope.runTest {
-            whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
-            whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
-            whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
-            val stateActions =
-                PlaybackState.ACTION_PAUSE or
-                    PlaybackState.ACTION_SKIP_TO_PREVIOUS or
-                    PlaybackState.ACTION_SKIP_TO_NEXT
-            val stateBuilder =
-                PlaybackState.Builder()
-                    .setState(PlaybackState.STATE_PLAYING, 0, 10f)
-                    .setActions(stateActions)
-            whenever(controller.playbackState).thenReturn(stateBuilder.build())
-            val userEntries by collectLastValue(mediaFilterRepository.selectedUserEntries)
+    fun postWithPlaybackActions_drawablesNotReused() {
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
+        whenever(notificationLockscreenUserManager.isCurrentProfile(USER_ID)).thenReturn(true)
+        whenever(notificationLockscreenUserManager.isProfileAvailable(USER_ID)).thenReturn(true)
+        val stateActions =
+            PlaybackState.ACTION_PAUSE or
+                PlaybackState.ACTION_SKIP_TO_PREVIOUS or
+                PlaybackState.ACTION_SKIP_TO_NEXT
+        val stateBuilder =
+            PlaybackState.Builder()
+                .setState(PlaybackState.STATE_PLAYING, 0, 10f)
+                .setActions(stateActions)
+        whenever(controller.playbackState).thenReturn(stateBuilder.build())
+        val userEntries by testScope.collectLastValue(mediaFilterRepository.selectedUserEntries)
 
-            mediaDataProcessor.addInternalListener(mediaDataFilter)
-            mediaDataFilter.mediaDataProcessor = mediaDataProcessor
-            addNotificationAndLoad()
+        mediaDataProcessor.addInternalListener(mediaDataFilter)
+        mediaDataFilter.mediaDataProcessor = mediaDataProcessor
+        addNotificationAndLoad()
 
-            assertThat(userEntries).hasSize(1)
-            val firstSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
+        assertThat(userEntries).hasSize(1)
+        val firstSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
 
-            addNotificationAndLoad()
+        addNotificationAndLoad()
 
-            assertThat(userEntries).hasSize(1)
-            val secondSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
-
-            assertThat(secondSemanticActions.playOrPause?.icon)
-                .isNotEqualTo(firstSemanticActions.playOrPause?.icon)
-            assertThat(secondSemanticActions.playOrPause?.background)
-                .isNotEqualTo(firstSemanticActions.playOrPause?.background)
-            assertThat(secondSemanticActions.nextOrCustom?.icon)
-                .isNotEqualTo(firstSemanticActions.nextOrCustom?.icon)
-            assertThat(secondSemanticActions.prevOrCustom?.icon)
-                .isNotEqualTo(firstSemanticActions.prevOrCustom?.icon)
-        }
+        assertThat(userEntries).hasSize(1)
+        val secondSemanticActions = userEntries?.values?.toList()?.get(0)?.semanticActions!!
+        assertThat(secondSemanticActions.nextOrCustom?.icon)
+            .isNotEqualTo(firstSemanticActions.nextOrCustom?.icon)
+        assertThat(secondSemanticActions.prevOrCustom?.icon)
+            .isNotEqualTo(firstSemanticActions.prevOrCustom?.icon)
+    }
 
     @Test
     fun testPlaybackActions_reservedSpace() {
         val customDesc = arrayOf("custom 1", "custom 2", "custom 3", "custom 4")
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val stateActions = PlaybackState.ACTION_PLAY
         val stateBuilder = PlaybackState.Builder().setActions(stateActions)
         customDesc.forEach {
@@ -1927,7 +1917,7 @@
 
     @Test
     fun testPlaybackActions_playPause_hasButton() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val stateActions = PlaybackState.ACTION_PLAY_PAUSE
         val stateBuilder = PlaybackState.Builder().setActions(stateActions)
         whenever(controller.playbackState).thenReturn(stateBuilder.build())
@@ -1964,8 +1954,7 @@
 
         // update to remote cast
         mediaDataProcessor.onNotificationAdded(KEY, remoteCastNotification)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(logger)
             .logPlaybackLocationChange(
                 anyInt(),
@@ -2027,7 +2016,7 @@
 
     @Test
     fun testPlaybackState_PauseWhenFlagTrue_keyExists_callsListener() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         val state = PlaybackState.Builder().setState(PlaybackState.STATE_PAUSED, 0L, 1f).build()
         whenever(controller.playbackState).thenReturn(state)
 
@@ -2071,6 +2060,7 @@
             pendingIntent,
             PACKAGE_NAME
         )
+        testScope.runCurrent()
         backgroundExecutor.runAllReady()
         foregroundExecutor.runAllReady()
 
@@ -2149,7 +2139,7 @@
 
     @Test
     fun testRetain_notifPlayer_notifRemoved_setToResume() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
 
         // When a media control based on notification is added, times out, and then removed
         addNotificationAndLoad()
@@ -2179,7 +2169,7 @@
 
     @Test
     fun testRetain_notifPlayer_sessionDestroyed_doesNotChange() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
 
         // When a media control based on notification is added and times out
         addNotificationAndLoad()
@@ -2197,7 +2187,7 @@
 
     @Test
     fun testRetain_notifPlayer_removeWhileActive_fullyRemoved() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
 
         // When a media control based on notification is added and then removed, without timing out
         addNotificationAndLoad()
@@ -2214,7 +2204,7 @@
 
     @Test
     fun testRetain_canResume_removeWhileActive_setToResume() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
 
         // When a media control that supports resumption is added
         addNotificationAndLoad()
@@ -2246,8 +2236,8 @@
 
     @Test
     fun testRetain_sessionPlayer_notifRemoved_doesNotChange() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control with PlaybackState actions is added, times out,
@@ -2266,8 +2256,8 @@
 
     @Test
     fun testRetain_sessionPlayer_sessionDestroyed_setToResume() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control with PlaybackState actions is added, times out,
@@ -2300,8 +2290,8 @@
 
     @Test
     fun testRetain_sessionPlayer_destroyedWhileActive_noResume_fullyRemoved() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control using session actions is added, and then the session is destroyed
@@ -2320,8 +2310,8 @@
 
     @Test
     fun testRetain_sessionPlayer_canResume_destroyedWhileActive_setToResume() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control using session actions and that does allow resumption is added,
@@ -2354,7 +2344,7 @@
 
     @Test
     fun testSessionPlayer_sessionDestroyed_noResume_fullyRemoved() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control with PlaybackState actions is added, times out,
@@ -2381,7 +2371,7 @@
 
     @Test
     fun testSessionPlayer_destroyedWhileActive_noResume_fullyRemoved() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control using session actions is added, and then the session is destroyed
@@ -2400,7 +2390,7 @@
 
     @Test
     fun testSessionPlayer_canResume_destroyedWhileActive_setToResume() {
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
         addPlaybackStateAction()
 
         // When a media control using session actions and that does allow resumption is added,
@@ -2433,8 +2423,8 @@
 
     @Test
     fun testSessionDestroyed_noNotificationKey_stillRemoved() {
-        whenever(mediaFlags.isRetainingPlayersEnabled()).thenReturn(true)
-        whenever(mediaFlags.areMediaSessionActionsEnabled(any(), any())).thenReturn(true)
+        fakeFeatureFlags.set(MEDIA_RETAIN_SESSIONS, true)
+        fakeFeatureFlags.set(MEDIA_SESSION_ACTIONS, true)
 
         // When a notiifcation is added and then removed before it is fully processed
         mediaDataProcessor.onNotificationAdded(KEY, mediaNotification)
@@ -2505,6 +2495,23 @@
         assertThat(mediaDataCaptor.value.artwork).isNull()
     }
 
+    private fun TestScope.assertRunAllReady(foreground: Int = 0, background: Int = 0) {
+        runCurrent()
+        if (Flags.mediaLoadMetadataViaMediaDataLoader()) {
+            // It doesn't make much sense to count tasks when we use coroutines in loader
+            // so this check is skipped in that scenario.
+            backgroundExecutor.runAllReady()
+            foregroundExecutor.runAllReady()
+        } else {
+            if (background > 0) {
+                assertThat(backgroundExecutor.runAllReady()).isEqualTo(background)
+            }
+            if (foreground > 0) {
+                assertThat(foregroundExecutor.runAllReady()).isEqualTo(foreground)
+            }
+        }
+    }
+
     /** Helper function to add a basic media notification and capture the resulting MediaData */
     private fun addNotificationAndLoad() {
         addNotificationAndLoad(mediaNotification)
@@ -2513,8 +2520,7 @@
     /** Helper function to add the given notification and capture the resulting MediaData */
     private fun addNotificationAndLoad(sbn: StatusBarNotification) {
         mediaDataProcessor.onNotificationAdded(KEY, sbn)
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
         verify(listener)
             .onMediaDataLoaded(
                 eq(KEY),
@@ -2548,8 +2554,7 @@
             pendingIntent,
             packageName
         )
-        assertThat(backgroundExecutor.runAllReady()).isEqualTo(1)
-        assertThat(foregroundExecutor.runAllReady()).isEqualTo(1)
+        testScope.assertRunAllReady(foreground = 1, background = 1)
 
         verify(listener)
             .onMediaDataLoaded(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerTest.kt
index 6a66c40..0c8d880 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerTest.kt
@@ -54,6 +54,7 @@
 import com.android.systemui.media.muteawait.MediaMuteAwaitConnectionManagerFactory
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.policy.ConfigurationController
+import com.android.systemui.testKosmos
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
@@ -125,6 +126,8 @@
     private lateinit var mediaData: MediaData
     @JvmField @Rule val mockito = MockitoJUnit.rule()
 
+    private val kosmos = testKosmos()
+
     @Before
     fun setUp() {
         fakeFgExecutor = FakeExecutor(FakeSystemClock())
@@ -141,6 +144,7 @@
                 { localBluetoothManager },
                 fakeFgExecutor,
                 fakeBgExecutor,
+                kosmos.mediaDeviceLogger,
             )
         manager.addListener(listener)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaCarouselControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaCarouselControllerTest.kt
index 850916b..03667cf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaCarouselControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/controller/MediaCarouselControllerTest.kt
@@ -27,22 +27,28 @@
 import android.view.View
 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.internal.logging.InstanceId
 import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.keyguard.KeyguardUpdateMonitorCallback
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
 import com.android.systemui.dump.DumpManager
 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.fakeDeviceEntryFingerprintAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.keyguard.shared.model.KeyguardState
+import com.android.systemui.keyguard.shared.model.SuccessFingerprintAuthenticationStatus
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.kosmos.unconfinedTestDispatcher
 import com.android.systemui.media.controls.MediaTestUtils
 import com.android.systemui.media.controls.domain.pipeline.EMPTY_SMARTSPACE_MEDIA_DATA
 import com.android.systemui.media.controls.domain.pipeline.MediaDataManager
@@ -67,18 +73,19 @@
 import com.android.systemui.testKosmos
 import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.concurrency.FakeExecutor
-import com.android.systemui.util.settings.FakeSettings
 import com.android.systemui.util.settings.GlobalSettings
-import com.android.systemui.util.settings.SecureSettings
+import com.android.systemui.util.settings.unconfinedDispatcherFakeSettings
 import com.android.systemui.util.time.FakeSystemClock
+import com.google.common.truth.Truth.assertThat
 import java.util.Locale
 import javax.inject.Provider
 import junit.framework.Assert.assertEquals
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestDispatcher
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.After
 import org.junit.Before
@@ -106,11 +113,14 @@
 private const val PAUSED_LOCAL = "paused local"
 private const val PLAYING_LOCAL = "playing local"
 
+@ExperimentalCoroutinesApi
 @SmallTest
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 @RunWith(AndroidJUnit4::class)
 class MediaCarouselControllerTest : SysuiTestCase() {
-    val kosmos = testKosmos()
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.unconfinedTestDispatcher
+    private val secureSettings = kosmos.unconfinedDispatcherFakeSettings
 
     @Mock lateinit var mediaControlPanelFactory: Provider<MediaControlPanel>
     @Mock lateinit var mediaViewControllerFactory: Provider<MediaViewController>
@@ -132,7 +142,6 @@
     @Mock lateinit var mediaFlags: MediaFlags
     @Mock lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
     @Mock lateinit var globalSettings: GlobalSettings
-    private lateinit var secureSettings: SecureSettings
     private val transitionRepository = kosmos.fakeKeyguardTransitionRepository
     @Captor lateinit var listener: ArgumentCaptor<MediaDataManager.Listener>
     @Captor
@@ -144,7 +153,6 @@
 
     private val clock = FakeSystemClock()
     private lateinit var bgExecutor: FakeExecutor
-    private lateinit var testDispatcher: TestDispatcher
     private lateinit var mediaCarouselController: MediaCarouselController
 
     private var originalResumeSetting =
@@ -153,10 +161,8 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        secureSettings = FakeSettings()
         context.resources.configuration.setLocales(LocaleList(Locale.US, Locale.UK))
         bgExecutor = FakeExecutor(clock)
-        testDispatcher = UnconfinedTestDispatcher()
         mediaCarouselController =
             MediaCarouselController(
                 applicationScope = kosmos.applicationCoroutineScope,
@@ -183,7 +189,7 @@
                 secureSettings = secureSettings,
                 mediaCarouselViewModel = kosmos.mediaCarouselViewModel,
                 mediaViewControllerFactory = mediaViewControllerFactory,
-                sceneInteractor = kosmos.sceneInteractor,
+                deviceEntryInteractor = kosmos.deviceEntryInteractor,
             )
         verify(configurationController).addCallback(capture(configListener))
         verify(visualStabilityProvider)
@@ -868,7 +874,6 @@
     }
 
     @DisableSceneContainer
-    @ExperimentalCoroutinesApi
     @Test
     fun testKeyguardGone_showMediaCarousel() =
         kosmos.testScope.runTest {
@@ -892,7 +897,6 @@
         }
 
     @EnableSceneContainer
-    @ExperimentalCoroutinesApi
     @Test
     fun testKeyguardGone_showMediaCarousel_scene_container() =
         kosmos.testScope.runTest {
@@ -910,7 +914,6 @@
             job.cancel()
         }
 
-    @ExperimentalCoroutinesApi
     @Test
     fun keyguardShowing_notAllowedOnLockscreen_updateVisibility() {
         kosmos.testScope.runTest {
@@ -940,7 +943,6 @@
         }
     }
 
-    @ExperimentalCoroutinesApi
     @Test
     fun keyguardShowing_allowedOnLockscreen_updateVisibility() {
         kosmos.testScope.runTest {
@@ -970,6 +972,74 @@
         }
     }
 
+    @EnableSceneContainer
+    @Test
+    fun deviceEntered_mediaAllowed_notLockedAndHidden() {
+        kosmos.testScope.runTest {
+            val settingsJob =
+                mediaCarouselController.listenForLockscreenSettingChanges(
+                    kosmos.applicationCoroutineScope
+                )
+            secureSettings.putBool(Settings.Secure.MEDIA_CONTROLS_LOCK_SCREEN, true)
+            setDeviceEntered(true)
+
+            assertEquals(false, mediaCarouselController.isLockedAndHidden())
+
+            settingsJob.cancel()
+        }
+    }
+
+    @EnableSceneContainer
+    @Test
+    fun deviceEntered_mediaNotAllowed_notLockedAndHidden() {
+        kosmos.testScope.runTest {
+            val settingsJob =
+                mediaCarouselController.listenForLockscreenSettingChanges(
+                    kosmos.applicationCoroutineScope
+                )
+            secureSettings.putBool(Settings.Secure.MEDIA_CONTROLS_LOCK_SCREEN, false)
+            setDeviceEntered(true)
+
+            assertEquals(false, mediaCarouselController.isLockedAndHidden())
+
+            settingsJob.cancel()
+        }
+    }
+
+    @EnableSceneContainer
+    @Test
+    fun deviceNotEntered_mediaAllowed_notLockedAndHidden() {
+        kosmos.testScope.runTest {
+            val settingsJob =
+                mediaCarouselController.listenForLockscreenSettingChanges(
+                    kosmos.applicationCoroutineScope
+                )
+            secureSettings.putBool(Settings.Secure.MEDIA_CONTROLS_LOCK_SCREEN, true)
+            setDeviceEntered(false)
+
+            assertEquals(false, mediaCarouselController.isLockedAndHidden())
+
+            settingsJob.cancel()
+        }
+    }
+
+    @EnableSceneContainer
+    @Test
+    fun deviceNotEntered_mediaNotAllowed_lockedAndHidden() {
+        kosmos.testScope.runTest {
+            val settingsJob =
+                mediaCarouselController.listenForLockscreenSettingChanges(
+                    kosmos.applicationCoroutineScope
+                )
+            secureSettings.putBool(Settings.Secure.MEDIA_CONTROLS_LOCK_SCREEN, false)
+            setDeviceEntered(false)
+
+            assertEquals(true, mediaCarouselController.isLockedAndHidden())
+
+            settingsJob.cancel()
+        }
+    }
+
     @Test
     fun testInvisibleToUserAndExpanded_playersNotListening() {
         // Add players to carousel.
@@ -1129,4 +1199,30 @@
             mediaCarouselController.mediaCarouselScrollHandler.visibleMediaIndex
         )
     }
+
+    private fun TestScope.setDeviceEntered(isEntered: Boolean) {
+        if (isEntered) {
+            // Unlock the device, marking the device as entered
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                SuccessFingerprintAuthenticationStatus(0, true)
+            )
+            runCurrent()
+        }
+        setScene(
+            if (isEntered) {
+                Scenes.Gone
+            } else {
+                Scenes.Lockscreen
+            }
+        )
+        assertThat(kosmos.deviceEntryInteractor.isDeviceEntered.value).isEqualTo(isEntered)
+    }
+
+    private fun TestScope.setScene(key: SceneKey) {
+        kosmos.sceneInteractor.changeScene(key, "test")
+        kosmos.sceneInteractor.setTransitionState(
+            MutableStateFlow<ObservableTransitionState>(ObservableTransitionState.Idle(key))
+        )
+        runCurrent()
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
index 19735e2..848c8db 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/qs/tiles/ModesTileTest.kt
@@ -25,10 +25,13 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.internal.logging.MetricsLogger
-import com.android.settingslib.notification.data.repository.FakeZenModeRepository
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.classifier.FalsingManagerFake
+import com.android.systemui.common.shared.model.asIcon
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.ActivityStarter
+import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.qs.QSHost
 import com.android.systemui.qs.QsEventLogger
@@ -36,21 +39,21 @@
 import com.android.systemui.qs.tiles.base.actions.FakeQSTileIntentUserInputHandler
 import com.android.systemui.qs.tiles.impl.modes.domain.interactor.ModesTileDataInteractor
 import com.android.systemui.qs.tiles.impl.modes.domain.interactor.ModesTileUserActionInteractor
+import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import com.android.systemui.qs.tiles.impl.modes.ui.ModesTileMapper
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfigProvider
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfigTestBuilder
 import com.android.systemui.qs.tiles.viewmodel.QSTileUIConfig
 import com.android.systemui.res.R
-import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
-import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor
+import com.android.systemui.statusbar.policy.data.repository.zenModeRepository
+import com.android.systemui.statusbar.policy.domain.interactor.zenModeInteractor
 import com.android.systemui.statusbar.policy.ui.dialog.ModesDialogDelegate
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.settings.FakeSettings
 import com.android.systemui.util.settings.SecureSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.After
@@ -59,7 +62,6 @@
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
-import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -68,6 +70,9 @@
 @RunWith(AndroidJUnit4::class)
 @RunWithLooper(setAsMainLooper = true)
 class ModesTileTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val testDispatcher = kosmos.testDispatcher
 
     @Mock private lateinit var qsHost: QSHost
 
@@ -85,25 +90,13 @@
 
     @Mock private lateinit var dialogDelegate: ModesDialogDelegate
 
-    private val testDispatcher = UnconfinedTestDispatcher()
-    private val testScope = TestScope(testDispatcher)
-
     private val inputHandler = FakeQSTileIntentUserInputHandler()
-    private val zenModeRepository = FakeZenModeRepository()
+    private val zenModeRepository = kosmos.zenModeRepository
     private val tileDataInteractor =
-        ModesTileDataInteractor(
-            context,
-            ZenModeInteractor(context, zenModeRepository, mock<NotificationSettingsRepository>()),
-            testDispatcher
-        )
+        ModesTileDataInteractor(context, kosmos.zenModeInteractor, testDispatcher)
     private val mapper =
         ModesTileMapper(
-            context.orCreateTestableResources
-                .apply {
-                    addOverride(R.drawable.qs_dnd_icon_on, TestStubDrawable())
-                    addOverride(R.drawable.qs_dnd_icon_off, TestStubDrawable())
-                }
-                .resources,
+            context.resources,
             context.theme,
         )
 
@@ -127,7 +120,7 @@
                 QSTileConfigTestBuilder.build {
                     uiConfig =
                         QSTileUIConfig.Resource(
-                            iconRes = R.drawable.qs_dnd_icon_off,
+                            iconRes = ModesTile.ICON_RES_ID,
                             labelRes = R.string.quick_settings_modes_label,
                         )
                 }
@@ -179,4 +172,43 @@
 
             assertThat(underTest.state.state).isEqualTo(Tile.STATE_ACTIVE)
         }
+
+    @Test
+    fun handleUpdateState_withTileModel_updatesState() =
+        testScope.runTest {
+            val tileState =
+                QSTile.State().apply {
+                    state = Tile.STATE_INACTIVE
+                    secondaryLabel = "Old secondary label"
+                }
+            val model =
+                ModesTileModel(
+                    isActivated = true,
+                    activeModes = listOf("One", "Two"),
+                    icon = TestStubDrawable().asIcon()
+                )
+
+            underTest.handleUpdateState(tileState, model)
+
+            assertThat(tileState.state).isEqualTo(Tile.STATE_ACTIVE)
+            assertThat(tileState.secondaryLabel).isEqualTo("2 modes are active")
+        }
+
+    @Test
+    fun handleUpdateState_withNull_updatesState() =
+        testScope.runTest {
+            val tileState =
+                QSTile.State().apply {
+                    state = Tile.STATE_INACTIVE
+                    secondaryLabel = "Old secondary label"
+                }
+            zenModeRepository.addMode("One", active = true)
+            zenModeRepository.addMode("Two", active = true)
+            runCurrent()
+
+            underTest.handleUpdateState(tileState, null)
+
+            assertThat(tileState.state).isEqualTo(Tile.STATE_ACTIVE)
+            assertThat(tileState.secondaryLabel).isEqualTo("2 modes are active")
+        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/recordissue/CustomTraceStateTest.kt b/packages/SystemUI/tests/src/com/android/systemui/recordissue/CustomTraceStateTest.kt
new file mode 100644
index 0000000..d8618fa
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/recordissue/CustomTraceStateTest.kt
@@ -0,0 +1,84 @@
+/*
+ * 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.recordissue
+
+import android.testing.TestableLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.recordissue.RecordIssueModule.Companion.TILE_SPEC
+import com.android.traceur.PresetTraceConfigs
+import com.google.common.truth.Truth
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.MockitoAnnotations
+
+/**
+ * CustomTraceState is customized Traceur settings for power users. These settings determine what
+ * tracing is used during the Record Issue Quick Settings flow. This class tests that those features
+ * are persistently and accurately stored across sessions.
+ */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper(setAsMainLooper = true)
+class CustomTraceStateTest : SysuiTestCase() {
+
+    private lateinit var underTest: CustomTraceState
+
+    @Before
+    fun setup() {
+        underTest = CustomTraceState(context.getSharedPreferences(TILE_SPEC, 0))
+        MockitoAnnotations.initMocks(this)
+    }
+
+    @Test
+    fun trace_config_is_stored_accurately() {
+        val expected = PresetTraceConfigs.getUiConfig()
+
+        underTest.traceConfig = expected
+
+        val actual = underTest.traceConfig
+        Truth.assertThat(actual.longTrace).isEqualTo(expected.longTrace)
+        Truth.assertThat(actual.maxLongTraceSizeMb).isEqualTo(expected.maxLongTraceSizeMb)
+        Truth.assertThat(actual.maxLongTraceDurationMinutes)
+            .isEqualTo(expected.maxLongTraceDurationMinutes)
+        Truth.assertThat(actual.apps).isEqualTo(expected.apps)
+        Truth.assertThat(actual.winscope).isEqualTo(expected.winscope)
+        Truth.assertThat(actual.attachToBugreport).isEqualTo(expected.attachToBugreport)
+        Truth.assertThat(actual.bufferSizeKb).isEqualTo(expected.bufferSizeKb)
+        Truth.assertThat(actual.tags).isEqualTo(expected.tags)
+    }
+
+    @Test
+    fun trace_config_is_persistently_stored_between_instances() {
+        val expected = PresetTraceConfigs.getUiConfig()
+
+        underTest.traceConfig = expected
+
+        val actual = CustomTraceState(context.getSharedPreferences(TILE_SPEC, 0)).traceConfig
+        Truth.assertThat(actual.longTrace).isEqualTo(expected.longTrace)
+        Truth.assertThat(actual.maxLongTraceSizeMb).isEqualTo(expected.maxLongTraceSizeMb)
+        Truth.assertThat(actual.maxLongTraceDurationMinutes)
+            .isEqualTo(expected.maxLongTraceDurationMinutes)
+        Truth.assertThat(actual.apps).isEqualTo(expected.apps)
+        Truth.assertThat(actual.winscope).isEqualTo(expected.winscope)
+        Truth.assertThat(actual.attachToBugreport).isEqualTo(expected.attachToBugreport)
+        Truth.assertThat(actual.bufferSizeKb).isEqualTo(expected.bufferSizeKb)
+        Truth.assertThat(actual.tags).isEqualTo(expected.tags)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
index b31d21e..15da77d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/MessageContainerControllerTest.kt
@@ -2,8 +2,6 @@
 
 import android.graphics.drawable.Drawable
 import android.os.UserHandle
-import android.platform.test.annotations.DisableFlags
-import android.platform.test.annotations.EnableFlags
 import android.view.View
 import android.view.ViewGroup
 import android.widget.FrameLayout
@@ -90,20 +88,6 @@
     }
 
     @Test
-    @DisableFlags(com.android.systemui.Flags.FLAG_SCREENSHOT_PRIVATE_PROFILE_BEHAVIOR_FIX)
-    fun testOnScreenshotTakenUserHandle_withWorkProfileFirstRun() {
-        whenever(workProfileMessageController.onScreenshotTaken(eq(userHandle)))
-            .thenReturn(workProfileData)
-        messageContainer.onScreenshotTaken(screenshotData)
-
-        verify(workProfileMessageController)
-            .populateView(eq(workProfileFirstRunView), eq(workProfileData), any())
-        assertEquals(View.VISIBLE, workProfileFirstRunView.visibility)
-        assertEquals(View.GONE, detectionNoticeView.visibility)
-    }
-
-    @Test
-    @EnableFlags(com.android.systemui.Flags.FLAG_SCREENSHOT_PRIVATE_PROFILE_BEHAVIOR_FIX)
     fun testOnScreenshotTakenUserHandle_withProfileProfileFirstRun() = runTest {
         val profileData =
             ProfileMessageController.ProfileFirstRunData(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt
deleted file mode 100644
index 0847f01..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/RequestProcessorTest.kt
+++ /dev/null
@@ -1,184 +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.systemui.screenshot
-
-import android.content.ComponentName
-import android.graphics.Bitmap
-import android.graphics.ColorSpace
-import android.graphics.Insets
-import android.graphics.Rect
-import android.hardware.HardwareBuffer
-import android.os.UserHandle
-import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_OTHER
-import android.view.WindowManager.ScreenshotSource.SCREENSHOT_OTHER
-import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
-import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
-import com.android.internal.util.ScreenshotRequest
-import com.android.systemui.screenshot.ScreenshotPolicy.DisplayContentInfo
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.runBlocking
-import org.junit.Assert
-import org.junit.Test
-
-private const val USER_ID = 1
-private const val TASK_ID = 1
-
-class RequestProcessorTest {
-    private val imageCapture = FakeImageCapture()
-    private val component = ComponentName("android.test", "android.test.Component")
-    private val bounds = Rect(25, 25, 75, 75)
-
-    private val policy = FakeScreenshotPolicy()
-
-    @Test
-    fun testFullScreenshot() = runBlocking {
-        // Indicate that the primary content belongs to a normal user
-        policy.setManagedProfile(USER_ID, false)
-        policy.setDisplayContentInfo(
-            policy.getDefaultDisplayId(),
-            DisplayContentInfo(component, bounds, UserHandle.of(USER_ID), TASK_ID)
-        )
-
-        val request =
-            ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_OTHER).build()
-        val processor = RequestProcessor(imageCapture, policy)
-
-        val processedData = processor.process(ScreenshotData.fromRequest(request))
-
-        // Request has topComponent added, but otherwise unchanged.
-        assertThat(processedData.type).isEqualTo(TAKE_SCREENSHOT_FULLSCREEN)
-        assertThat(processedData.source).isEqualTo(SCREENSHOT_OTHER)
-        assertThat(processedData.topComponent).isEqualTo(component)
-    }
-
-    @Test
-    fun testFullScreenshot_managedProfile() = runBlocking {
-        // Provide a fake task bitmap when asked
-        val bitmap = makeHardwareBitmap(100, 100)
-        imageCapture.image = bitmap
-
-        // Indicate that the primary content belongs to a manged profile
-        policy.setManagedProfile(USER_ID, true)
-        policy.setDisplayContentInfo(
-            policy.getDefaultDisplayId(),
-            DisplayContentInfo(component, bounds, UserHandle.of(USER_ID), TASK_ID)
-        )
-
-        val request =
-            ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_OTHER).build()
-        val processor = RequestProcessor(imageCapture, policy)
-
-        val processedData = processor.process(ScreenshotData.fromRequest(request))
-
-        // Expect a task snapshot is taken, overriding the full screen mode
-        assertThat(processedData.type).isEqualTo(TAKE_SCREENSHOT_PROVIDED_IMAGE)
-        assertThat(processedData.bitmap).isEqualTo(bitmap)
-        assertThat(processedData.screenBounds).isEqualTo(bounds)
-        assertThat(processedData.insets).isEqualTo(Insets.NONE)
-        assertThat(processedData.taskId).isEqualTo(TASK_ID)
-        assertThat(imageCapture.requestedTaskId).isEqualTo(TASK_ID)
-    }
-
-    @Test
-    fun testFullScreenshot_managedProfile_nullBitmap() {
-        // Provide a null task bitmap when asked
-        imageCapture.image = null
-
-        // Indicate that the primary content belongs to a manged profile
-        policy.setManagedProfile(USER_ID, true)
-        policy.setDisplayContentInfo(
-            policy.getDefaultDisplayId(),
-            DisplayContentInfo(component, bounds, UserHandle.of(USER_ID), TASK_ID)
-        )
-
-        val request =
-            ScreenshotRequest.Builder(TAKE_SCREENSHOT_FULLSCREEN, SCREENSHOT_KEY_OTHER).build()
-        val processor = RequestProcessor(imageCapture, policy)
-
-        Assert.assertThrows(IllegalStateException::class.java) {
-            runBlocking { processor.process(ScreenshotData.fromRequest(request)) }
-        }
-    }
-
-    @Test
-    fun testProvidedImageScreenshot() = runBlocking {
-        val bounds = Rect(50, 50, 150, 150)
-        val processor = RequestProcessor(imageCapture, policy)
-
-        policy.setManagedProfile(USER_ID, false)
-
-        val bitmap = makeHardwareBitmap(100, 100)
-
-        val request =
-            ScreenshotRequest.Builder(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER)
-                .setTopComponent(component)
-                .setTaskId(TASK_ID)
-                .setUserId(USER_ID)
-                .setBitmap(bitmap)
-                .setBoundsOnScreen(bounds)
-                .setInsets(Insets.NONE)
-                .build()
-
-        val screenshotData = ScreenshotData.fromRequest(request)
-        val processedData = processor.process(screenshotData)
-
-        assertThat(processedData).isEqualTo(screenshotData)
-    }
-
-    @Test
-    fun testProvidedImageScreenshot_managedProfile() = runBlocking {
-        val bounds = Rect(50, 50, 150, 150)
-        val processor = RequestProcessor(imageCapture, policy)
-
-        // Indicate that the screenshot belongs to a manged profile
-        policy.setManagedProfile(USER_ID, true)
-
-        val bitmap = makeHardwareBitmap(100, 100)
-
-        val request =
-            ScreenshotRequest.Builder(TAKE_SCREENSHOT_PROVIDED_IMAGE, SCREENSHOT_OTHER)
-                .setTopComponent(component)
-                .setTaskId(TASK_ID)
-                .setUserId(USER_ID)
-                .setBitmap(bitmap)
-                .setBoundsOnScreen(bounds)
-                .setInsets(Insets.NONE)
-                .build()
-
-        val screenshotData = ScreenshotData.fromRequest(request)
-        val processedData = processor.process(screenshotData)
-
-        // Work profile, but already a task snapshot, so no changes
-        assertThat(processedData).isEqualTo(screenshotData)
-    }
-
-    private fun makeHardwareBitmap(width: Int, height: Int): Bitmap {
-        val buffer =
-            HardwareBuffer.create(
-                width,
-                height,
-                HardwareBuffer.RGBA_8888,
-                1,
-                HardwareBuffer.USAGE_GPU_SAMPLED_IMAGE
-            )
-        return Bitmap.wrapHardwareBuffer(buffer, ColorSpace.get(ColorSpace.Named.SRGB))!!
-    }
-
-    private fun Bitmap.equalsHardwareBitmap(bitmap: Bitmap): Boolean {
-        return bitmap.hardwareBuffer == this.hardwareBuffer && bitmap.colorSpace == this.colorSpace
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
index a8d5008..eb1a04d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/appclips/AppClipsActivityTest.java
@@ -308,6 +308,40 @@
         assertThat(backlinksData.getCompoundDrawablesRelative()[2]).isNotNull();
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_APP_CLIPS_BACKLINKS)
+    public void appClipsLaunched_backlinks_multipleBacklinksAvailable_duplicateName()
+            throws RemoteException {
+        // Set up mocking for multiple backlinks.
+        ResolveInfo resolveInfo1 = createBacklinksTaskResolveInfo();
+
+        ResolveInfo resolveInfo2 = createBacklinksTaskResolveInfo();
+        RunningTaskInfo runningTaskInfo2 = createTaskInfoForBacklinksTask();
+        int taskId2 = BACKLINKS_TASK_ID + 2;
+        runningTaskInfo2.taskId = taskId2;
+
+        when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
+                mDisplayIdCaptor.capture()))
+                .thenReturn(List.of(TASK_THAT_SUPPORTS_BACKLINKS, runningTaskInfo2));
+        when(mPackageManager.resolveActivity(any(Intent.class), anyInt())).thenReturn(resolveInfo1,
+                resolveInfo1, resolveInfo1, resolveInfo2, resolveInfo2, resolveInfo2);
+        when(mPackageManager.loadItemIcon(any(), any())).thenReturn(FAKE_DRAWABLE);
+
+        // Using same AssistContent data for both tasks.
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, BACKLINKS_TASK_ID);
+        mockForAssistContent(ASSIST_CONTENT_FOR_BACKLINKS_TASK, taskId2);
+
+        // Mocking complete, trigger backlinks.
+        launchActivity();
+        waitForIdleSync();
+
+        // Verify default backlink shown to user has the numerical suffix.
+        TextView backlinksData = mActivity.findViewById(R.id.backlinks_data);
+        assertThat(backlinksData.getText().toString()).isEqualTo(
+                mContext.getString(R.string.backlinks_duplicate_label_format,
+                        BACKLINKS_TASK_APP_NAME, 1));
+    }
+
     private void setUpMocksForBacklinks() throws RemoteException {
         when(mAtmService.getTasks(eq(Integer.MAX_VALUE), eq(false), eq(false),
                 mDisplayIdCaptor.capture()))
diff --git a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
index 774aa51..2e2ac3e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/settings/UserTrackerImplTest.kt
@@ -33,9 +33,11 @@
 import com.android.systemui.flags.Flags
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.TruthJUnit.assume
+import java.util.concurrent.Executor
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
@@ -54,10 +56,7 @@
 import org.mockito.Mock
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
-import java.util.concurrent.Executor
-
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -71,27 +70,19 @@
         fun isBackgroundUserTrackerEnabled(): Iterable<Boolean> = listOf(true, false)
     }
 
-    @Mock
-    private lateinit var context: Context
+    @Mock private lateinit var context: Context
 
-    @Mock
-    private lateinit var userManager: UserManager
+    @Mock private lateinit var userManager: UserManager
 
-    @Mock
-    private lateinit var iActivityManager: IActivityManager
+    @Mock private lateinit var iActivityManager: IActivityManager
 
-    @Mock
-    private lateinit var userSwitchingReply: IRemoteCallback
+    @Mock private lateinit var userSwitchingReply: IRemoteCallback
 
-    @Mock(stubOnly = true)
-    private lateinit var dumpManager: DumpManager
+    @Mock(stubOnly = true) private lateinit var dumpManager: DumpManager
 
-    @Mock(stubOnly = true)
-    private lateinit var handler: Handler
+    @Mock(stubOnly = true) private lateinit var handler: Handler
 
-    @Parameterized.Parameter
-    @JvmField
-    var isBackgroundUserTrackerEnabled: Boolean = false
+    @Parameterized.Parameter @JvmField var isBackgroundUserTrackerEnabled: Boolean = false
 
     private val testScope = TestScope()
     private val testDispatcher = StandardTestDispatcher(testScope.testScheduler)
@@ -104,365 +95,379 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        `when`(context.userId).thenReturn(UserHandle.USER_SYSTEM)
-        `when`(context.user).thenReturn(UserHandle.SYSTEM)
-        `when`(context.createContextAsUser(any(), anyInt())).thenAnswer { invocation ->
+        whenever(context.userId).thenReturn(UserHandle.USER_SYSTEM)
+        whenever(context.user).thenReturn(UserHandle.SYSTEM)
+        whenever(context.createContextAsUser(any(), anyInt())).thenAnswer { invocation ->
             val user = invocation.getArgument<UserHandle>(0)
-            `when`(context.user).thenReturn(user)
-            `when`(context.userId).thenReturn(user.identifier)
+            whenever(context.user).thenReturn(user)
+            whenever(context.userId).thenReturn(user.identifier)
             context
         }
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+        whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
             val info = UserInfo(invocation.getArgument<Int>(0), "", UserInfo.FLAG_FULL)
             listOf(info)
         }
 
         featureFlags.set(Flags.USER_TRACKER_BACKGROUND_CALLBACKS, isBackgroundUserTrackerEnabled)
         tracker =
-                UserTrackerImpl(
-                        context,
-                        { featureFlags },
-                        userManager,
-                        iActivityManager,
-                        dumpManager,
-                        testScope.backgroundScope,
-                        testDispatcher,
-                        handler,
-                )
+            UserTrackerImpl(
+                context,
+                { featureFlags },
+                userManager,
+                iActivityManager,
+                dumpManager,
+                testScope.backgroundScope,
+                testDispatcher,
+                handler,
+            )
     }
 
+    @Test fun testNotInitialized() = testScope.runTest { assertThat(tracker.initialized).isFalse() }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetUserIdBeforeInitThrowsException() = testScope.runTest { tracker.userId }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetUserHandleBeforeInitThrowsException() = testScope.runTest { tracker.userHandle }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetUserContextBeforeInitThrowsException() = testScope.runTest { tracker.userContext }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetUserContentResolverBeforeInitThrowsException() =
+        testScope.runTest { tracker.userContentResolver }
+
+    @Test(expected = IllegalStateException::class)
+    fun testGetUserProfilesBeforeInitThrowsException() = testScope.runTest { tracker.userProfiles }
+
     @Test
-    fun testNotInitialized() = testScope.runTest {
-        assertThat(tracker.initialized).isFalse()
-    }
+    fun testInitialize() =
+        testScope.runTest {
+            tracker.initialize(0)
 
-    @Test(expected = IllegalStateException::class)
-    fun testGetUserIdBeforeInitThrowsException() = testScope.runTest {
-        tracker.userId
-    }
-
-    @Test(expected = IllegalStateException::class)
-    fun testGetUserHandleBeforeInitThrowsException() = testScope.runTest {
-        tracker.userHandle
-    }
-
-    @Test(expected = IllegalStateException::class)
-    fun testGetUserContextBeforeInitThrowsException() = testScope.runTest {
-        tracker.userContext
-    }
-
-    @Test(expected = IllegalStateException::class)
-    fun testGetUserContentResolverBeforeInitThrowsException() = testScope.runTest {
-        tracker.userContentResolver
-    }
-
-    @Test(expected = IllegalStateException::class)
-    fun testGetUserProfilesBeforeInitThrowsException() = testScope.runTest {
-        tracker.userProfiles
-    }
+            assertThat(tracker.initialized).isTrue()
+        }
 
     @Test
-    fun testInitialize() = testScope.runTest {
-        tracker.initialize(0)
+    fun testReceiverRegisteredOnInitialize() =
+        testScope.runTest {
+            tracker.initialize(0)
 
-        assertThat(tracker.initialized).isTrue()
-    }
+            val captor = ArgumentCaptor.forClass(IntentFilter::class.java)
 
-    @Test
-    fun testReceiverRegisteredOnInitialize() = testScope.runTest {
-        tracker.initialize(0)
-
-        val captor = ArgumentCaptor.forClass(IntentFilter::class.java)
-
-        verify(context)
+            verify(context)
                 .registerReceiverForAllUsers(eq(tracker), capture(captor), isNull(), eq(handler))
-        with(captor.value) {
-            assertThat(countActions()).isEqualTo(11)
-            assertThat(hasAction(Intent.ACTION_LOCALE_CHANGED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_USER_INFO_CHANGED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)).isTrue()
-            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)).isTrue()
-            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_ADDED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_REMOVED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_PROFILE_ADDED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_PROFILE_REMOVED)).isTrue()
-            assertThat(hasAction(Intent.ACTION_PROFILE_AVAILABLE)).isTrue()
-            assertThat(hasAction(Intent.ACTION_PROFILE_UNAVAILABLE)).isTrue()
-        }
-    }
-
-    @Test
-    fun testInitialValuesSet() = testScope.runTest {
-        val testID = 4
-        tracker.initialize(testID)
-
-        verify(userManager).getProfiles(testID)
-
-        assertThat(tracker.userId).isEqualTo(testID)
-        assertThat(tracker.userHandle).isEqualTo(UserHandle.of(testID))
-        assertThat(tracker.userContext.userId).isEqualTo(testID)
-        assertThat(tracker.userContext.user).isEqualTo(UserHandle.of(testID))
-        assertThat(tracker.userProfiles).hasSize(1)
-
-        val info = tracker.userProfiles[0]
-        assertThat(info.id).isEqualTo(testID)
-    }
-
-    @Test
-    fun testUserSwitch() = testScope.runTest {
-        tracker.initialize(0)
-        val newID = 5
-
-        val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
-        verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-        captor.value.onBeforeUserSwitching(newID)
-        captor.value.onUserSwitching(newID, userSwitchingReply)
-        runCurrent()
-        verify(userSwitchingReply).sendResult(any())
-
-        verify(userManager).getProfiles(newID)
-
-        assertThat(tracker.userId).isEqualTo(newID)
-        assertThat(tracker.userHandle).isEqualTo(UserHandle.of(newID))
-        assertThat(tracker.userContext.userId).isEqualTo(newID)
-        assertThat(tracker.userContext.user).isEqualTo(UserHandle.of(newID))
-        assertThat(tracker.userProfiles).hasSize(1)
-
-        val info = tracker.userProfiles[0]
-        assertThat(info.id).isEqualTo(newID)
-    }
-
-    @Test
-    fun testManagedProfileAvailable() = testScope.runTest {
-        tracker.initialize(0)
-        val profileID = tracker.userId + 10
-
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            val id = invocation.getArgument<Int>(0)
-            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
-            val infoProfile = UserInfo(
-                    id + 10,
-                    "",
-                    "",
-                    UserInfo.FLAG_MANAGED_PROFILE,
-                    UserManager.USER_TYPE_PROFILE_MANAGED
-            )
-            infoProfile.profileGroupId = id
-            listOf(info, infoProfile)
+            with(captor.value) {
+                assertThat(countActions()).isEqualTo(11)
+                assertThat(hasAction(Intent.ACTION_LOCALE_CHANGED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_USER_INFO_CHANGED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)).isTrue()
+                assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)).isTrue()
+                assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_ADDED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_REMOVED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_PROFILE_ADDED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_PROFILE_REMOVED)).isTrue()
+                assertThat(hasAction(Intent.ACTION_PROFILE_AVAILABLE)).isTrue()
+                assertThat(hasAction(Intent.ACTION_PROFILE_UNAVAILABLE)).isTrue()
+            }
         }
 
-        val intent = Intent(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
-        tracker.onReceive(context, intent)
-
-        assertThat(tracker.userProfiles.map { it.id }).containsExactly(tracker.userId, profileID)
-    }
-
     @Test
-    fun testManagedProfileUnavailable() = testScope.runTest {
-        tracker.initialize(0)
-        val profileID = tracker.userId + 10
+    fun testInitialValuesSet() =
+        testScope.runTest {
+            val testID = 4
+            tracker.initialize(testID)
 
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            val id = invocation.getArgument<Int>(0)
-            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
-            val infoProfile = UserInfo(
-                    id + 10,
-                    "",
-                    "",
-                    UserInfo.FLAG_MANAGED_PROFILE or UserInfo.FLAG_QUIET_MODE,
-                    UserManager.USER_TYPE_PROFILE_MANAGED
-            )
-            infoProfile.profileGroupId = id
-            listOf(info, infoProfile)
+            verify(userManager).getProfiles(testID)
+
+            assertThat(tracker.userId).isEqualTo(testID)
+            assertThat(tracker.userHandle).isEqualTo(UserHandle.of(testID))
+            assertThat(tracker.userContext.userId).isEqualTo(testID)
+            assertThat(tracker.userContext.user).isEqualTo(UserHandle.of(testID))
+            assertThat(tracker.userProfiles).hasSize(1)
+
+            val info = tracker.userProfiles[0]
+            assertThat(info.id).isEqualTo(testID)
         }
 
-        val intent = Intent(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
-        tracker.onReceive(context, intent)
-
-        assertThat(tracker.userProfiles.map { it.id }).containsExactly(tracker.userId, profileID)
-    }
-
     @Test
-    fun testManagedProfileStartedAndRemoved() = testScope.runTest {
-        tracker.initialize(0)
-        val profileID = tracker.userId + 10
+    fun testUserSwitch() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val newID = 5
 
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            val id = invocation.getArgument<Int>(0)
-            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
-            val infoProfile = UserInfo(
-                    id + 10,
-                    "",
-                    "",
-                    UserInfo.FLAG_MANAGED_PROFILE,
-                    UserManager.USER_TYPE_PROFILE_MANAGED
-            )
-            infoProfile.profileGroupId = id
-            listOf(info, infoProfile)
+            val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
+            verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onUserSwitching(newID, userSwitchingReply)
+            runCurrent()
+            verify(userSwitchingReply).sendResult(any())
+
+            verify(userManager).getProfiles(newID)
+
+            assertThat(tracker.userId).isEqualTo(newID)
+            assertThat(tracker.userHandle).isEqualTo(UserHandle.of(newID))
+            assertThat(tracker.userContext.userId).isEqualTo(newID)
+            assertThat(tracker.userContext.user).isEqualTo(UserHandle.of(newID))
+            assertThat(tracker.userProfiles).hasSize(1)
+
+            val info = tracker.userProfiles[0]
+            assertThat(info.id).isEqualTo(newID)
         }
 
-        // Managed profile started
-        val intent = Intent(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
-        tracker.onReceive(context, intent)
+    @Test
+    fun testManagedProfileAvailable() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val profileID = tracker.userId + 10
 
-        assertThat(tracker.userProfiles.map { it.id }).containsExactly(tracker.userId, profileID)
+            whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+                val id = invocation.getArgument<Int>(0)
+                val info = UserInfo(id, "", UserInfo.FLAG_FULL)
+                val infoProfile =
+                    UserInfo(
+                        id + 10,
+                        "",
+                        "",
+                        UserInfo.FLAG_MANAGED_PROFILE,
+                        UserManager.USER_TYPE_PROFILE_MANAGED
+                    )
+                infoProfile.profileGroupId = id
+                listOf(info, infoProfile)
+            }
 
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            listOf(UserInfo(invocation.getArgument(0), "", UserInfo.FLAG_FULL))
+            val intent =
+                Intent(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+            tracker.onReceive(context, intent)
+
+            assertThat(tracker.userProfiles.map { it.id })
+                .containsExactly(tracker.userId, profileID)
         }
 
-        val intent2 = Intent(Intent.ACTION_MANAGED_PROFILE_REMOVED)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
-        tracker.onReceive(context, intent2)
-
-        assertThat(tracker.userProfiles.map { it.id }).containsExactly(tracker.userId)
-    }
-
     @Test
-    fun testCallbackNotCalledOnAdd() = testScope.runTest {
-        tracker.initialize(0)
-        val callback = TestCallback()
+    fun testManagedProfileUnavailable() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val profileID = tracker.userId + 10
 
-        tracker.addCallback(callback, executor)
+            whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+                val id = invocation.getArgument<Int>(0)
+                val info = UserInfo(id, "", UserInfo.FLAG_FULL)
+                val infoProfile =
+                    UserInfo(
+                        id + 10,
+                        "",
+                        "",
+                        UserInfo.FLAG_MANAGED_PROFILE or UserInfo.FLAG_QUIET_MODE,
+                        UserManager.USER_TYPE_PROFILE_MANAGED
+                    )
+                infoProfile.profileGroupId = id
+                listOf(info, infoProfile)
+            }
 
-        assertThat(callback.calledOnProfilesChanged).isEqualTo(0)
-        assertThat(callback.calledOnUserChanged).isEqualTo(0)
-    }
+            val intent =
+                Intent(Intent.ACTION_MANAGED_PROFILE_UNAVAILABLE)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+            tracker.onReceive(context, intent)
 
-    @Test
-    fun testCallbackCalledOnUserChanging() = testScope.runTest {
-        tracker.initialize(0)
-        val callback = TestCallback()
-        tracker.addCallback(callback, executor)
-
-        val newID = 5
-
-        val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
-        verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-        captor.value.onBeforeUserSwitching(newID)
-        captor.value.onUserSwitching(newID, userSwitchingReply)
-        runCurrent()
-
-        verify(userSwitchingReply).sendResult(any())
-        assertThat(callback.calledOnUserChanging).isEqualTo(1)
-        assertThat(callback.lastUser).isEqualTo(newID)
-        assertThat(callback.lastUserContext?.userId).isEqualTo(newID)
-    }
-
-    @Test
-    fun testAsyncCallbackWaitsUserToChange() = testScope.runTest {
-        // Skip this test for CountDownLatch variation. The problem is that there would be a
-        // deadlock if the callbacks processing runs on the same thread as the callback (which
-        // is blocked by the latch). Before the change it works because the callbacks are
-        // processed on a binder thread which is always distinct.
-        // This is the issue that this feature addresses.
-        assume().that(isBackgroundUserTrackerEnabled).isTrue()
-
-        tracker.initialize(0)
-        val callback = TestCallback()
-        val callbackExecutor = FakeExecutor(FakeSystemClock())
-        tracker.addCallback(callback, callbackExecutor)
-
-        val newID = 5
-
-        val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
-        verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-        captor.value.onUserSwitching(newID, userSwitchingReply)
-
-        assertThat(callback.calledOnUserChanging).isEqualTo(0)
-        verify(userSwitchingReply, never()).sendResult(any())
-
-        FakeExecutor.exhaustExecutors(callbackExecutor)
-        runCurrent()
-        FakeExecutor.exhaustExecutors(callbackExecutor)
-        runCurrent()
-
-        assertThat(callback.calledOnUserChanging).isEqualTo(1)
-        verify(userSwitchingReply).sendResult(any())
-    }
-
-    @Test
-    fun testCallbackCalledOnUserChanged() = testScope.runTest {
-        tracker.initialize(0)
-        val callback = TestCallback()
-        tracker.addCallback(callback, executor)
-
-        val newID = 5
-
-        val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
-        verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-        captor.value.onBeforeUserSwitching(newID)
-        captor.value.onUserSwitchComplete(newID)
-        runCurrent()
-
-        assertThat(callback.calledOnUserChanged).isEqualTo(1)
-        assertThat(callback.lastUser).isEqualTo(newID)
-        assertThat(callback.lastUserContext?.userId).isEqualTo(newID)
-        assertThat(callback.calledOnProfilesChanged).isEqualTo(1)
-        assertThat(callback.lastUserProfiles.map { it.id }).containsExactly(newID)
-    }
-
-    @Test
-    fun testCallbackCalledOnUserInfoChanged() = testScope.runTest {
-        tracker.initialize(0)
-        val callback = TestCallback()
-        tracker.addCallback(callback, executor)
-        val profileID = tracker.userId + 10
-
-        `when`(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
-            val id = invocation.getArgument<Int>(0)
-            val info = UserInfo(id, "", UserInfo.FLAG_FULL)
-            val infoProfile = UserInfo(
-                    id + 10,
-                    "",
-                    "",
-                    UserInfo.FLAG_MANAGED_PROFILE,
-                    UserManager.USER_TYPE_PROFILE_MANAGED
-            )
-            infoProfile.profileGroupId = id
-            listOf(info, infoProfile)
+            assertThat(tracker.userProfiles.map { it.id })
+                .containsExactly(tracker.userId, profileID)
         }
 
-        val intent = Intent(Intent.ACTION_USER_INFO_CHANGED)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+    @Test
+    fun testManagedProfileStartedAndRemoved() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val profileID = tracker.userId + 10
 
-        tracker.onReceive(context, intent)
+            whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+                val id = invocation.getArgument<Int>(0)
+                val info = UserInfo(id, "", UserInfo.FLAG_FULL)
+                val infoProfile =
+                    UserInfo(
+                        id + 10,
+                        "",
+                        "",
+                        UserInfo.FLAG_MANAGED_PROFILE,
+                        UserManager.USER_TYPE_PROFILE_MANAGED
+                    )
+                infoProfile.profileGroupId = id
+                listOf(info, infoProfile)
+            }
 
-        assertThat(callback.calledOnUserChanged).isEqualTo(0)
-        assertThat(callback.calledOnProfilesChanged).isEqualTo(1)
-        assertThat(callback.lastUserProfiles.map { it.id }).containsExactly(0, profileID)
-    }
+            // Managed profile started
+            val intent =
+                Intent(Intent.ACTION_MANAGED_PROFILE_UNLOCKED)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+            tracker.onReceive(context, intent)
+
+            assertThat(tracker.userProfiles.map { it.id })
+                .containsExactly(tracker.userId, profileID)
+
+            whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+                listOf(UserInfo(invocation.getArgument(0), "", UserInfo.FLAG_FULL))
+            }
+
+            val intent2 =
+                Intent(Intent.ACTION_MANAGED_PROFILE_REMOVED)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+            tracker.onReceive(context, intent2)
+
+            assertThat(tracker.userProfiles.map { it.id }).containsExactly(tracker.userId)
+        }
 
     @Test
-    fun testCallbackRemoved() = testScope.runTest {
-        tracker.initialize(0)
-        val newID = 5
-        val profileID = newID + 10
+    fun testCallbackNotCalledOnAdd() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val callback = TestCallback()
 
-        val callback = TestCallback()
-        tracker.addCallback(callback, executor)
-        tracker.removeCallback(callback)
+            tracker.addCallback(callback, executor)
 
-        val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
-        verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
-        captor.value.onUserSwitching(newID, userSwitchingReply)
-        runCurrent()
-        verify(userSwitchingReply).sendResult(any())
-        captor.value.onUserSwitchComplete(newID)
+            assertThat(callback.calledOnProfilesChanged).isEqualTo(0)
+            assertThat(callback.calledOnUserChanged).isEqualTo(0)
+        }
 
-        val intentProfiles = Intent(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
-                .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+    @Test
+    fun testCallbackCalledOnUserChanging() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val callback = TestCallback()
+            tracker.addCallback(callback, executor)
 
-        tracker.onReceive(context, intentProfiles)
+            val newID = 5
 
-        assertThat(callback.calledOnUserChanging).isEqualTo(0)
-        assertThat(callback.calledOnUserChanged).isEqualTo(0)
-        assertThat(callback.calledOnProfilesChanged).isEqualTo(0)
-    }
+            val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
+            verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onUserSwitching(newID, userSwitchingReply)
+            runCurrent()
+
+            verify(userSwitchingReply).sendResult(any())
+            assertThat(callback.calledOnUserChanging).isEqualTo(1)
+            assertThat(callback.lastUser).isEqualTo(newID)
+            assertThat(callback.lastUserContext?.userId).isEqualTo(newID)
+        }
+
+    @Test
+    fun testAsyncCallbackWaitsUserToChange() =
+        testScope.runTest {
+            // Skip this test for CountDownLatch variation. The problem is that there would be a
+            // deadlock if the callbacks processing runs on the same thread as the callback (which
+            // is blocked by the latch). Before the change it works because the callbacks are
+            // processed on a binder thread which is always distinct.
+            // This is the issue that this feature addresses.
+            assume().that(isBackgroundUserTrackerEnabled).isTrue()
+
+            tracker.initialize(0)
+            val callback = TestCallback()
+            val callbackExecutor = FakeExecutor(FakeSystemClock())
+            tracker.addCallback(callback, callbackExecutor)
+
+            val newID = 5
+
+            val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
+            verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onUserSwitching(newID, userSwitchingReply)
+
+            assertThat(callback.calledOnUserChanging).isEqualTo(0)
+            verify(userSwitchingReply, never()).sendResult(any())
+
+            FakeExecutor.exhaustExecutors(callbackExecutor)
+            runCurrent()
+            FakeExecutor.exhaustExecutors(callbackExecutor)
+            runCurrent()
+
+            assertThat(callback.calledOnUserChanging).isEqualTo(1)
+            verify(userSwitchingReply).sendResult(any())
+        }
+
+    @Test
+    fun testCallbackCalledOnUserChanged() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val callback = TestCallback()
+            tracker.addCallback(callback, executor)
+
+            val newID = 5
+
+            val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
+            verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onBeforeUserSwitching(newID)
+            captor.value.onUserSwitchComplete(newID)
+            runCurrent()
+
+            assertThat(callback.calledOnUserChanged).isEqualTo(1)
+            assertThat(callback.lastUser).isEqualTo(newID)
+            assertThat(callback.lastUserContext?.userId).isEqualTo(newID)
+            assertThat(callback.calledOnProfilesChanged).isEqualTo(1)
+            assertThat(callback.lastUserProfiles.map { it.id }).containsExactly(newID)
+        }
+
+    @Test
+    fun testCallbackCalledOnUserInfoChanged() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val callback = TestCallback()
+            tracker.addCallback(callback, executor)
+            val profileID = tracker.userId + 10
+
+            whenever(userManager.getProfiles(anyInt())).thenAnswer { invocation ->
+                val id = invocation.getArgument<Int>(0)
+                val info = UserInfo(id, "", UserInfo.FLAG_FULL)
+                val infoProfile =
+                    UserInfo(
+                        id + 10,
+                        "",
+                        "",
+                        UserInfo.FLAG_MANAGED_PROFILE,
+                        UserManager.USER_TYPE_PROFILE_MANAGED
+                    )
+                infoProfile.profileGroupId = id
+                listOf(info, infoProfile)
+            }
+
+            val intent =
+                Intent(Intent.ACTION_USER_INFO_CHANGED)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+
+            tracker.onReceive(context, intent)
+
+            assertThat(callback.calledOnUserChanged).isEqualTo(0)
+            assertThat(callback.calledOnProfilesChanged).isEqualTo(1)
+            assertThat(callback.lastUserProfiles.map { it.id }).containsExactly(0, profileID)
+        }
+
+    @Test
+    fun testCallbackRemoved() =
+        testScope.runTest {
+            tracker.initialize(0)
+            val newID = 5
+            val profileID = newID + 10
+
+            val callback = TestCallback()
+            tracker.addCallback(callback, executor)
+            tracker.removeCallback(callback)
+
+            val captor = ArgumentCaptor.forClass(IUserSwitchObserver::class.java)
+            verify(iActivityManager).registerUserSwitchObserver(capture(captor), anyString())
+            captor.value.onUserSwitching(newID, userSwitchingReply)
+            runCurrent()
+            verify(userSwitchingReply).sendResult(any())
+            captor.value.onUserSwitchComplete(newID)
+
+            val intentProfiles =
+                Intent(Intent.ACTION_MANAGED_PROFILE_AVAILABLE)
+                    .putExtra(Intent.EXTRA_USER, UserHandle.of(profileID))
+
+            tracker.onReceive(context, intentProfiles)
+
+            assertThat(callback.calledOnUserChanging).isEqualTo(0)
+            assertThat(callback.calledOnUserChanged).isEqualTo(0)
+            assertThat(callback.calledOnProfilesChanged).isEqualTo(0)
+        }
 
     private class TestCallback : UserTracker.Callback {
         var calledOnUserChanging = 0
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
index 5a5cdcd..3ba1447e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/GlanceableHubContainerControllerTest.kt
@@ -61,7 +61,6 @@
 import com.android.systemui.media.controls.controller.keyguardMediaController
 import com.android.systemui.res.R
 import com.android.systemui.scene.shared.model.sceneDataSourceDelegator
-import com.android.systemui.shade.data.repository.fakeShadeRepository
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 import com.android.systemui.statusbar.lockscreen.lockscreenSmartspaceController
 import com.android.systemui.statusbar.notification.stack.notificationStackScrollLayoutController
@@ -417,13 +416,17 @@
                 // Communal is open.
                 goToScene(CommunalScenes.Communal)
 
-                // Shade shows up.
-                shadeTestUtil.setQsExpansion(0.5f)
-                testableLooper.processAllMessages()
+                // Touch starts and ends.
                 assertThat(underTest.onTouchEvent(DOWN_EVENT)).isTrue()
                 assertThat(underTest.onTouchEvent(CANCEL_EVENT)).isTrue()
+
+                // Up event is no longer processed
                 assertThat(underTest.onTouchEvent(UP_EVENT)).isFalse()
+
+                // Move event can still be processed
                 assertThat(underTest.onTouchEvent(MOVE_EVENT)).isTrue()
+                assertThat(underTest.onTouchEvent(MOVE_EVENT)).isTrue()
+                assertThat(underTest.onTouchEvent(UP_EVENT)).isTrue()
             }
         }
 
@@ -702,7 +705,7 @@
                 verify(containerView).onTouchEvent(DOWN_EVENT)
 
                 // User is interacting with shade on lockscreen.
-                fakeShadeRepository.setLegacyLockscreenShadeTracking(true)
+                shadeTestUtil.setLockscreenShadeTracking(true)
                 testableLooper.processAllMessages()
 
                 // A move event is ignored while the user is already interacting.
@@ -716,6 +719,30 @@
         }
 
     @Test
+    fun onTouchEvent_shadeExpanding_touchesNotDispatched() =
+        with(kosmos) {
+            testScope.runTest {
+                // On lockscreen.
+                goToScene(CommunalScenes.Blank)
+                whenever(
+                        notificationStackScrollLayoutController.isBelowLastNotification(
+                            any(),
+                            any()
+                        )
+                    )
+                    .thenReturn(true)
+
+                // Shade is open slightly.
+                shadeTestUtil.setShadeExpansion(0.01f)
+                testableLooper.processAllMessages()
+
+                // Touches are not consumed.
+                assertThat(underTest.onTouchEvent(DOWN_EVENT)).isFalse()
+                verify(containerView, never()).onTouchEvent(DOWN_EVENT)
+            }
+        }
+
+    @Test
     fun onTouchEvent_bouncerInteracting_movesNotDispatched() =
         with(kosmos) {
             testScope.runTest {
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 523d15c..9481e5a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerBaseTest.java
@@ -155,7 +155,6 @@
 import com.android.systemui.statusbar.notification.NotificationWakeUpCoordinatorLogger;
 import com.android.systemui.statusbar.notification.data.repository.NotificationsKeyguardViewStateRepository;
 import com.android.systemui.statusbar.notification.domain.interactor.ActiveNotificationsInteractor;
-import com.android.systemui.statusbar.notification.domain.interactor.HeadsUpNotificationInteractor;
 import com.android.systemui.statusbar.notification.domain.interactor.NotificationsKeyguardInteractor;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
 import com.android.systemui.statusbar.notification.stack.AmbientState;
@@ -163,7 +162,6 @@
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController;
 import com.android.systemui.statusbar.notification.stack.NotificationStackSizeCalculator;
-import com.android.systemui.statusbar.notification.stack.data.repository.FakeHeadsUpNotificationRepository;
 import com.android.systemui.statusbar.notification.stack.domain.interactor.SharedNotificationContainerInteractor;
 import com.android.systemui.statusbar.phone.CentralSurfaces;
 import com.android.systemui.statusbar.phone.ConfigurationControllerImpl;
@@ -365,13 +363,6 @@
     protected TestScope mTestScope = mKosmos.getTestScope();
     protected ShadeInteractor mShadeInteractor;
     protected PowerInteractor mPowerInteractor;
-    protected FakeHeadsUpNotificationRepository mFakeHeadsUpNotificationRepository =
-            new FakeHeadsUpNotificationRepository();
-    protected NotificationsKeyguardViewStateRepository mNotificationsKeyguardViewStateRepository =
-            new NotificationsKeyguardViewStateRepository();
-    protected NotificationsKeyguardInteractor mNotificationsKeyguardInteractor =
-            new NotificationsKeyguardInteractor(mNotificationsKeyguardViewStateRepository);
-    protected HeadsUpNotificationInteractor mHeadsUpNotificationInteractor;
     protected NotificationPanelViewController.TouchHandler mTouchHandler;
     protected ConfigurationController mConfigurationController;
     protected SysuiStatusBarStateController mStatusBarStateController;
@@ -689,12 +680,6 @@
         when(longPressHandlingView.getResources()).thenReturn(longPressHandlingViewRes);
         when(longPressHandlingViewRes.getString(anyInt())).thenReturn("");
 
-
-        mHeadsUpNotificationInteractor =
-                new HeadsUpNotificationInteractor(mFakeHeadsUpNotificationRepository,
-                        mDeviceEntryFaceAuthInteractor, mKeyguardTransitionInteractor,
-                        mNotificationsKeyguardInteractor, mShadeInteractor);
-
         mNotificationPanelViewController = new NotificationPanelViewController(
                 mView,
                 mMainHandler,
@@ -769,7 +754,6 @@
                 mActivityStarter,
                 mSharedNotificationContainerInteractor,
                 mActiveNotificationsInteractor,
-                mHeadsUpNotificationInteractor,
                 mShadeAnimationInteractor,
                 mKeyguardViewConfigurator,
                 mDeviceEntryFaceAuthInteractor,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
index 905cc4c..a7fd160 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerTest.java
@@ -58,13 +58,13 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.DejankUtils;
+import com.android.systemui.flags.DisableSceneContainer;
 import com.android.systemui.flags.Flags;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.power.domain.interactor.PowerInteractor;
 import com.android.systemui.res.R;
 import com.android.systemui.statusbar.notification.row.ExpandableView;
 import com.android.systemui.statusbar.notification.row.ExpandableView.OnHeightChangedListener;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.notification.stack.AmbientState;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController;
 import com.android.systemui.statusbar.phone.KeyguardClockPositionAlgorithm;
@@ -1375,7 +1375,7 @@
     }
 
     @Test
-    @DisableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @DisableSceneContainer
     public void shadeExpanded_whenHunIsPresent() {
         when(mHeadsUpManager.hasPinnedHeadsUp()).thenReturn(true);
         assertThat(mNotificationPanelViewController.isExpanded()).isTrue();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
index 64eadb7..90655c3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/NotificationPanelViewControllerWithCoroutinesTest.kt
@@ -19,7 +19,6 @@
 package com.android.systemui.shade
 
 import android.platform.test.annotations.DisableFlags
-import android.platform.test.annotations.EnableFlags
 import android.testing.TestableLooper
 import android.view.HapticFeedbackConstants
 import android.view.View
@@ -35,7 +34,6 @@
 import com.android.systemui.statusbar.StatusBarState.SHADE
 import com.android.systemui.statusbar.StatusBarState.SHADE_LOCKED
 import com.android.systemui.statusbar.notification.data.repository.FakeHeadsUpRowRepository
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
@@ -243,36 +241,4 @@
         val bottomAreaAlpha by collectLastValue(mFakeKeyguardRepository.bottomAreaAlpha)
         assertThat(bottomAreaAlpha).isEqualTo(1f)
     }
-
-    @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
-    fun shadeExpanded_whenHunIsPresent() = runTest {
-        launch(mainDispatcher) {
-            givenViewAttached()
-
-            // WHEN a pinned heads up is present
-            mFakeHeadsUpNotificationRepository.setNotifications(
-                FakeHeadsUpRowRepository("key", isPinned = true)
-            )
-        }
-        advanceUntilIdle()
-
-        // THEN the panel should be visible
-        assertThat(mNotificationPanelViewController.isExpanded).isTrue()
-    }
-
-    @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
-    fun shadeExpanded_whenHunIsAnimatingAway() = runTest {
-        launch(mainDispatcher) {
-            givenViewAttached()
-
-            // WHEN a heads up is animating away
-            mFakeHeadsUpNotificationRepository.isHeadsUpAnimatingAway.value = true
-        }
-        advanceUntilIdle()
-
-        // THEN the panel should be visible
-        assertThat(mNotificationPanelViewController.isExpanded).isTrue()
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
index fadb1d7..b65a902 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/shade/domain/interactor/ShadeInteractorSceneContainerImplTest.kt
@@ -156,6 +156,25 @@
         }
 
     @Test
+    fun qsFullscreen_falseWhenIdleSplitShadeQs() =
+        testScope.runTest {
+            val actual by collectLastValue(underTest.isQsFullscreen)
+
+            // WHEN split shade is enabled and Idle on QuickSettings scene
+            shadeTestUtil.setSplitShade(true)
+            keyguardRepository.setStatusBarState(StatusBarState.SHADE)
+            val transitionState =
+                MutableStateFlow<ObservableTransitionState>(
+                    ObservableTransitionState.Idle(Scenes.QuickSettings)
+                )
+            sceneInteractor.setTransitionState(transitionState)
+            runCurrent()
+
+            // THEN QS is not fullscreen
+            Truth.assertThat(actual).isFalse()
+        }
+
+    @Test
     fun qsFullscreen_trueWhenIdleQS() =
         testScope.runTest {
             val actual by collectLastValue(underTest.isQsFullscreen)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/OperatorNameViewControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/OperatorNameViewControllerTest.kt
index 396d017..d6b3b91 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/OperatorNameViewControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/OperatorNameViewControllerTest.kt
@@ -28,7 +28,7 @@
 import com.android.systemui.plugins.DarkIconDispatcher
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeSubscriptionManagerProxy
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.tuner.TunerService
@@ -71,7 +71,6 @@
 
     private val airplaneModeRepository = FakeAirplaneModeRepository()
     private val connectivityRepository = FakeConnectivityRepository()
-    private val mobileConnectionsRepository = FakeMobileConnectionsRepository()
 
     @Before
     fun setup() {
@@ -81,7 +80,7 @@
             AirplaneModeInteractor(
                 airplaneModeRepository,
                 connectivityRepository,
-                mobileConnectionsRepository,
+                kosmos.fakeMobileConnectionsRepository,
             )
 
         underTest =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarSignalPolicyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarSignalPolicyTest.kt
new file mode 100644
index 0000000..593f873
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/StatusBarSignalPolicyTest.kt
@@ -0,0 +1,173 @@
+/*
+ * 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
+
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.R
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.Flags.FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testCase
+import com.android.systemui.statusbar.connectivity.IconState
+import com.android.systemui.statusbar.connectivity.NetworkController
+import com.android.systemui.statusbar.phone.StatusBarSignalPolicy
+import com.android.systemui.statusbar.phone.StatusBarSignalPolicy_Factory
+import com.android.systemui.statusbar.phone.ui.StatusBarIconController
+import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.airplaneModeInteractor
+import com.android.systemui.statusbar.policy.SecurityController
+import com.android.systemui.tuner.TunerService
+import com.android.systemui.util.CarrierConfigTracker
+import com.android.systemui.util.kotlin.JavaAdapter
+import com.android.systemui.util.mockito.mock
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.runner.RunWith
+import org.mockito.Mockito.verify
+import org.mockito.kotlin.clearInvocations
+import org.mockito.kotlin.verifyZeroInteractions
+import kotlin.test.Test
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class StatusBarSignalPolicyTest : SysuiTestCase() {
+    private val kosmos = Kosmos().also { it.testCase = this }
+
+    private lateinit var underTest: StatusBarSignalPolicy
+
+    private val testScope = TestScope()
+
+    private val javaAdapter = JavaAdapter(testScope.backgroundScope)
+    private val airplaneModeInteractor = kosmos.airplaneModeInteractor
+
+    private val securityController = mock<SecurityController>()
+    private val tunerService = mock<TunerService>()
+    private val statusBarIconController = mock<StatusBarIconController>()
+    private val networkController = mock<NetworkController>()
+    private val carrierConfigTracker = mock<CarrierConfigTracker>()
+
+    private var slotAirplane: String? = null
+
+    @Before
+    fun setup() {
+        underTest =
+            StatusBarSignalPolicy_Factory.newInstance(
+                mContext,
+                statusBarIconController,
+                carrierConfigTracker,
+                networkController,
+                securityController,
+                tunerService,
+                javaAdapter,
+                airplaneModeInteractor,
+            )
+
+        slotAirplane = mContext.getString(R.string.status_bar_airplane)
+    }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun airplaneModeViaInteractor_statusBarSignalPolicyRefactorFlagEnabled_iconUpdated() =
+        testScope.runTest {
+            underTest.start()
+            airplaneModeInteractor.setIsAirplaneMode(true)
+            runCurrent()
+            verify(statusBarIconController).setIconVisibility(slotAirplane, true)
+
+            airplaneModeInteractor.setIsAirplaneMode(false)
+            runCurrent()
+            verify(statusBarIconController).setIconVisibility(slotAirplane, false)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun airplaneModeViaSignalCallback_statusBarSignalPolicyRefactorFlagEnabled_iconNotUpdated() =
+        testScope.runTest {
+            underTest.start()
+            runCurrent()
+            clearInvocations(statusBarIconController)
+
+            // Make sure the legacy code path does not change airplane mode when the refactor
+            // flag is enabled.
+            underTest.setIsAirplaneMode(IconState(true, TelephonyIcons.FLIGHT_MODE_ICON, ""))
+            runCurrent()
+            verifyZeroInteractions(statusBarIconController)
+
+            underTest.setIsAirplaneMode(IconState(false, TelephonyIcons.FLIGHT_MODE_ICON, ""))
+            runCurrent()
+            verifyZeroInteractions(statusBarIconController)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun statusBarSignalPolicyInitialization_statusBarSignalPolicyRefactorFlagEnabled_initNoOp() =
+        testScope.runTest {
+            // Make sure StatusBarSignalPolicy.init does no initialization when
+            // the refactor flag is disabled.
+            underTest.init()
+            verifyZeroInteractions(securityController, networkController, tunerService)
+        }
+
+    @Test
+    @DisableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun airplaneModeViaSignalCallback_statusBarSignalPolicyRefactorFlagDisabled_iconUpdated() =
+        testScope.runTest {
+            underTest.init()
+
+            underTest.setIsAirplaneMode(IconState(true, TelephonyIcons.FLIGHT_MODE_ICON, ""))
+            runCurrent()
+            verify(statusBarIconController).setIconVisibility(slotAirplane, true)
+
+            underTest.setIsAirplaneMode(IconState(false, TelephonyIcons.FLIGHT_MODE_ICON, ""))
+            runCurrent()
+            verify(statusBarIconController).setIconVisibility(slotAirplane, false)
+        }
+
+    @Test
+    @DisableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun airplaneModeViaInteractor_statusBarSignalPolicyRefactorFlagDisabled_iconNotUpdated() =
+        testScope.runTest {
+            underTest.init()
+
+            // Make sure changing airplane mode from airplaneModeRepository does nothing
+            // if the StatusBarSignalPolicyRefactor is not enabled.
+            airplaneModeInteractor.setIsAirplaneMode(true)
+            runCurrent()
+            verifyZeroInteractions(statusBarIconController)
+
+            airplaneModeInteractor.setIsAirplaneMode(false)
+            runCurrent()
+            verifyZeroInteractions(statusBarIconController)
+        }
+
+    @Test
+    @DisableFlags(FLAG_STATUS_BAR_SIGNAL_POLICY_REFACTOR)
+    fun statusBarSignalPolicyInitialization_statusBarSignalPolicyRefactorFlagDisabled_startNoOp() =
+        testScope.runTest {
+            // Make sure StatusBarSignalPolicy.start does no initialization when
+            // the refactor flag is disabled.
+            underTest.start()
+            verifyZeroInteractions(securityController, networkController, tunerService)
+        }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModelTest.kt
index ce79fbd..7bc6d4a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/call/ui/viewmodel/CallChipViewModelTest.kt
@@ -132,10 +132,10 @@
             )
 
             assertThat((latest as OngoingActivityChipModel.Shown).icon)
-                .isInstanceOf(OngoingActivityChipModel.ChipIcon.Basic::class.java)
+                .isInstanceOf(OngoingActivityChipModel.ChipIcon.SingleColorIcon::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(com.android.internal.R.drawable.ic_phone)
             assertThat(icon.contentDescription).isNotNull()
@@ -170,10 +170,10 @@
             )
 
             assertThat((latest as OngoingActivityChipModel.Shown).icon)
-                .isInstanceOf(OngoingActivityChipModel.ChipIcon.Basic::class.java)
+                .isInstanceOf(OngoingActivityChipModel.ChipIcon.SingleColorIcon::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(com.android.internal.R.drawable.ic_phone)
             assertThat(icon.contentDescription).isNotNull()
@@ -206,10 +206,10 @@
             repo.setOngoingCallState(inCallModel(startTimeMs = 1000, notificationIcon = null))
 
             assertThat((latest as OngoingActivityChipModel.Shown).icon)
-                .isInstanceOf(OngoingActivityChipModel.ChipIcon.Basic::class.java)
+                .isInstanceOf(OngoingActivityChipModel.ChipIcon.SingleColorIcon::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(com.android.internal.R.drawable.ic_phone)
             assertThat(icon.contentDescription).isNotNull()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
index a8d2c5b..77992db 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/casttootherdevice/ui/viewmodel/CastToOtherDeviceChipViewModelTest.kt
@@ -127,7 +127,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_cast_connected)
             assertThat((icon.contentDescription as ContentDescription.Resource).res)
@@ -146,7 +146,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_cast_connected)
             assertThat((icon.contentDescription as ContentDescription.Resource).res)
@@ -184,7 +184,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.IconOnly::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_cast_connected)
             // This content description is just generic "Casting", not "Casting screen"
@@ -214,7 +214,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_cast_connected)
             // MediaProjection == screen casting, so this content description reflects that we're
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelTest.kt
new file mode 100644
index 0000000..118dea6
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelTest.kt
@@ -0,0 +1,146 @@
+/*
+ * 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.chips.ron.demo.ui.viewmodel
+
+import android.content.packageManager
+import android.graphics.drawable.BitmapDrawable
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_STATUS_BAR_RON_CHIPS
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.statusbar.chips.ui.model.ColorsModel
+import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
+import com.android.systemui.statusbar.commandline.CommandRegistry
+import com.android.systemui.statusbar.commandline.commandRegistry
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import java.io.PrintWriter
+import java.io.StringWriter
+import kotlin.test.Test
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
+
+@SmallTest
+class DemoRonChipViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+    private val commandRegistry = kosmos.commandRegistry
+    private val pw = PrintWriter(StringWriter())
+
+    private val underTest = kosmos.demoRonChipViewModel
+
+    @Before
+    fun setUp() {
+        underTest.start()
+        whenever(kosmos.packageManager.getApplicationIcon(any<String>()))
+            .thenReturn(BitmapDrawable())
+    }
+
+    @Test
+    @DisableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_flagOff_hidden() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            addDemoRonChip()
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Hidden::class.java)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_noPackage_hidden() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            commandRegistry.onShellCommand(pw, arrayOf("demo-ron"))
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Hidden::class.java)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_hasPackage_shown() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            commandRegistry.onShellCommand(pw, arrayOf("demo-ron", "-p", "com.android.systemui"))
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_hasText_shownWithText() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            commandRegistry.onShellCommand(
+                pw,
+                arrayOf("demo-ron", "-p", "com.android.systemui", "-t", "test")
+            )
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Text::class.java)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_supportsColor() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            commandRegistry.onShellCommand(
+                pw,
+                arrayOf("demo-ron", "-p", "com.android.systemui", "-c", "#434343")
+            )
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
+            assertThat((latest as OngoingActivityChipModel.Shown).colors)
+                .isInstanceOf(ColorsModel.Custom::class.java)
+        }
+
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
+    fun chip_hasHideArg_hidden() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.chip)
+
+            // First, show a chip
+            addDemoRonChip()
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
+
+            // Then, hide the chip
+            commandRegistry.onShellCommand(pw, arrayOf("demo-ron", "--hide"))
+
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Hidden::class.java)
+        }
+
+    private fun addDemoRonChip() {
+        Companion.addDemoRonChip(commandRegistry, pw)
+    }
+
+    companion object {
+        fun addDemoRonChip(commandRegistry: CommandRegistry, pw: PrintWriter) {
+            commandRegistry.onShellCommand(pw, arrayOf("demo-ron", "-p", "com.android.systemui"))
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
index 804eb5c..16101bf 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/screenrecord/ui/viewmodel/ScreenRecordChipViewModelTest.kt
@@ -150,7 +150,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_screenrecord)
             assertThat(icon.contentDescription).isNotNull()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
index a2ef599..791a21d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/sharetoapp/ui/viewmodel/ShareToAppChipViewModelTest.kt
@@ -135,7 +135,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_present_to_all)
             assertThat(icon.contentDescription).isNotNull()
@@ -152,7 +152,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown.Timer::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_present_to_all)
             assertThat(icon.contentDescription).isNotNull()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/ChipTransitionHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/ChipTransitionHelperTest.kt
index a724cfaa..4977c548 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/ChipTransitionHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/ChipTransitionHelperTest.kt
@@ -154,5 +154,7 @@
         }
 
     private fun createIcon(@DrawableRes drawable: Int) =
-        OngoingActivityChipModel.ChipIcon.Basic(Icon.Resource(drawable, contentDescription = null))
+        OngoingActivityChipModel.ChipIcon.SingleColorIcon(
+            Icon.Resource(drawable, contentDescription = null)
+        )
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelTest.kt
index 556ec6a..bd5df07 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelTest.kt
@@ -19,9 +19,12 @@
 import android.content.DialogInterface
 import android.content.packageManager
 import android.content.pm.PackageManager
+import android.graphics.drawable.BitmapDrawable
+import android.platform.test.annotations.EnableFlags
 import android.view.View
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_STATUS_BAR_RON_CHIPS
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.coroutines.collectLastValue
@@ -36,8 +39,11 @@
 import com.android.systemui.screenrecord.data.repository.screenRecordRepository
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.NORMAL_PACKAGE
 import com.android.systemui.statusbar.chips.mediaprojection.domain.interactor.MediaProjectionChipInteractorTest.Companion.setUpPackageManagerForMediaProjection
+import com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel.DemoRonChipViewModelTest.Companion.addDemoRonChip
+import com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel.demoRonChipViewModel
 import com.android.systemui.statusbar.chips.ui.model.OngoingActivityChipModel
 import com.android.systemui.statusbar.chips.ui.view.ChipBackgroundContainer
+import com.android.systemui.statusbar.commandline.commandRegistry
 import com.android.systemui.statusbar.phone.SystemUIDialog
 import com.android.systemui.statusbar.phone.mockSystemUIDialogFactory
 import com.android.systemui.statusbar.phone.ongoingcall.data.repository.ongoingCallRepository
@@ -45,6 +51,8 @@
 import com.android.systemui.statusbar.phone.ongoingcall.shared.model.inCallModel
 import com.android.systemui.util.time.fakeSystemClock
 import com.google.common.truth.Truth.assertThat
+import java.io.PrintWriter
+import java.io.StringWriter
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
@@ -68,11 +76,14 @@
     private val kosmos = Kosmos().also { it.testCase = this }
     private val testScope = kosmos.testScope
     private val systemClock = kosmos.fakeSystemClock
+    private val commandRegistry = kosmos.commandRegistry
 
     private val screenRecordState = kosmos.screenRecordRepository.screenRecordState
     private val mediaProjectionState = kosmos.fakeMediaProjectionRepository.mediaProjectionState
     private val callRepo = kosmos.ongoingCallRepository
 
+    private val pw = PrintWriter(StringWriter())
+
     private val mockSystemUIDialog = mock<SystemUIDialog>()
     private val chipBackgroundView = mock<ChipBackgroundContainer>()
     private val chipView =
@@ -90,6 +101,9 @@
     @Before
     fun setUp() {
         setUpPackageManagerForMediaProjection(kosmos)
+        kosmos.demoRonChipViewModel.start()
+        whenever(kosmos.packageManager.getApplicationIcon(any<String>()))
+            .thenReturn(BitmapDrawable())
     }
 
     @Test
@@ -169,15 +183,24 @@
         }
 
     @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
     fun chip_higherPriorityChipAdded_lowerPriorityChipReplaced() =
         testScope.runTest {
-            // Start with just the lower priority call chip
-            callRepo.setOngoingCallState(inCallModel(startTimeMs = 34))
+            // Start with just the lowest priority chip shown
+            addDemoRonChip(commandRegistry, pw)
+            // And everything else hidden
+            callRepo.setOngoingCallState(OngoingCallModel.NoCall)
             mediaProjectionState.value = MediaProjectionState.NotProjecting
             screenRecordState.value = ScreenRecordModel.DoingNothing
 
             val latest by collectLastValue(underTest.chip)
 
+            assertIsDemoRonChip(latest)
+
+            // WHEN the higher priority call chip is added
+            callRepo.setOngoingCallState(inCallModel(startTimeMs = 34))
+
+            // THEN the higher priority call chip is used
             assertIsCallChip(latest)
 
             // WHEN the higher priority media projection chip is added
@@ -199,14 +222,15 @@
         }
 
     @Test
+    @EnableFlags(FLAG_STATUS_BAR_RON_CHIPS)
     fun chip_highestPriorityChipRemoved_showsNextPriorityChip() =
         testScope.runTest {
             // WHEN all chips are active
             screenRecordState.value = ScreenRecordModel.Recording
             mediaProjectionState.value =
                 MediaProjectionState.Projecting.EntireScreen(NORMAL_PACKAGE)
-
             callRepo.setOngoingCallState(inCallModel(startTimeMs = 34))
+            addDemoRonChip(commandRegistry, pw)
 
             val latest by collectLastValue(underTest.chip)
 
@@ -224,6 +248,12 @@
 
             // THEN the lower priority call is used
             assertIsCallChip(latest)
+
+            // WHEN the higher priority call is removed
+            callRepo.setOngoingCallState(OngoingCallModel.NoCall)
+
+            // THEN the lower priority demo RON is used
+            assertIsDemoRonChip(latest)
         }
 
     /** Regression test for b/347726238. */
@@ -338,7 +368,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_screenrecord)
         }
@@ -347,7 +377,7 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(R.drawable.ic_present_to_all)
         }
@@ -356,9 +386,15 @@
             assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
             val icon =
                 (((latest as OngoingActivityChipModel.Shown).icon)
-                        as OngoingActivityChipModel.ChipIcon.Basic)
+                        as OngoingActivityChipModel.ChipIcon.SingleColorIcon)
                     .impl as Icon.Resource
             assertThat(icon.res).isEqualTo(com.android.internal.R.drawable.ic_phone)
         }
+
+        fun assertIsDemoRonChip(latest: OngoingActivityChipModel?) {
+            assertThat(latest).isInstanceOf(OngoingActivityChipModel.Shown::class.java)
+            assertThat((latest as OngoingActivityChipModel.Shown).icon)
+                .isInstanceOf(OngoingActivityChipModel.ChipIcon.FullColorAppIcon::class.java)
+        }
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
index 8cf7473..1efb3f7 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/commandline/ValueParserTest.kt
@@ -40,6 +40,15 @@
     }
 
     @Test
+    fun parseColor() {
+        assertThat(Type.Color.parseValue("#434343").isSuccess).isTrue()
+        assertThat(Type.Color.parseValue("#aa123456").isSuccess).isTrue()
+        assertThat(Type.Color.parseValue("red").isSuccess).isTrue()
+
+        assertThat(Type.Color.parseValue("not a color").isFailure).isTrue()
+    }
+
+    @Test
     fun mapToComplexType() {
         val parseSquare = Type.Int.map { Rect(it, it, it, it) }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
index e4945fc..1a1af2e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
@@ -342,7 +342,7 @@
         val stackTop = 200f
         val stackHeight = 800f
         whenever(ambientState.stackTop).thenReturn(stackTop)
-        whenever(ambientState.stackHeight).thenReturn(stackHeight)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(stackHeight)
         val shelfTop = stackTop + stackHeight - shelf.height
         val stackScrollAlgorithmState = StackScrollAlgorithmState()
         val viewInShelf = mock(ExpandableView::class.java)
@@ -378,7 +378,7 @@
         val stackTop = 200f
         val stackHeight = 800f
         whenever(ambientState.stackTop).thenReturn(stackTop)
-        whenever(ambientState.stackHeight).thenReturn(stackHeight)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(stackHeight)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         whenever(ambientState.isShadeExpanded).thenReturn(true)
@@ -404,7 +404,7 @@
     fun updateState_withNullLastVisibleBackgroundChild_hideShelf() {
         // GIVEN
         whenever(ambientState.stackY).thenReturn(100f)
-        whenever(ambientState.stackHeight).thenReturn(100f)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(100f)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         val endOfStack = 200f + paddingBetweenElements
@@ -433,7 +433,7 @@
         val stackTop = 200f
         val stackHeight = 800f
         whenever(ambientState.stackTop).thenReturn(stackTop)
-        whenever(ambientState.stackHeight).thenReturn(stackHeight)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(stackHeight)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         whenever(ambientState.isShadeExpanded).thenReturn(true)
@@ -459,7 +459,7 @@
     fun updateState_withNullFirstViewInShelf_hideShelf() {
         // GIVEN
         whenever(ambientState.stackY).thenReturn(100f)
-        whenever(ambientState.stackHeight).thenReturn(100f)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(100f)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         val endOfStack = 200f + paddingBetweenElements
@@ -488,7 +488,7 @@
         val stackTop = 200f
         val stackHeight = 800f
         whenever(ambientState.stackTop).thenReturn(stackTop)
-        whenever(ambientState.stackHeight).thenReturn(stackHeight)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(stackHeight)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         val lastVisibleBackgroundChild = mock<ExpandableView>()
@@ -514,7 +514,7 @@
     fun updateState_withCollapsedShade_hideShelf() {
         // GIVEN
         whenever(ambientState.stackY).thenReturn(100f)
-        whenever(ambientState.stackHeight).thenReturn(100f)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(100f)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         val endOfStack = 200f + paddingBetweenElements
@@ -543,7 +543,7 @@
         val stackTop = 200f
         val stackHeight = 800f
         whenever(ambientState.stackTop).thenReturn(stackTop)
-        whenever(ambientState.stackHeight).thenReturn(stackHeight)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(stackHeight)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         whenever(ambientState.isShadeExpanded).thenReturn(true)
@@ -583,7 +583,7 @@
     fun updateState_withHiddenSectionBeforeShelf_hideShelf() {
         // GIVEN
         whenever(ambientState.stackY).thenReturn(100f)
-        whenever(ambientState.stackHeight).thenReturn(100f)
+        whenever(ambientState.interpolatedStackHeight).thenReturn(100f)
         val paddingBetweenElements =
             context.resources.getDimensionPixelSize(R.dimen.notification_divider_height)
         val endOfStack = 200f + paddingBetweenElements
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 1717f4c..a18de68 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
@@ -95,7 +95,6 @@
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.ExpandableView;
 import com.android.systemui.statusbar.notification.shared.NotificationThrottleHun;
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor;
 import com.android.systemui.statusbar.phone.KeyguardBypassController;
 import com.android.systemui.statusbar.phone.ScreenOffAnimationController;
 import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
@@ -246,26 +245,12 @@
         when(mStackSizeCalculator.computeHeight(eq(mStackScroller), anyInt(), anyFloat()))
                 .thenReturn((float) stackHeight);
 
-        mStackScroller.updateContentHeight();
+        mStackScroller.updateStackHeight();
 
         assertThat(mStackScroller.getIntrinsicStackHeight()).isEqualTo(stackHeight);
     }
 
     @Test
-    @DisableSceneContainer
-    public void testIntrinsicStackHeight_includesTopScrimPadding() {
-        int stackHeight = 300;
-        int topScrimPadding = px(R.dimen.notification_side_paddings);
-        when(mStackSizeCalculator.computeHeight(eq(mStackScroller), anyInt(), anyFloat()))
-                .thenReturn((float) stackHeight);
-
-        mStackScroller.updateContentHeight();
-
-        assertThat(mStackScroller.getIntrinsicStackHeight())
-                .isEqualTo(stackHeight + topScrimPadding);
-    }
-
-    @Test
     @DisableSceneContainer // TODO(b/312473478): address disabled test
     public void testUpdateStackHeight_qsExpansionZero() {
         final float expansionFraction = 0.2f;
@@ -290,8 +275,8 @@
                 endHeight * StackScrollAlgorithm.START_FRACTION,
                 endHeight, expansionFraction);
 
-        mStackScroller.updateStackHeight(endHeight, expansionFraction);
-        assertThat(mAmbientState.getStackHeight()).isEqualTo(expected);
+        mStackScroller.updateInterpolatedStackHeight(endHeight, expansionFraction);
+        assertThat(mAmbientState.getInterpolatedStackHeight()).isEqualTo(expected);
     }
 
     @Test
@@ -310,7 +295,7 @@
 
         // THEN stackHeight and stackEndHeight are the same
         verify(mAmbientState).setStackEndHeight(stackEndHeight);
-        verify(mAmbientState).setStackHeight(stackEndHeight);
+        verify(mAmbientState).setInterpolatedStackHeight(stackEndHeight);
     }
 
     @Test
@@ -330,7 +315,7 @@
 
         // THEN stackHeight is changed by the expansion frac
         verify(mAmbientState).setStackEndHeight(stackEndHeight);
-        verify(mAmbientState).setStackHeight(stackEndHeight * 0.75f);
+        verify(mAmbientState).setInterpolatedStackHeight(stackEndHeight * 0.75f);
     }
 
     @Test
@@ -350,7 +335,7 @@
 
         // THEN stackHeight is measured from the stack top
         verify(mAmbientState).setStackEndHeight(stackEndHeight);
-        verify(mAmbientState).setStackHeight(stackEndHeight);
+        verify(mAmbientState).setInterpolatedStackHeight(stackEndHeight);
     }
 
     @Test
@@ -364,7 +349,7 @@
         clearInvocations(mAmbientState);
         mStackScroller.updateStackEndHeightAndStackHeight(1f);
 
-        verify(mAmbientState).setStackHeight(eq(300f));
+        verify(mAmbientState).setInterpolatedStackHeight(eq(300f));
     }
 
     @Test
@@ -377,7 +362,7 @@
         clearInvocations(mAmbientState);
         mStackScroller.updateStackEndHeightAndStackHeight(expansionFraction);
         verify(mAmbientState, never()).setStackEndHeight(anyFloat());
-        verify(mAmbientState).setStackHeight(anyFloat());
+        verify(mAmbientState).setInterpolatedStackHeight(anyFloat());
     }
 
     @Test
@@ -392,14 +377,14 @@
         clearInvocations(mAmbientState);
         mStackScroller.updateStackEndHeightAndStackHeight(expansionFraction);
         verify(mAmbientState, never()).setStackEndHeight(anyFloat());
-        verify(mAmbientState).setStackHeight(anyFloat());
+        verify(mAmbientState).setInterpolatedStackHeight(anyFloat());
 
         // Validate that when the animation ends the stackEndHeight is recalculated immediately
         clearInvocations(mAmbientState);
         mStackScroller.setPanelFlinging(false);
         verify(mAmbientState).setFlinging(eq(false));
         verify(mAmbientState).setStackEndHeight(anyFloat());
-        verify(mAmbientState).setStackHeight(anyFloat());
+        verify(mAmbientState).setInterpolatedStackHeight(anyFloat());
     }
 
     @Test
@@ -441,6 +426,86 @@
     }
 
     @Test
+    @EnableSceneContainer
+    public void setExpandFraction_fullyCollapsed() {
+        // Given: NSSL has a height
+        when(mStackScroller.getHeight()).thenReturn(1200);
+        // And: stack bounds are set
+        float expandFraction = 0.0f;
+        float stackTop = 100;
+        float stackCutoff = 1100;
+        float stackHeight = stackCutoff - stackTop;
+        mStackScroller.setStackTop(stackTop);
+        mStackScroller.setStackCutoff(stackCutoff);
+
+        // When: panel is fully collapsed
+        mStackScroller.setExpandFraction(expandFraction);
+
+        // Then
+        assertThat(mAmbientState.getExpansionFraction()).isEqualTo(expandFraction);
+        assertThat(mAmbientState.isExpansionChanging()).isFalse();
+        assertThat(mAmbientState.getStackEndHeight()).isEqualTo(stackHeight);
+        assertThat(mAmbientState.getInterpolatedStackHeight()).isEqualTo(
+                stackHeight * StackScrollAlgorithm.START_FRACTION);
+        assertThat(mAmbientState.isShadeExpanded()).isFalse();
+        assertThat(mStackScroller.getExpandedHeight()).isZero();
+    }
+
+    @Test
+    @EnableSceneContainer
+    public void setExpandFraction_expanding() {
+        // Given: NSSL has a height
+        when(mStackScroller.getHeight()).thenReturn(1200);
+        // And: stack bounds are set
+        float expandFraction = 0.6f;
+        float stackTop = 100;
+        float stackCutoff = 1100;
+        float stackHeight = stackCutoff - stackTop;
+        mStackScroller.setStackTop(stackTop);
+        mStackScroller.setStackCutoff(stackCutoff);
+
+        // When: panel is expanding
+        mStackScroller.setExpandFraction(expandFraction);
+
+        // Then
+        assertThat(mAmbientState.getExpansionFraction()).isEqualTo(expandFraction);
+        assertThat(mAmbientState.isExpansionChanging()).isTrue();
+        assertThat(mAmbientState.getStackEndHeight()).isEqualTo(stackHeight);
+        assertThat(mAmbientState.getInterpolatedStackHeight()).isGreaterThan(
+                stackHeight * StackScrollAlgorithm.START_FRACTION);
+        assertThat(mAmbientState.getInterpolatedStackHeight()).isLessThan(stackHeight);
+        assertThat(mStackScroller.getExpandedHeight()).isGreaterThan(0f);
+        assertThat(mAmbientState.isShadeExpanded()).isTrue();
+    }
+
+    @Test
+    @EnableSceneContainer
+    public void setExpandFraction_fullyExpanded() {
+        // Given: NSSL has a height
+        int viewHeight = 1200;
+        when(mStackScroller.getHeight()).thenReturn(viewHeight);
+        // And: stack bounds are set
+        float expandFraction = 1.0f;
+        float stackTop = 100;
+        float stackCutoff = 1100;
+        float stackHeight = stackCutoff - stackTop;
+        mStackScroller.setStackTop(stackTop);
+        mStackScroller.setStackCutoff(stackCutoff);
+
+        // When: panel is fully expanded
+        mStackScroller.setExpandFraction(expandFraction);
+
+        // Then
+        assertThat(mAmbientState.getExpansionFraction()).isEqualTo(expandFraction);
+        assertThat(mAmbientState.isExpansionChanging()).isFalse();
+        assertThat(mAmbientState.getStackEndHeight()).isEqualTo(stackHeight);
+        assertThat(mAmbientState.getInterpolatedStackHeight()).isEqualTo(stackHeight);
+        assertThat(mStackScroller.getExpandedHeight()).isEqualTo(viewHeight);
+        assertThat(mAmbientState.isShadeExpanded()).isTrue();
+    }
+
+    @Test
+    @DisableSceneContainer
     public void testSetExpandedHeight_listenerReceivedCallbacks() {
         final float expectedHeight = 0f;
 
@@ -467,6 +532,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testSetExpandedHeight_withSplitShade_doesntInterpolateStackHeight() {
         mTestableResources
                 .addOverride(R.bool.config_use_split_notification_shade, /* value= */ true);
@@ -1206,7 +1272,7 @@
 
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     public void testGenerateHeadsUpDisappearEvent_setsHeadsUpAnimatingAway() {
         // GIVEN NSSL is ready for HUN animations
         Consumer<Boolean> headsUpAnimatingAwayListener = mock(BooleanConsumer.class);
@@ -1222,7 +1288,7 @@
     }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     public void testGenerateHeadsUpDisappearEvent_stackExpanded_headsUpAnimatingAwayNotSet() {
         // GIVEN NSSL would be ready for HUN animations, BUT it is expanded
         Consumer<Boolean> headsUpAnimatingAwayListener = mock(BooleanConsumer.class);
@@ -1241,7 +1307,7 @@
     }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     public void testGenerateHeadsUpDisappearEvent_pendingAppearEvent_headsUpAnimatingAwayNotSet() {
         // GIVEN NSSL is ready for HUN animations
         Consumer<Boolean> headsUpAnimatingAwayListener = mock(BooleanConsumer.class);
@@ -1259,7 +1325,7 @@
     }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     public void testGenerateHeadsUpAppearEvent_headsUpAnimatingAwayNotSet() {
         // GIVEN NSSL is ready for HUN animations
         Consumer<Boolean> headsUpAnimatingAwayListener = mock(BooleanConsumer.class);
@@ -1295,7 +1361,7 @@
     }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     public void testOnChildAnimationsFinished_resetsheadsUpAnimatingAway() {
         // GIVEN NSSL is ready for HUN animations
         Consumer<Boolean> headsUpAnimatingAwayListener = mock(BooleanConsumer.class);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
index 7e79019..3e8bf47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithmTest.kt
@@ -1114,6 +1114,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     fun shadeOpened_hunDoesNotOverlapQQS_hunShouldHaveNoShadow() {
         // Given: shade is opened, yTranslation of HUN is equal to QQS Panel's height,
         // the height of HUN is equal to the height of QQS Panel,
@@ -1144,6 +1145,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     fun shadeClosed_hunShouldHaveFullShadow() {
         // Given: shade is closed, ambientState.stackTranslation == -ambientState.topPadding,
         // the height of HUN is equal to the height of QQS Panel,
@@ -1172,6 +1174,7 @@
     }
 
     @Test
+    @DisableSceneContainer
     fun draggingHunToOpenShade_hunShouldHavePartialShadow() {
         // Given: shade is closed when HUN pops up,
         // now drags down the HUN to open shade
@@ -1447,7 +1450,7 @@
         // set stackEndHeight and stackHeight
         // ExpansionFractionWithoutShelf == stackHeight / stackEndHeight
         ambientState.stackEndHeight = 100f
-        ambientState.stackHeight = ambientState.stackEndHeight * fraction
+        ambientState.interpolatedStackHeight = ambientState.stackEndHeight * fraction
     }
 
     private fun resetViewStates_hunYTranslationIs(expected: Float) {
@@ -1531,7 +1534,7 @@
         // shade is fully open
         ambientState.expansionFraction = 1.0f
         with(fullStackHeight) {
-            ambientState.stackHeight = this
+            ambientState.interpolatedStackHeight = this
             ambientState.stackEndHeight = this
         }
         stackScrollAlgorithm.setIsExpanded(true)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
index 2f81027..c7919df 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
@@ -20,6 +20,7 @@
 import android.os.UserManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.keyguard.KeyguardUpdateMonitor
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.settings.UserTracker
 import com.android.systemui.util.concurrency.FakeExecutor
@@ -35,6 +36,7 @@
 import org.mockito.Mockito.verify
 import org.mockito.Mockito.`when`
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.never
 
 @SmallTest
 @RunWith(AndroidJUnit4::class)
@@ -46,12 +48,20 @@
 
     @Mock private lateinit var userTracker: UserTracker
     @Mock private lateinit var userManager: UserManager
+    @Mock private lateinit var keyguardUpdateMonitor: KeyguardUpdateMonitor
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
 
-        controller = ManagedProfileControllerImpl(context, mainExecutor, userTracker, userManager)
+        controller =
+            ManagedProfileControllerImpl(
+                context,
+                mainExecutor,
+                userTracker,
+                userManager,
+                keyguardUpdateMonitor
+            )
     }
 
     @Test
@@ -107,6 +117,24 @@
         captor.value.onProfilesChanged(userManager.getEnabledProfiles(1))
     }
 
+    @Test
+    fun hasWorkingProfile_setWorkModeEnabled_callsAwakenFromDream() {
+        `when`(userTracker.userId).thenReturn(1)
+        setupWorkingProfile(1)
+        `when`(userManager.requestQuietModeEnabled(any(), any())).thenReturn(false)
+        controller.hasActiveProfile()
+
+        controller.isWorkModeEnabled = true
+
+        verify(keyguardUpdateMonitor).awakenFromDream()
+    }
+
+    @Test
+    fun noWorkingProfile_setWorkModeEnabled_NoAwakenFromDreamCall() {
+        controller.isWorkModeEnabled = true
+        verify(keyguardUpdateMonitor, never()).awakenFromDream()
+    }
+
     private fun setupWorkingProfile(userId: Int) {
         `when`(userManager.getEnabledProfiles(userId))
             .thenReturn(
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
index 76dc65c..2ed3473 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
@@ -34,6 +34,7 @@
 import android.testing.TestableLooper.RunWithLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
@@ -41,6 +42,7 @@
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.PendingDisplay
 import com.android.systemui.display.domain.interactor.ConnectedDisplayInteractor.State
+import com.android.systemui.kosmos.testScope
 import com.android.systemui.privacy.PrivacyItemController
 import com.android.systemui.privacy.logging.PrivacyLogger
 import com.android.systemui.screenrecord.RecordingController
@@ -71,9 +73,7 @@
 import com.android.systemui.util.time.FakeSystemClock
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -145,7 +145,7 @@
     private lateinit var alarmCallbackCaptor:
         ArgumentCaptor<NextAlarmController.NextAlarmChangeCallback>
 
-    private val testScope = TestScope(UnconfinedTestDispatcher())
+    private val testScope = kosmos.testScope
     private val fakeConnectedDisplayStateProvider = FakeConnectedDisplayStateProvider()
     private val zenModeController = FakeZenModeController()
 
@@ -249,7 +249,7 @@
             statusBarPolicy.init()
             clearInvocations(iconController)
 
-            fakeConnectedDisplayStateProvider.emit(State.CONNECTED)
+            fakeConnectedDisplayStateProvider.setState(State.CONNECTED)
             runCurrent()
 
             verify(iconController).setIconVisibility(CONNECTED_DISPLAY_SLOT, true)
@@ -261,7 +261,8 @@
             statusBarPolicy.init()
             clearInvocations(iconController)
 
-            fakeConnectedDisplayStateProvider.emit(State.DISCONNECTED)
+            fakeConnectedDisplayStateProvider.setState(State.DISCONNECTED)
+            runCurrent()
 
             verify(iconController).setIconVisibility(CONNECTED_DISPLAY_SLOT, false)
         }
@@ -272,9 +273,12 @@
             statusBarPolicy.init()
             clearInvocations(iconController)
 
-            fakeConnectedDisplayStateProvider.emit(State.CONNECTED)
-            fakeConnectedDisplayStateProvider.emit(State.DISCONNECTED)
-            fakeConnectedDisplayStateProvider.emit(State.CONNECTED)
+            fakeConnectedDisplayStateProvider.setState(State.CONNECTED)
+            runCurrent()
+            fakeConnectedDisplayStateProvider.setState(State.DISCONNECTED)
+            runCurrent()
+            fakeConnectedDisplayStateProvider.setState(State.CONNECTED)
+            runCurrent()
 
             inOrder(iconController).apply {
                 verify(iconController).setIconVisibility(CONNECTED_DISPLAY_SLOT, true)
@@ -289,7 +293,8 @@
             statusBarPolicy.init()
             clearInvocations(iconController)
 
-            fakeConnectedDisplayStateProvider.emit(State.CONNECTED_SECURE)
+            fakeConnectedDisplayStateProvider.setState(State.CONNECTED_SECURE)
+            runCurrent()
 
             verify(iconController).setIconVisibility(CONNECTED_DISPLAY_SLOT, true)
         }
@@ -390,7 +395,7 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_MODES_UI_ICONS)
+    @EnableFlags(android.app.Flags.FLAG_MODES_UI, android.app.Flags.FLAG_MODES_UI_ICONS)
     fun zenModeInteractorActiveModeChanged_showsModeIcon() =
         testScope.runTest {
             statusBarPolicy.init()
@@ -403,8 +408,8 @@
                         .setName("Bedtime Mode")
                         .setType(AutomaticZenRule.TYPE_BEDTIME)
                         .setActive(true)
-                        .setPackage("some.package")
-                        .setIconResId(123)
+                        .setPackage(mContext.packageName)
+                        .setIconResId(android.R.drawable.ic_lock_lock)
                         .build(),
                     TestModeBuilder()
                         .setId("other")
@@ -412,7 +417,7 @@
                         .setType(AutomaticZenRule.TYPE_OTHER)
                         .setActive(true)
                         .setPackage(SystemZenRules.PACKAGE_ANDROID)
-                        .setIconResId(456)
+                        .setIconResId(android.R.drawable.ic_media_play)
                         .build(),
                 )
             )
@@ -422,17 +427,25 @@
             verify(iconController)
                 .setResourceIcon(
                     eq(ZEN_SLOT),
-                    eq("some.package"),
-                    eq(123),
-                    eq(null),
-                    eq("Bedtime Mode")
+                    eq(mContext.packageName),
+                    eq(android.R.drawable.ic_lock_lock),
+                    any(), // non-null
+                    eq("Bedtime Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
                 )
 
             zenModeRepository.deactivateMode("bedtime")
             runCurrent()
 
             verify(iconController)
-                .setResourceIcon(eq(ZEN_SLOT), eq(null), eq(456), eq(null), eq("Other Mode"))
+                .setResourceIcon(
+                    eq(ZEN_SLOT),
+                    eq(null),
+                    eq(android.R.drawable.ic_media_play),
+                    any(), // non-null
+                    eq("Other Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
+                )
 
             zenModeRepository.deactivateMode("other")
             runCurrent()
@@ -441,7 +454,7 @@
         }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_MODES_UI_ICONS)
+    @EnableFlags(android.app.Flags.FLAG_MODES_UI, android.app.Flags.FLAG_MODES_UI_ICONS)
     fun zenModeControllerOnGlobalZenChanged_doesNotUpdateDndIcon() {
         statusBarPolicy.init()
         reset(iconController)
@@ -450,7 +463,8 @@
 
         verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
         verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
-        verify(iconController, never()).setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+        verify(iconController, never())
+            .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
     }
 
     @Test
@@ -466,7 +480,7 @@
             verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
             verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
             verify(iconController, never())
-                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
         }
 
     @Test
@@ -529,9 +543,11 @@
     }
 
     private class FakeConnectedDisplayStateProvider : ConnectedDisplayInteractor {
-        private val flow = MutableSharedFlow<State>()
+        private val flow = MutableStateFlow(State.DISCONNECTED)
 
-        suspend fun emit(value: State) = flow.emit(value)
+        fun setState(value: State) {
+            flow.value = value
+        }
 
         override val connectedDisplayState: Flow<State>
             get() = flow
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
index b75ac2b..3e3c046 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManagerTest.java
@@ -83,10 +83,10 @@
 import com.android.systemui.dreams.DreamOverlayStateController;
 import com.android.systemui.flags.DisableSceneContainer;
 import com.android.systemui.flags.EnableSceneContainer;
+import com.android.systemui.keyguard.DismissCallbackRegistry;
+import com.android.systemui.keyguard.domain.interactor.KeyguardDismissTransitionInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardDismissActionInteractor;
-import com.android.systemui.keyguard.domain.interactor.KeyguardSurfaceBehindInteractor;
 import com.android.systemui.keyguard.domain.interactor.KeyguardTransitionInteractor;
-import com.android.systemui.keyguard.domain.interactor.WindowManagerLockscreenVisibilityInteractor;
 import com.android.systemui.keyguard.shared.model.KeyguardState;
 import com.android.systemui.keyguard.shared.model.TransitionState;
 import com.android.systemui.keyguard.shared.model.TransitionStep;
@@ -171,6 +171,7 @@
     @Mock private SelectedUserInteractor mSelectedUserInteractor;
     @Mock private DeviceEntryInteractor mDeviceEntryInteractor;
     @Mock private SceneInteractor mSceneInteractor;
+    @Mock private DismissCallbackRegistry mDismissCallbackRegistry;
 
     private StatusBarKeyguardViewManager mStatusBarKeyguardViewManager;
     private PrimaryBouncerCallbackInteractor.PrimaryBouncerExpansionCallback
@@ -233,16 +234,16 @@
                         mUdfpsOverlayInteractor,
                         mActivityStarter,
                         mKeyguardTransitionInteractor,
+                        mock(KeyguardDismissTransitionInteractor.class),
                         StandardTestDispatcher(null, null),
-                        () -> mock(WindowManagerLockscreenVisibilityInteractor.class),
                         () -> mock(KeyguardDismissActionInteractor.class),
                         mSelectedUserInteractor,
-                        () -> mock(KeyguardSurfaceBehindInteractor.class),
                         mock(JavaAdapter.class),
                         () -> mSceneInteractor,
                         mock(StatusBarKeyguardViewManagerInteractor.class),
                         mExecutor,
-                        () -> mDeviceEntryInteractor) {
+                        () -> mDeviceEntryInteractor,
+                        mDismissCallbackRegistry) {
                     @Override
                     public ViewRootImpl getViewRootImpl() {
                         return mViewRootImpl;
@@ -756,16 +757,16 @@
                         mUdfpsOverlayInteractor,
                         mActivityStarter,
                         mock(KeyguardTransitionInteractor.class),
+                        mock(KeyguardDismissTransitionInteractor.class),
                         StandardTestDispatcher(null, null),
-                        () -> mock(WindowManagerLockscreenVisibilityInteractor.class),
                         () -> mock(KeyguardDismissActionInteractor.class),
                         mSelectedUserInteractor,
-                        () -> mock(KeyguardSurfaceBehindInteractor.class),
                         mock(JavaAdapter.class),
                         () -> mSceneInteractor,
                         mock(StatusBarKeyguardViewManagerInteractor.class),
                         mExecutor,
-                        () -> mDeviceEntryInteractor) {
+                        () -> mDeviceEntryInteractor,
+                        mDismissCallbackRegistry) {
                     @Override
                     public ViewRootImpl getViewRootImpl() {
                         return mViewRootImpl;
@@ -777,7 +778,11 @@
     }
 
     @Test
+    @DisableSceneContainer
     public void testResetHideBouncerWhenShowing_alternateBouncerHides() {
+        reset(mDismissCallbackRegistry);
+        reset(mPrimaryBouncerInteractor);
+
         // GIVEN the keyguard is showing
         reset(mAlternateBouncerInteractor);
         when(mKeyguardStateController.isShowing()).thenReturn(true);
@@ -785,8 +790,10 @@
         // WHEN SBKV is reset with hideBouncerWhenShowing=true
         mStatusBarKeyguardViewManager.reset(true);
 
-        // THEN alternate bouncer is hidden
+        // THEN alternate bouncer is hidden and dismiss actions reset
         verify(mAlternateBouncerInteractor).hide();
+        verify(mDismissCallbackRegistry).notifyDismissCancelled();
+        verify(mPrimaryBouncerInteractor).setDismissAction(eq(null), eq(null));
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
new file mode 100644
index 0000000..90732d01
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
@@ -0,0 +1,122 @@
+/*
+ * 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.ui
+
+import android.app.Flags
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.view.ViewGroup
+import android.widget.ImageView
+import android.widget.LinearLayout
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.connectivity.ui.MobileContextProvider
+import com.android.systemui.statusbar.phone.StatusBarLocation
+import com.android.systemui.statusbar.pipeline.mobile.ui.MobileUiAdapter
+import com.android.systemui.statusbar.pipeline.wifi.ui.WifiUiAdapter
+import com.android.systemui.util.Assert
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.RETURNS_DEEP_STUBS
+import org.mockito.kotlin.mock
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class IconManagerTest : SysuiTestCase() {
+
+    private lateinit var underTest: IconManager
+    private lateinit var viewGroup: ViewGroup
+
+    @Before
+    fun setUp() {
+        Assert.setTestThread(Thread.currentThread())
+        viewGroup = LinearLayout(context)
+        underTest =
+            IconManager(
+                viewGroup,
+                StatusBarLocation.HOME,
+                mock<WifiUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileContextProvider>(defaultAnswer = RETURNS_DEEP_STUBS),
+            )
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeWrapContent_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.WRAP_CONTENT)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeFixedSpace_addsIconViewWithFixedWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isNotEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.layoutParams.width).isEqualTo(iconView.layoutParams.height)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.FIT_CENTER)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_iconsFlagOff_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    private fun newStatusBarIcon(shape: StatusBarIcon.Shape) =
+        StatusBarIcon(
+            UserHandle.CURRENT,
+            context.packageName,
+            Icon.createWithResource(context, android.R.drawable.ic_media_next),
+            0,
+            0,
+            "",
+            StatusBarIcon.Type.ResourceIcon,
+            shape,
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
index 26a57e4..50a13b9 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
@@ -424,7 +424,14 @@
     @EnableFlags(android.app.Flags.FLAG_MODES_UI, android.app.Flags.FLAG_MODES_UI_ICONS)
     fun setResourceIcon_setsIconAndPreloadedIconInHolder() {
         val drawable = ColorDrawable(1)
-        underTest.setResourceIcon("slot", "some.package", 123, drawable, "description")
+        underTest.setResourceIcon(
+            "slot",
+            "some.package",
+            123,
+            drawable,
+            "description",
+            StatusBarIcon.Shape.FIXED_SPACE
+        )
 
         val iconHolder = iconList.getIconHolder("slot", 0)
         assertThat(iconHolder).isNotNull()
@@ -432,6 +439,7 @@
         assertThat(iconHolder?.icon?.icon?.resId).isEqualTo(123)
         assertThat(iconHolder?.icon?.icon?.resPackage).isEqualTo("some.package")
         assertThat(iconHolder?.icon?.contentDescription).isEqualTo("description")
+        assertThat(iconHolder?.icon?.shape).isEqualTo(StatusBarIcon.Shape.FIXED_SPACE)
         assertThat(iconHolder?.icon?.preloadedIcon).isEqualTo(drawable)
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
index db3e533..7901f47 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorTest.kt
@@ -19,13 +19,11 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.launchIn
@@ -39,9 +37,9 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class AirplaneModeInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
-    private val mobileConnectionsRepository =
-        FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), mock<TableLogBuffer> {})
+    private val mobileConnectionsRepository = kosmos.fakeMobileConnectionsRepository
     private val airplaneModeRepository = FakeAirplaneModeRepository()
     private val connectivityRepository = FakeConnectivityRepository()
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/ui/viewmodel/AirplaneModeViewModelImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/ui/viewmodel/AirplaneModeViewModelImplTest.kt
index b823333..8beed01 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/ui/viewmodel/AirplaneModeViewModelImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/airplane/ui/viewmodel/AirplaneModeViewModelImplTest.kt
@@ -19,12 +19,13 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -35,7 +36,6 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
 import org.mockito.MockitoAnnotations
 
 @SmallTest
@@ -43,10 +43,11 @@
 @Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
 @RunWith(AndroidJUnit4::class)
 class AirplaneModeViewModelImplTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var underTest: AirplaneModeViewModelImpl
 
-    @Mock private lateinit var logger: TableLogBuffer
+    private val logger = logcatTableLogBuffer(kosmos, "AirplaneModeViewModelImplTest")
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
     private lateinit var connectivityRepository: FakeConnectivityRepository
     private lateinit var interactor: AirplaneModeInteractor
@@ -61,7 +62,7 @@
             AirplaneModeInteractor(
                 airplaneModeRepository,
                 connectivityRepository,
-                FakeMobileConnectionsRepository(),
+                kosmos.fakeMobileConnectionsRepository,
             )
         scope = CoroutineScope(IMMEDIATE)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
index 7d586cd..36f5236 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
@@ -27,7 +27,7 @@
 import com.android.systemui.demomode.DemoModeController
 import com.android.systemui.dump.DumpManager
 import com.android.systemui.log.table.TableLogBuffer
-import com.android.systemui.log.table.TableLogBufferFactory
+import com.android.systemui.log.table.tableLogBufferFactory
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.MobileInputLogger
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
@@ -42,11 +42,11 @@
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.FakeWifiRepository
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.demo.DemoModeWifiDataSource
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.kotlinArgumentCaptor
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -56,7 +56,6 @@
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
@@ -75,12 +74,13 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class MobileRepositorySwitcherTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: MobileRepositorySwitcher
     private lateinit var realRepo: MobileConnectionsRepositoryImpl
     private lateinit var demoRepo: DemoMobileConnectionsRepository
     private lateinit var mobileDataSource: DemoModeMobileConnectionDataSource
     private lateinit var wifiDataSource: DemoModeWifiDataSource
-    private lateinit var logFactory: TableLogBufferFactory
     private lateinit var wifiRepository: FakeWifiRepository
     private lateinit var connectivityRepository: ConnectivityRepository
 
@@ -95,16 +95,12 @@
     private val mobileMappings = FakeMobileMappingsProxy()
     private val subscriptionManagerProxy = FakeSubscriptionManagerProxy()
 
-    private val testDispatcher = UnconfinedTestDispatcher()
     private val scope = CoroutineScope(IMMEDIATE)
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
-        logFactory =
-            TableLogBufferFactory(dumpManager, FakeSystemClock(), mock(), testDispatcher, scope)
-
         // Never start in demo mode
         whenever(demoModeController.isInDemoMode).thenReturn(false)
 
@@ -147,7 +143,7 @@
                 wifiDataSource = wifiDataSource,
                 scope = scope,
                 context = context,
-                logFactory = logFactory,
+                logFactory = kosmos.tableLogBufferFactory,
             )
 
         underTest =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
index db6f5927..7d32021 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
@@ -23,16 +23,16 @@
 import com.android.settingslib.SignalIcon
 import com.android.settingslib.mobile.TelephonyIcons
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.log.table.TableLogBufferFactory
+import com.android.systemui.log.table.tableLogBufferFactory
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
 import com.android.systemui.statusbar.pipeline.mobile.data.model.NetworkNameModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
 import com.android.systemui.statusbar.pipeline.shared.data.model.toMobileDataActivityModel
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.demo.DemoModeWifiDataSource
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.demo.model.FakeWifiEventModel
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
@@ -43,12 +43,11 @@
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.runTest
 import org.junit.After
-import platform.test.runner.parameterized.ParameterizedAndroidJunit4
-import platform.test.runner.parameterized.Parameters
-import platform.test.runner.parameterized.Parameter
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
 
 /**
  * Parameterized test for all of the common values of [FakeNetworkEventModel]. This test simply
@@ -60,19 +59,11 @@
 @RunWith(ParameterizedAndroidJunit4::class)
 internal class DemoMobileConnectionParameterizedTest(private val testCase: TestCase) :
     SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private val testDispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(testDispatcher)
 
-    private val logFactory =
-        TableLogBufferFactory(
-            mock(),
-            FakeSystemClock(),
-            mock(),
-            testDispatcher,
-            testScope.backgroundScope,
-        )
-
     private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
     private val fakeWifiEventFlow = MutableStateFlow<FakeWifiEventModel?>(null)
 
@@ -99,7 +90,7 @@
                 wifiDataSource = mockWifiDataSource,
                 scope = testScope.backgroundScope,
                 context = context,
-                logFactory = logFactory,
+                logFactory = kosmos.tableLogBufferFactory,
             )
 
         connectionsRepo.startProcessingCommands()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
index 5e0d2fb0..5017dda 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
@@ -24,8 +24,7 @@
 import com.android.settingslib.SignalIcon
 import com.android.settingslib.mobile.TelephonyIcons.THREE_G
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.dump.DumpManager
-import com.android.systemui.log.table.TableLogBufferFactory
+import com.android.systemui.log.table.tableLogBufferFactory
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
 import com.android.systemui.statusbar.pipeline.mobile.data.model.NetworkNameModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
@@ -34,9 +33,9 @@
 import com.android.systemui.statusbar.pipeline.shared.data.model.toMobileDataActivityModel
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.demo.DemoModeWifiDataSource
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.demo.model.FakeWifiEventModel
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -56,21 +55,13 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class DemoMobileConnectionsRepositoryTest : SysuiTestCase() {
-    private val dumpManager: DumpManager = mock()
+    private val kosmos = testKosmos()
 
     private val testDispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(testDispatcher)
 
     private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
     private val fakeWifiEventFlow = MutableStateFlow<FakeWifiEventModel?>(null)
-    private val logFactory =
-        TableLogBufferFactory(
-            dumpManager,
-            FakeSystemClock(),
-            mock(),
-            testDispatcher,
-            testScope.backgroundScope,
-        )
 
     private lateinit var underTest: DemoMobileConnectionsRepository
     private lateinit var mobileDataSource: DemoModeMobileConnectionDataSource
@@ -94,7 +85,7 @@
                 wifiDataSource = wifiDataSource,
                 scope = testScope.backgroundScope,
                 context = context,
-                logFactory = logFactory,
+                logFactory = kosmos.tableLogBufferFactory,
             )
 
         underTest.startProcessingCommands()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/FullMobileConnectionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/FullMobileConnectionRepositoryTest.kt
index fd4c370..c029850 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/FullMobileConnectionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/FullMobileConnectionRepositoryTest.kt
@@ -29,8 +29,8 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags.ROAMING_INDICATOR_VIA_DISPLAY_INFO
-import com.android.systemui.log.table.TableLogBuffer
-import com.android.systemui.log.table.TableLogBufferFactory
+import com.android.systemui.log.table.logcatTableLogBuffer
+import com.android.systemui.log.table.tableLogBufferFactory
 import com.android.systemui.statusbar.pipeline.mobile.data.model.NetworkNameModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfig
@@ -42,11 +42,11 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileTelephonyHelpers.getTelephonyCallbackForType
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.FakeWifiRepository
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiNetworkModel
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.io.PrintWriter
 import java.io.StringWriter
@@ -73,23 +73,16 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class FullMobileConnectionRepositoryTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: FullMobileConnectionRepository
 
     private val flags =
         FakeFeatureFlagsClassic().also { it.set(ROAMING_INDICATOR_VIA_DISPLAY_INFO, true) }
 
-    private val systemClock = FakeSystemClock()
     private val testDispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(testDispatcher)
-    private val tableLogBuffer =
-        TableLogBuffer(
-            maxSize = 100,
-            name = "TestName",
-            systemClock,
-            mock(),
-            testDispatcher,
-            testScope.backgroundScope,
-        )
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "TestName")
     private val mobileFactory = mock<MobileConnectionRepositoryImpl.Factory>()
     private val carrierMergedFactory = mock<CarrierMergedConnectionRepository.Factory>()
     private val connectivityManager = mock<ConnectivityManager>()
@@ -372,19 +365,10 @@
     @Test
     fun factory_reusesLogBuffersForSameConnection() =
         testScope.runTest {
-            val realLoggerFactory =
-                TableLogBufferFactory(
-                    mock(),
-                    FakeSystemClock(),
-                    mock(),
-                    testDispatcher,
-                    testScope.backgroundScope,
-                )
-
             val factory =
                 FullMobileConnectionRepository.Factory(
                     scope = testScope.backgroundScope,
-                    realLoggerFactory,
+                    kosmos.tableLogBufferFactory,
                     mobileFactory,
                     carrierMergedFactory,
                 )
@@ -416,19 +400,10 @@
     @Test
     fun factory_reusesLogBuffersForSameSubIDevenIfCarrierMerged() =
         testScope.runTest {
-            val realLoggerFactory =
-                TableLogBufferFactory(
-                    mock(),
-                    FakeSystemClock(),
-                    mock(),
-                    testDispatcher,
-                    testScope.backgroundScope,
-                )
-
             val factory =
                 FullMobileConnectionRepository.Factory(
                     scope = testScope.backgroundScope,
-                    realLoggerFactory,
+                    kosmos.tableLogBufferFactory,
                     mobileFactory,
                     carrierMergedFactory,
                 )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
index 171520f..fe408e3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
@@ -88,7 +88,6 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfig
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfigTest.Companion.configWithOverride
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfigTest.Companion.createTestConfig
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository.Companion.DEFAULT_NUM_LEVELS
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileTelephonyHelpers.signalStrength
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileTelephonyHelpers.telephonyDisplayInfo
@@ -121,7 +120,6 @@
 @RunWith(AndroidJUnit4::class)
 class MobileConnectionRepositoryTest : SysuiTestCase() {
     private lateinit var underTest: MobileConnectionRepositoryImpl
-    private lateinit var connectionsRepo: FakeMobileConnectionsRepository
 
     private val flags =
         FakeFeatureFlagsClassic().also { it.set(ROAMING_INDICATOR_VIA_DISPLAY_INFO, true) }
@@ -156,8 +154,6 @@
         MockitoAnnotations.initMocks(this)
         whenever(telephonyManager.subscriptionId).thenReturn(SUB_1_ID)
 
-        connectionsRepo = FakeMobileConnectionsRepository(mobileMappings, tableLogger)
-
         underTest =
             MobileConnectionRepositoryImpl(
                 SUB_1_ID,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionTelephonySmokeTests.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionTelephonySmokeTests.kt
index 2ab8c0a..0d82c79 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionTelephonySmokeTests.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionTelephonySmokeTests.kt
@@ -42,7 +42,6 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfig
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SystemUiCarrierConfigTest
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileTelephonyHelpers.getTelephonyCallbackForType
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileTelephonyHelpers.signalStrength
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
@@ -97,7 +96,6 @@
 @SmallTest
 class MobileConnectionTelephonySmokeTests : SysuiTestCase() {
     private lateinit var underTest: MobileConnectionRepositoryImpl
-    private lateinit var connectionsRepo: FakeMobileConnectionsRepository
 
     private val flags =
         FakeFeatureFlagsClassic().also { it.set(Flags.ROAMING_INDICATOR_VIA_DISPLAY_INFO, true) }
@@ -123,12 +121,6 @@
         MockitoAnnotations.initMocks(this)
         whenever(telephonyManager.subscriptionId).thenReturn(SUB_1_ID)
 
-        connectionsRepo =
-            FakeMobileConnectionsRepository(
-                mobileMappings,
-                tableLogger,
-            )
-
         underTest =
             MobileConnectionRepositoryImpl(
                 SUB_1_ID,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
index 76982ae..6de2caa 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
@@ -28,7 +28,6 @@
 import android.net.vcn.VcnTransportInfo
 import android.net.wifi.WifiInfo
 import android.net.wifi.WifiManager
-import android.os.Bundle
 import android.os.ParcelUuid
 import android.telephony.CarrierConfigManager
 import android.telephony.ServiceState
@@ -56,7 +55,6 @@
 import com.android.systemui.statusbar.connectivity.WifiPickerTrackerFactory
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.MobileInputLogger
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.CarrierConfigRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
@@ -74,7 +72,6 @@
 import com.android.systemui.util.mockito.capture
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.time.FakeSystemClock
 import com.android.wifitrackerlib.MergedCarrierEntry
 import com.android.wifitrackerlib.WifiEntry
@@ -98,6 +95,7 @@
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.whenever
 
 @Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -602,47 +600,85 @@
 
     @SuppressLint("UnspecifiedRegisterReceiverFlag")
     @Test
-    fun testDeviceServiceStateFromBroadcast_eagerlyWatchesBroadcast() =
+    fun testDeviceEmergencyCallState_eagerlyChecksState() =
         testScope.runTest {
-            // Value starts out empty (null)
-            assertThat(underTest.deviceServiceState.value).isNull()
+            // Value starts out false
+            assertThat(underTest.isDeviceEmergencyCallCapable.value).isFalse()
+            whenever(telephonyManager.activeModemCount).thenReturn(1)
+            whenever(telephonyManager.getServiceStateForSlot(any())).thenAnswer { _ ->
+                ServiceState().apply { isEmergencyOnly = true }
+            }
 
             // WHEN an appropriate intent gets sent out
-            val intent = serviceStateIntent(subId = -1, emergencyOnly = false)
+            val intent = serviceStateIntent(subId = -1)
             fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
                 context,
                 intent,
             )
             runCurrent()
 
-            // THEN the repo's state is updated
-            val expected = ServiceStateModel(isEmergencyOnly = false)
-            assertThat(underTest.deviceServiceState.value).isEqualTo(expected)
+            // THEN the repo's state is updated despite no listeners
+            assertThat(underTest.isDeviceEmergencyCallCapable.value).isEqualTo(true)
         }
 
     @Test
-    fun testDeviceServiceStateFromBroadcast_followsSubIdNegativeOne() =
+    fun testDeviceEmergencyCallState_aggregatesAcrossSlots_oneTrue() =
         testScope.runTest {
-            // device based state tracks -1
-            val intent = serviceStateIntent(subId = -1, emergencyOnly = false)
+            val latest by collectLastValue(underTest.isDeviceEmergencyCallCapable)
+
+            // GIVEN there are multiple slots
+            whenever(telephonyManager.activeModemCount).thenReturn(4)
+            // GIVEN only one of them reports ECM
+            whenever(telephonyManager.getServiceStateForSlot(any())).thenAnswer { invocation ->
+                when (invocation.getArgument(0) as Int) {
+                    0 -> ServiceState().apply { isEmergencyOnly = false }
+                    1 -> ServiceState().apply { isEmergencyOnly = false }
+                    2 -> ServiceState().apply { isEmergencyOnly = true }
+                    3 -> ServiceState().apply { isEmergencyOnly = false }
+                    else -> null
+                }
+            }
+
+            // GIVEN a broadcast goes out for the appropriate subID
+            val intent = serviceStateIntent(subId = -1)
             fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
                 context,
                 intent,
             )
             runCurrent()
 
-            val deviceBasedState = ServiceStateModel(isEmergencyOnly = false)
-            assertThat(underTest.deviceServiceState.value).isEqualTo(deviceBasedState)
+            // THEN the device is in ECM, because one of the service states is
+            assertThat(latest).isTrue()
+        }
 
-            // ... and ignores any other subId
-            val intent2 = serviceStateIntent(subId = 1, emergencyOnly = true)
+    @Test
+    fun testDeviceEmergencyCallState_aggregatesAcrossSlots_allFalse() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.isDeviceEmergencyCallCapable)
+
+            // GIVEN there are multiple slots
+            whenever(telephonyManager.activeModemCount).thenReturn(4)
+            // GIVEN only one of them reports ECM
+            whenever(telephonyManager.getServiceStateForSlot(any())).thenAnswer { invocation ->
+                when (invocation.getArgument(0) as Int) {
+                    0 -> ServiceState().apply { isEmergencyOnly = false }
+                    1 -> ServiceState().apply { isEmergencyOnly = false }
+                    2 -> ServiceState().apply { isEmergencyOnly = false }
+                    3 -> ServiceState().apply { isEmergencyOnly = false }
+                    else -> null
+                }
+            }
+
+            // GIVEN a broadcast goes out for the appropriate subID
+            val intent = serviceStateIntent(subId = -1)
             fakeBroadcastDispatcher.sendIntentToMatchingReceiversOnly(
                 context,
-                intent2,
+                intent,
             )
             runCurrent()
 
-            assertThat(underTest.deviceServiceState.value).isEqualTo(deviceBasedState)
+            // THEN the device is in ECM, because one of the service states is
+            assertThat(latest).isFalse()
         }
 
     @Test
@@ -1549,15 +1585,8 @@
          */
         private fun serviceStateIntent(
             subId: Int,
-            emergencyOnly: Boolean = false,
         ): Intent {
-            val serviceState = ServiceState().apply { isEmergencyOnly = emergencyOnly }
-
-            val bundle = Bundle()
-            serviceState.fillInNotifierBundle(bundle)
-
             return Intent(Intent.ACTION_SERVICE_STATE).apply {
-                putExtras(bundle)
                 putExtra(SubscriptionManager.EXTRA_SUBSCRIPTION_INDEX, subId)
             }
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
index e439aff..4fd830d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconInteractorTest.kt
@@ -18,7 +18,6 @@
 
 import android.platform.test.annotations.EnableFlags
 import android.telephony.CellSignalStrength
-import android.telephony.SubscriptionManager.PROFILE_CLASS_UNSET
 import android.telephony.TelephonyManager.NETWORK_TYPE_UNKNOWN
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -27,12 +26,12 @@
 import com.android.settingslib.mobile.TelephonyIcons
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
 import com.android.systemui.statusbar.pipeline.mobile.data.model.NetworkNameModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType.CarrierMergedNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType.DefaultNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType.OverrideNetworkType
-import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionRepository
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.FakeMobileIconsInteractor.Companion.FIVE_G_OVERRIDE
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.FakeMobileIconsInteractor.Companion.FOUR_G
@@ -40,12 +39,12 @@
 import com.android.systemui.statusbar.pipeline.mobile.domain.model.NetworkTypeIconModel
 import com.android.systemui.statusbar.pipeline.mobile.domain.model.SignalIconModel
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.test.TestScope
@@ -61,21 +60,18 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class MobileIconInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: MobileIconInteractor
     private val mobileMappingsProxy = FakeMobileMappingsProxy()
     private val mobileIconsInteractor = FakeMobileIconsInteractor(mobileMappingsProxy, mock())
 
-    private val subscriptionModel =
-        MutableStateFlow(
-            SubscriptionModel(
-                subscriptionId = SUB_1_ID,
-                carrierName = DEFAULT_NAME,
-                profileClass = PROFILE_CLASS_UNSET,
-            )
+    private val connectionRepository =
+        FakeMobileConnectionRepository(
+            SUB_1_ID,
+            logcatTableLogBuffer(kosmos, "MobileIconInteractorTest"),
         )
 
-    private val connectionRepository = FakeMobileConnectionRepository(SUB_1_ID, mock())
-
     private val testDispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(testDispatcher)
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
index cc0eae7..f6d439a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/domain/interactor/MobileIconsInteractorTest.kt
@@ -28,8 +28,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
-import com.android.systemui.log.table.TableLogBuffer
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
@@ -37,10 +36,10 @@
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.statusbar.policy.data.repository.FakeUserSetupRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.CarrierConfigTracker
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
-import com.android.systemui.util.time.FakeSystemClock
 import com.google.common.truth.Truth.assertThat
 import java.util.UUID
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -59,6 +58,8 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class MobileIconsInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: MobileIconsInteractor
     private lateinit var connectivityRepository: FakeConnectivityRepository
     private lateinit var connectionsRepository: FakeMobileConnectionsRepository
@@ -72,15 +73,7 @@
     private val testDispatcher = StandardTestDispatcher()
     private val testScope = TestScope(testDispatcher)
 
-    private val tableLogBuffer =
-        TableLogBuffer(
-            8,
-            "MobileIconsInteractorTest",
-            FakeSystemClock(),
-            mock(),
-            testDispatcher,
-            testScope.backgroundScope,
-        )
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "MobileIconsInteractorTest")
 
     @Mock private lateinit var carrierConfigTracker: CarrierConfigTracker
 
@@ -897,13 +890,11 @@
         testScope.runTest {
             val latest by collectLastValue(underTest.isDeviceInEmergencyCallsOnlyMode)
 
-            connectionsRepository.deviceServiceState.value =
-                ServiceStateModel(isEmergencyOnly = true)
+            connectionsRepository.isDeviceEmergencyCallCapable.value = true
 
             assertThat(latest).isTrue()
 
-            connectionsRepository.deviceServiceState.value =
-                ServiceStateModel(isEmergencyOnly = false)
+            connectionsRepository.isDeviceEmergencyCallCapable.value = false
 
             assertThat(latest).isFalse()
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/view/ModernStatusBarMobileViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/view/ModernStatusBarMobileViewTest.kt
index 42cb660..84846a1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/view/ModernStatusBarMobileViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/view/ModernStatusBarMobileViewTest.kt
@@ -28,12 +28,12 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.StatusBarIconView
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.FakeMobileIconInteractor
 import com.android.systemui.statusbar.pipeline.mobile.ui.MobileViewLogger
 import com.android.systemui.statusbar.pipeline.mobile.ui.viewmodel.LocationBasedMobileViewModel
@@ -41,6 +41,7 @@
 import com.android.systemui.statusbar.pipeline.mobile.ui.viewmodel.QsMobileIconViewModel
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -58,13 +59,13 @@
 @RunWithLooper(setAsMainLooper = true)
 @OptIn(ExperimentalCoroutinesApi::class)
 class ModernStatusBarMobileViewTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var testableLooper: TestableLooper
     private val testDispatcher = UnconfinedTestDispatcher()
     private val testScope = TestScope(testDispatcher)
     private val flags = FakeFeatureFlagsClassic().also { it.set(Flags.NEW_NETWORK_SLICE_UI, false) }
 
-    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var viewLogger: MobileViewLogger
     @Mock private lateinit var constants: ConnectivityConstants
     private lateinit var interactor: FakeMobileIconInteractor
@@ -88,10 +89,11 @@
             AirplaneModeInteractor(
                 airplaneModeRepository,
                 FakeConnectivityRepository(),
-                FakeMobileConnectionsRepository(),
+                kosmos.fakeMobileConnectionsRepository,
             )
 
-        interactor = FakeMobileIconInteractor(tableLogBuffer)
+        interactor =
+            FakeMobileIconInteractor(logcatTableLogBuffer(kosmos, "ModernStatusBarMobileViewTest"))
         createViewModel()
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
index deb9fcf..f99fcac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/LocationBasedMobileIconViewModelTest.kt
@@ -21,24 +21,23 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.statusbar.connectivity.MobileIconCarrierIdOverridesFake
-import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
 import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconInteractor
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconInteractorImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractor
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractorImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.model.SignalIconModel
-import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.statusbar.policy.data.repository.FakeUserSetupRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.CarrierConfigTracker
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
@@ -59,13 +58,15 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class LocationBasedMobileIconViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var commonImpl: MobileIconViewModelCommon
     private lateinit var homeIcon: HomeMobileIconViewModel
     private lateinit var qsIcon: QsMobileIconViewModel
     private lateinit var keyguardIcon: KeyguardMobileIconViewModel
     private lateinit var iconsInteractor: MobileIconsInteractor
     private lateinit var interactor: MobileIconInteractor
-    private lateinit var connectionsRepository: FakeMobileConnectionsRepository
+    private val connectionsRepository = kosmos.fakeMobileConnectionsRepository
     private lateinit var repository: FakeMobileConnectionRepository
     private lateinit var airplaneModeInteractor: AirplaneModeInteractor
 
@@ -76,9 +77,9 @@
             it.set(Flags.FILTER_PROVISIONING_NETWORK_SUBSCRIPTIONS, true)
         }
 
-    @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var constants: ConnectivityConstants
-    @Mock private lateinit var tableLogBuffer: TableLogBuffer
+    private val tableLogBuffer =
+        logcatTableLogBuffer(kosmos, "LocationBasedMobileIconViewModelTest")
     @Mock private lateinit var carrierConfigTracker: CarrierConfigTracker
 
     private val testDispatcher = UnconfinedTestDispatcher()
@@ -91,10 +92,8 @@
             AirplaneModeInteractor(
                 FakeAirplaneModeRepository(),
                 FakeConnectivityRepository(),
-                FakeMobileConnectionsRepository(),
+                connectionsRepository,
             )
-        connectionsRepository =
-            FakeMobileConnectionsRepository(FakeMobileMappingsProxy(), tableLogBuffer)
         repository =
             FakeMobileConnectionRepository(SUB_1_ID, tableLogBuffer).apply {
                 isInService.value = true
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
index e510924..4c7cdfa 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconViewModelTest.kt
@@ -34,7 +34,7 @@
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
 import com.android.systemui.flags.Flags.NEW_NETWORK_SLICE_UI
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.connectivity.MobileIconCarrierIdOverridesFake
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
@@ -42,6 +42,7 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconInteractorImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractorImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.model.SignalIconModel
@@ -51,6 +52,7 @@
 import com.android.systemui.statusbar.pipeline.shared.data.model.DataActivityModel
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.statusbar.policy.data.repository.FakeUserSetupRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.CarrierConfigTracker
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
@@ -74,6 +76,8 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class MobileIconViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private var connectivityRepository = FakeConnectivityRepository()
 
     private lateinit var underTest: MobileIconViewModel
@@ -84,7 +88,7 @@
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
     private lateinit var airplaneModeInteractor: AirplaneModeInteractor
     @Mock private lateinit var constants: ConnectivityConstants
-    @Mock private lateinit var tableLogBuffer: TableLogBuffer
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "MobileIconViewModelTest")
     @Mock private lateinit var carrierConfigTracker: CarrierConfigTracker
 
     private val flags =
@@ -118,7 +122,7 @@
             AirplaneModeInteractor(
                 airplaneModeRepository,
                 connectivityRepository,
-                FakeMobileConnectionsRepository(),
+                kosmos.fakeMobileConnectionsRepository,
             )
 
         iconsInteractor =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
index 47899a6..31ba837 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/ui/viewmodel/MobileIconsViewModelTest.kt
@@ -24,11 +24,10 @@
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
 import com.android.systemui.statusbar.phone.StatusBarLocation
-import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.FakeMobileIconsInteractor
 import com.android.systemui.statusbar.pipeline.mobile.domain.model.NetworkTypeIconModel
 import com.android.systemui.statusbar.pipeline.mobile.ui.MobileViewLogger
@@ -36,6 +35,7 @@
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
 import junit.framework.Assert.assertFalse
@@ -58,12 +58,13 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class MobileIconsViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: MobileIconsViewModel
     private val interactor = FakeMobileIconsInteractor(FakeMobileMappingsProxy(), mock())
     private val flags = FakeFeatureFlagsClassic().also { it.set(Flags.NEW_NETWORK_SLICE_UI, false) }
 
     private lateinit var airplaneModeInteractor: AirplaneModeInteractor
-    @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var constants: ConnectivityConstants
     @Mock private lateinit var logger: MobileViewLogger
     @Mock private lateinit var verboseLogger: VerboseMobileViewLogger
@@ -79,7 +80,7 @@
             AirplaneModeInteractor(
                 FakeAirplaneModeRepository(),
                 FakeConnectivityRepository(),
-                FakeMobileConnectionsRepository(),
+                kosmos.fakeMobileConnectionsRepository,
             )
 
         underTest =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepositoryImplTest.kt
index f486787..0945742 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/data/repository/ConnectivityRepositoryImplTest.kt
@@ -21,60 +21,61 @@
 import android.net.NetworkCapabilities
 import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.net.NetworkCapabilities.TRANSPORT_ETHERNET
+import android.net.NetworkCapabilities.TRANSPORT_VPN
 import android.net.NetworkCapabilities.TRANSPORT_WIFI
+import android.net.VpnTransportInfo
 import android.net.vcn.VcnTransportInfo
 import android.net.wifi.WifiInfo
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_STATUS_BAR_ALWAYS_CHECK_UNDERLYING_NETWORKS
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.dump.DumpManager
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.dump.dumpManager
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.log.core.FakeLogBuffer
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityInputLogger
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlots
-import com.android.systemui.statusbar.pipeline.shared.data.model.DefaultConnectionModel
 import com.android.systemui.statusbar.pipeline.shared.data.repository.ConnectivityRepositoryImpl.Companion.DEFAULT_HIDDEN_ICONS_RESOURCE
 import com.android.systemui.statusbar.pipeline.shared.data.repository.ConnectivityRepositoryImpl.Companion.HIDDEN_ICONS_TUNABLE_KEY
 import com.android.systemui.statusbar.pipeline.shared.data.repository.ConnectivityRepositoryImpl.Companion.getMainOrUnderlyingWifiInfo
+import com.android.systemui.testKosmos
 import com.android.systemui.tuner.TunerService
-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.mock
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
-import kotlinx.coroutines.yield
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class ConnectivityRepositoryImplTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var underTest: ConnectivityRepositoryImpl
 
-    @Mock private lateinit var connectivityManager: ConnectivityManager
-    @Mock private lateinit var connectivitySlots: ConnectivitySlots
-    @Mock private lateinit var dumpManager: DumpManager
-    @Mock private lateinit var logger: ConnectivityInputLogger
-    private lateinit var testScope: TestScope
-    @Mock private lateinit var tunerService: TunerService
+    private val connectivityManager = mock<ConnectivityManager>()
+    private val connectivitySlots = mock<ConnectivitySlots>()
+    private val dumpManager = kosmos.dumpManager
+    private val logger = ConnectivityInputLogger(FakeLogBuffer.Factory.create())
+    private val testScope = kosmos.testScope
+    private val tunerService = mock<TunerService>()
 
     @Before
     fun setUp() {
-        MockitoAnnotations.initMocks(this)
-        testScope = TestScope(UnconfinedTestDispatcher())
         createAndSetRepo()
     }
 
@@ -89,12 +90,10 @@
             // config_statusBarIconsToExclude when it's first constructed
             createAndSetRepo()
 
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             assertThat(latest).containsExactly(ConnectivitySlot.ETHERNET, ConnectivitySlot.WIFI)
-
-            job.cancel()
         }
 
     @Test
@@ -102,14 +101,12 @@
         testScope.runTest {
             setUpEthernetWifiMobileSlotNames()
 
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, SLOT_MOBILE)
 
             assertThat(latest).containsExactly(ConnectivitySlot.MOBILE)
-
-            job.cancel()
         }
 
     @Test
@@ -117,19 +114,16 @@
         testScope.runTest {
             setUpEthernetWifiMobileSlotNames()
 
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, SLOT_MOBILE)
 
             // WHEN onTuningChanged with the wrong key
             getTunable().onTuningChanged("wrongKey", SLOT_WIFI)
-            yield()
 
             // THEN we didn't update our value and still have the old one
             assertThat(latest).containsExactly(ConnectivitySlot.MOBILE)
-
-            job.cancel()
         }
 
     @Test
@@ -143,8 +137,8 @@
             // config_statusBarIconsToExclude when it's first constructed
             createAndSetRepo()
 
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             // First, update the slots
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, SLOT_MOBILE)
@@ -152,19 +146,16 @@
 
             // WHEN we update to a null value
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, null)
-            yield()
 
             // THEN we go back to our default value
             assertThat(latest).containsExactly(ConnectivitySlot.ETHERNET, ConnectivitySlot.WIFI)
-
-            job.cancel()
         }
 
     @Test
     fun forceHiddenSlots_someInvalidSlotNames_flowHasValidSlotsOnly() =
         testScope.runTest {
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             whenever(connectivitySlots.getSlotFromName(SLOT_WIFI)).thenReturn(ConnectivitySlot.WIFI)
             whenever(connectivitySlots.getSlotFromName(SLOT_MOBILE)).thenReturn(null)
@@ -172,8 +163,6 @@
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, "$SLOT_WIFI,$SLOT_MOBILE")
 
             assertThat(latest).containsExactly(ConnectivitySlot.WIFI)
-
-            job.cancel()
         }
 
     @Test
@@ -181,23 +170,21 @@
         testScope.runTest {
             setUpEthernetWifiMobileSlotNames()
 
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             // WHEN there's empty and blank slot names
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, "$SLOT_MOBILE,  ,,$SLOT_WIFI")
 
             // THEN we skip that slot but still process the other ones
             assertThat(latest).containsExactly(ConnectivitySlot.WIFI, ConnectivitySlot.MOBILE)
-
-            job.cancel()
         }
 
     @Test
     fun forceHiddenSlots_allInvalidOrEmptySlotNames_flowHasEmpty() =
         testScope.runTest {
-            var latest: Set<ConnectivitySlot>? = null
-            val job = underTest.forceHiddenSlots.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             whenever(connectivitySlots.getSlotFromName(SLOT_WIFI)).thenReturn(null)
             whenever(connectivitySlots.getSlotFromName(SLOT_ETHERNET)).thenReturn(null)
@@ -210,8 +197,6 @@
                 )
 
             assertThat(latest).isEmpty()
-
-            job.cancel()
         }
 
     @Test
@@ -219,29 +204,25 @@
         testScope.runTest {
             setUpEthernetWifiMobileSlotNames()
 
-            var latest1: Set<ConnectivitySlot>? = null
-            val job1 = underTest.forceHiddenSlots.onEach { latest1 = it }.launchIn(this)
+            val latest1 by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             getTunable().onTuningChanged(HIDDEN_ICONS_TUNABLE_KEY, "$SLOT_WIFI,$SLOT_ETHERNET")
 
             assertThat(latest1).containsExactly(ConnectivitySlot.WIFI, ConnectivitySlot.ETHERNET)
 
             // WHEN we add a second subscriber after having already emitted a value
-            var latest2: Set<ConnectivitySlot>? = null
-            val job2 = underTest.forceHiddenSlots.onEach { latest2 = it }.launchIn(this)
+            val latest2 by collectLastValue(underTest.forceHiddenSlots)
+            runCurrent()
 
             // THEN the second subscribe receives the already-emitted value
             assertThat(latest2).containsExactly(ConnectivitySlot.WIFI, ConnectivitySlot.ETHERNET)
-
-            job1.cancel()
-            job2.cancel()
         }
 
     @Test
     fun defaultConnections_noTransports_nothingIsDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -256,15 +237,12 @@
             assertThat(latest!!.wifi.isDefault).isFalse()
             assertThat(latest!!.ethernet.isDefault).isFalse()
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_cellularTransport_mobileIsDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -279,15 +257,12 @@
             assertThat(latest!!.wifi.isDefault).isFalse()
             assertThat(latest!!.ethernet.isDefault).isFalse()
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_wifiTransport_wifiIsDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -302,15 +277,12 @@
             assertThat(latest!!.ethernet.isDefault).isFalse()
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
             assertThat(latest!!.mobile.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_ethernetTransport_ethernetIsDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -325,15 +297,12 @@
             assertThat(latest!!.wifi.isDefault).isFalse()
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
             assertThat(latest!!.mobile.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_carrierMergedViaWifi_wifiAndCarrierMergedDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(true) }
@@ -350,15 +319,12 @@
             assertThat(latest!!.wifi.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.mobile.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_carrierMergedViaMobile_mobileCarrierMergedWifiDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(true) }
@@ -375,15 +341,12 @@
             assertThat(latest!!.mobile.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.wifi.isDefault).isTrue()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_carrierMergedViaWifiWithVcnTransport_wifiAndCarrierMergedDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(true) }
@@ -400,15 +363,13 @@
             assertThat(latest!!.wifi.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.mobile.isDefault).isFalse()
-
-            job.cancel()
         }
 
+    /** VCN over W+ (aka VCN over carrier merged). See b/352162710#comment27 scenario #1. */
     @Test
     fun defaultConnections_carrierMergedViaMobileWithVcnTransport_mobileCarrierMergedWifiDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(true) }
@@ -425,15 +386,48 @@
             assertThat(latest!!.mobile.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.wifi.isDefault).isTrue()
+        }
 
-            job.cancel()
+    /** VPN over W+ (aka VPN over carrier merged). See b/352162710#comment27 scenario #2. */
+    @Test
+    @EnableFlags(FLAG_STATUS_BAR_ALWAYS_CHECK_UNDERLYING_NETWORKS)
+    fun defaultConnections_vpnOverCarrierMerged_carrierMergedDefault() =
+        testScope.runTest {
+            val latest by collectLastValue(underTest.defaultConnections)
+
+            // Underlying carrier merged network
+            val underlyingCarrierMergedNetwork = mock<Network>()
+            val carrierMergedInfo =
+                mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(true) }
+            val underlyingCapabilities =
+                mock<NetworkCapabilities>().also {
+                    whenever(it.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
+                    whenever(it.transportInfo).thenReturn(carrierMergedInfo)
+                }
+            whenever(connectivityManager.getNetworkCapabilities(underlyingCarrierMergedNetwork))
+                .thenReturn(underlyingCapabilities)
+
+            val mainCapabilities =
+                mock<NetworkCapabilities>().also {
+                    whenever(it.hasTransport(TRANSPORT_ETHERNET)).thenReturn(false)
+                    // Transports are WIFI|VPN, *not* CELLULAR.
+                    whenever(it.hasTransport(TRANSPORT_CELLULAR)).thenReturn(false)
+                    whenever(it.hasTransport(TRANSPORT_WIFI)).thenReturn(true)
+                    whenever(it.hasTransport(TRANSPORT_VPN)).thenReturn(true)
+                    whenever(it.transportInfo).thenReturn(VpnTransportInfo(0, null, false, false))
+                    whenever(it.underlyingNetworks)
+                        .thenReturn(listOf(underlyingCarrierMergedNetwork))
+                }
+
+            getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, mainCapabilities)
+
+            assertThat(latest!!.carrierMerged.isDefault).isTrue()
         }
 
     @Test
     fun defaultConnections_notCarrierMergedViaWifi_carrierMergedNotDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(false) }
@@ -448,15 +442,12 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_notCarrierMergedViaMobile_carrierMergedNotDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val carrierMergedInfo =
                 mock<WifiInfo>().apply { whenever(this.isCarrierMerged).thenReturn(false) }
@@ -471,15 +462,12 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest!!.carrierMerged.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_transportInfoNotWifi_wifiNotDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -492,8 +480,6 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest!!.wifi.isDefault).isFalse()
-
-            job.cancel()
         }
 
     @Test
@@ -531,8 +517,7 @@
     @Test
     fun defaultConnections_cellular_underlyingCarrierMergedViaWifi_allDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             // Underlying carrier merged network
             val underlyingCarrierMergedNetwork = mock<Network>()
@@ -560,16 +545,17 @@
             assertThat(latest!!.mobile.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.wifi.isDefault).isTrue()
-
-            job.cancel()
         }
 
-    /** Test for b/225902574. */
+    /**
+     * Test for b/225902574: VPN over VCN over W+ (aka VPN over VCN over carrier merged).
+     *
+     * Also see b/352162710#comment27 scenario #3 and b/352162710#comment30.
+     */
     @Test
     fun defaultConnections_cellular_underlyingCarrierMergedViaMobileWithVcnTransport_allDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             // Underlying carrier merged network
             val underlyingCarrierMergedNetwork = mock<Network>()
@@ -587,6 +573,7 @@
             val mainCapabilities =
                 mock<NetworkCapabilities>().also {
                     whenever(it.hasTransport(TRANSPORT_CELLULAR)).thenReturn(true)
+                    whenever(it.hasTransport(TRANSPORT_VPN)).thenReturn(true)
                     whenever(it.transportInfo).thenReturn(null)
                     whenever(it.underlyingNetworks)
                         .thenReturn(listOf(underlyingCarrierMergedNetwork))
@@ -597,15 +584,12 @@
             assertThat(latest!!.mobile.isDefault).isTrue()
             assertThat(latest!!.carrierMerged.isDefault).isTrue()
             assertThat(latest!!.wifi.isDefault).isTrue()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_multipleTransports_multipleDefault() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -619,15 +603,12 @@
             assertThat(latest!!.mobile.isDefault).isTrue()
             assertThat(latest!!.ethernet.isDefault).isTrue()
             assertThat(latest!!.wifi.isDefault).isTrue()
-
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_hasValidated_isValidatedTrue() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -638,14 +619,12 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest!!.isValidated).isTrue()
-            job.cancel()
         }
 
     @Test
     fun defaultConnections_noValidated_isValidatedFalse() =
         testScope.runTest {
-            var latest: DefaultConnectionModel? = null
-            val job = underTest.defaultConnections.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.defaultConnections)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -656,7 +635,6 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest!!.isValidated).isFalse()
-            job.cancel()
         }
 
     @Test
@@ -669,8 +647,7 @@
         testScope.runTest {
             val vcnInfo = VcnTransportInfo(SUB_1_ID)
 
-            var latest: Int? = null
-            val job = underTest.vcnSubId.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.vcnSubId)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -681,7 +658,6 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest).isEqualTo(SUB_1_ID)
-            job.cancel()
         }
 
     @Test
@@ -689,8 +665,7 @@
         testScope.runTest {
             val vcnInfo = VcnTransportInfo(INVALID_SUBSCRIPTION_ID)
 
-            var latest: Int? = null
-            val job = underTest.vcnSubId.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.vcnSubId)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -701,14 +676,12 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest).isNull()
-            job.cancel()
         }
 
     @Test
     fun vcnSubId_nullIfNoTransportInfo() =
         testScope.runTest {
-            var latest: Int? = null
-            val job = underTest.vcnSubId.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.vcnSubId)
 
             val capabilities =
                 mock<NetworkCapabilities>().also {
@@ -719,7 +692,6 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest).isNull()
-            job.cancel()
         }
 
     @Test
@@ -728,8 +700,7 @@
             // If the underlying network of the VCN is a WiFi network, then there is no subId that
             // could disagree with telephony's active data subscription id.
 
-            var latest: Int? = null
-            val job = underTest.vcnSubId.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.vcnSubId)
 
             val wifiInfo = mock<WifiInfo>()
             val vcnInfo = VcnTransportInfo(wifiInfo)
@@ -742,14 +713,12 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest).isNull()
-            job.cancel()
         }
 
     @Test
     fun vcnSubId_changingVcnInfoIsTracked() =
         testScope.runTest {
-            var latest: Int? = null
-            val job = underTest.vcnSubId.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.vcnSubId)
 
             val wifiInfo = mock<WifiInfo>()
             val wifiVcnInfo = VcnTransportInfo(wifiInfo)
@@ -788,8 +757,6 @@
             getDefaultNetworkCallback().onCapabilitiesChanged(NETWORK, capabilities)
 
             assertThat(latest).isNull()
-
-            job.cancel()
         }
 
     @Test
@@ -862,6 +829,7 @@
     }
 
     @Test
+    @DisableFlags(FLAG_STATUS_BAR_ALWAYS_CHECK_UNDERLYING_NETWORKS)
     fun getMainOrUnderlyingWifiInfo_notCellular_underlyingWifi_noInfo() {
         val underlyingNetwork = mock<Network>()
         val underlyingWifiInfo = mock<WifiInfo>()
@@ -916,6 +884,7 @@
     }
 
     @Test
+    @DisableFlags(FLAG_STATUS_BAR_ALWAYS_CHECK_UNDERLYING_NETWORKS)
     fun getMainOrUnderlyingWifiInfo_notCellular_underlyingVcnWithWifi_noInfo() {
         val underlyingNetwork = mock<Network>()
         val underlyingVcnInfo = VcnTransportInfo(mock<WifiInfo>())
@@ -1076,12 +1045,13 @@
                 testScope.backgroundScope,
                 tunerService,
             )
+        testScope.runCurrent()
     }
 
     private fun getTunable(): TunerService.Tunable {
         val callbackCaptor = argumentCaptor<TunerService.Tunable>()
         verify(tunerService).addTunable(callbackCaptor.capture(), any())
-        return callbackCaptor.value!!
+        return callbackCaptor.firstValue
     }
 
     private fun setUpEthernetWifiMobileSlotNames() {
@@ -1094,7 +1064,7 @@
     private fun getDefaultNetworkCallback(): ConnectivityManager.NetworkCallback {
         val callbackCaptor = argumentCaptor<ConnectivityManager.NetworkCallback>()
         verify(connectivityManager).registerDefaultNetworkCallback(callbackCaptor.capture())
-        return callbackCaptor.value!!
+        return callbackCaptor.firstValue
     }
 
     private companion object {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
index 2238bff..50f262c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/shared/ui/viewmodel/InternetTileViewModelTest.kt
@@ -26,7 +26,7 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.Flags
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.qs.tileimpl.QSTileImpl.ResourceIcon
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.connectivity.WifiIcons
@@ -49,6 +49,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiScanEntry
 import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.policy.data.repository.FakeUserSetupRepository
+import com.android.systemui.testKosmos
 import com.android.systemui.util.CarrierConfigTracker
 import com.android.systemui.util.mockito.mock
 import com.google.common.truth.Truth.assertThat
@@ -61,6 +62,8 @@
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class InternetTileViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
     private lateinit var underTest: InternetTileViewModel
     private lateinit var mobileIconsInteractor: MobileIconsInteractor
 
@@ -73,7 +76,7 @@
     private val wifiInteractor =
         WifiInteractorImpl(connectivityRepository, wifiRepository, testScope.backgroundScope)
 
-    private val tableLogBuffer: TableLogBuffer = mock()
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "InternetTileViewModelTest")
     private val carrierConfigTracker: CarrierConfigTracker = mock()
 
     private val mobileConnectionsRepository =
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
index 80b10c0..161c4f5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
@@ -26,7 +26,7 @@
 import android.widget.ImageView
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logcatTableLogBuffer
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.StatusBarIconView.STATE_DOT
 import com.android.systemui.statusbar.StatusBarIconView.STATE_HIDDEN
@@ -36,7 +36,7 @@
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModel
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModelImpl
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 import com.android.systemui.statusbar.pipeline.wifi.data.repository.FakeWifiRepository
@@ -47,6 +47,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel.Companion.viewModelForLocation
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.WifiViewModel
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -61,10 +62,11 @@
 @RunWith(AndroidTestingRunner::class)
 @RunWithLooper(setAsMainLooper = true)
 class ModernStatusBarWifiViewTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var testableLooper: TestableLooper
 
-    @Mock private lateinit var tableLogBuffer: TableLogBuffer
+    private val tableLogBuffer = logcatTableLogBuffer(kosmos, "ModernStatusBarWifiViewTest")
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -91,7 +93,7 @@
                 AirplaneModeInteractor(
                     airplaneModeRepository,
                     connectivityRepository,
-                    FakeMobileConnectionsRepository(),
+                    kosmos.fakeMobileConnectionsRepository,
                 ),
                 tableLogBuffer,
                 scope,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
index 489319e..d2a4bf3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
@@ -31,7 +31,7 @@
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModel
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModelImpl
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityConstants
 import com.android.systemui.statusbar.pipeline.shared.data.model.ConnectivitySlot
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
@@ -42,6 +42,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiNetworkModel
 import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon.Companion.NO_INTERNET
+import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -64,6 +65,7 @@
 @RunWith(Parameterized::class)
 internal class WifiViewModelIconParameterizedTest(private val testCase: TestCase) :
     SysuiTestCase() {
+    private val kosmos = testKosmos()
 
     private lateinit var underTest: WifiViewModel
 
@@ -91,7 +93,7 @@
                 AirplaneModeInteractor(
                     airplaneModeRepository,
                     connectivityRepository,
-                    FakeMobileConnectionsRepository(),
+                    kosmos.fakeMobileConnectionsRepository,
                 ),
                 tableLogBuffer,
                 scope,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
index 12cfdcf..e396b56 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/ui/viewmodel/KeyguardStatusBarViewModelTest.kt
@@ -16,11 +16,11 @@
 
 package com.android.systemui.statusbar.ui.viewmodel
 
-import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.FlagsParameterization
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.flags.EnableSceneContainer
 import com.android.systemui.flags.andSceneContainer
 import com.android.systemui.keyguard.data.repository.fakeDeviceEntryFaceAuthRepository
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
@@ -33,7 +33,6 @@
 import com.android.systemui.scene.shared.model.Scenes
 import com.android.systemui.statusbar.domain.interactor.keyguardStatusBarInteractor
 import com.android.systemui.statusbar.notification.data.repository.FakeHeadsUpRowRepository
-import com.android.systemui.statusbar.notification.shared.NotificationsHeadsUpRefactor
 import com.android.systemui.statusbar.notification.stack.data.repository.headsUpNotificationRepository
 import com.android.systemui.statusbar.notification.stack.domain.interactor.headsUpNotificationInteractor
 import com.android.systemui.statusbar.policy.BatteryController
@@ -127,7 +126,7 @@
         }
 
     @Test
-    @EnableFlags(NotificationsHeadsUpRefactor.FLAG_NAME)
+    @EnableSceneContainer
     fun isVisible_headsUpStatusBarShown_false() =
         testScope.runTest {
             val latest by collectLastValue(underTest.isVisible)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
index 37a73cf..c235954 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/user/data/repository/UserRepositoryImplTest.kt
@@ -24,21 +24,21 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.kosmos.unconfinedTestDispatcher
+import com.android.systemui.kosmos.unconfinedTestScope
 import com.android.systemui.settings.FakeUserTracker
+import com.android.systemui.testKosmos
 import com.android.systemui.user.data.model.SelectedUserModel
 import com.android.systemui.user.data.model.SelectionStatus
 import com.android.systemui.user.data.model.UserSwitcherSettingsModel
-import com.android.systemui.util.settings.FakeGlobalSettings
+import com.android.systemui.util.settings.unconfinedDispatcherFakeGlobalSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.test.TestCoroutineScope
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -52,143 +52,153 @@
 @RunWith(AndroidJUnit4::class)
 class UserRepositoryImplTest : SysuiTestCase() {
 
+    private val kosmos = testKosmos()
+    private val testDispatcher = kosmos.unconfinedTestDispatcher
+    private val testScope = kosmos.unconfinedTestScope
+    private val globalSettings = kosmos.unconfinedDispatcherFakeGlobalSettings
+
     @Mock private lateinit var manager: UserManager
 
     private lateinit var underTest: UserRepositoryImpl
 
-    private lateinit var globalSettings: FakeGlobalSettings
     private lateinit var tracker: FakeUserTracker
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
-
-        globalSettings = FakeGlobalSettings()
         tracker = FakeUserTracker()
     }
 
     @Test
-    fun userSwitcherSettings() = runSelfCancelingTest {
-        setUpGlobalSettings(
-            isSimpleUserSwitcher = true,
-            isAddUsersFromLockscreen = true,
-            isUserSwitcherEnabled = true,
-        )
-        underTest = create(this)
+    fun userSwitcherSettings() =
+        testScope.runTest {
+            setUpGlobalSettings(
+                isSimpleUserSwitcher = true,
+                isAddUsersFromLockscreen = true,
+                isUserSwitcherEnabled = true,
+            )
+            underTest = create(testScope.backgroundScope)
+            var value: UserSwitcherSettingsModel? = null
+            val job = underTest.userSwitcherSettings.onEach { value = it }.launchIn(this)
 
-        var value: UserSwitcherSettingsModel? = null
-        underTest.userSwitcherSettings.onEach { value = it }.launchIn(this)
+            assertUserSwitcherSettings(
+                model = value,
+                expectedSimpleUserSwitcher = true,
+                expectedAddUsersFromLockscreen = true,
+                expectedUserSwitcherEnabled = true,
+            )
 
-        assertUserSwitcherSettings(
-            model = value,
-            expectedSimpleUserSwitcher = true,
-            expectedAddUsersFromLockscreen = true,
-            expectedUserSwitcherEnabled = true,
-        )
-
-        setUpGlobalSettings(
-            isSimpleUserSwitcher = false,
-            isAddUsersFromLockscreen = true,
-            isUserSwitcherEnabled = true,
-        )
-        assertUserSwitcherSettings(
-            model = value,
-            expectedSimpleUserSwitcher = false,
-            expectedAddUsersFromLockscreen = true,
-            expectedUserSwitcherEnabled = true,
-        )
-    }
+            setUpGlobalSettings(
+                isSimpleUserSwitcher = false,
+                isAddUsersFromLockscreen = true,
+                isUserSwitcherEnabled = true,
+            )
+            assertUserSwitcherSettings(
+                model = value,
+                expectedSimpleUserSwitcher = false,
+                expectedAddUsersFromLockscreen = true,
+                expectedUserSwitcherEnabled = true,
+            )
+            job.cancel()
+        }
 
     @Test
-    fun userSwitcherSettings_isUserSwitcherEnabled_notInitialized() = runSelfCancelingTest {
-        underTest = create(this)
+    fun userSwitcherSettings_isUserSwitcherEnabled_notInitialized() =
+        testScope.runTest {
+            underTest = create(testScope.backgroundScope)
 
-        var value: UserSwitcherSettingsModel? = null
-        underTest.userSwitcherSettings.onEach { value = it }.launchIn(this)
+            var value: UserSwitcherSettingsModel? = null
+            val job = underTest.userSwitcherSettings.onEach { value = it }.launchIn(this)
 
-        assertUserSwitcherSettings(
-            model = value,
-            expectedSimpleUserSwitcher = false,
-            expectedAddUsersFromLockscreen = false,
-            expectedUserSwitcherEnabled =
-                context.resources.getBoolean(
-                    com.android.internal.R.bool.config_showUserSwitcherByDefault
-                ),
-        )
-    }
+            assertUserSwitcherSettings(
+                model = value,
+                expectedSimpleUserSwitcher = false,
+                expectedAddUsersFromLockscreen = false,
+                expectedUserSwitcherEnabled =
+                    context.resources.getBoolean(
+                        com.android.internal.R.bool.config_showUserSwitcherByDefault
+                    ),
+            )
+            job.cancel()
+        }
 
     @Test
-    fun refreshUsers() = runSelfCancelingTest {
-        val mainUserId = 10
-        val mainUser = mock(UserHandle::class.java)
-        whenever(manager.mainUser).thenReturn(mainUser)
-        whenever(mainUser.identifier).thenReturn(mainUserId)
+    fun refreshUsers() =
+        testScope.runTest {
+            val mainUserId = 10
+            val mainUser = mock(UserHandle::class.java)
+            whenever(manager.mainUser).thenReturn(mainUser)
+            whenever(mainUser.identifier).thenReturn(mainUserId)
 
-        underTest = create(this)
-        val initialExpectedValue =
-            setUpUsers(
-                count = 3,
-                selectedIndex = 0,
-            )
-        var userInfos: List<UserInfo>? = null
-        var selectedUserInfo: UserInfo? = null
-        underTest.userInfos.onEach { userInfos = it }.launchIn(this)
-        underTest.selectedUserInfo.onEach { selectedUserInfo = it }.launchIn(this)
+            underTest = create(testScope.backgroundScope)
+            val initialExpectedValue =
+                setUpUsers(
+                    count = 3,
+                    selectedIndex = 0,
+                )
+            var userInfos: List<UserInfo>? = null
+            var selectedUserInfo: UserInfo? = null
+            val job1 = underTest.userInfos.onEach { userInfos = it }.launchIn(this)
+            val job2 = underTest.selectedUserInfo.onEach { selectedUserInfo = it }.launchIn(this)
 
-        underTest.refreshUsers()
-        assertThat(userInfos).isEqualTo(initialExpectedValue)
-        assertThat(selectedUserInfo).isEqualTo(initialExpectedValue[0])
-        assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedUserInfo?.id)
+            underTest.refreshUsers()
+            assertThat(userInfos).isEqualTo(initialExpectedValue)
+            assertThat(selectedUserInfo).isEqualTo(initialExpectedValue[0])
+            assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedUserInfo?.id)
 
-        val secondExpectedValue =
-            setUpUsers(
-                count = 4,
-                selectedIndex = 1,
-            )
-        underTest.refreshUsers()
-        assertThat(userInfos).isEqualTo(secondExpectedValue)
-        assertThat(selectedUserInfo).isEqualTo(secondExpectedValue[1])
-        assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedUserInfo?.id)
+            val secondExpectedValue =
+                setUpUsers(
+                    count = 4,
+                    selectedIndex = 1,
+                )
+            underTest.refreshUsers()
+            assertThat(userInfos).isEqualTo(secondExpectedValue)
+            assertThat(selectedUserInfo).isEqualTo(secondExpectedValue[1])
+            assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedUserInfo?.id)
 
-        val selectedNonGuestUserId = selectedUserInfo?.id
-        val thirdExpectedValue =
-            setUpUsers(
-                count = 2,
-                isLastGuestUser = true,
-                selectedIndex = 1,
-            )
-        underTest.refreshUsers()
-        assertThat(userInfos).isEqualTo(thirdExpectedValue)
-        assertThat(selectedUserInfo).isEqualTo(thirdExpectedValue[1])
-        assertThat(selectedUserInfo?.isGuest).isTrue()
-        assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedNonGuestUserId)
-        assertThat(underTest.mainUserId).isEqualTo(mainUserId)
-    }
+            val selectedNonGuestUserId = selectedUserInfo?.id
+            val thirdExpectedValue =
+                setUpUsers(
+                    count = 2,
+                    isLastGuestUser = true,
+                    selectedIndex = 1,
+                )
+            underTest.refreshUsers()
+            assertThat(userInfos).isEqualTo(thirdExpectedValue)
+            assertThat(selectedUserInfo).isEqualTo(thirdExpectedValue[1])
+            assertThat(selectedUserInfo?.isGuest).isTrue()
+            assertThat(underTest.lastSelectedNonGuestUserId).isEqualTo(selectedNonGuestUserId)
+            assertThat(underTest.mainUserId).isEqualTo(mainUserId)
+            job1.cancel()
+            job2.cancel()
+        }
 
     @Test
-    fun refreshUsers_sortsByCreationTime_guestUserLast() = runSelfCancelingTest {
-        underTest = create(this)
-        val unsortedUsers =
-            setUpUsers(
-                count = 3,
-                selectedIndex = 0,
-                isLastGuestUser = true,
-            )
-        unsortedUsers[0].creationTime = 999
-        unsortedUsers[1].creationTime = 900
-        unsortedUsers[2].creationTime = 950
-        val expectedUsers =
-            listOf(
-                unsortedUsers[1],
-                unsortedUsers[0],
-                unsortedUsers[2], // last because this is the guest
-            )
-        var userInfos: List<UserInfo>? = null
-        underTest.userInfos.onEach { userInfos = it }.launchIn(this)
+    fun refreshUsers_sortsByCreationTime_guestUserLast() =
+        testScope.runTest {
+            underTest = create(testScope.backgroundScope)
+            val unsortedUsers =
+                setUpUsers(
+                    count = 3,
+                    selectedIndex = 0,
+                    isLastGuestUser = true,
+                )
+            unsortedUsers[0].creationTime = 999
+            unsortedUsers[1].creationTime = 900
+            unsortedUsers[2].creationTime = 950
+            val expectedUsers =
+                listOf(
+                    unsortedUsers[1],
+                    unsortedUsers[0],
+                    unsortedUsers[2], // last because this is the guest
+                )
+            var userInfos: List<UserInfo>? = null
+            val job = underTest.userInfos.onEach { userInfos = it }.launchIn(this)
 
-        underTest.refreshUsers()
-        assertThat(userInfos).isEqualTo(expectedUsers)
-    }
+            underTest.refreshUsers()
+            assertThat(userInfos).isEqualTo(expectedUsers)
+            job.cancel()
+        }
 
     private fun setUpUsers(
         count: Int,
@@ -206,58 +216,68 @@
         tracker.set(userInfos, selectedIndex)
         return userInfos
     }
-    @Test
-    fun userTrackerCallback_updatesSelectedUserInfo() = runSelfCancelingTest {
-        underTest = create(this)
-        var selectedUserInfo: UserInfo? = null
-        underTest.selectedUserInfo.onEach { selectedUserInfo = it }.launchIn(this)
-        setUpUsers(
-            count = 2,
-            selectedIndex = 0,
-        )
-        tracker.onProfileChanged()
-        assertThat(selectedUserInfo?.id).isEqualTo(0)
-        setUpUsers(
-            count = 2,
-            selectedIndex = 1,
-        )
-        tracker.onProfileChanged()
-        assertThat(selectedUserInfo?.id).isEqualTo(1)
-    }
 
     @Test
-    fun userTrackerCallback_updatesSelectionStatus() = runSelfCancelingTest {
-        underTest = create(this)
-        var selectedUser: SelectedUserModel? = null
-        underTest.selectedUser.onEach { selectedUser = it }.launchIn(this)
-        setUpUsers(count = 2, selectedIndex = 1)
+    fun userTrackerCallback_updatesSelectedUserInfo() =
+        testScope.runTest {
+            underTest = create(testScope.backgroundScope)
+            var selectedUserInfo: UserInfo? = null
+            val job = underTest.selectedUserInfo.onEach { selectedUserInfo = it }.launchIn(this)
 
-        // WHEN the user is changing
-        tracker.onUserChanging(userId = 1)
+            setUpUsers(
+                count = 2,
+                selectedIndex = 0,
+            )
+            tracker.onProfileChanged()
+            assertThat(selectedUserInfo?.id).isEqualTo(0)
+            setUpUsers(
+                count = 2,
+                selectedIndex = 1,
+            )
+            tracker.onProfileChanged()
+            assertThat(selectedUserInfo?.id).isEqualTo(1)
+            job.cancel()
+        }
 
-        // THEN the selection status is IN_PROGRESS
-        assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
+    @Test
+    fun userTrackerCallback_updatesSelectionStatus() =
+        testScope.runTest {
+            underTest = create(testScope.backgroundScope)
+            var selectedUser: SelectedUserModel? = null
+            val job = underTest.selectedUser.onEach { selectedUser = it }.launchIn(this)
 
-        // WHEN the user has finished changing
-        tracker.onUserChanged(userId = 1)
+            setUpUsers(count = 2, selectedIndex = 1)
 
-        // THEN the selection status is COMPLETE
-        assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_COMPLETE)
+            // WHEN the user is changing
+            tracker.onUserChanging(userId = 1)
 
-        tracker.onProfileChanged()
-        assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_COMPLETE)
+            // THEN the selection status is IN_PROGRESS
+            assertThat(selectedUser!!.selectionStatus)
+                .isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
 
-        setUpUsers(count = 2, selectedIndex = 0)
+            // WHEN the user has finished changing
+            tracker.onUserChanged(userId = 1)
 
-        tracker.onUserChanging(userId = 0)
-        assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
+            // THEN the selection status is COMPLETE
+            assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_COMPLETE)
 
-        // WHEN a profile change occurs while a user is changing
-        tracker.onProfileChanged()
+            tracker.onProfileChanged()
+            assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_COMPLETE)
 
-        // THEN the selection status remains as IN_PROGRESS
-        assertThat(selectedUser!!.selectionStatus).isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
-    }
+            setUpUsers(count = 2, selectedIndex = 0)
+
+            tracker.onUserChanging(userId = 0)
+            assertThat(selectedUser!!.selectionStatus)
+                .isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
+
+            // WHEN a profile change occurs while a user is changing
+            tracker.onProfileChanged()
+
+            // THEN the selection status remains as IN_PROGRESS
+            assertThat(selectedUser!!.selectionStatus)
+                .isEqualTo(SelectionStatus.SELECTION_IN_PROGRESS)
+            job.cancel()
+        }
 
     private fun createUserInfo(
         id: Int,
@@ -308,26 +328,13 @@
         assertThat(model.isUserSwitcherEnabled).isEqualTo(expectedUserSwitcherEnabled)
     }
 
-    /**
-     * Executes the given block of execution within the scope of a dedicated [CoroutineScope] which
-     * is then automatically canceled and cleaned-up.
-     */
-    private fun runSelfCancelingTest(
-        block: suspend CoroutineScope.() -> Unit,
-    ) =
-        runBlocking(Dispatchers.Main.immediate) {
-            val scope = CoroutineScope(coroutineContext + Job())
-            block(scope)
-            scope.cancel()
-        }
-
-    private fun create(scope: CoroutineScope = TestCoroutineScope()): UserRepositoryImpl {
+    private fun create(scope: CoroutineScope): UserRepositoryImpl {
         return UserRepositoryImpl(
             appContext = context,
             manager = manager,
             applicationScope = scope,
-            mainDispatcher = IMMEDIATE,
-            backgroundDispatcher = IMMEDIATE,
+            mainDispatcher = testDispatcher,
+            backgroundDispatcher = testDispatcher,
             globalSettings = globalSettings,
             tracker = tracker,
         )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/settings/repository/UserAwareSecureSettingsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/util/settings/repository/UserAwareSecureSettingsRepositoryTest.kt
index 88b2630..a0cfab4d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/settings/repository/UserAwareSecureSettingsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/settings/repository/UserAwareSecureSettingsRepositoryTest.kt
@@ -23,12 +23,13 @@
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.testKosmos
 import com.android.systemui.user.data.repository.fakeUserRepository
-import com.android.systemui.util.settings.FakeSettings
+import com.android.systemui.util.settings.fakeSettings
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
-import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -40,9 +41,10 @@
 @RunWith(AndroidJUnit4::class)
 class UserAwareSecureSettingsRepositoryTest : SysuiTestCase() {
 
-    private val dispatcher = StandardTestDispatcher()
-    private val testScope = TestScope(dispatcher)
-    private val secureSettings = FakeSettings()
+    private val kosmos = testKosmos()
+    private val dispatcher = kosmos.testDispatcher
+    private val testScope = kosmos.testScope
+    private val secureSettings = kosmos.fakeSettings
     private val userRepository = Kosmos().fakeUserRepository
     private lateinit var repository: UserAwareSecureSettingsRepository
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerAdapterTest.kt
similarity index 86%
rename from packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerAdapterTest.kt
index dd78e4a..c140364 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerCollectorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeControllerAdapterTest.kt
@@ -19,16 +19,17 @@
 import android.media.IVolumeController
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.settingslib.media.data.repository.VolumeControllerEvent
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
 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.data.repository.audioRepository
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.kotlin.eq
@@ -38,14 +39,20 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 @RunWith(AndroidJUnit4::class)
 @SmallTest
-class VolumeControllerCollectorTest : SysuiTestCase() {
+class VolumeControllerAdapterTest : SysuiTestCase() {
 
     private val kosmos = testKosmos()
     private val eventsFlow = MutableStateFlow<VolumeControllerEvent?>(null)
-    private val underTest = VolumeControllerCollector(kosmos.applicationCoroutineScope)
+    private val underTest =
+        with(kosmos) { VolumeControllerAdapter(applicationCoroutineScope, audioRepository) }
 
     private val volumeController = mock<IVolumeController> {}
 
+    @Before
+    fun setUp() {
+        kosmos.audioRepository.init()
+    }
+
     @Test
     fun volumeControllerEvent_volumeChanged_callsMethod() =
         testEvent(VolumeControllerEvent.VolumeChanged(3, 0)) {
@@ -90,7 +97,8 @@
 
     private fun testEvent(event: VolumeControllerEvent, verify: () -> Unit) =
         kosmos.testScope.runTest {
-            underTest.collectToController(eventsFlow.filterNotNull(), volumeController)
+            kosmos.audioRepository.sendVolumeControllerEvent(event)
+            underTest.collectToController(volumeController)
 
             eventsFlow.value = event
             runCurrent()
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
index 4ea1a0c..f62beeb 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTest.java
@@ -48,9 +48,11 @@
 
 import com.android.settingslib.flags.Flags;
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.SysuiTestCaseExtKt;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
+import com.android.systemui.kosmos.Kosmos;
 import com.android.systemui.plugins.VolumeDialogController;
 import com.android.systemui.settings.UserTracker;
 import com.android.systemui.statusbar.VibratorHelper;
@@ -78,6 +80,8 @@
 @TestableLooper.RunWithLooper
 public class VolumeDialogControllerImplTest extends SysuiTestCase {
 
+    private final Kosmos mKosmos = SysuiTestCaseExtKt.testKosmos(this);
+
     TestableVolumeDialogControllerImpl mVolumeController;
     VolumeDialogControllerImpl.C mCallback;
     @Mock
@@ -146,6 +150,7 @@
                         mNotificationManager,
                         mVibrator,
                         mIAudioService,
+                        VolumeControllerAdapterKosmosKt.getVolumeControllerAdapter(mKosmos),
                         mAccessibilityManager,
                         mPackageManager,
                         mWakefullnessLifcycle,
@@ -323,6 +328,7 @@
                 NotificationManager notificationManager,
                 VibratorHelper optionalVibrator,
                 IAudioService iAudioService,
+                VolumeControllerAdapter volumeControllerAdapter,
                 AccessibilityManager accessibilityManager,
                 PackageManager packageManager,
                 WakefulnessLifecycle wakefulnessLifecycle,
@@ -342,6 +348,7 @@
                     notificationManager,
                     optionalVibrator,
                     iAudioService,
+                    volumeControllerAdapter,
                     accessibilityManager,
                     packageManager,
                     wakefulnessLifecycle,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt
new file mode 100644
index 0000000..98cea9d
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogControllerImplTestKt.kt
@@ -0,0 +1,156 @@
+/*
+ * 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
+
+import android.app.activityManager
+import android.app.keyguardManager
+import android.content.applicationContext
+import android.content.packageManager
+import android.media.AudioManager
+import android.media.IVolumeController
+import android.os.Handler
+import android.os.looper
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.testing.TestableLooper
+import android.view.accessibility.accessibilityManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
+import com.android.systemui.Flags
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.dump.dumpManager
+import com.android.systemui.keyguard.WakefulnessLifecycle
+import com.android.systemui.keyguard.wakefulnessLifecycle
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.plugins.VolumeDialogController
+import com.android.systemui.testKosmos
+import com.android.systemui.util.RingerModeLiveData
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.concurrency.FakeThreadFactory
+import com.android.systemui.util.time.fakeSystemClock
+import com.android.systemui.volume.data.repository.audioRepository
+import com.android.systemui.volume.domain.interactor.audioSharingInteractor
+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.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+@TestableLooper.RunWithLooper
+class VolumeDialogControllerImplTestKt : SysuiTestCase() {
+
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
+    private val kosmos: Kosmos = testKosmos()
+    private val audioManager: AudioManager = mock {}
+    private val callbacks: VolumeDialogController.Callbacks = mock {}
+
+    private lateinit var threadFactory: FakeThreadFactory
+    private lateinit var underTest: VolumeDialogControllerImpl
+
+    @Before
+    fun setUp() =
+        with(kosmos) {
+            audioRepository.init()
+            threadFactory =
+                FakeThreadFactory(FakeExecutor(fakeSystemClock)).apply { setLooper(looper) }
+            underTest =
+                VolumeDialogControllerImpl(
+                        applicationContext,
+                        mock {},
+                        mock {
+                            on { ringerMode }.thenReturn(mock<RingerModeLiveData> {})
+                            on { ringerModeInternal }.thenReturn(mock<RingerModeLiveData> {})
+                        },
+                        threadFactory,
+                        audioManager,
+                        mock {},
+                        mock {},
+                        mock {},
+                        volumeControllerAdapter,
+                        accessibilityManager,
+                        packageManager,
+                        wakefulnessLifecycle,
+                        keyguardManager,
+                        activityManager,
+                        mock { on { userContext }.thenReturn(applicationContext) },
+                        dumpManager,
+                        audioSharingInteractor,
+                        mock {},
+                    )
+                    .apply {
+                        setEnableDialogs(true, true)
+                        addCallback(callbacks, Handler(looper))
+                    }
+        }
+
+    @Test
+    @EnableFlags(Flags.FLAG_USE_VOLUME_CONTROLLER)
+    fun useVolumeControllerEnabled_listensToVolumeController() =
+        testVolumeController { stream: Int, flags: Int ->
+            audioRepository.sendVolumeControllerEvent(
+                VolumeControllerEvent.VolumeChanged(streamType = stream, flags = flags)
+            )
+        }
+
+    @Test
+    @DisableFlags(Flags.FLAG_USE_VOLUME_CONTROLLER)
+    fun useVolumeControllerDisabled_listensToVolumeController() =
+        testVolumeController { stream: Int, flags: Int ->
+            audioManager.emitVolumeChange(stream, flags)
+        }
+
+    private fun testVolumeController(
+        emitVolumeChange: suspend Kosmos.(stream: Int, flags: Int) -> Unit
+    ) =
+        with(kosmos) {
+            testScope.runTest {
+                whenever(wakefulnessLifecycle.wakefulness)
+                    .thenReturn(WakefulnessLifecycle.WAKEFULNESS_AWAKE)
+                underTest.setVolumeController()
+                runCurrent()
+
+                emitVolumeChange(AudioManager.STREAM_SYSTEM, AudioManager.FLAG_SHOW_UI)
+                runCurrent()
+                TestableLooper.get(this@VolumeDialogControllerImplTestKt).processAllMessages()
+
+                verify(callbacks) { 1 * { onShowRequested(any(), any(), any()) } }
+            }
+        }
+
+    private companion object {
+
+        private fun AudioManager.emitVolumeChange(stream: Int, flags: Int = 0) {
+            val captor = argumentCaptor<IVolumeController>()
+            verify(this) { 1 * { volumeController = captor.capture() } }
+            captor.firstValue.volumeChanged(stream, flags)
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
index caa1779..1e2648b22 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/volume/VolumeDialogImplTest.java
@@ -257,13 +257,13 @@
 
     private State createShellState() {
         State state = new VolumeDialogController.State();
-        for (int i = AudioManager.STREAM_VOICE_CALL; i <= AudioManager.STREAM_ACCESSIBILITY; i++) {
+        for (int stream : STREAMS.keySet()) {
             VolumeDialogController.StreamState ss = new VolumeDialogController.StreamState();
-            ss.name = STREAMS.get(i);
+            ss.name = STREAMS.get(stream);
             ss.level = 1;
             ss.levelMin = 0;
             ss.levelMax = 25;
-            state.states.append(i, ss);
+            state.states.append(stream, ss);
         }
         return state;
     }
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 6e39365..3e7980d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -184,11 +184,11 @@
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.TaskStackListenerImpl;
-import com.android.wm.shell.common.bubbles.BubbleBarLocation;
-import com.android.wm.shell.common.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.shared.animation.PhysicsAnimatorTestUtils;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleBarUpdate;
 import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/settingslib/notification/modes/ZenIconLoaderKosmos.kt
similarity index 72%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
rename to packages/SystemUI/tests/utils/src/com/android/settingslib/notification/modes/ZenIconLoaderKosmos.kt
index 60d97d1..8541d77 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/settingslib/notification/modes/ZenIconLoaderKosmos.kt
@@ -14,9 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.bouncer.shared.flag
+package com.android.settingslib.notification.modes
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.google.common.util.concurrent.MoreExecutors
 
-var Kosmos.fakeComposeBouncerFlags by Kosmos.Fixture { FakeComposeBouncerFlags() }
-val Kosmos.composeBouncerFlags by Kosmos.Fixture<ComposeBouncerFlags> { fakeComposeBouncerFlags }
+val Kosmos.zenIconLoader by Fixture { ZenIconLoader(MoreExecutors.newDirectExecutorService()) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/FakeComposeBouncerFlags.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/FakeComposeBouncerFlags.kt
deleted file mode 100644
index 7482c0f..0000000
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/FakeComposeBouncerFlags.kt
+++ /dev/null
@@ -1,32 +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.bouncer.shared.flag
-
-import com.android.systemui.scene.shared.flag.SceneContainerFlag
-
-class FakeComposeBouncerFlags(var composeBouncerEnabled: Boolean = false) : ComposeBouncerFlags {
-    override fun isComposeBouncerOrSceneContainerEnabled(): Boolean {
-        return SceneContainerFlag.isEnabled || composeBouncerEnabled
-    }
-
-    @Deprecated(
-        "Avoid using this, this is meant to be used only by the glue code " +
-            "that includes compose bouncer in legacy keyguard.",
-        replaceWith = ReplaceWith("isComposeBouncerOrSceneContainerEnabled()")
-    )
-    override fun isOnlyComposeBouncerEnabled(): Boolean = composeBouncerEnabled
-}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelKosmos.kt
index e8612d08..5c5969d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/ui/viewmodel/BouncerMessageViewModelKosmos.kt
@@ -22,7 +22,6 @@
 import com.android.systemui.authentication.domain.interactor.authenticationInteractor
 import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
 import com.android.systemui.bouncer.domain.interactor.simBouncerInteractor
-import com.android.systemui.bouncer.shared.flag.composeBouncerFlags
 import com.android.systemui.deviceentry.domain.interactor.biometricMessageInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryBiometricsAllowedInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryFaceAuthInteractor
@@ -45,7 +44,6 @@
         faceAuthInteractor = deviceEntryFaceAuthInteractor,
         deviceUnlockedInteractor = deviceUnlockedInteractor,
         deviceEntryBiometricsAllowedInteractor = deviceEntryBiometricsAllowedInteractor,
-        flags = composeBouncerFlags,
     )
 }
 
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 e405d17..649e4e8 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
@@ -25,7 +25,6 @@
 import com.android.systemui.bouncer.domain.interactor.bouncerActionButtonInteractor
 import com.android.systemui.bouncer.domain.interactor.bouncerInteractor
 import com.android.systemui.bouncer.domain.interactor.simBouncerInteractor
-import com.android.systemui.bouncer.shared.flag.composeBouncerFlags
 import com.android.systemui.inputmethod.domain.interactor.inputMethodInteractor
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
@@ -34,16 +33,16 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.StateFlow
 
-val Kosmos.bouncerSceneActionsViewModel by Fixture {
-    BouncerSceneActionsViewModel(
+val Kosmos.bouncerUserActionsViewModel by Fixture {
+    BouncerUserActionsViewModel(
         bouncerInteractor = bouncerInteractor,
     )
 }
 
-val Kosmos.bouncerSceneActionsViewModelFactory by Fixture {
-    object : BouncerSceneActionsViewModel.Factory {
-        override fun create(): BouncerSceneActionsViewModel {
-            return bouncerSceneActionsViewModel
+val Kosmos.bouncerUserActionsViewModelFactory by Fixture {
+    object : BouncerUserActionsViewModel.Factory {
+        override fun create(): BouncerUserActionsViewModel {
+            return bouncerUserActionsViewModel
         }
     }
 }
@@ -55,7 +54,6 @@
         authenticationInteractor = authenticationInteractor,
         devicePolicyManager = devicePolicyManager,
         bouncerMessageViewModelFactory = bouncerMessageViewModelFactory,
-        flags = composeBouncerFlags,
         userSwitcher = userSwitcherViewModel,
         actionButtonInteractor = bouncerActionButtonInteractor,
         pinViewModelFactory = pinBouncerViewModelFactory,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/common/ui/ConfigurationStateKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/common/ui/ConfigurationStateKosmos.kt
index 86a8ae5..1ef3464 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/common/ui/ConfigurationStateKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/common/ui/ConfigurationStateKosmos.kt
@@ -17,11 +17,8 @@
 package com.android.systemui.common.ui
 
 import android.content.applicationContext
-import android.view.layoutInflater
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.statusbar.policy.configurationController
 
 val Kosmos.configurationState: ConfigurationState by
-    Kosmos.Fixture {
-        ConfigurationState(configurationController, applicationContext, layoutInflater)
-    }
+    Kosmos.Fixture { ConfigurationStateImpl(configurationController, applicationContext) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt
index 4ad046c..629fda6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/domain/interactor/CommunalInteractorKosmos.kt
@@ -36,6 +36,7 @@
 import com.android.systemui.plugins.activityStarter
 import com.android.systemui.scene.domain.interactor.sceneInteractor
 import com.android.systemui.settings.userTracker
+import com.android.systemui.statusbar.phone.fakeManagedProfileController
 import com.android.systemui.user.data.repository.fakeUserRepository
 import com.android.systemui.util.mockito.mock
 
@@ -61,6 +62,7 @@
         sceneInteractor = sceneInteractor,
         logBuffer = logcatLogBuffer("CommunalInteractor"),
         tableLogBuffer = mock(),
+        managedProfileController = fakeManagedProfileController
     )
 }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorKosmos.kt
index caa6e99..13116e7 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/deviceentry/domain/interactor/DeviceEntryInteractorKosmos.kt
@@ -19,6 +19,7 @@
 import com.android.systemui.authentication.domain.interactor.authenticationInteractor
 import com.android.systemui.bouncer.domain.interactor.alternateBouncerInteractor
 import com.android.systemui.deviceentry.data.repository.deviceEntryRepository
+import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.scene.domain.interactor.sceneInteractor
@@ -34,5 +35,6 @@
             sceneInteractor = sceneInteractor,
             deviceUnlockedInteractor = deviceUnlockedInteractor,
             alternateBouncerInteractor = alternateBouncerInteractor,
+            dismissCallbackRegistry = dismissCallbackRegistry,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
index 88ab170..7ccacb6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
@@ -16,14 +16,17 @@
 
 package com.android.systemui.education.domain.interactor
 
+import android.hardware.input.InputManager
 import com.android.systemui.education.data.repository.fakeEduClock
 import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
+import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
 import com.android.systemui.keyboard.data.repository.keyboardRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.touchpad.data.repository.touchpadRepository
 import com.android.systemui.user.data.repository.userRepository
+import org.mockito.kotlin.mock
 
 var Kosmos.keyboardTouchpadEduInteractor by
     Kosmos.Fixture {
@@ -37,14 +40,23 @@
                     touchpadRepository,
                     userRepository
                 ),
-            clock = fakeEduClock
+            clock = fakeEduClock,
+            inputManager = mockEduInputManager
         )
     }
 
+var Kosmos.mockEduInputManager by Kosmos.Fixture { mock<InputManager>() }
+
 var Kosmos.keyboardTouchpadEduStatsInteractor by
     Kosmos.Fixture {
         KeyboardTouchpadEduStatsInteractorImpl(
             backgroundScope = testScope.backgroundScope,
-            contextualEducationInteractor = contextualEducationInteractor
+            contextualEducationInteractor = contextualEducationInteractor,
+            inputDeviceRepository = mockUserInputDeviceRepository,
+            tutorialRepository = mockTutorialSchedulerRepository,
+            clock = fakeEduClock
         )
     }
+
+var mockUserInputDeviceRepository = mock<UserInputDeviceRepository>()
+var mockTutorialSchedulerRepository = mock<TutorialSchedulerRepository>()
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/EnableSceneContainer.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/EnableSceneContainer.kt
index 1107971..c252924 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/flags/EnableSceneContainer.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/flags/EnableSceneContainer.kt
@@ -22,7 +22,7 @@
 import com.android.systemui.Flags.FLAG_KEYGUARD_BOTTOM_AREA_REFACTOR
 import com.android.systemui.Flags.FLAG_KEYGUARD_WM_STATE_REFACTOR
 import com.android.systemui.Flags.FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT
-import com.android.systemui.Flags.FLAG_NOTIFICATIONS_HEADS_UP_REFACTOR
+import com.android.systemui.Flags.FLAG_NOTIFICATION_AVALANCHE_THROTTLE_HUN
 import com.android.systemui.Flags.FLAG_PREDICTIVE_BACK_SYSUI
 import com.android.systemui.Flags.FLAG_SCENE_CONTAINER
 
@@ -35,7 +35,7 @@
     FLAG_KEYGUARD_BOTTOM_AREA_REFACTOR,
     FLAG_KEYGUARD_WM_STATE_REFACTOR,
     FLAG_MIGRATE_CLOCKS_TO_BLUEPRINT,
-    FLAG_NOTIFICATIONS_HEADS_UP_REFACTOR,
+    FLAG_NOTIFICATION_AVALANCHE_THROTTLE_HUN,
     FLAG_PREDICTIVE_BACK_SYSUI,
     FLAG_SCENE_CONTAINER,
     FLAG_DEVICE_ENTRY_UDFPS_REFACTOR,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/VibratorHelperKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/VibratorHelperKosmos.kt
index 434953f..ff71f2f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/VibratorHelperKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/haptics/VibratorHelperKosmos.kt
@@ -17,5 +17,7 @@
 package com.android.systemui.haptics
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.statusbar.VibratorHelper
 
-var Kosmos.vibratorHelper by Kosmos.Fixture { FakeVibratorHelper() }
+var Kosmos.vibratorHelper: VibratorHelper by Kosmos.Fixture { fakeVibratorHelper }
+val Kosmos.fakeVibratorHelper by Kosmos.Fixture { FakeVibratorHelper() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeDeviceEntryFaceAuthRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeDeviceEntryFaceAuthRepository.kt
index e96aeada..5753c6c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeDeviceEntryFaceAuthRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeDeviceEntryFaceAuthRepository.kt
@@ -27,7 +27,6 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.filterNotNull
 
 @SysUISingleton
@@ -37,38 +36,41 @@
     private val _authenticationStatus = MutableStateFlow<FaceAuthenticationStatus?>(null)
     override val authenticationStatus: Flow<FaceAuthenticationStatus> =
         _authenticationStatus.filterNotNull()
+
     fun setAuthenticationStatus(status: FaceAuthenticationStatus) {
         _authenticationStatus.value = status
     }
+
     private val _detectionStatus = MutableStateFlow<FaceDetectionStatus?>(null)
     override val detectionStatus: Flow<FaceDetectionStatus>
         get() = _detectionStatus.filterNotNull()
+
     fun setDetectionStatus(status: FaceDetectionStatus) {
         _detectionStatus.value = status
     }
 
     private val _isLockedOut = MutableStateFlow(false)
     override val isLockedOut = _isLockedOut
-    private val _runningAuthRequest = MutableStateFlow<Pair<FaceAuthUiEvent, Boolean>?>(null)
-    val runningAuthRequest: StateFlow<Pair<FaceAuthUiEvent, Boolean>?> =
-        _runningAuthRequest.asStateFlow()
+    val runningAuthRequest: MutableStateFlow<Pair<FaceAuthUiEvent, Boolean>?> =
+        MutableStateFlow(null)
 
     private val _isAuthRunning = MutableStateFlow(false)
     override val isAuthRunning: StateFlow<Boolean> = _isAuthRunning
 
     override val isBypassEnabled = MutableStateFlow(false)
+
     override fun setLockedOut(isLockedOut: Boolean) {
         _isLockedOut.value = isLockedOut
     }
 
     override fun requestAuthenticate(uiEvent: FaceAuthUiEvent, fallbackToDetection: Boolean) {
-        _runningAuthRequest.value = uiEvent to fallbackToDetection
+        runningAuthRequest.value = uiEvent to fallbackToDetection
         _isAuthRunning.value = true
     }
 
     override fun cancel() {
         _isAuthRunning.value = false
-        _runningAuthRequest.value = null
+        runningAuthRequest.value = null
     }
 }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
index 616f2b6..a73c184 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/data/repository/FakeKeyguardTransitionRepository.kt
@@ -102,6 +102,45 @@
     }
 
     /**
+     * Sends the provided [step] and makes sure that all previous [TransitionState]'s are sent when
+     * [fillInSteps] is true. e.g. when a step FINISHED is provided, a step with STARTED and RUNNING
+     * is also sent.
+     */
+    suspend fun sendTransitionSteps(
+        step: TransitionStep,
+        testScope: TestScope,
+        fillInSteps: Boolean = true,
+    ) {
+        if (fillInSteps && step.transitionState != TransitionState.STARTED) {
+            sendTransitionStep(
+                step =
+                    TransitionStep(
+                        transitionState = TransitionState.STARTED,
+                        from = step.from,
+                        to = step.to,
+                        value = 0f,
+                    )
+            )
+            testScope.testScheduler.runCurrent()
+
+            if (step.transitionState != TransitionState.RUNNING) {
+                sendTransitionStep(
+                    step =
+                        TransitionStep(
+                            transitionState = TransitionState.RUNNING,
+                            from = step.from,
+                            to = step.to,
+                            value = 0.6f,
+                        )
+                )
+                testScope.testScheduler.runCurrent()
+            }
+        }
+        sendTransitionStep(step = step)
+        testScope.testScheduler.runCurrent()
+    }
+
+    /**
      * Sends TransitionSteps between [from] and [to], calling [runCurrent] after each step.
      *
      * By default, sends steps through FINISHED (STARTED, RUNNING, FINISHED) but can be halted part
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/DismissKeyguardInteractor.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/DismissKeyguardInteractor.kt
new file mode 100644
index 0000000..82a5311
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/DismissKeyguardInteractor.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2023 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 com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
+import com.android.systemui.kosmos.Kosmos
+
+val Kosmos.keyguardDismissTransitionInteractor: KeyguardDismissTransitionInteractor by
+    Kosmos.Fixture {
+        KeyguardDismissTransitionInteractor(
+            repository = keyguardTransitionRepository,
+            fromLockscreenTransitionInteractor = fromLockscreenTransitionInteractor,
+            fromPrimaryBouncerTransitionInteractor = fromPrimaryBouncerTransitionInteractor,
+            fromAodTransitionInteractor = fromAodTransitionInteractor,
+            fromAlternateBouncerTransitionInteractor = fromAlternateBouncerTransitionInteractor,
+            fromDozingTransitionInteractor = fromDozingTransitionInteractor,
+            fromOccludedTransitionInteractor = fromOccludedTransitionInteractor,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractorKosmos.kt
index c6b5ed0..007d229 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardEnabledInteractorKosmos.kt
@@ -27,7 +27,7 @@
             applicationCoroutineScope,
             keyguardRepository,
             biometricSettingsRepository,
-            keyguardTransitionInteractor,
+            keyguardDismissTransitionInteractor,
             internalTransitionInteractor = internalKeyguardTransitionInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
index b68d6a0..aa94c36 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/domain/interactor/KeyguardTransitionInteractorKosmos.kt
@@ -16,7 +16,6 @@
 
 package com.android.systemui.keyguard.domain.interactor
 
-import com.android.systemui.keyguard.data.repository.keyguardRepository
 import com.android.systemui.keyguard.data.repository.keyguardTransitionRepository
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
@@ -27,13 +26,6 @@
         KeyguardTransitionInteractor(
             scope = applicationCoroutineScope,
             repository = keyguardTransitionRepository,
-            keyguardRepository = keyguardRepository,
-            fromLockscreenTransitionInteractor = { fromLockscreenTransitionInteractor },
-            fromPrimaryBouncerTransitionInteractor = { fromPrimaryBouncerTransitionInteractor },
-            fromAodTransitionInteractor = { fromAodTransitionInteractor },
-            fromAlternateBouncerTransitionInteractor = { fromAlternateBouncerTransitionInteractor },
-            fromDozingTransitionInteractor = { fromDozingTransitionInteractor },
-            fromOccludedTransitionInteractor = { fromOccludedTransitionInteractor },
             sceneInteractor = sceneInteractor
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
index 2958315..f1d87fe 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/AlternateBouncerViewModelKosmos.kt
@@ -19,6 +19,7 @@
 package com.android.systemui.keyguard.ui.viewmodel
 
 import com.android.systemui.bouncer.domain.interactor.alternateBouncerInteractor
+import com.android.systemui.bouncer.domain.interactor.primaryBouncerInteractor
 import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.kosmos.Kosmos
@@ -32,5 +33,6 @@
         keyguardTransitionInteractor = keyguardTransitionInteractor,
         dismissCallbackRegistry = dismissCallbackRegistry,
         alternateBouncerInteractor = { alternateBouncerInteractor },
+        primaryBouncerInteractor = primaryBouncerInteractor,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
index b9443bc..7cf4083 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/KeyguardRootViewModelKosmos.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.applicationCoroutineScope
 import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.shade.ui.viewmodel.notificationShadeWindowModel
 import com.android.systemui.statusbar.notification.stack.domain.interactor.notificationsKeyguardInteractor
 import com.android.systemui.statusbar.phone.dozeParameters
 import com.android.systemui.statusbar.phone.screenOffAnimationController
@@ -39,6 +40,7 @@
         communalInteractor = communalInteractor,
         keyguardTransitionInteractor = keyguardTransitionInteractor,
         notificationsKeyguardInteractor = notificationsKeyguardInteractor,
+        notificationShadeWindowModel = notificationShadeWindowModel,
         alternateBouncerToAodTransitionViewModel = alternateBouncerToAodTransitionViewModel,
         alternateBouncerToGoneTransitionViewModel = alternateBouncerToGoneTransitionViewModel,
         alternateBouncerToLockscreenTransitionViewModel =
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModelKosmos.kt
index a05e606..4196e54 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/keyguard/ui/viewmodel/OccludedToDozingTransitionViewModelKosmos.kt
@@ -18,6 +18,7 @@
 
 package com.android.systemui.keyguard.ui.viewmodel
 
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryUdfpsInteractor
 import com.android.systemui.keyguard.ui.keyguardTransitionAnimationFlow
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
@@ -25,6 +26,7 @@
 
 val Kosmos.occludedToDozingTransitionViewModel by Fixture {
     OccludedToDozingTransitionViewModel(
+        deviceEntryUdfpsInteractor = deviceEntryUdfpsInteractor,
         animationFlow = keyguardTransitionAnimationFlow,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
index b34681a..f8df707 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/GeneralKosmos.kt
@@ -5,9 +5,12 @@
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
 
 var Kosmos.testDispatcher by Fixture { StandardTestDispatcher() }
+var Kosmos.unconfinedTestDispatcher by Fixture { UnconfinedTestDispatcher() }
 var Kosmos.testScope by Fixture { TestScope(testDispatcher) }
+var Kosmos.unconfinedTestScope by Fixture { TestScope(unconfinedTestDispatcher) }
 var Kosmos.applicationCoroutineScope by Fixture { testScope.backgroundScope }
 var Kosmos.testCase: SysuiTestCase by Fixture()
 var Kosmos.backgroundCoroutineContext: CoroutineContext by Fixture {
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
index 953363d..851a378 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/KosmosJavaAdapter.kt
@@ -36,6 +36,7 @@
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryUdfpsInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
 import com.android.systemui.globalactions.domain.interactor.globalActionsInteractor
+import com.android.systemui.haptics.msdl.msdlPlayer
 import com.android.systemui.haptics.qs.qsLongPressEffect
 import com.android.systemui.jank.interactionJankMonitor
 import com.android.systemui.keyguard.data.repository.fakeKeyguardRepository
@@ -154,4 +155,5 @@
     val scrimController by lazy { kosmos.scrimController }
     val scrimStartable by lazy { kosmos.scrimStartable }
     val sceneContainerOcclusionInteractor by lazy { kosmos.sceneContainerOcclusionInteractor }
+    val msdlPlayer by lazy { kosmos.msdlPlayer }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/lifecycle/FakeSysUiViewModel.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/lifecycle/FakeSysUiViewModel.kt
index 1652462..2eb7ce6 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/lifecycle/FakeSysUiViewModel.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/lifecycle/FakeSysUiViewModel.kt
@@ -26,17 +26,25 @@
 class FakeSysUiViewModel(
     private val onActivation: () -> Unit = {},
     private val onDeactivation: () -> Unit = {},
-    private val upstreamFlow: Flow<Boolean> = flowOf(true),
-    private val upstreamStateFlow: StateFlow<Boolean> = MutableStateFlow(true).asStateFlow(),
-) : SysUiViewModel, ExclusiveActivatable() {
+    upstreamFlow: Flow<Boolean> = flowOf(true),
+    upstreamStateFlow: StateFlow<Boolean> = MutableStateFlow(true).asStateFlow(),
+) : ExclusiveActivatable() {
 
     var activationCount = 0
     var cancellationCount = 0
 
-    private val hydrator = Hydrator()
+    private val hydrator = Hydrator("test")
     val stateBackedByFlow: Boolean by
-        hydrator.hydratedStateOf(initialValue = true, source = upstreamFlow)
-    val stateBackedByStateFlow: Boolean by hydrator.hydratedStateOf(source = upstreamStateFlow)
+        hydrator.hydratedStateOf(
+            traceName = "test",
+            initialValue = true,
+            source = upstreamFlow,
+        )
+    val stateBackedByStateFlow: Boolean by
+        hydrator.hydratedStateOf(
+            traceName = "test",
+            source = upstreamStateFlow,
+        )
 
     override suspend fun onActivated(): Nothing {
         activationCount++
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferFactoryKosmos.kt
similarity index 61%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferFactoryKosmos.kt
index 5d70ed6..c55dc6a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferFactoryKosmos.kt
@@ -14,15 +14,18 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.ui.viewmodel
+package com.android.systemui.log.table
 
+import com.android.systemui.dump.dumpManager
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.log.LogcatEchoTrackerAlways
+import com.android.systemui.util.time.fakeSystemClock
 
-val Kosmos.quickSettingsShadeSceneActionsViewModel: QuickSettingsShadeSceneActionsViewModel by
+val Kosmos.tableLogBufferFactory: TableLogBufferFactory by
     Kosmos.Fixture {
-        QuickSettingsShadeSceneActionsViewModel(
-            shadeInteractor = shadeInteractor,
-            quickSettingsContainerViewModel = quickSettingsContainerViewModel,
+        TableLogBufferFactory(
+            dumpManager = dumpManager,
+            systemClock = fakeSystemClock,
+            logcatEchoTracker = LogcatEchoTrackerAlways(),
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferHelper.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferHelper.kt
new file mode 100644
index 0000000..64c3974
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/log/table/TableLogBufferHelper.kt
@@ -0,0 +1,36 @@
+/*
+ * 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.log.table
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.log.LogcatEchoTrackerAlways
+import com.android.systemui.util.time.SystemClock
+import com.android.systemui.util.time.fakeSystemClock
+
+/**
+ * Creates a [TableLogBuffer] that will echo everything to logcat, which is useful for debugging
+ * tests.
+ */
+fun logcatTableLogBuffer(kosmos: Kosmos, name: String = "EchoToLogcatTableLogBuffer") =
+    logcatTableLogBuffer(kosmos.fakeSystemClock, name)
+
+/**
+ * Creates a [TableLogBuffer] that will echo everything to logcat, which is useful for debugging
+ * tests.
+ */
+fun logcatTableLogBuffer(systemClock: SystemClock, name: String = "EchoToLogcatTableLogBuffer") =
+    TableLogBuffer(maxSize = 50, name, systemClock, logcatEchoTracker = LogcatEchoTrackerAlways())
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorKosmos.kt
index 2127a88..632436a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDataProcessorKosmos.kt
@@ -18,7 +18,6 @@
 
 import android.app.smartspace.SmartspaceManager
 import android.content.applicationContext
-import android.os.fakeExecutorHandler
 import com.android.keyguard.keyguardUpdateMonitor
 import com.android.systemui.broadcast.broadcastDispatcher
 import com.android.systemui.concurrency.fakeExecutor
@@ -45,7 +44,7 @@
             backgroundExecutor = fakeExecutor,
             uiExecutor = fakeExecutor,
             foregroundExecutor = fakeExecutor,
-            handler = fakeExecutorHandler,
+            mainDispatcher = testDispatcher,
             mediaControllerFactory = fakeMediaControllerFactory,
             broadcastDispatcher = broadcastDispatcher,
             dumpManager = dumpManager,
@@ -60,5 +59,6 @@
             smartspaceManager = SmartspaceManager(applicationContext),
             keyguardUpdateMonitor = keyguardUpdateMonitor,
             mediaDataRepository = mediaDataRepository,
+            mediaDataLoader = { mediaDataLoader },
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLoggerKosmos.kt
similarity index 74%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLoggerKosmos.kt
index d60f14c..76d71dd 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceLoggerKosmos.kt
@@ -14,10 +14,10 @@
  * limitations under the License.
  */
 
-package com.android.systemui.volume
+package com.android.systemui.media.controls.domain.pipeline
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.log.logcatLogBuffer
 
-val Kosmos.volumeControllerCollector by
-    Kosmos.Fixture { VolumeControllerCollector(applicationCoroutineScope) }
+var Kosmos.mediaDeviceLogger by
+    Kosmos.Fixture { MediaDeviceLogger(logcatLogBuffer("MediaDeviceLoggerKosmos")) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerKosmos.kt
index c479ce6..11408d8 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/media/controls/domain/pipeline/MediaDeviceManagerKosmos.kt
@@ -41,5 +41,6 @@
             },
             fgExecutor = fakeExecutor,
             bgExecutor = fakeExecutor,
+            logger = mediaDeviceLogger,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
similarity index 71%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
index 60d97d1..8fc40e4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/bouncer/shared/flag/ComposeBouncerFlagsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayActionsViewModelKosmos.kt
@@ -14,9 +14,12 @@
  * limitations under the License.
  */
 
-package com.android.systemui.bouncer.shared.flag
+package com.android.systemui.qs.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
 
-var Kosmos.fakeComposeBouncerFlags by Kosmos.Fixture { FakeComposeBouncerFlags() }
-val Kosmos.composeBouncerFlags by Kosmos.Fixture<ComposeBouncerFlags> { fakeComposeBouncerFlags }
+val Kosmos.quickSettingsShadeOverlayActionsViewModel:
+    QuickSettingsShadeOverlayActionsViewModel by Fixture {
+    QuickSettingsShadeOverlayActionsViewModel()
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelKosmos.kt
similarity index 67%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelKosmos.kt
index 5d70ed6..ff8b478 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeOverlayContentViewModelKosmos.kt
@@ -17,12 +17,14 @@
 package com.android.systemui.qs.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.shade.ui.viewmodel.shadeHeaderViewModelFactory
 
-val Kosmos.quickSettingsShadeSceneActionsViewModel: QuickSettingsShadeSceneActionsViewModel by
+val Kosmos.quickSettingsShadeOverlayContentViewModel: QuickSettingsShadeOverlayContentViewModel by
     Kosmos.Fixture {
-        QuickSettingsShadeSceneActionsViewModel(
-            shadeInteractor = shadeInteractor,
+        QuickSettingsShadeOverlayContentViewModel(
+            sceneInteractor = sceneInteractor,
+            shadeHeaderViewModelFactory = shadeHeaderViewModelFactory,
             quickSettingsContainerViewModel = quickSettingsContainerViewModel,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModelKosmos.kt
index 5ad5cb2..cd1704c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneContentViewModelKosmos.kt
@@ -14,18 +14,13 @@
  * limitations under the License.
  */
 
-@file:OptIn(ExperimentalCoroutinesApi::class)
-
 package com.android.systemui.qs.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.ui.viewmodel.overlayShadeViewModelFactory
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 
 val Kosmos.quickSettingsShadeSceneContentViewModel: QuickSettingsShadeSceneContentViewModel by
     Kosmos.Fixture {
         QuickSettingsShadeSceneContentViewModel(
-            overlayShadeViewModelFactory = overlayShadeViewModelFactory,
             quickSettingsContainerViewModel = quickSettingsContainerViewModel,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelKosmos.kt
similarity index 76%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelKosmos.kt
index 5d70ed6..06592b1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeUserActionsViewModelKosmos.kt
@@ -17,12 +17,10 @@
 package com.android.systemui.qs.ui.viewmodel
 
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.domain.interactor.shadeInteractor
 
-val Kosmos.quickSettingsShadeSceneActionsViewModel: QuickSettingsShadeSceneActionsViewModel by
+val Kosmos.quickSettingsShadeUserActionsViewModel: QuickSettingsShadeUserActionsViewModel by
     Kosmos.Fixture {
-        QuickSettingsShadeSceneActionsViewModel(
-            shadeInteractor = shadeInteractor,
+        QuickSettingsShadeUserActionsViewModel(
             quickSettingsContainerViewModel = quickSettingsContainerViewModel,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
index dd93141..b3664e1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/SceneKosmos.kt
@@ -2,9 +2,10 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.scene.shared.model.FakeScene
+import com.android.systemui.scene.shared.model.Overlays
 import com.android.systemui.scene.shared.model.SceneContainerConfig
 import com.android.systemui.scene.shared.model.Scenes
+import com.android.systemui.scene.ui.FakeOverlay
 
 var Kosmos.sceneKeys by Fixture {
     listOf(
@@ -17,11 +18,19 @@
     )
 }
 
-val Kosmos.fakeScenes by Fixture { sceneKeys.map { key -> FakeScene(key) }.toSet() }
-
-val Kosmos.scenes by Fixture { fakeScenes }
-
 val Kosmos.initialSceneKey by Fixture { Scenes.Lockscreen }
+
+var Kosmos.overlayKeys by Fixture {
+    listOf(
+        Overlays.NotificationsShade,
+        Overlays.QuickSettingsShade,
+    )
+}
+
+val Kosmos.fakeOverlays by Fixture { overlayKeys.map { key -> FakeOverlay(key) }.toSet() }
+
+val Kosmos.overlays by Fixture { fakeOverlays }
+
 var Kosmos.sceneContainerConfig by Fixture {
     val navigationDistances =
         mapOf(
@@ -32,5 +41,11 @@
             Scenes.QuickSettings to 3,
             Scenes.Bouncer to 4,
         )
-    SceneContainerConfig(sceneKeys, initialSceneKey, navigationDistances)
+
+    SceneContainerConfig(
+        sceneKeys = sceneKeys,
+        initialSceneKey = initialSceneKey,
+        overlayKeys = overlayKeys,
+        navigationDistances = navigationDistances,
+    )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
index 53d3c01..c95b2dc 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
@@ -19,6 +19,7 @@
 import com.android.compose.animation.scene.ObservableTransitionState
 import com.android.compose.animation.scene.SceneKey
 import com.android.systemui.keyguard.data.repository.fakeKeyguardTransitionRepository
+import com.android.systemui.keyguard.shared.model.KeyguardState
 import com.android.systemui.keyguard.shared.model.TransitionStep
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
@@ -28,7 +29,6 @@
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.runCurrent
 
 private val mutableTransitionState =
     MutableStateFlow<ObservableTransitionState>(ObservableTransitionState.Idle(Scenes.Lockscreen))
@@ -36,20 +36,26 @@
 suspend fun Kosmos.setTransition(
     sceneTransition: ObservableTransitionState,
     stateTransition: TransitionStep? = null,
+    fillInStateSteps: Boolean = true,
     scope: TestScope = testScope,
     repository: SceneContainerRepository = sceneContainerRepository
 ) {
+    var state: TransitionStep? = stateTransition
     if (SceneContainerFlag.isEnabled) {
         setSceneTransition(sceneTransition, scope, repository)
-    } else {
-        if (stateTransition == null) throw IllegalArgumentException("No transitionStep provided")
-        fakeKeyguardTransitionRepository.sendTransitionSteps(
-            from = stateTransition.from,
-            to = stateTransition.to,
-            testScope = scope,
-            throughTransitionState = stateTransition.transitionState
-        )
+
+        if (state != null) {
+            state = getStateWithUndefined(sceneTransition, state)
+        }
     }
+
+    if (state == null) return
+    fakeKeyguardTransitionRepository.sendTransitionSteps(
+        step = state,
+        testScope = scope,
+        fillInSteps = fillInStateSteps,
+    )
+    scope.testScheduler.runCurrent()
 }
 
 fun Kosmos.setSceneTransition(
@@ -59,7 +65,7 @@
 ) {
     repository.setTransitionState(mutableTransitionState)
     mutableTransitionState.value = transition
-    scope.runCurrent()
+    scope.testScheduler.runCurrent()
 }
 
 fun Transition(
@@ -87,3 +93,43 @@
 fun Idle(currentScene: SceneKey): ObservableTransitionState.Idle {
     return ObservableTransitionState.Idle(currentScene)
 }
+
+private fun getStateWithUndefined(
+    sceneTransition: ObservableTransitionState,
+    state: TransitionStep
+): TransitionStep {
+    return when (sceneTransition) {
+        is ObservableTransitionState.Idle -> {
+            TransitionStep(
+                from = state.from,
+                to =
+                    if (sceneTransition.currentScene != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.to
+                    },
+                value = state.value,
+                transitionState = state.transitionState
+            )
+        }
+        is ObservableTransitionState.Transition -> {
+            TransitionStep(
+                from =
+                    if (sceneTransition.fromContent != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.from
+                    },
+                to =
+                    if (sceneTransition.toContent != Scenes.Lockscreen) {
+                        KeyguardState.UNDEFINED
+                    } else {
+                        state.from
+                    },
+                value = state.value,
+                transitionState = state.transitionState
+            )
+        }
+        else -> state
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/startable/SceneContainerStartableKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/startable/SceneContainerStartableKosmos.kt
index 53b6a2e..b612a8b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/startable/SceneContainerStartableKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/domain/startable/SceneContainerStartableKosmos.kt
@@ -24,8 +24,10 @@
 import com.android.systemui.classifier.falsingCollector
 import com.android.systemui.classifier.falsingManager
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryFaceAuthInteractor
+import com.android.systemui.deviceentry.domain.interactor.deviceEntryHapticsInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceEntryInteractor
 import com.android.systemui.deviceentry.domain.interactor.deviceUnlockedInteractor
+import com.android.systemui.haptics.vibratorHelper
 import com.android.systemui.keyguard.dismissCallbackRegistry
 import com.android.systemui.keyguard.domain.interactor.keyguardEnabledInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardInteractor
@@ -55,6 +57,7 @@
         applicationScope = testScope.backgroundScope,
         sceneInteractor = sceneInteractor,
         deviceEntryInteractor = deviceEntryInteractor,
+        deviceEntryHapticsInteractor = deviceEntryHapticsInteractor,
         deviceUnlockedInteractor = deviceUnlockedInteractor,
         bouncerInteractor = bouncerInteractor,
         keyguardInteractor = keyguardInteractor,
@@ -82,5 +85,6 @@
         dismissCallbackRegistry = dismissCallbackRegistry,
         statusBarStateController = sysuiStatusBarStateController,
         alternateBouncerInteractor = alternateBouncerInteractor,
+        vibratorHelper = vibratorHelper,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeScene.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeScene.kt
deleted file mode 100644
index 78358f5..0000000
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeScene.kt
+++ /dev/null
@@ -1,49 +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.scene.shared.model
-
-import com.android.compose.animation.scene.SceneKey
-import com.android.compose.animation.scene.UserAction
-import com.android.compose.animation.scene.UserActionResult
-import com.android.systemui.lifecycle.ExclusiveActivatable
-import kotlinx.coroutines.awaitCancellation
-import kotlinx.coroutines.channels.Channel
-import kotlinx.coroutines.flow.onCompletion
-import kotlinx.coroutines.flow.onStart
-import kotlinx.coroutines.flow.receiveAsFlow
-
-class FakeScene(
-    override val key: SceneKey,
-) : ExclusiveActivatable(), Scene {
-    var isDestinationScenesBeingCollected = false
-
-    private val destinationScenesChannel = Channel<Map<UserAction, UserActionResult>>()
-
-    override val destinationScenes =
-        destinationScenesChannel
-            .receiveAsFlow()
-            .onStart { isDestinationScenesBeingCollected = true }
-            .onCompletion { isDestinationScenesBeingCollected = false }
-
-    override suspend fun onActivated(): Nothing {
-        awaitCancellation()
-    }
-
-    suspend fun setDestinationScenes(value: Map<UserAction, UserActionResult>) {
-        destinationScenesChannel.send(value)
-    }
-}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
index 957a60f..f52572a 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/shared/model/FakeSceneDataSource.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.scene.shared.model
 
+import com.android.compose.animation.scene.OverlayKey
 import com.android.compose.animation.scene.SceneKey
 import com.android.compose.animation.scene.TransitionKey
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -29,11 +30,18 @@
     private val _currentScene = MutableStateFlow(initialSceneKey)
     override val currentScene: StateFlow<SceneKey> = _currentScene.asStateFlow()
 
+    private val _currentOverlays = MutableStateFlow<Set<OverlayKey>>(emptySet())
+    override val currentOverlays: StateFlow<Set<OverlayKey>> = _currentOverlays.asStateFlow()
+
     var isPaused = false
         private set
+
     var pendingScene: SceneKey? = null
         private set
 
+    var pendingOverlays: Set<OverlayKey>? = null
+        private set
+
     override fun changeScene(toScene: SceneKey, transitionKey: TransitionKey?) {
         if (isPaused) {
             pendingScene = toScene
@@ -46,10 +54,32 @@
         changeScene(toScene)
     }
 
+    override fun showOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        if (isPaused) {
+            pendingOverlays = (pendingOverlays ?: currentOverlays.value) + overlay
+        } else {
+            _currentOverlays.value += overlay
+        }
+    }
+
+    override fun hideOverlay(overlay: OverlayKey, transitionKey: TransitionKey?) {
+        if (isPaused) {
+            pendingOverlays = (pendingOverlays ?: currentOverlays.value) - overlay
+        } else {
+            _currentOverlays.value -= overlay
+        }
+    }
+
+    override fun replaceOverlay(from: OverlayKey, to: OverlayKey, transitionKey: TransitionKey?) {
+        hideOverlay(from, transitionKey)
+        showOverlay(to, transitionKey)
+    }
+
     /**
-     * Pauses scene changes.
+     * Pauses scene and overlay changes.
      *
-     * Any following calls to [changeScene] will be conflated and the last one will be remembered.
+     * Any following calls to [changeScene] or overlay changing functions will be conflated and the
+     * last one will be remembered.
      */
     fun pause() {
         check(!isPaused) { "Can't pause what's already paused!" }
@@ -58,11 +88,14 @@
     }
 
     /**
-     * Unpauses scene changes.
+     * Unpauses scene and overlay changes.
      *
      * If there were any calls to [changeScene] since [pause] was called, the latest of the bunch
      * will be replayed.
      *
+     * If there were any calls to show, hide or replace overlays since [pause] was called, they will
+     * all be applied at once.
+     *
      * If [force] is `true`, there will be no check that [isPaused] is true.
      *
      * If [expectedScene] is provided, will assert that it's indeed the latest called.
@@ -76,6 +109,8 @@
         isPaused = false
         pendingScene?.let { _currentScene.value = it }
         pendingScene = null
+        pendingOverlays?.let { _currentOverlays.value = it }
+        pendingOverlays = null
 
         check(expectedScene == null || currentScene.value == expectedScene) {
             """
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt
new file mode 100644
index 0000000..f4f30cd
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/ui/FakeOverlay.kt
@@ -0,0 +1,36 @@
+/*
+ * 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.scene.ui
+
+import androidx.compose.runtime.Composable
+import androidx.compose.ui.Modifier
+import com.android.compose.animation.scene.ContentScope
+import com.android.compose.animation.scene.OverlayKey
+import com.android.systemui.lifecycle.ExclusiveActivatable
+import com.android.systemui.scene.ui.composable.Overlay
+import kotlinx.coroutines.awaitCancellation
+
+class FakeOverlay(
+    override val key: OverlayKey,
+) : ExclusiveActivatable(), Overlay {
+
+    @Composable override fun ContentScope.Content(modifier: Modifier) = Unit
+
+    override suspend fun onActivated(): Nothing {
+        awaitCancellation()
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
index 4660337..4a86fd5 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/data/repository/FakeShadeRepository.kt
@@ -72,10 +72,6 @@
     @Deprecated("Use ShadeInteractor.isUserInteractingWithShade instead")
     override val legacyLockscreenShadeTracking = MutableStateFlow(false)
 
-    private var _isDualShadeAlignedToBottom = false
-    override val isDualShadeAlignedToBottom
-        get() = _isDualShadeAlignedToBottom
-
     private var _isShadeLayoutWide = MutableStateFlow(false)
     override val isShadeLayoutWide: StateFlow<Boolean> = _isShadeLayoutWide.asStateFlow()
 
@@ -155,10 +151,6 @@
         _legacyShadeExpansion.value = expandedFraction
     }
 
-    fun setDualShadeAlignedToBottom(isAlignedToBottom: Boolean) {
-        _isDualShadeAlignedToBottom = isAlignedToBottom
-    }
-
     override fun setShadeLayoutWide(isShadeLayoutWide: Boolean) {
         _isShadeLayoutWide.value = isShadeLayoutWide
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayActionsViewModelKosmos.kt
similarity index 73%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayActionsViewModelKosmos.kt
index 9bf4756..1e00ac4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayActionsViewModelKosmos.kt
@@ -18,12 +18,9 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeSceneActionsViewModel
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeOverlayActionsViewModel
 
-val Kosmos.notificationsShadeSceneActionsViewModel:
-    NotificationsShadeSceneActionsViewModel by Fixture {
-    NotificationsShadeSceneActionsViewModel(
-        shadeInteractor = shadeInteractor,
-    )
+val Kosmos.notificationsShadeOverlayActionsViewModel:
+    NotificationsShadeOverlayActionsViewModel by Fixture {
+    NotificationsShadeOverlayActionsViewModel()
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayContentViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayContentViewModelKosmos.kt
new file mode 100644
index 0000000..9cdd519
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeOverlayContentViewModelKosmos.kt
@@ -0,0 +1,32 @@
+/*
+ * 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.shade.ui.viewmodel
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeOverlayContentViewModel
+import com.android.systemui.scene.domain.interactor.sceneInteractor
+import com.android.systemui.statusbar.notification.stack.ui.viewmodel.notificationsPlaceholderViewModelFactory
+
+val Kosmos.notificationsShadeOverlayContentViewModel:
+    NotificationsShadeOverlayContentViewModel by Fixture {
+    NotificationsShadeOverlayContentViewModel(
+        shadeHeaderViewModelFactory = shadeHeaderViewModelFactory,
+        notificationsPlaceholderViewModelFactory = notificationsPlaceholderViewModelFactory,
+        sceneInteractor = sceneInteractor,
+    )
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeUserActionsViewModelKosmos.kt
similarity index 73%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeUserActionsViewModelKosmos.kt
index 9bf4756..6345c40 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeUserActionsViewModelKosmos.kt
@@ -18,12 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeSceneActionsViewModel
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeUserActionsViewModel
 
-val Kosmos.notificationsShadeSceneActionsViewModel:
-    NotificationsShadeSceneActionsViewModel by Fixture {
-    NotificationsShadeSceneActionsViewModel(
-        shadeInteractor = shadeInteractor,
-    )
-}
+val Kosmos.notificationsShadeUserActionsViewModel:
+    NotificationsShadeUserActionsViewModel by Fixture { NotificationsShadeUserActionsViewModel() }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeWindowModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeWindowModelKosmos.kt
index 4b42e07..9742595 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeWindowModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeWindowModelKosmos.kt
@@ -16,12 +16,18 @@
 
 package com.android.systemui.shade.ui.viewmodel
 
+import com.android.systemui.authentication.domain.interactor.authenticationInteractor
+import com.android.systemui.bouncer.domain.interactor.primaryBouncerInteractor
 import com.android.systemui.keyguard.domain.interactor.keyguardTransitionInteractor
 import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.scene.domain.interactor.sceneInteractor
 
 val Kosmos.notificationShadeWindowModel: NotificationShadeWindowModel by
     Kosmos.Fixture {
         NotificationShadeWindowModel(
             keyguardTransitionInteractor,
+            sceneInteractor = { sceneInteractor },
+            authenticationInteractor = { authenticationInteractor },
+            primaryBouncerInteractor = primaryBouncerInteractor,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt
deleted file mode 100644
index 00f1526..0000000
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/OverlayShadeViewModelKosmos.kt
+++ /dev/null
@@ -1,38 +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.shade.ui.viewmodel
-
-import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.scene.domain.interactor.sceneInteractor
-import com.android.systemui.shade.domain.interactor.shadeInteractor
-
-val Kosmos.overlayShadeViewModel: OverlayShadeViewModel by
-    Kosmos.Fixture {
-        OverlayShadeViewModel(
-            sceneInteractor = sceneInteractor,
-            shadeInteractor = shadeInteractor,
-        )
-    }
-
-val Kosmos.overlayShadeViewModelFactory: OverlayShadeViewModel.Factory by
-    Kosmos.Fixture {
-        object : OverlayShadeViewModel.Factory {
-            override fun create(): OverlayShadeViewModel {
-                return overlayShadeViewModel
-            }
-        }
-    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelKosmos.kt
similarity index 89%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelKosmos.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelKosmos.kt
index 2387aa8..48c5121 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/ShadeUserActionsViewModelKosmos.kt
@@ -21,8 +21,8 @@
 import com.android.systemui.qs.ui.adapter.qsSceneAdapter
 import com.android.systemui.shade.domain.interactor.shadeInteractor
 
-val Kosmos.shadeSceneActionsViewModel: ShadeSceneActionsViewModel by Fixture {
-    ShadeSceneActionsViewModel(
+val Kosmos.shadeUserActionsViewModel: ShadeUserActionsViewModel by Fixture {
+    ShadeUserActionsViewModel(
         qsSceneAdapter = qsSceneAdapter,
         shadeInteractor = shadeInteractor,
     )
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelKosmos.kt
similarity index 60%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelKosmos.kt
index 5d70ed6..c0d65a0 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ron/demo/ui/viewmodel/DemoRonChipViewModelKosmos.kt
@@ -14,15 +14,18 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.ui.viewmodel
+package com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel
 
+import android.content.packageManager
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.statusbar.commandline.commandRegistry
+import com.android.systemui.util.time.fakeSystemClock
 
-val Kosmos.quickSettingsShadeSceneActionsViewModel: QuickSettingsShadeSceneActionsViewModel by
+val Kosmos.demoRonChipViewModel: DemoRonChipViewModel by
     Kosmos.Fixture {
-        QuickSettingsShadeSceneActionsViewModel(
-            shadeInteractor = shadeInteractor,
-            quickSettingsContainerViewModel = quickSettingsContainerViewModel,
+        DemoRonChipViewModel(
+            commandRegistry = commandRegistry,
+            packageManager = packageManager,
+            systemClock = fakeSystemClock,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelKosmos.kt
index 16e288f..5382c1c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/chips/ui/viewmodel/OngoingActivityChipsViewModelKosmos.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.statusbar.chips.call.ui.viewmodel.callChipViewModel
 import com.android.systemui.statusbar.chips.casttootherdevice.ui.viewmodel.castToOtherDeviceChipViewModel
+import com.android.systemui.statusbar.chips.ron.demo.ui.viewmodel.demoRonChipViewModel
 import com.android.systemui.statusbar.chips.screenrecord.ui.viewmodel.screenRecordChipViewModel
 import com.android.systemui.statusbar.chips.sharetoapp.ui.viewmodel.shareToAppChipViewModel
 import com.android.systemui.statusbar.chips.statusBarChipsLogger
@@ -32,6 +33,7 @@
             shareToAppChipViewModel = shareToAppChipViewModel,
             castToOtherDeviceChipViewModel = castToOtherDeviceChipViewModel,
             callChipViewModel = callChipViewModel,
+            demoRonChipViewModel = demoRonChipViewModel,
             logger = statusBarChipsLogger,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/commandline/CommandRegistryKosmos.kt
similarity index 64%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/commandline/CommandRegistryKosmos.kt
index 5d70ed6..14777b4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/commandline/CommandRegistryKosmos.kt
@@ -14,15 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.ui.viewmodel
+package com.android.systemui.statusbar.commandline
 
+import android.content.applicationContext
 import com.android.systemui.kosmos.Kosmos
-import com.android.systemui.shade.domain.interactor.shadeInteractor
 
-val Kosmos.quickSettingsShadeSceneActionsViewModel: QuickSettingsShadeSceneActionsViewModel by
+val Kosmos.commandRegistry: CommandRegistry by
     Kosmos.Fixture {
-        QuickSettingsShadeSceneActionsViewModel(
-            shadeInteractor = shadeInteractor,
-            quickSettingsContainerViewModel = quickSettingsContainerViewModel,
+        CommandRegistry(
+            context = applicationContext,
+            // Immediately run anything that comes in
+            mainExecutor = { command -> command.run() },
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
index 4dd3ae7..2eb1573 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
@@ -35,7 +35,9 @@
 import com.android.systemui.flags.FakeFeatureFlagsClassic
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.media.controls.util.MediaFeatureFlag
 import com.android.systemui.media.dialog.MediaOutputDialogManager
+import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.plugins.statusbar.StatusBarStateController
 import com.android.systemui.shared.system.ActivityManagerWrapper
 import com.android.systemui.shared.system.DevicePolicyManagerWrapper
@@ -46,6 +48,7 @@
 import com.android.systemui.statusbar.RankingBuilder
 import com.android.systemui.statusbar.SmartReplyController
 import com.android.systemui.statusbar.notification.ColorUpdateLogger
+import com.android.systemui.statusbar.notification.ConversationNotificationManager
 import com.android.systemui.statusbar.notification.ConversationNotificationProcessor
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
@@ -69,6 +72,7 @@
 import com.android.systemui.statusbar.notification.row.shared.NotificationRowContentBinderRefactor
 import com.android.systemui.statusbar.notification.stack.NotificationChildrenContainerLogger
 import com.android.systemui.statusbar.phone.KeyguardBypassController
+import com.android.systemui.statusbar.phone.KeyguardDismissUtil
 import com.android.systemui.statusbar.policy.HeadsUpManager
 import com.android.systemui.statusbar.policy.SmartActionInflaterImpl
 import com.android.systemui.statusbar.policy.SmartReplyConstants
@@ -84,6 +88,7 @@
 import com.android.systemui.wmshell.BubblesManager
 import java.util.Optional
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.Executor
 import java.util.concurrent.TimeUnit
 import kotlinx.coroutines.test.TestScope
 import org.junit.Assert.assertTrue
@@ -128,19 +133,19 @@
         dependency.injectMockDependency(NotificationShadeWindowController::class.java)
         dependency.injectMockDependency(MediaOutputDialogManager::class.java)
 
-        mMockLogger = Mockito.mock(ExpandableNotificationRowLogger::class.java)
-        mStatusBarStateController = Mockito.mock(StatusBarStateController::class.java)
-        mKeyguardBypassController = Mockito.mock(KeyguardBypassController::class.java)
+        mMockLogger = Mockito.mock(ExpandableNotificationRowLogger::class.java, STUB_ONLY)
+        mStatusBarStateController = Mockito.mock(StatusBarStateController::class.java, STUB_ONLY)
+        mKeyguardBypassController = Mockito.mock(KeyguardBypassController::class.java, STUB_ONLY)
         mGroupMembershipManager = GroupMembershipManagerImpl()
-        mSmartReplyController = Mockito.mock(SmartReplyController::class.java)
+        mSmartReplyController = Mockito.mock(SmartReplyController::class.java, STUB_ONLY)
 
         val dumpManager = DumpManager()
         mGroupExpansionManager = GroupExpansionManagerImpl(dumpManager, mGroupMembershipManager)
-        mHeadsUpManager = Mockito.mock(HeadsUpManager::class.java)
+        mHeadsUpManager = Mockito.mock(HeadsUpManager::class.java, STUB_ONLY)
         mIconManager =
             IconManager(
-                Mockito.mock(CommonNotifCollection::class.java),
-                Mockito.mock(LauncherApps::class.java),
+                Mockito.mock(CommonNotifCollection::class.java, STUB_ONLY),
+                Mockito.mock(LauncherApps::class.java, STUB_ONLY),
                 IconBuilder(context),
                 mTestScope,
                 mBgCoroutineContext,
@@ -173,7 +178,7 @@
                 }
             )
         val remoteViewsFactories = getNotifRemoteViewsFactoryContainer(featureFlags)
-        val remoteInputManager = Mockito.mock(NotificationRemoteInputManager::class.java)
+        val remoteInputManager = Mockito.mock(NotificationRemoteInputManager::class.java, STUB_ONLY)
         val smartReplyStateInflater =
             SmartReplyStateInflaterImpl(
                 constants = mSmartReplyConstants,
@@ -183,7 +188,8 @@
                 smartRepliesInflater =
                     SmartReplyInflaterImpl(
                         constants = mSmartReplyConstants,
-                        keyguardDismissUtil = mock(),
+                        keyguardDismissUtil =
+                            Mockito.mock(KeyguardDismissUtil::class.java, STUB_ONLY),
                         remoteInputManager = remoteInputManager,
                         smartReplyController = mSmartReplyController,
                         context = context
@@ -191,7 +197,7 @@
                 smartActionsInflater =
                     SmartActionInflaterImpl(
                         constants = mSmartReplyConstants,
-                        activityStarter = mock(),
+                        activityStarter = Mockito.mock(ActivityStarter::class.java, STUB_ONLY),
                         smartReplyController = mSmartReplyController,
                         headsUpManager = mHeadsUpManager
                     )
@@ -206,41 +212,42 @@
             }
         val conversationProcessor =
             ConversationNotificationProcessor(
-                mock(),
-                mock(),
+                Mockito.mock(LauncherApps::class.java, STUB_ONLY),
+                Mockito.mock(ConversationNotificationManager::class.java, STUB_ONLY),
             )
+
         mContentBinder =
             if (NotificationRowContentBinderRefactor.isEnabled)
                 NotificationRowContentBinderImpl(
-                    mock(),
+                    Mockito.mock(NotifRemoteViewCache::class.java, STUB_ONLY),
                     remoteInputManager,
                     conversationProcessor,
-                    mock(),
-                    mock(),
-                    mock(),
+                    Mockito.mock(RichOngoingNotificationContentExtractor::class.java, STUB_ONLY),
+                    Mockito.mock(RichOngoingNotificationViewInflater::class.java, STUB_ONLY),
+                    Mockito.mock(Executor::class.java, STUB_ONLY),
                     smartReplyStateInflater,
                     notifLayoutInflaterFactoryProvider,
-                    mock(),
-                    mock(),
+                    Mockito.mock(HeadsUpStyleProvider::class.java, STUB_ONLY),
+                    Mockito.mock(NotificationRowContentBinderLogger::class.java, STUB_ONLY),
                 )
             else
                 NotificationContentInflater(
-                    mock(),
+                    Mockito.mock(NotifRemoteViewCache::class.java, STUB_ONLY),
                     remoteInputManager,
                     conversationProcessor,
-                    mock(),
-                    mock(),
+                    Mockito.mock(MediaFeatureFlag::class.java, STUB_ONLY),
+                    Mockito.mock(Executor::class.java, STUB_ONLY),
                     smartReplyStateInflater,
                     notifLayoutInflaterFactoryProvider,
-                    mock(),
-                    mock(),
+                    Mockito.mock(HeadsUpStyleProvider::class.java, STUB_ONLY),
+                    Mockito.mock(NotificationRowContentBinderLogger::class.java, STUB_ONLY),
                 )
         mContentBinder.setInflateSynchronously(true)
         mBindStage =
             RowContentBindStage(
                 mContentBinder,
-                mock(),
-                mock(),
+                Mockito.mock(NotifInflationErrorManager::class.java, STUB_ONLY),
+                Mockito.mock(RowContentBindStageLogger::class.java, STUB_ONLY),
             )
 
         val collection = Mockito.mock(CommonNotifCollection::class.java)
@@ -248,7 +255,7 @@
         mBindPipeline =
             NotifBindPipeline(
                 collection,
-                Mockito.mock(NotifBindPipelineLogger::class.java),
+                Mockito.mock(NotifBindPipelineLogger::class.java, STUB_ONLY),
                 NotificationEntryProcessorFactoryExecutorImpl(mMainExecutor),
             )
         mBindPipeline.setStage(mBindStage)
@@ -256,9 +263,11 @@
         val collectionListenerCaptor = ArgumentCaptor.forClass(NotifCollectionListener::class.java)
         Mockito.verify(collection).addCollectionListener(collectionListenerCaptor.capture())
         mBindPipelineEntryListener = collectionListenerCaptor.value
-        mPeopleNotificationIdentifier = Mockito.mock(PeopleNotificationIdentifier::class.java)
+        mPeopleNotificationIdentifier =
+            Mockito.mock(PeopleNotificationIdentifier::class.java, STUB_ONLY)
         mOnUserInteractionCallback = Mockito.mock(OnUserInteractionCallback::class.java)
-        mDismissibilityProvider = Mockito.mock(NotificationDismissibilityProvider::class.java)
+        mDismissibilityProvider =
+            Mockito.mock(NotificationDismissibilityProvider::class.java, STUB_ONLY)
         val mFutureDismissalRunnable = Mockito.mock(Runnable::class.java)
         whenever(
                 mOnUserInteractionCallback.registerFutureDismissal(
@@ -320,7 +329,10 @@
         //  set, but we do not want to override an existing value that is needed by a specific test.
 
         val rowInflaterTask =
-            RowInflaterTask(mFakeSystemClock, Mockito.mock(RowInflaterTaskLogger::class.java))
+            RowInflaterTask(
+                mFakeSystemClock,
+                Mockito.mock(RowInflaterTaskLogger::class.java, STUB_ONLY)
+            )
         val row = rowInflaterTask.inflateSynchronously(context, null, entry)
 
         entry.row = row
@@ -329,7 +341,7 @@
         mBindPipeline.manageRow(entry, row)
         row.initialize(
             entry,
-            Mockito.mock(RemoteInputViewSubcomponent.Factory::class.java),
+            Mockito.mock(RemoteInputViewSubcomponent.Factory::class.java, STUB_ONLY),
             APP_NAME,
             entry.key,
             mMockLogger,
@@ -338,23 +350,23 @@
             mGroupExpansionManager,
             mHeadsUpManager,
             mBindStage,
-            Mockito.mock(OnExpandClickListener::class.java),
-            Mockito.mock(CoordinateOnClickListener::class.java),
+            Mockito.mock(OnExpandClickListener::class.java, STUB_ONLY),
+            Mockito.mock(CoordinateOnClickListener::class.java, STUB_ONLY),
             FalsingManagerFake(),
             mStatusBarStateController,
             mPeopleNotificationIdentifier,
             mOnUserInteractionCallback,
-            Optional.of(Mockito.mock(BubblesManager::class.java)),
-            Mockito.mock(NotificationGutsManager::class.java),
+            Optional.of(Mockito.mock(BubblesManager::class.java, STUB_ONLY)),
+            Mockito.mock(NotificationGutsManager::class.java, STUB_ONLY),
             mDismissibilityProvider,
-            Mockito.mock(MetricsLogger::class.java),
-            Mockito.mock(NotificationChildrenContainerLogger::class.java),
-            Mockito.mock(ColorUpdateLogger::class.java),
+            Mockito.mock(MetricsLogger::class.java, STUB_ONLY),
+            Mockito.mock(NotificationChildrenContainerLogger::class.java, STUB_ONLY),
+            Mockito.mock(ColorUpdateLogger::class.java, STUB_ONLY),
             mSmartReplyConstants,
             mSmartReplyController,
             featureFlags,
-            Mockito.mock(IStatusBarService::class.java),
-            Mockito.mock(UiEventLogger::class.java)
+            Mockito.mock(IStatusBarService::class.java, STUB_ONLY),
+            Mockito.mock(UiEventLogger::class.java, STUB_ONLY)
         )
         row.setAboveShelfChangedListener { aboveShelf: Boolean -> }
         mBindStage.getStageParams(entry).requireContentViews(extraInflationFlags)
@@ -381,6 +393,8 @@
         private val Notification.isConversationStyleNotification
             get() = extras.getBoolean(IS_CONVERSATION_FLAG, false)
 
+        private val STUB_ONLY = Mockito.withSettings().stubOnly()
+
         fun markAsConversation(builder: Notification.Builder) {
             builder.addExtras(bundleOf(IS_CONVERSATION_FLAG to true))
         }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModelKosmos.kt
index 634354b..8bfc390 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/stack/ui/viewmodel/NotificationsPlaceholderViewModelKosmos.kt
@@ -35,3 +35,11 @@
         dumpManager = dumpManager,
     )
 }
+
+val Kosmos.notificationsPlaceholderViewModelFactory by Fixture {
+    object : NotificationsPlaceholderViewModel.Factory {
+        override fun create(): NotificationsPlaceholderViewModel {
+            return notificationsPlaceholderViewModel
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/phone/ManagedProfileControllerKosmos.kt
similarity index 62%
copy from packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
copy to packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/phone/ManagedProfileControllerKosmos.kt
index 9bf4756..ef04b9d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/shade/ui/viewmodel/NotificationsShadeSceneActionsViewModelKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/phone/ManagedProfileControllerKosmos.kt
@@ -14,16 +14,14 @@
  * limitations under the License.
  */
 
-package com.android.systemui.shade.ui.viewmodel
+@file:OptIn(ExperimentalCoroutinesApi::class)
 
+package com.android.systemui.statusbar.phone
+
+import android.testing.LeakCheck
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.notifications.ui.viewmodel.NotificationsShadeSceneActionsViewModel
-import com.android.systemui.shade.domain.interactor.shadeInteractor
+import com.android.systemui.utils.leaks.FakeManagedProfileController
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 
-val Kosmos.notificationsShadeSceneActionsViewModel:
-    NotificationsShadeSceneActionsViewModel by Fixture {
-    NotificationsShadeSceneActionsViewModel(
-        shadeInteractor = shadeInteractor,
-    )
-}
+val Kosmos.fakeManagedProfileController by Fixture { FakeManagedProfileController(LeakCheck()) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorKosmos.kt
index 386c7c5..d76defe 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/airplane/domain/interactor/AirplaneModeInteractorKosmos.kt
@@ -18,7 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.fakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.data.repository.FakeConnectivityRepository
 
 val Kosmos.airplaneModeInteractor: AirplaneModeInteractor by
@@ -26,6 +26,6 @@
         AirplaneModeInteractor(
             FakeAirplaneModeRepository(),
             FakeConnectivityRepository(),
-            FakeMobileConnectionsRepository(),
+            fakeMobileConnectionsRepository,
         )
     }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/FakeMobileConnectionsRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/FakeMobileConnectionsRepository.kt
index 8229575..de73d33 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/FakeMobileConnectionsRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/FakeMobileConnectionsRepository.kt
@@ -23,18 +23,16 @@
 import com.android.settingslib.mobile.MobileMappings
 import com.android.settingslib.mobile.TelephonyIcons
 import com.android.systemui.log.table.TableLogBuffer
-import com.android.systemui.statusbar.pipeline.mobile.data.model.ServiceStateModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.SubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy
 import com.android.systemui.statusbar.pipeline.mobile.util.MobileMappingsProxy
-import com.android.systemui.util.mockito.mock
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 
 // TODO(b/261632894): remove this in favor of the real impl or DemoMobileConnectionsRepository
 class FakeMobileConnectionsRepository(
     mobileMappings: MobileMappingsProxy = FakeMobileMappingsProxy(),
-    val tableLogBuffer: TableLogBuffer = mock<TableLogBuffer> {},
+    val tableLogBuffer: TableLogBuffer,
 ) : MobileConnectionsRepository {
 
     val GSM_KEY = mobileMappings.toIconKey(GSM)
@@ -94,9 +92,10 @@
     private val _defaultMobileIconGroup = MutableStateFlow(DEFAULT_ICON)
     override val defaultMobileIconGroup = _defaultMobileIconGroup
 
-    override val deviceServiceState = MutableStateFlow<ServiceStateModel?>(null)
+    override val isDeviceEmergencyCallCapable = MutableStateFlow(false)
 
     override val isAnySimSecure = MutableStateFlow(false)
+
     override fun getIsAnySimSecure(): Boolean = isAnySimSecure.value
 
     private var isInEcmMode: Boolean = false
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt
index 9d62ea5..cd22f1d 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryKosmos.kt
@@ -18,10 +18,12 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
-import com.android.systemui.util.mockito.mock
+import com.android.systemui.log.table.logcatTableLogBuffer
 
 val Kosmos.fakeMobileConnectionsRepository by Fixture {
-    FakeMobileConnectionsRepository(tableLogBuffer = mock())
+    FakeMobileConnectionsRepository(
+        tableLogBuffer = logcatTableLogBuffer(this, "FakeMobileConnectionsRepository"),
+    )
 }
 
 val Kosmos.mobileConnectionsRepository by
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
index 66be7e7..61b53c9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorKosmos.kt
@@ -17,8 +17,10 @@
 package com.android.systemui.statusbar.policy.domain.interactor
 
 import android.content.testableContext
+import com.android.settingslib.notification.modes.zenIconLoader
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
+import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.shared.notifications.data.repository.notificationSettingsRepository
 import com.android.systemui.statusbar.policy.data.repository.zenModeRepository
 
@@ -27,5 +29,7 @@
         context = testableContext,
         zenModeRepository = zenModeRepository,
         notificationSettingsRepository = notificationSettingsRepository,
+        bgDispatcher = testDispatcher,
+        iconLoader = zenIconLoader,
     )
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
index 35fa2af..73d423c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeGlobalSettingsKosmos.kt
@@ -19,5 +19,10 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.unconfinedTestDispatcher
 
 val Kosmos.fakeGlobalSettings: FakeGlobalSettings by Fixture { FakeGlobalSettings(testDispatcher) }
+
+val Kosmos.unconfinedDispatcherFakeGlobalSettings: FakeGlobalSettings by Fixture {
+    FakeGlobalSettings(unconfinedTestDispatcher)
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
index 76ef202..e1daf9b 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/util/settings/FakeSettingsKosmos.kt
@@ -19,8 +19,13 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.Kosmos.Fixture
 import com.android.systemui.kosmos.testDispatcher
+import com.android.systemui.kosmos.unconfinedTestDispatcher
 import com.android.systemui.settings.userTracker
 
 val Kosmos.fakeSettings: FakeSettings by Fixture {
     FakeSettings(testDispatcher) { userTracker.userId }
 }
+
+val Kosmos.unconfinedDispatcherFakeSettings: FakeSettings by Fixture {
+    FakeSettings(unconfinedTestDispatcher) { userTracker.userId }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
index 2dbac67..0089199 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
@@ -64,7 +64,8 @@
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage, int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription) {
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape) {
     }
 
     @Override
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerAdapterKosmos.kt
similarity index 79%
rename from packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerAdapterKosmos.kt
index d60f14c..4045135b9 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerCollectorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/VolumeControllerAdapterKosmos.kt
@@ -18,6 +18,7 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.volume.data.repository.audioRepository
 
-val Kosmos.volumeControllerCollector by
-    Kosmos.Fixture { VolumeControllerCollector(applicationCoroutineScope) }
+val Kosmos.volumeControllerAdapter by
+    Kosmos.Fixture { VolumeControllerAdapter(applicationCoroutineScope, audioRepository) }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioRepository.kt
index 135cb14..1fa6c3f 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/data/repository/FakeAudioRepository.kt
@@ -18,12 +18,16 @@
 
 import android.media.AudioDeviceInfo
 import android.media.AudioManager
+import com.android.settingslib.volume.data.model.VolumeControllerEvent
 import com.android.settingslib.volume.data.repository.AudioRepository
 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 kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.update
 
@@ -39,10 +43,26 @@
     override val communicationDevice: StateFlow<AudioDeviceInfo?> =
         mutableCommunicationDevice.asStateFlow()
 
+    private val mutableVolumeControllerEvents = MutableSharedFlow<VolumeControllerEvent>(replay = 1)
+    override val volumeControllerEvents: Flow<VolumeControllerEvent>
+        get() = mutableVolumeControllerEvents.asSharedFlow()
+
     private val models: MutableMap<AudioStream, MutableStateFlow<AudioStreamModel>> = mutableMapOf()
     private val lastAudibleVolumes: MutableMap<AudioStream, Int> = mutableMapOf()
     private val deviceCategories: MutableMap<String, Int> = mutableMapOf()
 
+    private val mutableIsVolumeControllerVisible = MutableStateFlow(false)
+    val isVolumeControllerVisible: StateFlow<Boolean>
+        get() = mutableIsVolumeControllerVisible.asStateFlow()
+
+    private var mutableIsInitialized: Boolean = false
+    val isInitialized: Boolean
+        get() = mutableIsInitialized
+
+    override fun init() {
+        mutableIsInitialized = true
+    }
+
     private fun getAudioStreamModelState(
         audioStream: AudioStream
     ): MutableStateFlow<AudioStreamModel> =
@@ -111,4 +131,16 @@
     override suspend fun getBluetoothAudioDeviceCategory(bluetoothAddress: String): Int {
         return deviceCategories[bluetoothAddress] ?: AudioManager.AUDIO_DEVICE_CATEGORY_UNKNOWN
     }
+
+    suspend fun sendVolumeControllerEvent(event: VolumeControllerEvent) {
+        if (isInitialized) {
+            mutableVolumeControllerEvents.emit(event)
+        }
+    }
+
+    override suspend fun notifyVolumeControllerVisible(isVisible: Boolean) {
+        if (isInitialized) {
+            mutableIsVolumeControllerVisible.value = isVisible
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/TestMediaDevicesFactory.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/TestMediaDevicesFactory.kt
index 83adc79..ad1292e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/TestMediaDevicesFactory.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/TestMediaDevicesFactory.kt
@@ -38,6 +38,7 @@
     ): MediaDevice = mock {
         whenever(name).thenReturn(deviceName)
         whenever(icon).thenReturn(deviceIcon)
+        whenever(deviceType).thenReturn(MediaDevice.MediaDeviceType.TYPE_PHONE_DEVICE)
     }
 
     fun wiredMediaDevice(
@@ -77,6 +78,18 @@
             whenever(name).thenReturn(deviceName)
             whenever(icon).thenReturn(deviceIcon)
             whenever(cachedDevice).thenReturn(cachedBluetoothDevice)
+            whenever(deviceType).thenReturn(MediaDevice.MediaDeviceType.TYPE_BLUETOOTH_DEVICE)
+        }
+    }
+
+    fun remoteMediaDevice(
+        deviceName: String = "remote_media",
+        deviceIcon: Drawable? = TestStubDrawable(),
+    ): MediaDevice {
+        return mock<MediaDevice> {
+            whenever(name).thenReturn(deviceName)
+            whenever(icon).thenReturn(deviceIcon)
+            whenever(deviceType).thenReturn(MediaDevice.MediaDeviceType.TYPE_CAST_DEVICE)
         }
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractorKosmos.kt
index 63386d0..dd5bbf3 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/volume/panel/component/volume/domain/interactor/AudioSlidersInteractorKosmos.kt
@@ -18,7 +18,6 @@
 
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.applicationCoroutineScope
-import com.android.systemui.volume.data.repository.audioRepository
 import com.android.systemui.volume.domain.interactor.audioModeInteractor
 import com.android.systemui.volume.mediaOutputInteractor
 
@@ -27,7 +26,6 @@
         AudioSlidersInteractor(
             applicationCoroutineScope,
             mediaOutputInteractor,
-            audioRepository,
             audioModeInteractor,
         )
     }
diff --git a/ravenwood/.gitignore b/ravenwood/.gitignore
new file mode 100644
index 0000000..751553b
--- /dev/null
+++ b/ravenwood/.gitignore
@@ -0,0 +1 @@
+*.bak
diff --git a/ravenwood/Android.bp b/ravenwood/Android.bp
index be4cd76..9b0c8e5 100644
--- a/ravenwood/Android.bp
+++ b/ravenwood/Android.bp
@@ -160,6 +160,7 @@
         "ravenwood-framework",
         "services.core.ravenwood",
         "junit",
+        "framework-annotations-lib",
     ],
     sdk_version: "core_current",
     visibility: ["//frameworks/base"],
diff --git a/ravenwood/OWNERS b/ravenwood/OWNERS
index badfca0..f7b13d1 100644
--- a/ravenwood/OWNERS
+++ b/ravenwood/OWNERS
@@ -1,8 +1,8 @@
 set noparent
 
-jsharkey@google.com
-omakoto@google.com
-dplotnikov@google.com
+topjohnwu@google.com
+hackbod@google.com  #{LAST_RESORT_SUGGESTION}
+
 
 per-file ravenwood-annotation-allowed-classes.txt = dplotnikov@google.com
 per-file texts/ravenwood-annotation-allowed-classes.txt = dplotnikov@google.com
diff --git a/ravenwood/TEST_MAPPING b/ravenwood/TEST_MAPPING
index 7e2ee3e..469759b 100644
--- a/ravenwood/TEST_MAPPING
+++ b/ravenwood/TEST_MAPPING
@@ -1,17 +1,12 @@
-// Keep the following two TEST_MAPPINGs in sync:
-// frameworks/base/ravenwood/TEST_MAPPING
-// frameworks/base/tools/hoststubgen/TEST_MAPPING
 {
   "presubmit": [
     { "name": "tiny-framework-dump-test" },
     { "name": "hoststubgentest" },
+    { "name": "hoststubgen-test-tiny-test" },
     { "name": "hoststubgen-invoke-test" },
-    {
-      "name": "RavenwoodMockitoTest_device"
-    },
-    {
-      "name": "RavenwoodBivalentTest_device"
-    },
+    { "name": "RavenwoodMockitoTest_device" },
+    { "name": "RavenwoodBivalentTest_device" },
+
     // The sysui tests should match vendor/unbundled_google/packages/SystemUIGoogle/TEST_MAPPING
     {
       "name": "SystemUIGoogleTests",
@@ -39,6 +34,113 @@
     }
   ],
   "ravenwood-presubmit": [
+    // AUTO-GENERATED-START
+    // DO NOT MODIFY MANUALLY
+    // Use scripts/update-test-mapping.sh to update it.
+    {
+      "name": "AdServicesSharedLibrariesUnitTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "android.test.mock.ravenwood.tests",
+      "host": true
+    },
+    {
+      "name": "CarLibHostUnitTest",
+      "host": true
+    },
+    {
+      "name": "CarServiceHostUnitTest",
+      "host": true
+    },
+    {
+      "name": "CarSystemUIRavenTests",
+      "host": true
+    },
+    {
+      "name": "CtsAccountManagerTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsAppTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsContentTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsDatabaseTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsGraphicsTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsIcuTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsInputMethodTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsOsTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsProtoTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsResourcesTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsTextTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "CtsUtilTestCasesRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksCoreSystemPropertiesTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksCoreTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksInputMethodSystemServerTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksMockingServicesTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksServicesTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "FrameworksUtilTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "InternalTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "PowerStatsTestsRavenwood",
+      "host": true
+    },
+    {
+      "name": "RavenwoodBivalentTest",
+      "host": true
+    },
     {
       "name": "RavenwoodMinimumTest",
       "host": true
@@ -48,16 +150,21 @@
       "host": true
     },
     {
-      "name": "CtsUtilTestCasesRavenwood",
-      "host": true
-    },
-    {
       "name": "RavenwoodResApkTest",
       "host": true
     },
     {
-      "name": "RavenwoodBivalentTest",
+      "name": "RavenwoodRuntimeTest",
+      "host": true
+    },
+    {
+      "name": "RavenwoodServicesTest",
+      "host": true
+    },
+    {
+      "name": "SystemUiRavenTests",
       "host": true
     }
+    // AUTO-GENERATED-END
   ]
 }
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodClassRuleDeviceOnlyTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodClassRuleDeviceOnlyTest.java
index 3a24c0e..e8f59db 100644
--- a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodClassRuleDeviceOnlyTest.java
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/RavenwoodClassRuleDeviceOnlyTest.java
@@ -26,6 +26,10 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+/**
+ * Test to ensure @DisabledOnRavenwood works. Note, now the DisabledOnRavenwood annotation
+ * is handled by the test runner, so it won't really need the class rule.
+ */
 @RunWith(AndroidJUnit4.class)
 @DisabledOnRavenwood
 public class RavenwoodClassRuleDeviceOnlyTest {
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/CallTracker.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/CallTracker.java
index 8dadd39..09a0aa8 100644
--- a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/CallTracker.java
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/CallTracker.java
@@ -64,7 +64,7 @@
     /**
      * Check the number of calls stored in {@link #mNumCalled}.
      */
-    protected void assertCalls(Object... methodNameAndCountPairs) {
+    public void assertCalls(Object... methodNameAndCountPairs) {
         // Create a local copy
         HashMap<String, Integer> counts = new HashMap<>(mNumCalled);
         for (int i = 0; i < methodNameAndCountPairs.length - 1; i += 2) {
@@ -95,7 +95,7 @@
      * Same as {@link #assertCalls(Object...)} but it kills the process if it fails.
      * Only use in @AfterClass.
      */
-    protected void assertCallsOrDie(Object... methodNameAndCountPairs) {
+    public void assertCallsOrDie(Object... methodNameAndCountPairs) {
         try {
             assertCalls(methodNameAndCountPairs);
         } catch (Throwable th) {
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodImplicitClassRuleDeviceOnlyTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodImplicitClassRuleDeviceOnlyTest.java
index 0f8be0e..7ef672e 100644
--- a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodImplicitClassRuleDeviceOnlyTest.java
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodImplicitClassRuleDeviceOnlyTest.java
@@ -34,7 +34,12 @@
 
     @BeforeClass
     public static void beforeClass() {
-        Assert.assertFalse(RavenwoodRule.isOnRavenwood());
+        // This method shouldn't be called -- unless RUN_DISABLED_TESTS is enabled.
+
+        // If we're doing RUN_DISABLED_TESTS, don't throw here, because that'd confuse junit.
+        if (!RavenwoodRule.private$ravenwood().isRunningDisabledTests()) {
+            Assert.assertFalse(RavenwoodRule.isOnRavenwood());
+        }
     }
 
     @Test
@@ -46,7 +51,10 @@
     public static void afterClass() {
         if (RavenwoodRule.isOnRavenwood()) {
             Log.e(TAG, "Even @AfterClass shouldn't be executed!");
-            System.exit(1);
+
+            if (!RavenwoodRule.private$ravenwood().isRunningDisabledTests()) {
+                System.exit(1);
+            }
         }
     }
 }
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodNoRavenizerTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodNoRavenizerTest.java
new file mode 100644
index 0000000..9d878f4
--- /dev/null
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodNoRavenizerTest.java
@@ -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.ravenwoodtest.bivalenttest.ravenizer;
+
+import android.platform.test.annotations.NoRavenizer;
+import android.platform.test.ravenwood.RavenwoodAwareTestRunner.RavenwoodTestRunnerInitializing;
+
+import org.junit.Test;
+
+/**
+ * Test for {@link android.platform.test.annotations.NoRavenizer}
+ */
+@NoRavenizer
+public class RavenwoodNoRavenizerTest {
+    public static final String TAG = "RavenwoodNoRavenizerTest";
+
+    private static final CallTracker sCallTracker = new CallTracker();
+
+    /**
+     * With @NoRavenizer, this method shouldn't be called.
+     */
+    @RavenwoodTestRunnerInitializing
+    public static void ravenwoodRunnerInitializing() {
+        sCallTracker.incrementMethodCallCount();
+    }
+
+    /**
+     * Make sure ravenwoodRunnerInitializing() wasn't called.
+     */
+    @Test
+    public void testNotRavenized() {
+        sCallTracker.assertCalls(
+                "ravenwoodRunnerInitializing", 0
+        );
+    }
+}
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsReallyDisabledTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsReallyDisabledTest.java
new file mode 100644
index 0000000..c77841b
--- /dev/null
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsReallyDisabledTest.java
@@ -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.ravenwoodtest.bivalenttest.ravenizer;
+
+import static org.junit.Assert.fail;
+
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.ravenwood.RavenwoodAwareTestRunner.RavenwoodTestRunnerInitializing;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.Log;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.AfterClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/**
+ * Test for "RAVENWOOD_RUN_DISABLED_TESTS" with "REALLY_DISABLED" set.
+ *
+ * This test is only executed on Ravenwood.
+ */
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodRunDisabledTestsReallyDisabledTest {
+    private static final String TAG = "RavenwoodRunDisabledTestsTest";
+
+    private static final CallTracker sCallTracker = new CallTracker();
+
+    @RavenwoodTestRunnerInitializing
+    public static void ravenwoodRunnerInitializing() {
+        RavenwoodRule.private$ravenwood().overrideRunDisabledTest(true,
+                "\\#testReallyDisabled$");
+    }
+
+    /**
+     * This test gets to run with RAVENWOOD_RUN_DISABLED_TESTS set.
+     */
+    @Test
+    @DisabledOnRavenwood
+    public void testDisabledTestGetsToRun() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        sCallTracker.incrementMethodCallCount();
+
+        fail("This test won't pass on Ravenwood.");
+    }
+
+    /**
+     * This will still not be executed due to the "really disabled" pattern.
+     */
+    @Test
+    @DisabledOnRavenwood
+    public void testReallyDisabled() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        sCallTracker.incrementMethodCallCount();
+
+        fail("This test won't pass on Ravenwood.");
+    }
+
+    @AfterClass
+    public static void afterClass() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        Log.i(TAG, "afterClass called");
+
+        sCallTracker.assertCallsOrDie(
+                "testDisabledTestGetsToRun", 1,
+                "testReallyDisabled", 0
+        );
+    }
+}
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsTest.java
new file mode 100644
index 0000000..ea1a29d
--- /dev/null
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodRunDisabledTestsTest.java
@@ -0,0 +1,87 @@
+/*
+ * 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.ravenwoodtest.bivalenttest.ravenizer;
+
+import static org.junit.Assert.fail;
+
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.ravenwood.RavenwoodAwareTestRunner.RavenwoodTestRunnerInitializing;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.Log;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import org.junit.AfterClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.runner.RunWith;
+
+/**
+ * Test for "RAVENWOOD_RUN_DISABLED_TESTS". (with no "REALLY_DISABLED" set.)
+ *
+ * This test is only executed on Ravenwood.
+ */
+@RunWith(AndroidJUnit4.class)
+public class RavenwoodRunDisabledTestsTest {
+    private static final String TAG = "RavenwoodRunDisabledTestsTest";
+
+    @Rule
+    public ExpectedException mExpectedException = ExpectedException.none();
+
+    private static final CallTracker sCallTracker = new CallTracker();
+
+    @RavenwoodTestRunnerInitializing
+    public static void ravenwoodRunnerInitializing() {
+        RavenwoodRule.private$ravenwood().overrideRunDisabledTest(true, null);
+    }
+
+    @Test
+    @DisabledOnRavenwood
+    public void testDisabledTestGetsToRun() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        sCallTracker.incrementMethodCallCount();
+
+        fail("This test won't pass on Ravenwood.");
+    }
+
+    @Test
+    @DisabledOnRavenwood
+    public void testDisabledButPass() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        sCallTracker.incrementMethodCallCount();
+
+        // When a @DisabledOnRavenwood actually passed, the runner should make fail().
+        mExpectedException.expectMessage("it actually passed under Ravenwood");
+    }
+
+    @AfterClass
+    public static void afterClass() {
+        if (!RavenwoodRule.isOnRavenwood()) {
+            return;
+        }
+        Log.i(TAG, "afterClass called");
+
+        sCallTracker.assertCallsOrDie(
+                "testDisabledTestGetsToRun", 1,
+                "testDisabledButPass", 1
+        );
+    }
+}
diff --git a/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodSuiteTest.java b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodSuiteTest.java
new file mode 100644
index 0000000..7e396c2
--- /dev/null
+++ b/ravenwood/bivalenttest/test/com/android/ravenwoodtest/bivalenttest/ravenizer/RavenwoodSuiteTest.java
@@ -0,0 +1,69 @@
+/*
+ * 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.ravenwoodtest.bivalenttest.ravenizer;
+
+import android.util.Log;
+
+import org.junit.AfterClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Suite;
+
+/**
+ * Test to make sure {@link Suite} works with the ravenwood test runner.
+ */
+@RunWith(Suite.class)
+@Suite.SuiteClasses({
+        RavenwoodSuiteTest.Test1.class,
+        RavenwoodSuiteTest.Test2.class
+})
+public class RavenwoodSuiteTest {
+    public static final String TAG = "RavenwoodSuiteTest";
+
+    private static final CallTracker sCallTracker = new CallTracker();
+
+    @AfterClass
+    public static void afterClass() {
+        Log.i(TAG, "afterClass called");
+
+        sCallTracker.assertCallsOrDie(
+                "test1", 1,
+                "test2", 1
+        );
+    }
+
+    /**
+     * Workaround for the issue where tradefed won't think a class is a test class
+     * if it has a @RunWith but no @Test methods, even if it is a Suite.
+     */
+    @Test
+    public void testEmpty() {
+    }
+
+    public static class Test1 {
+        @Test
+        public void test1() {
+            sCallTracker.incrementMethodCallCount();
+        }
+    }
+
+    public static class Test2 {
+        @Test
+        public void test2() {
+            sCallTracker.incrementMethodCallCount();
+        }
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
index 03600ad..f237ba9 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
@@ -17,14 +17,16 @@
 
 import static com.android.ravenwood.common.RavenwoodCommonUtils.RAVENWOOD_VERSION_JAVA_SYSPROP;
 
+import static org.junit.Assert.fail;
+
 import android.os.Bundle;
 import android.platform.test.ravenwood.RavenwoodAwareTestRunner.Order;
 import android.platform.test.ravenwood.RavenwoodAwareTestRunner.Scope;
+import android.platform.test.ravenwood.RavenwoodTestStats.Result;
+import android.util.Log;
 
 import androidx.test.platform.app.InstrumentationRegistry;
 
-import com.android.ravenwood.common.RavenwoodCommonUtils;
-
 import org.junit.runner.Description;
 import org.junit.runner.Runner;
 import org.junit.runners.model.TestClass;
@@ -38,12 +40,25 @@
     private RavenwoodAwareTestRunnerHook() {
     }
 
-    private static void log(String message) {
-        RavenwoodCommonUtils.log(TAG, message);
+    private static RavenwoodTestStats sStats; // lazy initialization.
+    private static Description sCurrentClassDescription;
+
+    private static RavenwoodTestStats getStats() {
+        if (sStats == null) {
+            // We don't want to throw in the static initializer, because tradefed may not report
+            // it properly, so we initialize it here.
+            sStats = new RavenwoodTestStats();
+        }
+        return sStats;
     }
 
+    /**
+     * Called when a runner starts, before the inner runner gets a chance to run.
+     */
     public static void onRunnerInitializing(Runner runner, TestClass testClass) {
-        log("onRunnerStart: testClass=" + testClass + " runner=" + runner);
+        // This log call also ensures the framework JNI is loaded.
+        Log.i(TAG, "onRunnerInitializing: testClass=" + testClass.getJavaClass()
+                + " runner=" + runner);
 
         // TODO: Move the initialization code to a better place.
 
@@ -52,26 +67,97 @@
                 "androidx.test.internal.runner.junit4.AndroidJUnit4ClassRunner");
         System.setProperty(RAVENWOOD_VERSION_JAVA_SYSPROP, "1");
 
+
         // This is needed to make AndroidJUnit4ClassRunner happy.
         InstrumentationRegistry.registerInstance(null, Bundle.EMPTY);
     }
 
+    /**
+     * Called when a whole test class is skipped.
+     */
+    public static void onClassSkipped(Description description) {
+        Log.i(TAG, "onClassSkipped: description=" + description);
+        getStats().onClassSkipped(description);
+    }
+
+    /**
+     * Called before a test / class.
+     *
+     * Return false if it should be skipped.
+     */
     public static boolean onBefore(RavenwoodAwareTestRunner runner, Description description,
             Scope scope, Order order) {
-        log("onBefore: description=" + description + ", " + scope + ", " + order);
+        Log.i(TAG, "onBefore: description=" + description + ", " + scope + ", " + order);
+
+        if (scope == Scope.Class && order == Order.First) {
+            // Keep track of the current class.
+            sCurrentClassDescription = description;
+        }
 
         // Class-level annotations are checked by the runner already, so we only check
         // method-level annotations here.
         if (scope == Scope.Instance && order == Order.First) {
-            if (!RavenwoodRule.shouldEnableOnRavenwood(description)) {
+            if (!RavenwoodEnablementChecker.shouldEnableOnRavenwood(
+                    description, true)) {
+                getStats().onTestFinished(sCurrentClassDescription, description, Result.Skipped);
                 return false;
             }
         }
         return true;
     }
 
-    public static void onAfter(RavenwoodAwareTestRunner runner, Description description,
+    /**
+     * Called after a test / class.
+     *
+     * Return false if the exception should be ignored.
+     */
+    public static boolean onAfter(RavenwoodAwareTestRunner runner, Description description,
             Scope scope, Order order, Throwable th) {
-        log("onAfter: description=" + description + ", " + scope + ", " + order + ", " + th);
+        Log.i(TAG, "onAfter: description=" + description + ", " + scope + ", " + order + ", " + th);
+
+        if (scope == Scope.Instance && order == Order.First) {
+            getStats().onTestFinished(sCurrentClassDescription, description,
+                    th == null ? Result.Passed : Result.Failed);
+
+        } else if (scope == Scope.Class && order == Order.Last) {
+            getStats().onClassFinished(sCurrentClassDescription);
+        }
+
+        // If RUN_DISABLED_TESTS is set, and the method did _not_ throw, make it an error.
+        if (RavenwoodRule.private$ravenwood().isRunningDisabledTests()
+                && scope == Scope.Instance && order == Order.First) {
+
+            boolean isTestEnabled = RavenwoodEnablementChecker.shouldEnableOnRavenwood(
+                    description, false);
+            if (th == null) {
+                // Test passed. Is the test method supposed to be enabled?
+                if (isTestEnabled) {
+                    // Enabled and didn't throw, okay.
+                    return true;
+                } else {
+                    // Disabled and didn't throw. We should report it.
+                    fail("Test wasn't included under Ravenwood, but it actually "
+                            + "passed under Ravenwood; consider updating annotations");
+                    return true; // unreachable.
+                }
+            } else {
+                // Test failed.
+                if (isTestEnabled) {
+                    // Enabled but failed. We should throw the exception.
+                    return true;
+                } else {
+                    // Disabled and failed. Expected. Don't throw.
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Called by {@link RavenwoodAwareTestRunner} to see if it should run a test class or not.
+     */
+    public static boolean shouldRunClassOnRavenwood(Class<?> clazz) {
+        return RavenwoodEnablementChecker.shouldRunClassOnRavenwood(clazz, true);
     }
 }
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodEnablementChecker.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodEnablementChecker.java
new file mode 100644
index 0000000..77275c4
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodEnablementChecker.java
@@ -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 android.platform.test.ravenwood;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.platform.test.annotations.DisabledOnRavenwood;
+import android.platform.test.annotations.EnabledOnRavenwood;
+import android.platform.test.annotations.IgnoreUnderRavenwood;
+
+import org.junit.runner.Description;
+
+/**
+ * Calculates which tests need to be executed on Ravenwood.
+ */
+public class RavenwoodEnablementChecker {
+    private static final String TAG = "RavenwoodDisablementChecker";
+
+    private RavenwoodEnablementChecker() {
+    }
+
+    /**
+     * Determine if the given {@link Description} should be enabled when running on the
+     * Ravenwood test environment.
+     *
+     * A more specific method-level annotation always takes precedence over any class-level
+     * annotation, and an {@link EnabledOnRavenwood} annotation always takes precedence over
+     * an {@link DisabledOnRavenwood} annotation.
+     */
+    public static boolean shouldEnableOnRavenwood(Description description,
+            boolean takeIntoAccountRunDisabledTestsFlag) {
+        // First, consult any method-level annotations
+        if (description.isTest()) {
+            Boolean result = null;
+
+            // Stopgap for http://g/ravenwood/EPAD-N5ntxM
+            if (description.getMethodName().endsWith("$noRavenwood")) {
+                result = false;
+            } else if (description.getAnnotation(EnabledOnRavenwood.class) != null) {
+                result = true;
+            } else if (description.getAnnotation(DisabledOnRavenwood.class) != null) {
+                result = false;
+            } else if (description.getAnnotation(IgnoreUnderRavenwood.class) != null) {
+                result = false;
+            }
+            if (result != null) {
+                if (takeIntoAccountRunDisabledTestsFlag
+                        && RavenwoodRule.private$ravenwood().isRunningDisabledTests()) {
+                    result = !shouldStillIgnoreInProbeIgnoreMode(
+                            description.getTestClass(), description.getMethodName());
+                }
+            }
+            if (result != null) {
+                return result;
+            }
+        }
+
+        // Otherwise, consult any class-level annotations
+        return shouldRunClassOnRavenwood(description.getTestClass(),
+                takeIntoAccountRunDisabledTestsFlag);
+    }
+
+    public static boolean shouldRunClassOnRavenwood(@NonNull Class<?> testClass,
+            boolean takeIntoAccountRunDisabledTestsFlag) {
+        boolean result = true;
+        if (testClass.getAnnotation(EnabledOnRavenwood.class) != null) {
+            result = true;
+        } else if (testClass.getAnnotation(DisabledOnRavenwood.class) != null) {
+            result = false;
+        } else if (testClass.getAnnotation(IgnoreUnderRavenwood.class) != null) {
+            result = false;
+        }
+        if (!result) {
+            if (takeIntoAccountRunDisabledTestsFlag
+                    && RavenwoodRule.private$ravenwood().isRunningDisabledTests()) {
+                result = !shouldStillIgnoreInProbeIgnoreMode(testClass, null);
+            }
+        }
+        return result;
+    }
+
+    /**
+     * Check if a test should _still_ disabled even if {@code RUN_DISABLED_TESTS}
+     * is true, using {@code REALLY_DISABLED_PATTERN}.
+     *
+     * This only works on tests, not on classes.
+     */
+    static boolean shouldStillIgnoreInProbeIgnoreMode(
+            @NonNull Class<?> testClass, @Nullable String methodName) {
+        if (RavenwoodRule.private$ravenwood().getReallyDisabledPattern().pattern().isEmpty()) {
+            return false;
+        }
+
+        final var fullname = testClass.getName() + (methodName != null ? "#" + methodName : "");
+
+        System.out.println("XXX=" + fullname);
+
+        if (RavenwoodRule.private$ravenwood().getReallyDisabledPattern().matcher(fullname).find()) {
+            System.out.println("Still ignoring " + fullname);
+            return true;
+        }
+        return false;
+    }
+}
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
index 7b4c173..a2088fd 100644
--- a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
@@ -40,7 +40,6 @@
 import com.android.server.LocalServices;
 
 import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
 
 import java.io.File;
 import java.io.IOException;
@@ -226,11 +225,6 @@
         }
     }
 
-    public static void validate(Statement base, Description description,
-            boolean enableOptionalValidation) {
-        // Nothing to check, for now.
-    }
-
     /**
      * Set the current configuration to the actual SystemProperties.
      */
diff --git a/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodTestStats.java b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodTestStats.java
new file mode 100644
index 0000000..631f68f
--- /dev/null
+++ b/ravenwood/junit-impl-src/android/platform/test/ravenwood/RavenwoodTestStats.java
@@ -0,0 +1,163 @@
+/*
+ * 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.test.ravenwood;
+
+import android.util.Log;
+
+import org.junit.runner.Description;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.time.LocalDateTime;
+import java.time.format.DateTimeFormatter;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Creats a "stats" CSV file containing the test results.
+ *
+ * The output file is created as `/tmp/Ravenwood-stats_[TEST-MODULE=NAME]_[TIMESTAMP].csv`.
+ * A symlink to the latest result will be created as
+ * `/tmp/Ravenwood-stats_[TEST-MODULE=NAME]_latest.csv`.
+ */
+public class RavenwoodTestStats {
+    private static final String TAG = "RavenwoodTestStats";
+    private static final String HEADER = "Module,Class,ClassDesc,Passed,Failed,Skipped";
+
+    public enum Result {
+        Passed,
+        Failed,
+        Skipped,
+    }
+
+    private final File mOutputFile;
+    private final PrintWriter mOutputWriter;
+    private final String mTestModuleName;
+
+    public final Map<Description, Map<Description, Result>> mStats = new HashMap<>();
+
+    /** Ctor */
+    public RavenwoodTestStats() {
+        mTestModuleName = guessTestModuleName();
+
+        var basename = "Ravenwood-stats_" + mTestModuleName + "_";
+
+        // Get the current time
+        LocalDateTime now = LocalDateTime.now();
+        DateTimeFormatter fmt = DateTimeFormatter.ofPattern("yyyyMMdd-HHmmss");
+
+        var tmpdir = System.getProperty("java.io.tmpdir");
+        mOutputFile = new File(tmpdir, basename + now.format(fmt) + ".csv");
+
+        try {
+            mOutputWriter = new PrintWriter(mOutputFile);
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to crete logfile. File=" + mOutputFile, e);
+        }
+
+        // Crete the "latest" symlink.
+        Path symlink = Paths.get(tmpdir, basename + "latest.csv");
+        try {
+            if (Files.exists(symlink)) {
+                Files.delete(symlink);
+            }
+            Files.createSymbolicLink(symlink, Paths.get(mOutputFile.getName()));
+
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to crete logfile. File=" + mOutputFile, e);
+        }
+
+        Log.i(TAG, "Test result stats file: " + mOutputFile);
+
+        // Print the header.
+        mOutputWriter.println(HEADER);
+        mOutputWriter.flush();
+    }
+
+    private String guessTestModuleName() {
+        // Assume the current directory name is the test module name.
+        File cwd;
+        try {
+            cwd = new File(".").getCanonicalFile();
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to get the current directory", e);
+        }
+        return cwd.getName();
+    }
+
+    private void addResult(Description classDescription, Description methodDescription,
+            Result result) {
+        mStats.compute(classDescription, (classDesc, value) -> {
+            if (value == null) {
+                value = new HashMap<>();
+            }
+            value.put(methodDescription, result);
+            return value;
+        });
+    }
+
+    public void onClassSkipped(Description classDescription) {
+        addResult(classDescription, Description.EMPTY, Result.Skipped);
+        onClassFinished(classDescription);
+    }
+
+    public void onTestFinished(Description classDescription, Description testDescription,
+            Result result) {
+        addResult(classDescription, testDescription, result);
+    }
+
+    public void onClassFinished(Description classDescription) {
+        int passed = 0;
+        int skipped = 0;
+        int failed = 0;
+        for (var e : mStats.get(classDescription).values()) {
+            switch (e) {
+                case Passed: passed++; break;
+                case Skipped: skipped++; break;
+                case Failed: failed++; break;
+            }
+        }
+
+        var testClass = extractTestClass(classDescription);
+
+        mOutputWriter.printf("%s,%s,%s,%d,%d,%d\n",
+                mTestModuleName, (testClass == null ? "?" : testClass.getCanonicalName()),
+                classDescription, passed, failed, skipped);
+        mOutputWriter.flush();
+    }
+
+    /**
+     * Try to extract the class from a description, which is needed because
+     * ParameterizedAndroidJunit4's description doesn't contain a class.
+     */
+    private Class<?> extractTestClass(Description desc) {
+        if (desc.getTestClass() != null) {
+            return desc.getTestClass();
+        }
+        // Look into the children.
+        for (var child : desc.getChildren()) {
+            var fromChild = extractTestClass(child);
+            if (fromChild != null) {
+                return fromChild;
+            }
+        }
+        return null;
+    }
+}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/LargeTest.java b/ravenwood/junit-src/android/platform/test/annotations/NoRavenizer.java
similarity index 62%
rename from tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/LargeTest.java
rename to ravenwood/junit-src/android/platform/test/annotations/NoRavenizer.java
index 76bbcad..a84f16f 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/LargeTest.java
+++ b/ravenwood/junit-src/android/platform/test/annotations/NoRavenizer.java
@@ -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.
@@ -13,16 +13,24 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.hoststubgen.test.tinyframework;
+
+package android.platform.test.annotations;
 
 import java.lang.annotation.ElementType;
+import java.lang.annotation.Inherited;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 /**
- * We don't want to use any android classes in this module, so we create our own copy of it here.
+ * Disable the ravenizer preprocessor for a class. This should be only used for testing
+ * ravenizer itself, or to workaround issues with the preprocessor. A test class probably won't run
+ * properly if it's not preprocessed.
+ *
+ * @hide
  */
+@Inherited
+@Target({ElementType.TYPE})
 @Retention(RetentionPolicy.RUNTIME)
-@Target({ElementType.METHOD, ElementType.TYPE})
-public @interface LargeTest {}
+public @interface NoRavenizer {
+}
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
index a4fa41a..7d99166 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodAwareTestRunner.java
@@ -15,18 +15,19 @@
  */
 package android.platform.test.ravenwood;
 
-import static android.platform.test.ravenwood.RavenwoodRule.shouldRunCassOnRavenwood;
-
 import static com.android.ravenwood.common.RavenwoodCommonUtils.ensureIsPublicVoidMethod;
 import static com.android.ravenwood.common.RavenwoodCommonUtils.isOnRavenwood;
 
 import static java.lang.annotation.ElementType.METHOD;
 import static java.lang.annotation.ElementType.TYPE;
 
+import android.util.Log;
+
 import com.android.ravenwood.common.RavenwoodCommonUtils;
 import com.android.ravenwood.common.SneakyThrow;
 
 import org.junit.Assume;
+import org.junit.internal.builders.AllDefaultPossibilitiesBuilder;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runner.Runner;
@@ -38,8 +39,10 @@
 import org.junit.runner.manipulation.Orderer;
 import org.junit.runner.manipulation.Sortable;
 import org.junit.runner.manipulation.Sorter;
+import org.junit.runner.notification.Failure;
 import org.junit.runner.notification.RunNotifier;
 import org.junit.runners.BlockJUnit4ClassRunner;
+import org.junit.runners.model.RunnerBuilder;
 import org.junit.runners.model.Statement;
 import org.junit.runners.model.TestClass;
 
@@ -53,8 +56,6 @@
 /**
  * A test runner used for Ravenwood.
  *
- * TODO: Handle ENABLE_PROBE_IGNORED
- *
  * It will delegate to another runner specified with {@link InnerRunner}
  * (default = {@link BlockJUnit4ClassRunner}) with the following features.
  * - Add a {@link RavenwoodAwareTestRunnerHook#onRunnerInitializing} hook, which is called before
@@ -136,12 +137,15 @@
         return runner;
     }
 
-    private final TestClass mTestClsas;
-    private final Runner mRealRunner;
+    private TestClass mTestClass = null;
+    private Runner mRealRunner = null;
+    private Description mDescription = null;
+    private Throwable mExceptionInConstructor = null;
 
     /** Simple logging method. */
     private void log(String message) {
-        RavenwoodCommonUtils.log(TAG, "[" + getTestClass() + "  @" + this + "] " + message);
+        RavenwoodCommonUtils.log(TAG, "[" + getTestClass().getJavaClass() + "  @" + this + "] "
+                + message);
     }
 
     private Error logAndFail(String message, Throwable innerException) {
@@ -151,44 +155,76 @@
     }
 
     public TestClass getTestClass() {
-        return mTestClsas;
+        return mTestClass;
     }
 
     /**
      * Constructor.
      */
     public RavenwoodAwareTestRunner(Class<?> testClass) {
-        mTestClsas = new TestClass(testClass);
-
-        /*
-         * If the class has @DisabledOnRavenwood, then we'll delegate to ClassSkippingTestRunner,
-         * which simply skips it.
-         */
-        if (isOnRavenwood() && !shouldRunCassOnRavenwood(mTestClsas.getJavaClass())) {
-            mRealRunner = new ClassSkippingTestRunner(mTestClsas);
-            return;
-        }
-
-        // Find the real runner.
-        final Class<? extends Runner> realRunner;
-        final InnerRunner innerRunnerAnnotation = mTestClsas.getAnnotation(InnerRunner.class);
-        if (innerRunnerAnnotation != null) {
-            realRunner = innerRunnerAnnotation.value();
-        } else {
-            // Default runner.
-            realRunner = BlockJUnit4ClassRunner.class;
-        }
-
-        onRunnerInitializing();
-
         try {
-            log("Initializing the inner runner: " + realRunner);
+            mTestClass = new TestClass(testClass);
 
-            mRealRunner = realRunner.getConstructor(Class.class).newInstance(testClass);
+            /*
+             * If the class has @DisabledOnRavenwood, then we'll delegate to
+             * ClassSkippingTestRunner, which simply skips it.
+             */
+            if (isOnRavenwood() && !RavenwoodAwareTestRunnerHook.shouldRunClassOnRavenwood(
+                    mTestClass.getJavaClass())) {
+                mRealRunner = new ClassSkippingTestRunner(mTestClass);
+                mDescription = mRealRunner.getDescription();
+                return;
+            }
 
-        } catch (InstantiationException | IllegalAccessException
-                 | InvocationTargetException | NoSuchMethodException e) {
-            throw logAndFail("Failed to instantiate " + realRunner, e);
+            // Find the real runner.
+            final Class<? extends Runner> realRunnerClass;
+            final InnerRunner innerRunnerAnnotation = mTestClass.getAnnotation(InnerRunner.class);
+            if (innerRunnerAnnotation != null) {
+                realRunnerClass = innerRunnerAnnotation.value();
+            } else {
+                // Default runner.
+                realRunnerClass = BlockJUnit4ClassRunner.class;
+            }
+
+            onRunnerInitializing();
+
+            try {
+                log("Initializing the inner runner: " + realRunnerClass);
+
+                mRealRunner = instantiateRealRunner(realRunnerClass, testClass);
+                mDescription = mRealRunner.getDescription();
+
+            } catch (InstantiationException | IllegalAccessException
+                     | InvocationTargetException | NoSuchMethodException e) {
+                throw logAndFail("Failed to instantiate " + realRunnerClass, e);
+            }
+        } catch (Throwable th) {
+            // If we throw in the constructor, Tradefed may not report it and just ignore the class,
+            // so record it and throw it when the test actually started.
+            log("Fatal: Exception detected in constructor: " + th.getMessage() + "\n"
+                    + Log.getStackTraceString(th));
+            mExceptionInConstructor = new RuntimeException("Exception detected in constructor",
+                    th);
+            mDescription = Description.createTestDescription(testClass, "Constructor");
+
+            // This is for testing if tradefed is fixed.
+            if ("1".equals(System.getenv("RAVENWOOD_THROW_EXCEPTION_IN_TEST_RUNNER"))) {
+                throw th;
+            }
+        }
+    }
+
+    private static Runner instantiateRealRunner(
+            Class<? extends Runner> realRunnerClass,
+            Class<?> testClass)
+            throws NoSuchMethodException, InvocationTargetException, InstantiationException,
+            IllegalAccessException {
+        try {
+            return realRunnerClass.getConstructor(Class.class).newInstance(testClass);
+        } catch (NoSuchMethodException e) {
+            var runnerBuilder = new AllDefaultPossibilitiesBuilder();
+            return realRunnerClass.getConstructor(Class.class,
+                    RunnerBuilder.class).newInstance(testClass, runnerBuilder);
         }
     }
 
@@ -203,7 +239,7 @@
 
         log("onRunnerInitializing");
 
-        RavenwoodAwareTestRunnerHook.onRunnerInitializing(this, mTestClsas);
+        RavenwoodAwareTestRunnerHook.onRunnerInitializing(this, mTestClass);
 
         // Hook point to allow more customization.
         runAnnotatedMethodsOnRavenwood(RavenwoodTestRunnerInitializing.class, null);
@@ -231,13 +267,18 @@
 
     @Override
     public Description getDescription() {
-        return mRealRunner.getDescription();
+        return mDescription;
     }
 
     @Override
     public void run(RunNotifier notifier) {
         if (mRealRunner instanceof ClassSkippingTestRunner) {
             mRealRunner.run(notifier);
+            RavenwoodAwareTestRunnerHook.onClassSkipped(getDescription());
+            return;
+        }
+
+        if (maybeReportExceptionFromConstructor(notifier)) {
             return;
         }
 
@@ -250,6 +291,18 @@
         }
     }
 
+    /** Throw the exception detected in the constructor, if any. */
+    private boolean maybeReportExceptionFromConstructor(RunNotifier notifier) {
+        if (mExceptionInConstructor == null) {
+            return false;
+        }
+        notifier.fireTestStarted(mDescription);
+        notifier.fireTestFailure(new Failure(mDescription, mExceptionInConstructor));
+        notifier.fireTestFinished(mDescription);
+
+        return true;
+    }
+
     @Override
     public void filter(Filter filter) throws NoTestsRemainException {
         if (mRealRunner instanceof Filterable r) {
@@ -294,19 +347,23 @@
     }
 
     private void runWithHooks(Description description, Scope scope, Order order, Statement s) {
-        Throwable th = null;
         if (isOnRavenwood()) {
             Assume.assumeTrue(
                     RavenwoodAwareTestRunnerHook.onBefore(this, description, scope, order));
         }
         try {
             s.evaluate();
-        } catch (Throwable t) {
-            th = t;
-            SneakyThrow.sneakyThrow(t);
-        } finally {
             if (isOnRavenwood()) {
-                RavenwoodAwareTestRunnerHook.onAfter(this, description, scope, order, th);
+                RavenwoodAwareTestRunnerHook.onAfter(this, description, scope, order, null);
+            }
+        } catch (Throwable t) {
+            boolean shouldThrow = true;
+            if (isOnRavenwood()) {
+                shouldThrow = RavenwoodAwareTestRunnerHook.onAfter(
+                        this, description, scope, order, t);
+            }
+            if (shouldThrow) {
+                SneakyThrow.sneakyThrow(t);
             }
         }
     }
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodClassRule.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodClassRule.java
index 6c8d96a..85297fe 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodClassRule.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodClassRule.java
@@ -16,37 +16,20 @@
 
 package android.platform.test.ravenwood;
 
-import static android.platform.test.ravenwood.RavenwoodRule.ENABLE_PROBE_IGNORED;
-import static android.platform.test.ravenwood.RavenwoodRule.IS_ON_RAVENWOOD;
-import static android.platform.test.ravenwood.RavenwoodRule.shouldEnableOnRavenwood;
-import static android.platform.test.ravenwood.RavenwoodRule.shouldStillIgnoreInProbeIgnoreMode;
-
-import android.platform.test.annotations.DisabledOnRavenwood;
-import android.platform.test.annotations.EnabledOnRavenwood;
-
-import org.junit.Assume;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
 /**
- * {@code @ClassRule} that respects Ravenwood-specific class annotations. This rule has no effect
- * when tests are run on non-Ravenwood test environments.
+ * No longer needed.
  *
- * By default, all tests are executed on Ravenwood, but annotations such as
- * {@link DisabledOnRavenwood} and {@link EnabledOnRavenwood} can be used at both the method
- * and class level to "ignore" tests that may not be ready.
+ * @deprecated this class used to be used to handle the class level annotation, which
+ * is now done by the test runner, so this class is not needed.
  */
+@Deprecated
 public class RavenwoodClassRule implements TestRule {
     @Override
     public Statement apply(Statement base, Description description) {
-        if (!IS_ON_RAVENWOOD) {
-            // No check on a real device.
-        } else if (ENABLE_PROBE_IGNORED) {
-            Assume.assumeFalse(shouldStillIgnoreInProbeIgnoreMode(description));
-        } else {
-            Assume.assumeTrue(shouldEnableOnRavenwood(description));
-        }
         return base;
     }
 }
diff --git a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
index 75faafb..d569896 100644
--- a/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
+++ b/ravenwood/junit-src/android/platform/test/ravenwood/RavenwoodRule.java
@@ -20,22 +20,20 @@
 import static android.os.Process.SYSTEM_UID;
 import static android.os.UserHandle.SYSTEM;
 
-import static org.junit.Assert.fail;
+import static com.android.ravenwood.common.RavenwoodCommonUtils.log;
 
+import android.annotation.Nullable;
 import android.app.Instrumentation;
 import android.content.Context;
 import android.platform.test.annotations.DisabledOnRavenwood;
 import android.platform.test.annotations.EnabledOnRavenwood;
-import android.platform.test.annotations.IgnoreUnderRavenwood;
 
 import com.android.ravenwood.common.RavenwoodCommonUtils;
 
-import org.junit.Assume;
 import org.junit.rules.TestRule;
 import org.junit.runner.Description;
 import org.junit.runners.model.Statement;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
@@ -56,18 +54,18 @@
  * before a test class is fully initialized.
  */
 public class RavenwoodRule implements TestRule {
+    private static final String TAG = "RavenwoodRule";
+
     static final boolean IS_ON_RAVENWOOD = RavenwoodCommonUtils.isOnRavenwood();
 
     /**
-     * When probing is enabled, all tests will be unconditionally run on Ravenwood to detect
+     * When this flag is enabled, all tests will be unconditionally run on Ravenwood to detect
      * cases where a test is able to pass despite being marked as {@link DisabledOnRavenwood}.
      *
      * This is typically helpful for internal maintainers discovering tests that had previously
      * been ignored, but now have enough Ravenwood-supported functionality to be enabled.
-     *
-     * TODO: Rename it to a more descriptive name.
      */
-    static final boolean ENABLE_PROBE_IGNORED = "1".equals(
+    private static final boolean RUN_DISABLED_TESTS = "1".equals(
             System.getenv("RAVENWOOD_RUN_DISABLED_TESTS"));
 
     /**
@@ -92,23 +90,17 @@
      *
      * Because we use a regex-find, setting "." would disable all tests.
      */
-    private static final Pattern REALLY_DISABLE_PATTERN = Pattern.compile(
-            Objects.requireNonNullElse(System.getenv("RAVENWOOD_REALLY_DISABLE"), ""));
+    private static final Pattern REALLY_DISABLED_PATTERN = Pattern.compile(
+            Objects.requireNonNullElse(System.getenv("RAVENWOOD_REALLY_DISABLED"), ""));
 
-    private static final boolean ENABLE_REALLY_DISABLE_PATTERN =
-            !REALLY_DISABLE_PATTERN.pattern().isEmpty();
-
-    /**
-     * If true, enable optional validation on running tests.
-     */
-    private static final boolean ENABLE_OPTIONAL_VALIDATION = "1".equals(
-            System.getenv("RAVENWOOD_OPTIONAL_VALIDATION"));
+    private static final boolean HAS_REALLY_DISABLE_PATTERN =
+            !REALLY_DISABLED_PATTERN.pattern().isEmpty();
 
     static {
-        if (ENABLE_PROBE_IGNORED) {
-            System.out.println("$RAVENWOOD_RUN_DISABLED_TESTS enabled: force running all tests");
-            if (ENABLE_REALLY_DISABLE_PATTERN) {
-                System.out.println("$RAVENWOOD_REALLY_DISABLE=" + REALLY_DISABLE_PATTERN.pattern());
+        if (RUN_DISABLED_TESTS) {
+            log(TAG, "$RAVENWOOD_RUN_DISABLED_TESTS enabled: force running all tests");
+            if (HAS_REALLY_DISABLE_PATTERN) {
+                log(TAG, "$RAVENWOOD_REALLY_DISABLED=" + REALLY_DISABLED_PATTERN.pattern());
             }
         }
     }
@@ -275,103 +267,18 @@
                 "Instrumentation is only available during @Test execution");
     }
 
-    /**
-     * Determine if the given {@link Description} should be enabled when running on the
-     * Ravenwood test environment.
-     *
-     * A more specific method-level annotation always takes precedence over any class-level
-     * annotation, and an {@link EnabledOnRavenwood} annotation always takes precedence over
-     * an {@link DisabledOnRavenwood} annotation.
-     */
-    public static boolean shouldEnableOnRavenwood(Description description) {
-        // First, consult any method-level annotations
-        if (description.isTest()) {
-            // Stopgap for http://g/ravenwood/EPAD-N5ntxM
-            if (description.getMethodName().endsWith("$noRavenwood")) {
-                return false;
-            }
-            if (description.getAnnotation(EnabledOnRavenwood.class) != null) {
-                return true;
-            }
-            if (description.getAnnotation(DisabledOnRavenwood.class) != null) {
-                return false;
-            }
-            if (description.getAnnotation(IgnoreUnderRavenwood.class) != null) {
-                return false;
-            }
-        }
-
-        // Otherwise, consult any class-level annotations
-        return shouldRunCassOnRavenwood(description.getTestClass());
-    }
-
-    public static boolean shouldRunCassOnRavenwood(Class<?> clazz) {
-        if (clazz != null) {
-            if (clazz.getAnnotation(EnabledOnRavenwood.class) != null) {
-                return true;
-            }
-            if (clazz.getAnnotation(DisabledOnRavenwood.class) != null) {
-                return false;
-            }
-            if (clazz.getAnnotation(IgnoreUnderRavenwood.class) != null) {
-                return false;
-            }
-        }
-        return true;
-    }
-
-    static boolean shouldStillIgnoreInProbeIgnoreMode(Description description) {
-        if (!ENABLE_REALLY_DISABLE_PATTERN) {
-            return false;
-        }
-
-        final var fullname = description.getTestClass().getName()
-                + (description.isTest() ? "#" + description.getMethodName() : "");
-
-        if (REALLY_DISABLE_PATTERN.matcher(fullname).find()) {
-            System.out.println("Still ignoring " + fullname);
-            return true;
-        }
-        return false;
-    }
 
     @Override
     public Statement apply(Statement base, Description description) {
-        // No special treatment when running outside Ravenwood; run tests as-is
-        if (!IS_ON_RAVENWOOD) {
-            return base;
-        }
-
-        if (ENABLE_PROBE_IGNORED) {
-            return applyProbeIgnored(base, description);
-        } else {
-            return applyDefault(base, description);
-        }
-    }
-
-    private void commonPrologue(Statement base, Description description) throws IOException {
-        RavenwoodRuleImpl.logTestRunner("started", description);
-        RavenwoodRuleImpl.validate(base, description, ENABLE_OPTIONAL_VALIDATION);
-        RavenwoodRuleImpl.init(RavenwoodRule.this);
-    }
-
-    /**
-     * Run the given {@link Statement} with no special treatment.
-     */
-    private Statement applyDefault(Statement base, Description description) {
+        // TODO: Here, we're calling init() / reset() once for each rule.
+        // That means if a test class has multiple rules -- even if they refer to the same
+        // rule instance -- we're calling them multiple times. We need to fix it.
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
-                Assume.assumeTrue(shouldEnableOnRavenwood(description));
-
-                commonPrologue(base, description);
+                RavenwoodRuleImpl.init(RavenwoodRule.this);
                 try {
                     base.evaluate();
-
-                    RavenwoodRuleImpl.logTestRunner("finished", description);
-                } catch (Throwable t) {
-                    RavenwoodRuleImpl.logTestRunner("failed", description);
-                    throw t;
                 } finally {
                     RavenwoodRuleImpl.reset(RavenwoodRule.this);
                 }
@@ -380,44 +287,6 @@
     }
 
     /**
-     * Run the given {@link Statement} with probing enabled. All tests will be unconditionally
-     * run on Ravenwood to detect cases where a test is able to pass despite being marked as
-     * {@code IgnoreUnderRavenwood}.
-     */
-    private Statement applyProbeIgnored(Statement base, Description description) {
-        return new Statement() {
-            @Override
-            public void evaluate() throws Throwable {
-                Assume.assumeFalse(shouldStillIgnoreInProbeIgnoreMode(description));
-
-                commonPrologue(base, description);
-                try {
-                    base.evaluate();
-                } catch (Throwable t) {
-                    // If the test isn't included, eat the exception and report the
-                    // assumption failure that test authors expect; otherwise throw
-                    Assume.assumeTrue(shouldEnableOnRavenwood(description));
-                    throw t;
-                } finally {
-                    RavenwoodRuleImpl.logTestRunner("finished", description);
-                    RavenwoodRuleImpl.reset(RavenwoodRule.this);
-                }
-
-                if (!shouldEnableOnRavenwood(description)) {
-                    fail("Test wasn't included under Ravenwood, but it actually "
-                            + "passed under Ravenwood; consider updating annotations");
-                }
-            }
-        };
-    }
-
-    public static class _$RavenwoodPrivate {
-        public static boolean isOptionalValidationEnabled() {
-            return ENABLE_OPTIONAL_VALIDATION;
-        }
-    }
-
-    /**
      * Returns the "real" result from {@link System#currentTimeMillis()}.
      *
      * Currently, it's the same thing as calling {@link System#currentTimeMillis()},
@@ -427,4 +296,47 @@
     public long realCurrentTimeMillis() {
         return System.currentTimeMillis();
     }
+
+    // Below are internal to ravenwood. Don't use them from normal tests...
+
+    public static class RavenwoodPrivate {
+        private RavenwoodPrivate() {
+        }
+
+        private volatile Boolean mRunDisabledTestsOverride = null;
+
+        private volatile Pattern mReallyDisabledPattern = null;
+
+        public boolean isRunningDisabledTests() {
+            if (mRunDisabledTestsOverride != null) {
+                return mRunDisabledTestsOverride;
+            }
+            return RUN_DISABLED_TESTS;
+        }
+
+        public Pattern getReallyDisabledPattern() {
+            if (mReallyDisabledPattern != null) {
+                return mReallyDisabledPattern;
+            }
+            return REALLY_DISABLED_PATTERN;
+        }
+
+        public void overrideRunDisabledTest(boolean runDisabledTests,
+                @Nullable String reallyDisabledPattern) {
+            mRunDisabledTestsOverride = runDisabledTests;
+            mReallyDisabledPattern =
+                    reallyDisabledPattern == null ? null : Pattern.compile(reallyDisabledPattern);
+        }
+
+        public void resetRunDisabledTest() {
+            mRunDisabledTestsOverride = null;
+            mReallyDisabledPattern = null;
+        }
+    }
+
+    private static final RavenwoodPrivate sRavenwoodPrivate = new  RavenwoodPrivate();
+
+    public static RavenwoodPrivate private$ravenwood() {
+        return sRavenwoodPrivate;
+    }
 }
diff --git a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
index 6b80e0c..1e4889c 100644
--- a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
+++ b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodAwareTestRunnerHook.java
@@ -23,27 +23,51 @@
 import org.junit.runners.model.TestClass;
 
 /**
- * Provide hook points created by {@link RavenwoodAwareTestRunner}.
+ * Provide hook points created by {@link RavenwoodAwareTestRunner}. This is a version
+ * that's used on a device side test.
+ *
+ * All methods are no-op in real device tests.
+ *
+ * TODO: Use some kind of factory to provide different implementation for the device test
+ * and the ravenwood test.
  */
 public class RavenwoodAwareTestRunnerHook {
     private RavenwoodAwareTestRunnerHook() {
     }
 
     /**
-     * Called when a runner starts, befre the inner runner gets a chance to run.
+     * Called when a runner starts, before the inner runner gets a chance to run.
      */
     public static void onRunnerInitializing(Runner runner, TestClass testClass) {
-        // No-op on a real device.
     }
 
+    /**
+     * Called when a whole test class is skipped.
+     */
+    public static void onClassSkipped(Description description) {
+    }
+
+    /**
+     * Called before a test / class.
+     *
+     * Return false if it should be skipped.
+     */
     public static boolean onBefore(RavenwoodAwareTestRunner runner, Description description,
             Scope scope, Order order) {
-        // No-op on a real device.
         return true;
     }
 
-    public static void onAfter(RavenwoodAwareTestRunner runner, Description description,
+    /**
+     * Called after a test / class.
+     *
+     * Return false if the exception should be ignored.
+     */
+    public static boolean onAfter(RavenwoodAwareTestRunner runner, Description description,
             Scope scope, Order order, Throwable th) {
-        // No-op on a real device.
+        return true;
+    }
+
+    public static boolean shouldRunClassOnRavenwood(Class<?> clazz) {
+        return true;
     }
 }
diff --git a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
index 483b98a..a470626 100644
--- a/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
+++ b/ravenwood/junit-stub-src/android/platform/test/ravenwood/RavenwoodRuleImpl.java
@@ -17,7 +17,6 @@
 package android.platform.test.ravenwood;
 
 import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
 
 public class RavenwoodRuleImpl {
     public static void init(RavenwoodRule rule) {
@@ -32,10 +31,6 @@
         // No-op when running on a real device
     }
 
-    public static void validate(Statement base, Description description,
-            boolean enableOptionalValidation) {
-    }
-
     public static long realCurrentTimeMillis() {
         return System.currentTimeMillis();
     }
diff --git a/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/RavenwoodEnvironment_host.java b/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/RavenwoodEnvironment_host.java
index f894b0e..58f6bbb 100644
--- a/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/RavenwoodEnvironment_host.java
+++ b/ravenwood/runtime-helper-src/framework/com/android/platform/test/ravenwood/nativesubstitution/RavenwoodEnvironment_host.java
@@ -36,7 +36,7 @@
     /**
      * Called from {@link RavenwoodEnvironment#ensureRavenwoodInitialized()}.
      */
-    public static void ensureRavenwoodInitialized() {
+    public static void nativeEnsureRavenwoodInitialized() {
 
         // TODO Unify it with the initialization code in RavenwoodAwareTestRunnerHook.
 
@@ -63,14 +63,14 @@
     /**
      * Called from {@link RavenwoodEnvironment#getRavenwoodRuntimePath()}.
      */
-    public static String getRavenwoodRuntimePath(RavenwoodEnvironment env) {
+    public static String nativeGetRavenwoodRuntimePath(RavenwoodEnvironment env) {
         return RavenwoodCommonUtils.getRavenwoodRuntimePath();
     }
 
     /**
      * Called from {@link RavenwoodEnvironment#fromAddress(long)}.
      */
-    public static <T> T fromAddress(RavenwoodEnvironment env, long address) {
+    public static <T> T nativeFromAddress(RavenwoodEnvironment env, long address) {
         return JvmWorkaround.getInstance().fromAddress(address);
     }
 }
diff --git a/ravenwood/runtime-helper-src/libcore-fake/libcore/util/FP16.java b/ravenwood/runtime-helper-src/libcore-fake/libcore/util/FP16.java
new file mode 100644
index 0000000..478503b
--- /dev/null
+++ b/ravenwood/runtime-helper-src/libcore-fake/libcore/util/FP16.java
@@ -0,0 +1,814 @@
+/*
+ * Copyright (C) 2019 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 libcore.util;
+
+/**
+ * <p>The {@code FP16} class is a wrapper and a utility class to manipulate half-precision 16-bit
+ * <a href="https://en.wikipedia.org/wiki/Half-precision_floating-point_format">IEEE 754</a>
+ * floating point data types (also called fp16 or binary16). A half-precision float can be
+ * created from or converted to single-precision floats, and is stored in a short data type.
+ *
+ * <p>The IEEE 754 standard specifies an fp16 as having the following format:</p>
+ * <ul>
+ * <li>Sign bit: 1 bit</li>
+ * <li>Exponent width: 5 bits</li>
+ * <li>Significand: 10 bits</li>
+ * </ul>
+ *
+ * <p>The format is laid out as follows:</p>
+ * <pre>
+ * 1   11111   1111111111
+ * ^   --^--   -----^----
+ * sign  |          |_______ significand
+ *       |
+ *       -- exponent
+ * </pre>
+ *
+ * <p>Half-precision floating points can be useful to save memory and/or
+ * bandwidth at the expense of range and precision when compared to single-precision
+ * floating points (fp32).</p>
+ * <p>To help you decide whether fp16 is the right storage type for you need, please
+ * refer to the table below that shows the available precision throughout the range of
+ * possible values. The <em>precision</em> column indicates the step size between two
+ * consecutive numbers in a specific part of the range.</p>
+ *
+ * <table summary="Precision of fp16 across the range">
+ *     <tr><th>Range start</th><th>Precision</th></tr>
+ *     <tr><td>0</td><td>1 &frasl; 16,777,216</td></tr>
+ *     <tr><td>1 &frasl; 16,384</td><td>1 &frasl; 16,777,216</td></tr>
+ *     <tr><td>1 &frasl; 8,192</td><td>1 &frasl; 8,388,608</td></tr>
+ *     <tr><td>1 &frasl; 4,096</td><td>1 &frasl; 4,194,304</td></tr>
+ *     <tr><td>1 &frasl; 2,048</td><td>1 &frasl; 2,097,152</td></tr>
+ *     <tr><td>1 &frasl; 1,024</td><td>1 &frasl; 1,048,576</td></tr>
+ *     <tr><td>1 &frasl; 512</td><td>1 &frasl; 524,288</td></tr>
+ *     <tr><td>1 &frasl; 256</td><td>1 &frasl; 262,144</td></tr>
+ *     <tr><td>1 &frasl; 128</td><td>1 &frasl; 131,072</td></tr>
+ *     <tr><td>1 &frasl; 64</td><td>1 &frasl; 65,536</td></tr>
+ *     <tr><td>1 &frasl; 32</td><td>1 &frasl; 32,768</td></tr>
+ *     <tr><td>1 &frasl; 16</td><td>1 &frasl; 16,384</td></tr>
+ *     <tr><td>1 &frasl; 8</td><td>1 &frasl; 8,192</td></tr>
+ *     <tr><td>1 &frasl; 4</td><td>1 &frasl; 4,096</td></tr>
+ *     <tr><td>1 &frasl; 2</td><td>1 &frasl; 2,048</td></tr>
+ *     <tr><td>1</td><td>1 &frasl; 1,024</td></tr>
+ *     <tr><td>2</td><td>1 &frasl; 512</td></tr>
+ *     <tr><td>4</td><td>1 &frasl; 256</td></tr>
+ *     <tr><td>8</td><td>1 &frasl; 128</td></tr>
+ *     <tr><td>16</td><td>1 &frasl; 64</td></tr>
+ *     <tr><td>32</td><td>1 &frasl; 32</td></tr>
+ *     <tr><td>64</td><td>1 &frasl; 16</td></tr>
+ *     <tr><td>128</td><td>1 &frasl; 8</td></tr>
+ *     <tr><td>256</td><td>1 &frasl; 4</td></tr>
+ *     <tr><td>512</td><td>1 &frasl; 2</td></tr>
+ *     <tr><td>1,024</td><td>1</td></tr>
+ *     <tr><td>2,048</td><td>2</td></tr>
+ *     <tr><td>4,096</td><td>4</td></tr>
+ *     <tr><td>8,192</td><td>8</td></tr>
+ *     <tr><td>16,384</td><td>16</td></tr>
+ *     <tr><td>32,768</td><td>32</td></tr>
+ * </table>
+ *
+ * <p>This table shows that numbers higher than 1024 lose all fractional precision.</p>
+ *
+ * @hide
+ */
+
+public final class FP16 {
+    /**
+     * The number of bits used to represent a half-precision float value.
+     *
+     * @hide
+     */
+    public static final int SIZE = 16;
+
+    /**
+     * Epsilon is the difference between 1.0 and the next value representable
+     * by a half-precision floating-point.
+     *
+     * @hide
+     */
+    public static final short EPSILON = (short) 0x1400;
+
+    /**
+     * Maximum exponent a finite half-precision float may have.
+     *
+     * @hide
+     */
+    public static final int MAX_EXPONENT = 15;
+    /**
+     * Minimum exponent a normalized half-precision float may have.
+     *
+     * @hide
+     */
+    public static final int MIN_EXPONENT = -14;
+
+    /**
+     * Smallest negative value a half-precision float may have.
+     *
+     * @hide
+     */
+    public static final short LOWEST_VALUE = (short) 0xfbff;
+    /**
+     * Maximum positive finite value a half-precision float may have.
+     *
+     * @hide
+     */
+    public static final short MAX_VALUE = (short) 0x7bff;
+    /**
+     * Smallest positive normal value a half-precision float may have.
+     *
+     * @hide
+     */
+    public static final short MIN_NORMAL = (short) 0x0400;
+    /**
+     * Smallest positive non-zero value a half-precision float may have.
+     *
+     * @hide
+     */
+    public static final short MIN_VALUE = (short) 0x0001;
+    /**
+     * A Not-a-Number representation of a half-precision float.
+     *
+     * @hide
+     */
+    public static final short NaN = (short) 0x7e00;
+    /**
+     * Negative infinity of type half-precision float.
+     *
+     * @hide
+     */
+    public static final short NEGATIVE_INFINITY = (short) 0xfc00;
+    /**
+     * Negative 0 of type half-precision float.
+     *
+     * @hide
+     */
+    public static final short NEGATIVE_ZERO = (short) 0x8000;
+    /**
+     * Positive infinity of type half-precision float.
+     *
+     * @hide
+     */
+    public static final short POSITIVE_INFINITY = (short) 0x7c00;
+    /**
+     * Positive 0 of type half-precision float.
+     *
+     * @hide
+     */
+    public static final short POSITIVE_ZERO = (short) 0x0000;
+
+    /**
+     * The offset to shift by to obtain the sign bit.
+     *
+     * @hide
+     */
+    public static final int SIGN_SHIFT                = 15;
+
+    /**
+     * The offset to shift by to obtain the exponent bits.
+     *
+     * @hide
+     */
+    public static final int EXPONENT_SHIFT            = 10;
+
+    /**
+     * The bitmask to AND a number with to obtain the sign bit.
+     *
+     * @hide
+     */
+    public static final int SIGN_MASK                 = 0x8000;
+
+    /**
+     * The bitmask to AND a number shifted by {@link #EXPONENT_SHIFT} right, to obtain exponent bits.
+     *
+     * @hide
+     */
+    public static final int SHIFTED_EXPONENT_MASK     = 0x1f;
+
+    /**
+     * The bitmask to AND a number with to obtain significand bits.
+     *
+     * @hide
+     */
+    public static final int SIGNIFICAND_MASK          = 0x3ff;
+
+    /**
+     * The bitmask to AND with to obtain exponent and significand bits.
+     *
+     * @hide
+     */
+    public static final int EXPONENT_SIGNIFICAND_MASK = 0x7fff;
+
+    /**
+     * The offset of the exponent from the actual value.
+     *
+     * @hide
+     */
+    public static final int EXPONENT_BIAS             = 15;
+
+    private static final int FP32_SIGN_SHIFT            = 31;
+    private static final int FP32_EXPONENT_SHIFT        = 23;
+    private static final int FP32_SHIFTED_EXPONENT_MASK = 0xff;
+    private static final int FP32_SIGNIFICAND_MASK      = 0x7fffff;
+    private static final int FP32_EXPONENT_BIAS         = 127;
+    private static final int FP32_QNAN_MASK             = 0x400000;
+    private static final int FP32_DENORMAL_MAGIC = 126 << 23;
+    private static final float FP32_DENORMAL_FLOAT = Float.intBitsToFloat(FP32_DENORMAL_MAGIC);
+
+    /** Hidden constructor to prevent instantiation. */
+    private FP16() {}
+
+    /**
+     * <p>Compares the two specified half-precision float values. The following
+     * conditions apply during the comparison:</p>
+     *
+     * <ul>
+     * <li>{@link #NaN} is considered by this method to be equal to itself and greater
+     * than all other half-precision float values (including {@code #POSITIVE_INFINITY})</li>
+     * <li>{@link #POSITIVE_ZERO} is considered by this method to be greater than
+     * {@link #NEGATIVE_ZERO}.</li>
+     * </ul>
+     *
+     * @param x The first half-precision float value to compare.
+     * @param y The second half-precision float value to compare
+     *
+     * @return  The value {@code 0} if {@code x} is numerically equal to {@code y}, a
+     *          value less than {@code 0} if {@code x} is numerically less than {@code y},
+     *          and a value greater than {@code 0} if {@code x} is numerically greater
+     *          than {@code y}
+     *
+     * @hide
+     */
+    public static int compare(short x, short y) {
+        if (less(x, y)) return -1;
+        if (greater(x, y)) return 1;
+
+        // Collapse NaNs, akin to halfToIntBits(), but we want to keep
+        // (signed) short value types to preserve the ordering of -0.0
+        // and +0.0
+        short xBits = isNaN(x) ? NaN : x;
+        short yBits = isNaN(y) ? NaN : y;
+
+        return (xBits == yBits ? 0 : (xBits < yBits ? -1 : 1));
+    }
+
+    /**
+     * Returns the closest integral half-precision float value to the specified
+     * half-precision float value. Special values are handled in the
+     * following ways:
+     * <ul>
+     * <li>If the specified half-precision float is NaN, the result is NaN</li>
+     * <li>If the specified half-precision float is infinity (negative or positive),
+     * the result is infinity (with the same sign)</li>
+     * <li>If the specified half-precision float is zero (negative or positive),
+     * the result is zero (with the same sign)</li>
+     * </ul>
+     *
+     * @param h A half-precision float value
+     * @return The value of the specified half-precision float rounded to the nearest
+     *         half-precision float value
+     *
+     * @hide
+     */
+    public static short rint(short h) {
+        int bits = h & 0xffff;
+        int abs = bits & EXPONENT_SIGNIFICAND_MASK;
+        int result = bits;
+
+        if (abs < 0x3c00) {
+            result &= SIGN_MASK;
+            if (abs > 0x3800){
+                result |= 0x3c00;
+            }
+        } else if (abs < 0x6400) {
+            int exp = 25 - (abs >> 10);
+            int mask = (1 << exp) - 1;
+            result += ((1 << (exp - 1)) - (~(abs >> exp) & 1));
+            result &= ~mask;
+        }
+        if (isNaN((short) result)) {
+            // if result is NaN mask with qNaN
+            // (i.e. mask the most significant mantissa bit with 1)
+            // to comply with hardware implementations (ARM64, Intel, etc).
+            result |= NaN;
+        }
+
+        return (short) result;
+    }
+
+    /**
+     * Returns the smallest half-precision float value toward negative infinity
+     * greater than or equal to the specified half-precision float value.
+     * Special values are handled in the following ways:
+     * <ul>
+     * <li>If the specified half-precision float is NaN, the result is NaN</li>
+     * <li>If the specified half-precision float is infinity (negative or positive),
+     * the result is infinity (with the same sign)</li>
+     * <li>If the specified half-precision float is zero (negative or positive),
+     * the result is zero (with the same sign)</li>
+     * </ul>
+     *
+     * @param h A half-precision float value
+     * @return The smallest half-precision float value toward negative infinity
+     *         greater than or equal to the specified half-precision float value
+     *
+     * @hide
+     */
+    public static short ceil(short h) {
+        int bits = h & 0xffff;
+        int abs = bits & EXPONENT_SIGNIFICAND_MASK;
+        int result = bits;
+
+        if (abs < 0x3c00) {
+            result &= SIGN_MASK;
+            result |= 0x3c00 & -(~(bits >> 15) & (abs != 0 ? 1 : 0));
+        } else if (abs < 0x6400) {
+            abs = 25 - (abs >> 10);
+            int mask = (1 << abs) - 1;
+            result += mask & ((bits >> 15) - 1);
+            result &= ~mask;
+        }
+        if (isNaN((short) result)) {
+            // if result is NaN mask with qNaN
+            // (i.e. mask the most significant mantissa bit with 1)
+            // to comply with hardware implementations (ARM64, Intel, etc).
+            result |= NaN;
+        }
+
+        return (short) result;
+    }
+
+    /**
+     * Returns the largest half-precision float value toward positive infinity
+     * less than or equal to the specified half-precision float value.
+     * Special values are handled in the following ways:
+     * <ul>
+     * <li>If the specified half-precision float is NaN, the result is NaN</li>
+     * <li>If the specified half-precision float is infinity (negative or positive),
+     * the result is infinity (with the same sign)</li>
+     * <li>If the specified half-precision float is zero (negative or positive),
+     * the result is zero (with the same sign)</li>
+     * </ul>
+     *
+     * @param h A half-precision float value
+     * @return The largest half-precision float value toward positive infinity
+     *         less than or equal to the specified half-precision float value
+     *
+     * @hide
+     */
+    public static short floor(short h) {
+        int bits = h & 0xffff;
+        int abs = bits & EXPONENT_SIGNIFICAND_MASK;
+        int result = bits;
+
+        if (abs < 0x3c00) {
+            result &= SIGN_MASK;
+            result |= 0x3c00 & (bits > 0x8000 ? 0xffff : 0x0);
+        } else if (abs < 0x6400) {
+            abs = 25 - (abs >> 10);
+            int mask = (1 << abs) - 1;
+            result += mask & -(bits >> 15);
+            result &= ~mask;
+        }
+        if (isNaN((short) result)) {
+            // if result is NaN mask with qNaN
+            // i.e. (Mask the most significant mantissa bit with 1)
+            result |= NaN;
+        }
+
+        return (short) result;
+    }
+
+    /**
+     * Returns the truncated half-precision float value of the specified
+     * half-precision float value. Special values are handled in the following ways:
+     * <ul>
+     * <li>If the specified half-precision float is NaN, the result is NaN</li>
+     * <li>If the specified half-precision float is infinity (negative or positive),
+     * the result is infinity (with the same sign)</li>
+     * <li>If the specified half-precision float is zero (negative or positive),
+     * the result is zero (with the same sign)</li>
+     * </ul>
+     *
+     * @param h A half-precision float value
+     * @return The truncated half-precision float value of the specified
+     *         half-precision float value
+     *
+     * @hide
+     */
+    public static short trunc(short h) {
+        int bits = h & 0xffff;
+        int abs = bits & EXPONENT_SIGNIFICAND_MASK;
+        int result = bits;
+
+        if (abs < 0x3c00) {
+            result &= SIGN_MASK;
+        } else if (abs < 0x6400) {
+            abs = 25 - (abs >> 10);
+            int mask = (1 << abs) - 1;
+            result &= ~mask;
+        }
+
+        return (short) result;
+    }
+
+    /**
+     * Returns the smaller of two half-precision float values (the value closest
+     * to negative infinity). Special values are handled in the following ways:
+     * <ul>
+     * <li>If either value is NaN, the result is NaN</li>
+     * <li>{@link #NEGATIVE_ZERO} is smaller than {@link #POSITIVE_ZERO}</li>
+     * </ul>
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     * @return The smaller of the two specified half-precision values
+     *
+     * @hide
+     */
+    public static short min(short x, short y) {
+        if (isNaN(x)) return NaN;
+        if (isNaN(y)) return NaN;
+
+        if ((x & EXPONENT_SIGNIFICAND_MASK) == 0 && (y & EXPONENT_SIGNIFICAND_MASK) == 0) {
+            return (x & SIGN_MASK) != 0 ? x : y;
+        }
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) <
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff) ? x : y;
+    }
+
+    /**
+     * Returns the larger of two half-precision float values (the value closest
+     * to positive infinity). Special values are handled in the following ways:
+     * <ul>
+     * <li>If either value is NaN, the result is NaN</li>
+     * <li>{@link #POSITIVE_ZERO} is greater than {@link #NEGATIVE_ZERO}</li>
+     * </ul>
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return The larger of the two specified half-precision values
+     *
+     * @hide
+     */
+    public static short max(short x, short y) {
+        if (isNaN(x)) return NaN;
+        if (isNaN(y)) return NaN;
+
+        if ((x & EXPONENT_SIGNIFICAND_MASK) == 0 && (y & EXPONENT_SIGNIFICAND_MASK) == 0) {
+            return (x & SIGN_MASK) != 0 ? y : x;
+        }
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) >
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff) ? x : y;
+    }
+
+    /**
+     * Returns true if the first half-precision float value is less (smaller
+     * toward negative infinity) than the second half-precision float value.
+     * If either of the values is NaN, the result is false.
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return True if x is less than y, false otherwise
+     *
+     * @hide
+     */
+    public static boolean less(short x, short y) {
+        if (isNaN(x)) return false;
+        if (isNaN(y)) return false;
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) <
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff);
+    }
+
+    /**
+     * Returns true if the first half-precision float value is less (smaller
+     * toward negative infinity) than or equal to the second half-precision
+     * float value. If either of the values is NaN, the result is false.
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return True if x is less than or equal to y, false otherwise
+     *
+     * @hide
+     */
+    public static boolean lessEquals(short x, short y) {
+        if (isNaN(x)) return false;
+        if (isNaN(y)) return false;
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) <=
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff);
+    }
+
+    /**
+     * Returns true if the first half-precision float value is greater (larger
+     * toward positive infinity) than the second half-precision float value.
+     * If either of the values is NaN, the result is false.
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return True if x is greater than y, false otherwise
+     *
+     * @hide
+     */
+    public static boolean greater(short x, short y) {
+        if (isNaN(x)) return false;
+        if (isNaN(y)) return false;
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) >
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff);
+    }
+
+    /**
+     * Returns true if the first half-precision float value is greater (larger
+     * toward positive infinity) than or equal to the second half-precision float
+     * value. If either of the values is NaN, the result is false.
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return True if x is greater than y, false otherwise
+     *
+     * @hide
+     */
+    public static boolean greaterEquals(short x, short y) {
+        if (isNaN(x)) return false;
+        if (isNaN(y)) return false;
+
+        return ((x & SIGN_MASK) != 0 ? 0x8000 - (x & 0xffff) : x & 0xffff) >=
+               ((y & SIGN_MASK) != 0 ? 0x8000 - (y & 0xffff) : y & 0xffff);
+    }
+
+    /**
+     * Returns true if the two half-precision float values are equal.
+     * If either of the values is NaN, the result is false. {@link #POSITIVE_ZERO}
+     * and {@link #NEGATIVE_ZERO} are considered equal.
+     *
+     * @param x The first half-precision value
+     * @param y The second half-precision value
+     *
+     * @return True if x is equal to y, false otherwise
+     *
+     * @hide
+     */
+    public static boolean equals(short x, short y) {
+        if (isNaN(x)) return false;
+        if (isNaN(y)) return false;
+
+        return x == y || ((x | y) & EXPONENT_SIGNIFICAND_MASK) == 0;
+    }
+
+    /**
+     * Returns true if the specified half-precision float value represents
+     * infinity, false otherwise.
+     *
+     * @param h A half-precision float value
+     * @return True if the value is positive infinity or negative infinity,
+     *         false otherwise
+     *
+     * @hide
+     */
+    public static boolean isInfinite(short h) {
+        return (h & EXPONENT_SIGNIFICAND_MASK) == POSITIVE_INFINITY;
+    }
+
+    /**
+     * Returns true if the specified half-precision float value represents
+     * a Not-a-Number, false otherwise.
+     *
+     * @param h A half-precision float value
+     * @return True if the value is a NaN, false otherwise
+     *
+     * @hide
+     */
+    public static boolean isNaN(short h) {
+        return (h & EXPONENT_SIGNIFICAND_MASK) > POSITIVE_INFINITY;
+    }
+
+    /**
+     * Returns true if the specified half-precision float value is normalized
+     * (does not have a subnormal representation). If the specified value is
+     * {@link #POSITIVE_INFINITY}, {@link #NEGATIVE_INFINITY},
+     * {@link #POSITIVE_ZERO}, {@link #NEGATIVE_ZERO}, NaN or any subnormal
+     * number, this method returns false.
+     *
+     * @param h A half-precision float value
+     * @return True if the value is normalized, false otherwise
+     *
+     * @hide
+     */
+    public static boolean isNormalized(short h) {
+        return (h & POSITIVE_INFINITY) != 0 && (h & POSITIVE_INFINITY) != POSITIVE_INFINITY;
+    }
+
+    /**
+     * <p>Converts the specified half-precision float value into a
+     * single-precision float value. The following special cases are handled:</p>
+     * <ul>
+     * <li>If the input is {@link #NaN}, the returned value is {@link Float#NaN}</li>
+     * <li>If the input is {@link #POSITIVE_INFINITY} or
+     * {@link #NEGATIVE_INFINITY}, the returned value is respectively
+     * {@link Float#POSITIVE_INFINITY} or {@link Float#NEGATIVE_INFINITY}</li>
+     * <li>If the input is 0 (positive or negative), the returned value is +/-0.0f</li>
+     * <li>Otherwise, the returned value is a normalized single-precision float value</li>
+     * </ul>
+     *
+     * @param h The half-precision float value to convert to single-precision
+     * @return A normalized single-precision float value
+     *
+     * @hide
+     */
+    public static float toFloat(short h) {
+        int bits = h & 0xffff;
+        int s = bits & SIGN_MASK;
+        int e = (bits >>> EXPONENT_SHIFT) & SHIFTED_EXPONENT_MASK;
+        int m = (bits                        ) & SIGNIFICAND_MASK;
+
+        int outE = 0;
+        int outM = 0;
+
+        if (e == 0) { // Denormal or 0
+            if (m != 0) {
+                // Convert denorm fp16 into normalized fp32
+                float o = Float.intBitsToFloat(FP32_DENORMAL_MAGIC + m);
+                o -= FP32_DENORMAL_FLOAT;
+                return s == 0 ? o : -o;
+            }
+        } else {
+            outM = m << 13;
+            if (e == 0x1f) { // Infinite or NaN
+                outE = 0xff;
+                if (outM != 0) { // SNaNs are quieted
+                    outM |= FP32_QNAN_MASK;
+                }
+            } else {
+                outE = e - EXPONENT_BIAS + FP32_EXPONENT_BIAS;
+            }
+        }
+
+        int out = (s << 16) | (outE << FP32_EXPONENT_SHIFT) | outM;
+        return Float.intBitsToFloat(out);
+    }
+
+    /**
+     * <p>Converts the specified single-precision float value into a
+     * half-precision float value. The following special cases are handled:</p>
+     * <ul>
+     * <li>If the input is NaN (see {@link Float#isNaN(float)}), the returned
+     * value is {@link #NaN}</li>
+     * <li>If the input is {@link Float#POSITIVE_INFINITY} or
+     * {@link Float#NEGATIVE_INFINITY}, the returned value is respectively
+     * {@link #POSITIVE_INFINITY} or {@link #NEGATIVE_INFINITY}</li>
+     * <li>If the input is 0 (positive or negative), the returned value is
+     * {@link #POSITIVE_ZERO} or {@link #NEGATIVE_ZERO}</li>
+     * <li>If the input is a less than {@link #MIN_VALUE}, the returned value
+     * is flushed to {@link #POSITIVE_ZERO} or {@link #NEGATIVE_ZERO}</li>
+     * <li>If the input is a less than {@link #MIN_NORMAL}, the returned value
+     * is a denorm half-precision float</li>
+     * <li>Otherwise, the returned value is rounded to the nearest
+     * representable half-precision float value</li>
+     * </ul>
+     *
+     * @param f The single-precision float value to convert to half-precision
+     * @return A half-precision float value
+     *
+     * @hide
+     */
+    public static short toHalf(float f) {
+        int bits = Float.floatToRawIntBits(f);
+        int s = (bits >>> FP32_SIGN_SHIFT    );
+        int e = (bits >>> FP32_EXPONENT_SHIFT) & FP32_SHIFTED_EXPONENT_MASK;
+        int m = (bits                        ) & FP32_SIGNIFICAND_MASK;
+
+        int outE = 0;
+        int outM = 0;
+
+        if (e == 0xff) { // Infinite or NaN
+            outE = 0x1f;
+            outM = m != 0 ? 0x200 : 0;
+        } else {
+            e = e - FP32_EXPONENT_BIAS + EXPONENT_BIAS;
+            if (e >= 0x1f) { // Overflow
+                outE = 0x1f;
+            } else if (e <= 0) { // Underflow
+                if (e < -10) {
+                    // The absolute fp32 value is less than MIN_VALUE, flush to +/-0
+                } else {
+                    // The fp32 value is a normalized float less than MIN_NORMAL,
+                    // we convert to a denorm fp16
+                    m = m | 0x800000;
+                    int shift = 14 - e;
+                    outM = m >> shift;
+
+                    int lowm = m & ((1 << shift) - 1);
+                    int hway = 1 << (shift - 1);
+                    // if above halfway or exactly halfway and outM is odd
+                    if (lowm + (outM & 1) > hway){
+                        // Round to nearest even
+                        // Can overflow into exponent bit, which surprisingly is OK.
+                        // This increment relies on the +outM in the return statement below
+                        outM++;
+                    }
+                }
+            } else {
+                outE = e;
+                outM = m >> 13;
+                // if above halfway or exactly halfway and outM is odd
+                if ((m & 0x1fff) + (outM & 0x1) > 0x1000) {
+                    // Round to nearest even
+                    // Can overflow into exponent bit, which surprisingly is OK.
+                    // This increment relies on the +outM in the return statement below
+                    outM++;
+                }
+            }
+        }
+        // The outM is added here as the +1 increments for outM above can
+        // cause an overflow in the exponent bit which is OK.
+        return (short) ((s << SIGN_SHIFT) | (outE << EXPONENT_SHIFT) + outM);
+    }
+
+    /**
+     * <p>Returns a hexadecimal string representation of the specified half-precision
+     * float value. If the value is a NaN, the result is <code>"NaN"</code>,
+     * otherwise the result follows this format:</p>
+     * <ul>
+     * <li>If the sign is positive, no sign character appears in the result</li>
+     * <li>If the sign is negative, the first character is <code>'-'</code></li>
+     * <li>If the value is inifinity, the string is <code>"Infinity"</code></li>
+     * <li>If the value is 0, the string is <code>"0x0.0p0"</code></li>
+     * <li>If the value has a normalized representation, the exponent and
+     * significand are represented in the string in two fields. The significand
+     * starts with <code>"0x1."</code> followed by its lowercase hexadecimal
+     * representation. Trailing zeroes are removed unless all digits are 0, then
+     * a single zero is used. The significand representation is followed by the
+     * exponent, represented by <code>"p"</code>, itself followed by a decimal
+     * string of the unbiased exponent</li>
+     * <li>If the value has a subnormal representation, the significand starts
+     * with <code>"0x0."</code> followed by its lowercase hexadecimal
+     * representation. Trailing zeroes are removed unless all digits are 0, then
+     * a single zero is used. The significand representation is followed by the
+     * exponent, represented by <code>"p-14"</code></li>
+     * </ul>
+     *
+     * @param h A half-precision float value
+     * @return A hexadecimal string representation of the specified value
+     *
+     * @hide
+     */
+    public static String toHexString(short h) {
+        StringBuilder o = new StringBuilder();
+
+        int bits = h & 0xffff;
+        int s = (bits >>> SIGN_SHIFT    );
+        int e = (bits >>> EXPONENT_SHIFT) & SHIFTED_EXPONENT_MASK;
+        int m = (bits                   ) & SIGNIFICAND_MASK;
+
+        if (e == 0x1f) { // Infinite or NaN
+            if (m == 0) {
+                if (s != 0) o.append('-');
+                o.append("Infinity");
+            } else {
+                o.append("NaN");
+            }
+        } else {
+            if (s == 1) o.append('-');
+            if (e == 0) {
+                if (m == 0) {
+                    o.append("0x0.0p0");
+                } else {
+                    o.append("0x0.");
+                    String significand = Integer.toHexString(m);
+                    o.append(significand.replaceFirst("0{2,}$", ""));
+                    o.append("p-14");
+                }
+            } else {
+                o.append("0x1.");
+                String significand = Integer.toHexString(m);
+                o.append(significand.replaceFirst("0{2,}$", ""));
+                o.append('p');
+                o.append(Integer.toString(e - EXPONENT_BIAS));
+            }
+        }
+
+        return o.toString();
+    }
+}
diff --git a/ravenwood/scripts/list-ravenwood-tests.sh b/ravenwood/scripts/list-ravenwood-tests.sh
index fb9b823..05f3fdf 100755
--- a/ravenwood/scripts/list-ravenwood-tests.sh
+++ b/ravenwood/scripts/list-ravenwood-tests.sh
@@ -15,4 +15,4 @@
 
 # List all the ravenwood test modules.
 
-jq -r 'to_entries[] | select( .value.compatibility_suites | index("ravenwood-tests") ) | .key' "$OUT/module-info.json"
+jq -r 'to_entries[] | select( .value.compatibility_suites | index("ravenwood-tests") ) | .key' "$OUT/module-info.json" | sort
diff --git a/ravenwood/scripts/remove-ravenizer-output.sh b/ravenwood/scripts/remove-ravenizer-output.sh
new file mode 100755
index 0000000..be15b71
--- /dev/null
+++ b/ravenwood/scripts/remove-ravenizer-output.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+# 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.
+
+# Delete all the ravenizer output jar files from Soong's intermediate directory.
+
+# `-a -prune` is needed because otherwise find would be confused if the directory disappears.
+
+find "${ANDROID_BUILD_TOP:?}/out/soong/.intermediates/" \
+    -type d \
+    -name 'ravenizer' \
+    -print \
+    -exec rm -fr \{\} \; \
+    -a -prune
diff --git a/ravenwood/scripts/update-test-mapping.sh b/ravenwood/scripts/update-test-mapping.sh
new file mode 100755
index 0000000..b6cf5b8
--- /dev/null
+++ b/ravenwood/scripts/update-test-mapping.sh
@@ -0,0 +1,82 @@
+#!/bin/bash
+# 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.
+
+# Update f/b/r/TEST_MAPPING with all the ravenwood tests as presubmit.
+#
+# Note, before running it, make sure module-info.json is up-to-date by running
+# (any) build.
+
+set -e
+
+main() {
+    local script_name="${0##*/}"
+    local script_dir="${0%/*}"
+    local test_mapping="$script_dir/../TEST_MAPPING"
+    local test_mapping_bak="$script_dir/../TEST_MAPPING.bak"
+
+    local header="$(sed -ne '1,/AUTO-GENERATED-START/p' "$test_mapping")"
+    local footer="$(sed -ne '/AUTO-GENERATED-END/,$p' "$test_mapping")"
+
+    echo "Getting all tests"
+    local tests=( $("$script_dir/list-ravenwood-tests.sh") )
+
+    local num_tests="${#tests[@]}"
+
+    if (( $num_tests == 0 )) ; then
+        echo "Something went wrong. No ravenwood tests detected." 1>&2
+        return 1
+    fi
+
+    echo "Tests: ${tests[@]}"
+
+    echo "Creating backup at $test_mapping_bak"
+    cp "$test_mapping" "$test_mapping_bak"
+
+    echo "Updating $test_mapping"
+    {
+        echo "$header"
+
+        echo "    // DO NOT MODIFY MANUALLY"
+        echo "    // Use scripts/$script_name to update it."
+
+        local i=0
+        while (( $i < $num_tests )) ; do
+            local comma=","
+            if (( $i == ($num_tests - 1) )); then
+                comma=""
+            fi
+            echo "    {"
+            echo "      \"name\": \"${tests[$i]}\","
+            echo "      \"host\": true"
+            echo "    }$comma"
+
+            i=$(( $i + 1 ))
+        done
+
+        echo "$footer"
+    } >"$test_mapping"
+
+    if cmp "$test_mapping_bak" "$test_mapping" ; then
+        echo "No change detecetd."
+        return 0
+    fi
+    echo "Updated $test_mapping"
+
+    # `|| true` is needed because of `set -e`.
+    diff -u "$test_mapping_bak" "$test_mapping" || true
+    return 0
+}
+
+main
diff --git a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
index d8366c5..9c86389 100644
--- a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
+++ b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
@@ -9,17 +9,13 @@
 com.android.internal.logging.testing.UiEventLoggerFake
 com.android.internal.os.AndroidPrintStream
 com.android.internal.os.BatteryStatsHistory
-com.android.internal.os.BatteryStatsHistory$TraceDelegate
-com.android.internal.os.BatteryStatsHistory$VarintParceler
 com.android.internal.os.BatteryStatsHistoryIterator
 com.android.internal.os.Clock
 com.android.internal.os.LongArrayMultiStateCounter
-com.android.internal.os.LongArrayMultiStateCounter$LongArrayContainer
 com.android.internal.os.LongMultiStateCounter
 com.android.internal.os.MonotonicClock
 com.android.internal.os.PowerProfile
 com.android.internal.os.PowerStats
-com.android.internal.os.PowerStats$Descriptor
 com.android.internal.os.RuntimeInit
 com.android.internal.power.EnergyConsumerStats
 com.android.internal.power.ModemPowerProfile
@@ -46,6 +42,7 @@
 android.util.EventLog
 android.util.FloatProperty
 android.util.FloatMath
+android.util.Half
 android.util.IndentingPrintWriter
 android.util.IntArray
 android.util.IntProperty
@@ -122,15 +119,9 @@
 android.os.BaseBundle
 android.os.BatteryConsumer
 android.os.BatteryStats
-android.os.BatteryStats$HistoryItem
-android.os.BatteryStats$HistoryStepDetails
-android.os.BatteryStats$HistoryTag
-android.os.BatteryStats$LongCounter
-android.os.BatteryStats$ProcessStateChange
 android.os.BatteryUsageStats
 android.os.BatteryUsageStatsQuery
 android.os.Binder
-android.os.Binder$IdentitySupplier
 android.os.BluetoothBatteryStats
 android.os.Broadcaster
 android.os.Build
@@ -141,7 +132,6 @@
 android.os.DeadObjectException
 android.os.DeadSystemException
 android.os.FileUtils
-android.os.FileUtils$MemoryPipe
 android.os.Handler
 android.os.HandlerExecutor
 android.os.HandlerThread
@@ -153,8 +143,6 @@
 android.os.PackageTagsList
 android.os.Parcel
 android.os.ParcelFileDescriptor
-android.os.ParcelFileDescriptor$AutoCloseInputStream
-android.os.ParcelFileDescriptor$AutoCloseOutputStream
 android.os.ParcelFormatException
 android.os.ParcelUuid
 android.os.Parcelable
@@ -168,7 +156,6 @@
 android.os.RemoteException
 android.os.ResultReceiver
 android.os.ServiceManager
-android.os.ServiceManager$ServiceNotFoundException
 android.os.ServiceSpecificException
 android.os.StrictMode
 android.os.SystemClock
@@ -179,18 +166,14 @@
 android.os.Trace
 android.os.TransactionTooLargeException
 android.os.UserBatteryConsumer
-android.os.UserBatteryConsumer$Builder
 android.os.UidBatteryConsumer
-android.os.UidBatteryConsumer$Builder
 android.os.UserHandle
 android.os.UserManager
 android.os.VibrationAttributes
-android.os.VibrationAttributes$Builder
 android.os.WakeLockStats
 android.os.WorkSource
 
 android.content.ClipData
-android.content.ClipData$Item
 android.content.ClipDescription
 android.content.ClipboardManager
 android.content.ComponentName
@@ -207,11 +190,7 @@
 android.content.pm.ComponentInfo
 android.content.pm.PackageInfo
 android.content.pm.PackageItemInfo
-android.content.pm.PackageManager$Flags
-android.content.pm.PackageManager$PackageInfoFlags
-android.content.pm.PackageManager$ApplicationInfoFlags
-android.content.pm.PackageManager$ComponentInfoFlags
-android.content.pm.PackageManager$ResolveInfoFlags
+android.content.pm.PackageManager
 android.content.pm.PathPermission
 android.content.pm.ProviderInfo
 android.content.pm.ResolveInfo
@@ -222,7 +201,6 @@
 android.content.res.ApkAssets
 android.content.res.AssetFileDescriptor
 android.content.res.AssetManager
-android.content.res.AssetManager$Builder
 android.content.res.ColorStateList
 android.content.res.ConfigurationBoundResourceCache
 android.content.res.Configuration
@@ -236,7 +214,6 @@
 android.content.res.FontScaleConverterImpl
 android.content.res.FontScaleConverterFactory
 android.content.res.Resources
-android.content.res.Resources$Theme
 android.content.res.ResourceId
 android.content.res.ResourcesImpl
 android.content.res.ResourcesKey
@@ -259,7 +236,6 @@
 android.database.DataSetObservable
 android.database.DataSetObserver
 android.database.MatrixCursor
-android.database.MatrixCursor$RowBuilder
 android.database.MergeCursor
 android.database.Observable
 android.database.SQLException
@@ -267,7 +243,6 @@
 android.database.sqlite.SQLiteException
 
 android.text.TextUtils
-android.text.TextUtils$SimpleStringSplitter
 
 android.accounts.Account
 
@@ -277,7 +252,11 @@
 android.graphics.Insets
 android.graphics.Interpolator
 android.graphics.Matrix
+android.graphics.Matrix44
+android.graphics.Outline
+android.graphics.ParcelableColorSpace
 android.graphics.Path
+android.graphics.PixelFormat
 android.graphics.Point
 android.graphics.PointF
 android.graphics.Rect
@@ -293,14 +272,11 @@
 android.app.Instrumentation
 android.app.LocaleConfig
 android.app.ResourcesManager
-android.app.ResourcesManager$UpdateHandler
 android.app.WindowConfiguration
 
 android.metrics.LogMaker
 
 android.view.Display
-android.view.Display$HdrCapabilities
-android.view.Display$Mode
 android.view.DisplayAdjustments
 android.view.DisplayInfo
 android.view.inputmethod.InputBinding
diff --git a/ravenwood/texts/ravenwood-standard-options.txt b/ravenwood/texts/ravenwood-standard-options.txt
index f64f26d..952ab82 100644
--- a/ravenwood/texts/ravenwood-standard-options.txt
+++ b/ravenwood/texts/ravenwood-standard-options.txt
@@ -6,8 +6,6 @@
 --default-throw
 
 # Uncomment below lines to enable each feature.
-# --enable-non-stub-method-check
---no-non-stub-method-check
 
 #--default-method-call-hook
 #    com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Exceptions.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Exceptions.kt
index 3a7fab3..0dcd271 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Exceptions.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Exceptions.kt
@@ -17,7 +17,14 @@
 
 package com.android.platform.test.ravenwood.ravenizer
 
+import com.android.hoststubgen.UserErrorException
+
 /**
  * Use it for internal exception that really shouldn't happen.
  */
 class RavenizerInternalException(message: String) : Exception(message)
+
+/**
+ * Thrown when an invalid test is detected in the target jar. (e.g. JUni3 tests)
+ */
+class RavenizerInvalidTestException(message: String) : Exception(message), UserErrorException
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
index e92ef72..a38512e 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Ravenizer.kt
@@ -44,6 +44,9 @@
     /** Time took to build [ClasNodes] */
     var loadStructureTime: Double = .0,
 
+    /** Time took to validate the classes */
+    var validationTime: Double = .0,
+
     /** Total real time spent for converting the jar file */
     var totalProcessTime: Double = .0,
 
@@ -67,6 +70,7 @@
             RavenizerStats{
               totalTime=$totalTime,
               loadStructureTime=$loadStructureTime,
+              validationTime=$validationTime,
               totalProcessTime=$totalProcessTime,
               totalConversionTime=$totalConversionTime,
               totalCopyTime=$totalCopyTime,
@@ -84,16 +88,44 @@
 class Ravenizer(val options: RavenizerOptions) {
     fun run() {
         val stats = RavenizerStats()
+
+        val fatalValidation = options.fatalValidation.get
+
         stats.totalTime = log.nTime {
-            process(options.inJar.get, options.outJar.get, stats)
+            process(
+                options.inJar.get,
+                options.outJar.get,
+                options.enableValidation.get,
+                fatalValidation,
+                stats,
+            )
         }
         log.i(stats.toString())
     }
 
-    private fun process(inJar: String, outJar: String, stats: RavenizerStats) {
+    private fun process(
+        inJar: String,
+        outJar: String,
+        enableValidation: Boolean,
+        fatalValidation: Boolean,
+        stats: RavenizerStats,
+    ) {
         var allClasses = ClassNodes.loadClassStructures(inJar) {
             time -> stats.loadStructureTime = time
         }
+        if (enableValidation) {
+            stats.validationTime = log.iTime("Validating classes") {
+                if (!validateClasses(allClasses)) {
+                    var message = "Invalid test class(es) detected." +
+                            " See error log for details."
+                    if (fatalValidation) {
+                        throw RavenizerInvalidTestException(message)
+                    } else {
+                        log.w("Warning: $message")
+                    }
+                }
+            }
+        }
 
         stats.totalProcessTime = log.iTime("$executableName processing $inJar") {
             ZipFile(inJar).use { inZip ->
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
index e85e3be..e8341e5 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/RavenizerOptions.kt
@@ -27,6 +27,12 @@
 
     /** Output jar file */
     var outJar: SetOnce<String> = SetOnce(""),
+
+    /** Whether to enable test validation. */
+    var enableValidation: SetOnce<Boolean> = SetOnce(true),
+
+    /** Whether the validation failure is fatal or not. */
+    var fatalValidation: SetOnce<Boolean> = SetOnce(false),
 ) {
     companion object {
         fun parseArgs(args: Array<String>): RavenizerOptions {
@@ -52,6 +58,12 @@
                         "--in-jar" -> ret.inJar.set(nextArg()).ensureFileExists()
                         "--out-jar" -> ret.outJar.set(nextArg())
 
+                        "--enable-validation" -> ret.enableValidation.set(true)
+                        "--disable-validation" -> ret.enableValidation.set(false)
+
+                        "--fatal-validation" -> ret.fatalValidation.set(true)
+                        "--no-fatal-validation" -> ret.fatalValidation.set(false)
+
                         else -> throw ArgumentsException("Unknown option: $arg")
                     }
                 } catch (e: SetOnce.SetMoreThanOnceException) {
@@ -74,6 +86,8 @@
             RavenizerOptions{
               inJar=$inJar,
               outJar=$outJar,
+              enableValidation=$enableValidation,
+              fatalValidation=$fatalValidation,
             }
             """.trimIndent()
     }
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Utils.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Utils.kt
index e026e7a..1aa70c08 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Utils.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Utils.kt
@@ -15,10 +15,12 @@
  */
 package com.android.platform.test.ravenwood.ravenizer
 
+import android.platform.test.annotations.NoRavenizer
 import android.platform.test.ravenwood.RavenwoodAwareTestRunner
 import com.android.hoststubgen.asm.ClassNodes
 import com.android.hoststubgen.asm.findAnyAnnotation
 import com.android.hoststubgen.asm.startsWithAny
+import com.android.hoststubgen.asm.toHumanReadableClassName
 import org.junit.rules.TestRule
 import org.junit.runner.RunWith
 import org.objectweb.asm.Type
@@ -30,6 +32,7 @@
     val desc = type.descriptor
     val descAsSet = setOf<String>(desc)
     val internlName = type.internalName
+    val humanReadableName = type.internalName.toHumanReadableClassName()
 }
 
 val testAnotType = TypeHolder(org.junit.Test::class.java)
@@ -37,6 +40,7 @@
 val classRuleAnotType = TypeHolder(org.junit.ClassRule::class.java)
 val runWithAnotType = TypeHolder(RunWith::class.java)
 val innerRunnerAnotType = TypeHolder(RavenwoodAwareTestRunner.InnerRunner::class.java)
+val noRavenizerAnotType = TypeHolder(NoRavenizer::class.java)
 
 val testRuleType = TypeHolder(TestRule::class.java)
 val ravenwoodTestRunnerType = TypeHolder(RavenwoodAwareTestRunner::class.java)
@@ -87,9 +91,12 @@
     return this.startsWithAny(
         "java/", // just in case...
         "javax/",
+        "junit/",
         "org/junit/",
         "org/mockito/",
         "kotlin/",
+        "androidx/",
+        "android/support/",
         // TODO -- anything else?
     )
 }
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Validator.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Validator.kt
new file mode 100644
index 0000000..27092d2
--- /dev/null
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/Validator.kt
@@ -0,0 +1,81 @@
+/*
+ * 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.platform.test.ravenwood.ravenizer
+
+import com.android.hoststubgen.asm.ClassNodes
+import com.android.hoststubgen.asm.startsWithAny
+import com.android.hoststubgen.asm.toHumanReadableClassName
+import com.android.hoststubgen.log
+import org.objectweb.asm.tree.ClassNode
+
+fun validateClasses(classes: ClassNodes): Boolean {
+    var allOk = true
+    classes.forEach { allOk = checkClass(it, classes) && allOk }
+
+    return allOk
+}
+
+/**
+ * Validate a class.
+ *
+ * - A test class shouldn't extend
+ *
+ */
+fun checkClass(cn: ClassNode, classes: ClassNodes): Boolean {
+    if (cn.name.shouldByBypassed()) {
+        // Class doesn't need to be checked.
+        return true
+    }
+    var allOk = true
+
+    // See if there's any class that extends a legacy base class.
+    // But ignore the base classes in android.test.
+    if (!cn.name.startsWithAny("android/test/")) {
+        allOk = checkSuperClass(cn, cn, classes) && allOk
+    }
+    return allOk
+}
+
+fun checkSuperClass(targetClass: ClassNode, currentClass: ClassNode, classes: ClassNodes): Boolean {
+    if (currentClass.superName == null || currentClass.superName == "java/lang/Object") {
+        return true // No parent class
+    }
+    if (currentClass.superName.isLegacyTestBaseClass()) {
+        log.e("Error: Class ${targetClass.name.toHumanReadableClassName()} extends"
+                + " a legacy test class ${currentClass.superName.toHumanReadableClassName()}.")
+        return false
+    }
+    classes.findClass(currentClass.superName)?.let {
+        return checkSuperClass(targetClass, it, classes)
+    }
+    // Super class not found.
+    // log.w("Class ${currentClass.superName} not found.")
+    return true
+}
+
+/**
+ * Check if a class internal name is a known legacy test base class.
+ */
+fun String.isLegacyTestBaseClass(): Boolean {
+    return this.startsWithAny(
+        "junit/framework/TestCase",
+
+        // In case the test doesn't statically include JUnit, we need
+        "android/test/AndroidTestCase",
+        "android/test/InstrumentationTestCase",
+        "android/test/InstrumentationTestSuite",
+    )
+}
diff --git a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
index 25cad02..eaef2cf 100644
--- a/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
+++ b/ravenwood/tools/ravenizer/src/com/android/platform/test/ravenwood/ravenizer/adapter/RunnerRewritingAdapter.kt
@@ -30,6 +30,7 @@
 import com.android.platform.test.ravenwood.ravenizer.classRuleAnotType
 import com.android.platform.test.ravenwood.ravenizer.isTestLookingClass
 import com.android.platform.test.ravenwood.ravenizer.innerRunnerAnotType
+import com.android.platform.test.ravenwood.ravenizer.noRavenizerAnotType
 import com.android.platform.test.ravenwood.ravenizer.ravenwoodTestRunnerType
 import com.android.platform.test.ravenwood.ravenizer.ruleAnotType
 import com.android.platform.test.ravenwood.ravenizer.runWithAnotType
@@ -183,7 +184,7 @@
             av.visit("value", ravenwoodTestRunnerType.type)
             av.visitEnd()
         }
-        log.d("Processed ${classInternalName.toHumanReadableClassName()}")
+        log.i("Update the @RunWith: ${classInternalName.toHumanReadableClassName()}")
     }
 
     /*
@@ -435,7 +436,19 @@
 
     companion object {
         fun shouldProcess(classes: ClassNodes, className: String): Boolean {
-            return isTestLookingClass(classes, className)
+            if (!isTestLookingClass(classes, className)) {
+                return false
+            }
+            // Don't process a class if it has a @NoRavenizer annotation.
+            classes.findClass(className)?.let { cn ->
+                if (cn.findAnyAnnotation(noRavenizerAnotType.descAsSet) != null) {
+                    log.w("Class ${className.toHumanReadableClassName()} has" +
+                        " @${noRavenizerAnotType.humanReadableName}. Skipping."
+                    )
+                    return false
+                }
+            }
+            return true
         }
 
         fun maybeApply(
diff --git a/services/accessibility/accessibility.aconfig b/services/accessibility/accessibility.aconfig
index 8e2e0ad..ffa4841 100644
--- a/services/accessibility/accessibility.aconfig
+++ b/services/accessibility/accessibility.aconfig
@@ -117,7 +117,14 @@
     name: "enable_magnification_follows_mouse"
     namespace: "accessibility"
     description: "Whether to enable mouse following for fullscreen magnification"
-    bug: "335494097"
+    bug: "354696546"
+}
+
+flag {
+    name: "enable_magnification_keyboard_control"
+    namespace: "accessibility"
+    description: "Whether to enable keyboard control for magnification"
+    bug: "355487062"
 }
 
 flag {
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilityInputFilter.java b/services/accessibility/java/com/android/server/accessibility/AccessibilityInputFilter.java
index d16a665..1b2447e 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilityInputFilter.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilityInputFilter.java
@@ -26,6 +26,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.graphics.Region;
+import android.hardware.input.InputManager;
 import android.os.Looper;
 import android.os.PowerManager;
 import android.os.SystemClock;
@@ -57,6 +58,7 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Objects;
 import java.util.StringJoiner;
 
 /**
@@ -748,6 +750,7 @@
 
         if ((mEnabledFeatures & FLAG_FEATURE_MOUSE_KEYS) != 0) {
             mMouseKeysInterceptor = new MouseKeysInterceptor(mAms,
+                    Objects.requireNonNull(mContext.getSystemService(InputManager.class)),
                     Looper.myLooper(),
                     Display.DEFAULT_DISPLAY);
             addFirstEventHandler(Display.DEFAULT_DISPLAY, mMouseKeysInterceptor);
diff --git a/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java b/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
index 7cbb97e..b541345 100644
--- a/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
+++ b/services/accessibility/java/com/android/server/accessibility/AccessibilityManagerService.java
@@ -163,6 +163,7 @@
 import android.view.accessibility.IAccessibilityManager;
 import android.view.accessibility.IAccessibilityManagerClient;
 import android.view.accessibility.IMagnificationConnection;
+import android.view.accessibility.IUserInitializationCompleteCallback;
 import android.view.inputmethod.EditorInfo;
 
 import com.android.internal.R;
@@ -366,6 +367,11 @@
     private final List<SendWindowStateChangedEventRunnable> mSendWindowStateChangedEventRunnables =
             new ArrayList<>();
 
+    @VisibleForTesting
+    final HashSet<IUserInitializationCompleteCallback>
+            mUserInitializationCompleteCallbacks =
+            new HashSet<IUserInitializationCompleteCallback>();
+
     @GuardedBy("mLock")
     private @UserIdInt int mCurrentUserId = UserHandle.USER_SYSTEM;
 
@@ -2034,6 +2040,17 @@
                         obtainMessage(AccessibilityManagerService::announceNewUserIfNeeded, this),
                         WAIT_FOR_USER_STATE_FULLY_INITIALIZED_MILLIS);
             }
+
+            for (IUserInitializationCompleteCallback callback
+                    : mUserInitializationCompleteCallbacks) {
+                try {
+                    callback.onUserInitializationComplete(mCurrentUserId);
+                } catch (RemoteException re) {
+                    Log.e("AccessibilityManagerService",
+                            "Error while dispatching userInitializationComplete callback: ",
+                            re);
+                }
+            }
         }
     }
 
@@ -4886,40 +4903,26 @@
     }
 
     @Override
-    @RequiresNoPermission
-    public boolean startFlashNotificationSequence(String opPkg,
-            @FlashNotificationReason int reason, IBinder token) {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            return mFlashNotificationsController.startFlashNotificationSequence(opPkg,
-                    reason, token);
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
+    @EnforcePermission(MANAGE_ACCESSIBILITY)
+    public boolean startFlashNotificationSequence(String opPkg, @FlashNotificationReason int reason,
+            IBinder token) {
+        startFlashNotificationSequence_enforcePermission();
+        return mFlashNotificationsController.startFlashNotificationSequence(opPkg, reason, token);
     }
 
     @Override
-    @RequiresNoPermission
+    @EnforcePermission(MANAGE_ACCESSIBILITY)
     public boolean stopFlashNotificationSequence(String opPkg) {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            return mFlashNotificationsController.stopFlashNotificationSequence(opPkg);
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
+        stopFlashNotificationSequence_enforcePermission();
+        return mFlashNotificationsController.stopFlashNotificationSequence(opPkg);
     }
 
     @Override
-    @RequiresNoPermission
-    public boolean startFlashNotificationEvent(String opPkg,
-            @FlashNotificationReason int reason, String reasonPkg) {
-        final long identity = Binder.clearCallingIdentity();
-        try {
-            return mFlashNotificationsController.startFlashNotificationEvent(opPkg,
-                    reason, reasonPkg);
-        } finally {
-            Binder.restoreCallingIdentity(identity);
-        }
+    @EnforcePermission(MANAGE_ACCESSIBILITY)
+    public boolean startFlashNotificationEvent(String opPkg, @FlashNotificationReason int reason,
+            String reasonPkg) {
+        startFlashNotificationEvent_enforcePermission();
+        return mFlashNotificationsController.startFlashNotificationEvent(opPkg, reason, reasonPkg);
     }
 
     @Override
@@ -6251,6 +6254,24 @@
     }
 
     @Override
+    @RequiresNoPermission
+    public void registerUserInitializationCompleteCallback(
+            IUserInitializationCompleteCallback callback) {
+        synchronized (mLock) {
+            mUserInitializationCompleteCallbacks.add(callback);
+        }
+    }
+
+    @Override
+    @RequiresNoPermission
+    public void unregisterUserInitializationCompleteCallback(
+            IUserInitializationCompleteCallback callback) {
+        synchronized (mLock) {
+            mUserInitializationCompleteCallbacks.remove(callback);
+        }
+    }
+
+    @Override
     @EnforcePermission(INJECT_EVENTS)
     public void injectInputEventToInputFilter(InputEvent event) {
         injectInputEventToInputFilter_enforcePermission();
diff --git a/services/accessibility/java/com/android/server/accessibility/MouseKeysInterceptor.java b/services/accessibility/java/com/android/server/accessibility/MouseKeysInterceptor.java
index 2ce5c2b..54368ca 100644
--- a/services/accessibility/java/com/android/server/accessibility/MouseKeysInterceptor.java
+++ b/services/accessibility/java/com/android/server/accessibility/MouseKeysInterceptor.java
@@ -23,6 +23,7 @@
 import android.annotation.RequiresPermission;
 import android.companion.virtual.VirtualDeviceManager;
 import android.companion.virtual.VirtualDeviceParams;
+import android.hardware.input.InputManager;
 import android.hardware.input.VirtualMouse;
 import android.hardware.input.VirtualMouseButtonEvent;
 import android.hardware.input.VirtualMouseConfig;
@@ -34,8 +35,11 @@
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.view.InputDevice;
 import android.view.KeyEvent;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.server.LocalServices;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 
@@ -60,7 +64,7 @@
  * mouse keys of each physical keyboard will control a single (global) mouse pointer.
  */
 public class MouseKeysInterceptor extends BaseEventStreamTransformation
-        implements Handler.Callback {
+        implements Handler.Callback, InputManager.InputDeviceListener {
     private static final String LOG_TAG = "MouseKeysInterceptor";
 
     // To enable these logs, run: 'adb shell setprop log.tag.MouseKeysInterceptor DEBUG'
@@ -77,10 +81,19 @@
 
     private final AccessibilityManagerService mAms;
     private final Handler mHandler;
+    private final InputManager mInputManager;
 
     /** Thread to wait for virtual mouse creation to complete */
     private final Thread mCreateVirtualMouseThread;
 
+    /**
+     * Map of device IDs to a map of key codes to their corresponding {@link MouseKeyEvent} values.
+     * To ensure thread safety for the map, all access and modification of the map
+     * should happen on the same thread, i.e., on the handler thread.
+     */
+    private final SparseArray<SparseArray<MouseKeyEvent>> mDeviceKeyCodeMap =
+            new SparseArray<>();
+
     VirtualDeviceManager.VirtualDevice mVirtualDevice = null;
 
     private VirtualMouse mVirtualMouse = null;
@@ -102,6 +115,21 @@
     /** Whether scroll toggle is on */
     private boolean mScrollToggleOn = false;
 
+    /** The ID of the input device that is currently active */
+    private int mActiveInputDeviceId = 0;
+
+    /**
+     * Enum representing different types of mouse key events, each associated with a specific
+     * key code.
+     *
+     * <p> These events correspond to various mouse actions such as directional movements,
+     * clicks, and scrolls, mapped to specific keys on the keyboard.
+     * The key codes here are the QWERTY key codes, and should be accessed via
+     * {@link MouseKeyEvent#getKeyCode(InputDevice)}
+     * so that it is mapped to the equivalent key on the keyboard layout of the keyboard device
+     * that is actually in use.
+     * </p>
+     */
     public enum MouseKeyEvent {
         DIAGONAL_UP_LEFT_MOVE(KeyEvent.KEYCODE_7),
         UP_MOVE_OR_SCROLL(KeyEvent.KEYCODE_8),
@@ -117,34 +145,64 @@
         RELEASE(KeyEvent.KEYCODE_COMMA),
         SCROLL_TOGGLE(KeyEvent.KEYCODE_PERIOD);
 
-        private final int mKeyCode;
+        private final int mLocationKeyCode;
         MouseKeyEvent(int enumValue) {
-            mKeyCode = enumValue;
+            mLocationKeyCode = enumValue;
         }
 
-        private static final SparseArray<MouseKeyEvent> VALUE_TO_ENUM_MAP = new SparseArray<>();
-
-        static {
-            for (MouseKeyEvent type : MouseKeyEvent.values()) {
-                VALUE_TO_ENUM_MAP.put(type.mKeyCode, type);
-            }
-        }
-
+        @VisibleForTesting
         public final int getKeyCodeValue() {
-            return mKeyCode;
+            return mLocationKeyCode;
         }
 
         /**
-         * Convert int value of the key code to corresponding MouseEvent enum. If no matching
-         * value is found, this will return {@code null}.
+         * Get the key code associated with the given MouseKeyEvent for the given keyboard
+         * input device, taking into account its layout.
+         * The default is to return the keycode for the default layout (QWERTY).
+         * We check if the input device has been generated using {@link InputDevice#getGeneration()}
+         * to test with the default {@link MouseKeyEvent} values in the unit tests.
+         */
+        public int getKeyCode(InputDevice inputDevice) {
+            if (inputDevice.getGeneration() == -1) {
+                return mLocationKeyCode;
+            }
+            return inputDevice.getKeyCodeForKeyLocation(mLocationKeyCode);
+        }
+
+        /**
+         * Convert int value of the key code to corresponding {@link MouseKeyEvent}
+         * enum for a particular device ID.
+         * If no matching value is found, this will return {@code null}.
          */
         @Nullable
-        public static MouseKeyEvent from(int value) {
-            return VALUE_TO_ENUM_MAP.get(value);
+        public static MouseKeyEvent from(int keyCode, int deviceId,
+                SparseArray<SparseArray<MouseKeyEvent>> deviceKeyCodeMap) {
+            SparseArray<MouseKeyEvent> keyCodeToEnumMap = deviceKeyCodeMap.get(deviceId);
+            if (keyCodeToEnumMap != null) {
+                return keyCodeToEnumMap.get(keyCode);
+            }
+            return null;
         }
     }
 
     /**
+     * Create a map of key codes to their corresponding {@link MouseKeyEvent} values
+     * for a specific input device.
+     * The key for {@code mDeviceKeyCodeMap} is the deviceId.
+     * The key for {@code keyCodeToEnumMap} is the keycode for each
+     * {@link MouseKeyEvent} according to the keyboard layout of the input device.
+     */
+    public void initializeDeviceToEnumMap(InputDevice inputDevice) {
+        int deviceId = inputDevice.getId();
+        SparseArray<MouseKeyEvent> keyCodeToEnumMap = new SparseArray<>();
+        for (MouseKeyEvent mouseKeyEventType : MouseKeyEvent.values()) {
+            int keyCode = mouseKeyEventType.getKeyCode(inputDevice);
+            keyCodeToEnumMap.put(keyCode, mouseKeyEventType);
+        }
+        mDeviceKeyCodeMap.put(deviceId, keyCodeToEnumMap);
+    }
+
+    /**
      * Construct a new MouseKeysInterceptor.
      *
      * @param service The service to notify of key events
@@ -152,8 +210,10 @@
      * @param displayId Display ID to send mouse events to
      */
     @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
-    public MouseKeysInterceptor(AccessibilityManagerService service, Looper looper, int displayId) {
+    public MouseKeysInterceptor(AccessibilityManagerService service,
+            InputManager inputManager, Looper looper, int displayId) {
         mAms = service;
+        mInputManager = inputManager;
         mHandler = new Handler(looper, this);
         // Create the virtual mouse on a separate thread since virtual device creation
         // should happen on an auxiliary thread, and not from the handler's thread.
@@ -163,6 +223,9 @@
             mVirtualMouse = createVirtualMouse(displayId);
         });
         mCreateVirtualMouseThread.start();
+        // Register an input device listener to watch when input devices are
+        // added, removed or reconfigured.
+        mInputManager.registerInputDeviceListener(this, mHandler);
     }
 
     /**
@@ -215,7 +278,8 @@
      */
     @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
     private void performMouseScrollAction(int keyCode) {
-        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(keyCode);
+        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(
+                keyCode, mActiveInputDeviceId, mDeviceKeyCodeMap);
         float y = switch (mouseKeyEvent) {
             case UP_MOVE_OR_SCROLL -> 1.0f;
             case DOWN_MOVE_OR_SCROLL -> -1.0f;
@@ -247,15 +311,18 @@
      */
     @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
     private void performMouseButtonAction(int keyCode) {
-        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(keyCode);
+        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(
+                keyCode, mActiveInputDeviceId, mDeviceKeyCodeMap);
         int buttonCode = switch (mouseKeyEvent) {
             case LEFT_CLICK -> VirtualMouseButtonEvent.BUTTON_PRIMARY;
             case RIGHT_CLICK -> VirtualMouseButtonEvent.BUTTON_SECONDARY;
             default -> VirtualMouseButtonEvent.BUTTON_UNKNOWN;
         };
         if (buttonCode != VirtualMouseButtonEvent.BUTTON_UNKNOWN) {
-            sendVirtualMouseButtonEvent(buttonCode, VirtualMouseButtonEvent.ACTION_BUTTON_PRESS);
-            sendVirtualMouseButtonEvent(buttonCode, VirtualMouseButtonEvent.ACTION_BUTTON_RELEASE);
+            sendVirtualMouseButtonEvent(buttonCode,
+                    VirtualMouseButtonEvent.ACTION_BUTTON_PRESS);
+            sendVirtualMouseButtonEvent(buttonCode,
+                    VirtualMouseButtonEvent.ACTION_BUTTON_RELEASE);
         }
         if (DEBUG) {
             if (buttonCode == VirtualMouseButtonEvent.BUTTON_UNKNOWN) {
@@ -293,7 +360,9 @@
     private void performMousePointerAction(int keyCode) {
         float x = 0f;
         float y = 0f;
-        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(keyCode);
+        MouseKeyEvent mouseKeyEvent = MouseKeyEvent.from(
+                keyCode, mActiveInputDeviceId, mDeviceKeyCodeMap);
+
         switch (mouseKeyEvent) {
             case DIAGONAL_DOWN_LEFT_MOVE -> {
                 x = -MOUSE_POINTER_MOVEMENT_STEP / sqrt(2);
@@ -339,18 +408,19 @@
         }
     }
 
-    private boolean isMouseKey(int keyCode) {
-        return MouseKeyEvent.VALUE_TO_ENUM_MAP.contains(keyCode);
+    private boolean isMouseKey(int keyCode, int deviceId) {
+        SparseArray<MouseKeyEvent> keyCodeToEnumMap = mDeviceKeyCodeMap.get(deviceId);
+        return keyCodeToEnumMap.contains(keyCode);
     }
 
-    private boolean isMouseButtonKey(int keyCode) {
-        return keyCode == MouseKeyEvent.LEFT_CLICK.getKeyCodeValue()
-                || keyCode == MouseKeyEvent.RIGHT_CLICK.getKeyCodeValue();
+    private boolean isMouseButtonKey(int keyCode, InputDevice inputDevice) {
+        return keyCode == MouseKeyEvent.LEFT_CLICK.getKeyCode(inputDevice)
+                || keyCode == MouseKeyEvent.RIGHT_CLICK.getKeyCode(inputDevice);
     }
 
-    private boolean isMouseScrollKey(int keyCode) {
-        return keyCode == MouseKeyEvent.UP_MOVE_OR_SCROLL.getKeyCodeValue()
-                || keyCode == MouseKeyEvent.DOWN_MOVE_OR_SCROLL.getKeyCodeValue();
+    private boolean isMouseScrollKey(int keyCode, InputDevice inputDevice) {
+        return keyCode == MouseKeyEvent.UP_MOVE_OR_SCROLL.getKeyCode(inputDevice)
+                || keyCode == MouseKeyEvent.DOWN_MOVE_OR_SCROLL.getKeyCode(inputDevice);
     }
 
     /**
@@ -373,7 +443,7 @@
     }
 
     /**
-     * Handles key events and forwards mouse key events to the virtual mouse.
+     * Handles key events and forwards mouse key events to the virtual mouse on the handler thread.
      *
      * @param event The key event to handle.
      * @param policyFlags The policy flags associated with the key event.
@@ -385,31 +455,45 @@
             mAms.getTraceManager().logTrace(LOG_TAG + ".onKeyEvent",
                     FLAGS_INPUT_FILTER, "event=" + event + ";policyFlags=" + policyFlags);
         }
+
+        mHandler.post(() -> {
+            onKeyEventInternal(event, policyFlags);
+        });
+    }
+
+    @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
+    private void onKeyEventInternal(KeyEvent event, int policyFlags) {
         boolean isDown = event.getAction() == KeyEvent.ACTION_DOWN;
         int keyCode = event.getKeyCode();
+        mActiveInputDeviceId = event.getDeviceId();
+        InputDevice inputDevice = mInputManager.getInputDevice(mActiveInputDeviceId);
 
-        if (!isMouseKey(keyCode)) {
+        if (!mDeviceKeyCodeMap.contains(mActiveInputDeviceId)) {
+            initializeDeviceToEnumMap(inputDevice);
+        }
+
+        if (!isMouseKey(keyCode, mActiveInputDeviceId)) {
             // Pass non-mouse key events to the next handler
             super.onKeyEvent(event, policyFlags);
         } else if (isDown) {
-            if (keyCode == MouseKeyEvent.SCROLL_TOGGLE.getKeyCodeValue()) {
+            if (keyCode == MouseKeyEvent.SCROLL_TOGGLE.getKeyCode(inputDevice)) {
                 mScrollToggleOn = !mScrollToggleOn;
                 if (DEBUG) {
                     Slog.d(LOG_TAG, "Scroll toggle " + (mScrollToggleOn ? "ON" : "OFF"));
                 }
-            } else if (keyCode == MouseKeyEvent.HOLD.getKeyCodeValue()) {
+            } else if (keyCode == MouseKeyEvent.HOLD.getKeyCode(inputDevice)) {
                 sendVirtualMouseButtonEvent(
                         VirtualMouseButtonEvent.BUTTON_PRIMARY,
                         VirtualMouseButtonEvent.ACTION_BUTTON_PRESS
                 );
-            } else if (keyCode == MouseKeyEvent.RELEASE.getKeyCodeValue()) {
+            } else if (keyCode == MouseKeyEvent.RELEASE.getKeyCode(inputDevice)) {
                 sendVirtualMouseButtonEvent(
                         VirtualMouseButtonEvent.BUTTON_PRIMARY,
                         VirtualMouseButtonEvent.ACTION_BUTTON_RELEASE
                 );
-            } else if (isMouseButtonKey(keyCode)) {
+            } else if (isMouseButtonKey(keyCode, inputDevice)) {
                 performMouseButtonAction(keyCode);
-            } else if (mScrollToggleOn && isMouseScrollKey(keyCode)) {
+            } else if (mScrollToggleOn && isMouseScrollKey(keyCode, inputDevice)) {
                 // If the scroll key is pressed down and no other key is active,
                 // set it as the active key and send a message to scroll the pointer
                 if (mActiveScrollKey == KEY_NOT_SET) {
@@ -439,7 +523,8 @@
                 mHandler.removeMessages(MESSAGE_SCROLL_MOUSE_POINTER);
             } else {
                 Slog.i(LOG_TAG, "Dropping event with key code: '" + keyCode
-                        + "', with no matching down event from deviceId = " + event.getDeviceId());
+                        + "', with no matching down event from deviceId = "
+                        + event.getDeviceId());
             }
         }
     }
@@ -503,12 +588,40 @@
     @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
     @Override
     public void onDestroy() {
-        // Clear mouse state
-        mActiveMoveKey = KEY_NOT_SET;
-        mActiveScrollKey = KEY_NOT_SET;
-        mLastTimeKeyActionPerformed = 0;
+        mHandler.post(() -> {
+            // Clear mouse state
+            mActiveMoveKey = KEY_NOT_SET;
+            mActiveScrollKey = KEY_NOT_SET;
+            mLastTimeKeyActionPerformed = 0;
+            mDeviceKeyCodeMap.clear();
+        });
 
         mHandler.removeCallbacksAndMessages(null);
         mVirtualDevice.close();
     }
+
+    @Override
+    public void onInputDeviceAdded(int deviceId) {
+    }
+
+    @Override
+    public void onInputDeviceRemoved(int deviceId) {
+        mDeviceKeyCodeMap.remove(deviceId);
+    }
+
+    /**
+     * The user can change the keyboard layout from settings at anytime, which would change
+     * key character map for that device. Hence, we should use this callback to
+     * update the key code to enum mapping if there is a change in the physical keyboard detected.
+     *
+     * @param deviceId The id of the input device that changed.
+     */
+    @Override
+    public void onInputDeviceChanged(int deviceId) {
+        InputDevice inputDevice = mInputManager.getInputDevice(deviceId);
+        // Update the enum mapping only if input device that changed is a keyboard
+        if (inputDevice.isFullKeyboard() && !mDeviceKeyCodeMap.contains(deviceId)) {
+            initializeDeviceToEnumMap(inputDevice);
+        }
+    }
 }
diff --git a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
index 83f57b2..950246f 100644
--- a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
+++ b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerManager.java
@@ -87,7 +87,7 @@
     public Set<AndroidAccessibilityCheckerResult> maybeRunA11yChecker(
             List<AccessibilityNodeInfo> nodes, @Nullable String sourceEventClassName,
             ComponentName a11yServiceComponentName, @UserIdInt int userId) {
-        if (!shouldRunA11yChecker()) {
+        if (!shouldRunA11yChecker() || nodes.isEmpty()) {
             return Set.of();
         }
 
@@ -95,24 +95,33 @@
         String defaultBrowserName = mPackageManager.getDefaultBrowserPackageNameAsUser(userId);
 
         try {
+            AndroidAccessibilityCheckerResult.Builder commonResultBuilder =
+                    AccessibilityCheckerUtils.getCommonResultBuilder(nodes.getFirst(),
+                            sourceEventClassName, mPackageManager, a11yServiceComponentName);
+            if (commonResultBuilder == null) {
+                return Set.of();
+            }
             for (AccessibilityNodeInfo nodeInfo : nodes) {
-                // Skip browser results because they are mostly related to web content and not the
-                // browser app itself.
+                // Skip browser results because they are mostly related to web content and
+                // not the browser app itself.
                 if (nodeInfo.getPackageName() == null
                         || nodeInfo.getPackageName().toString().equals(defaultBrowserName)) {
                     continue;
                 }
-                List<AccessibilityHierarchyCheckResult> checkResults = runChecksOnNode(nodeInfo);
+                List<AccessibilityHierarchyCheckResult> checkResults = runChecksOnNode(
+                        nodeInfo);
                 Set<AndroidAccessibilityCheckerResult> filteredResults =
                         AccessibilityCheckerUtils.processResults(nodeInfo, checkResults,
-                                sourceEventClassName, mPackageManager, a11yServiceComponentName);
+                                commonResultBuilder);
                 allResults.addAll(filteredResults);
             }
             mCachedResults.addAll(allResults);
+            return allResults;
+
         } catch (RuntimeException e) {
             Slog.e(LOG_TAG, "An unknown error occurred while running a11y checker.", e);
+            return Set.of();
         }
-        return allResults;
     }
 
     private List<AccessibilityHierarchyCheckResult> runChecksOnNode(
diff --git a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
index eb24b02..a739304 100644
--- a/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
+++ b/services/accessibility/java/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtils.java
@@ -91,45 +91,55 @@
                             AccessibilityCheckClass.TRAVERSAL_ORDER_CHECK));
     // LINT.ThenChange(/services/accessibility/java/com/android/server/accessibility/a11ychecker/proto/a11ychecker.proto)
 
-    static Set<AndroidAccessibilityCheckerResult> processResults(
+
+    /**
+     * Returns AccessibilityCheckResultReported.Builder with the common fields for all nodes
+     * belonging in the same cache pre-filled.
+     */
+    static @Nullable AndroidAccessibilityCheckerResult.Builder getCommonResultBuilder(
             AccessibilityNodeInfo nodeInfo,
-            List<AccessibilityHierarchyCheckResult> checkResults,
             @Nullable String activityClassName,
             PackageManager packageManager,
             ComponentName a11yServiceComponentName) {
-        String appPackageName = nodeInfo.getPackageName().toString();
-        String nodePath = AccessibilityNodePathBuilder.createNodePath(nodeInfo);
-        if (nodePath == null) {
-            return Set.of();
+        if (nodeInfo.getPackageName() == null) {
+            return null;
         }
-        AndroidAccessibilityCheckerResult.Builder commonBuilder;
+        String appPackageName = nodeInfo.getPackageName().toString();
         try {
-            commonBuilder = AndroidAccessibilityCheckerResult.newBuilder()
+            return AndroidAccessibilityCheckerResult.newBuilder()
                     .setPackageName(appPackageName)
                     .setAppVersionCode(getAppVersionCode(packageManager, appPackageName))
-                    .setUiElementPath(nodePath)
                     .setActivityName(
                             getActivityName(packageManager, appPackageName, activityClassName))
                     .setWindowTitle(getWindowTitle(nodeInfo))
                     .setSourceComponentName(a11yServiceComponentName)
-                    .setSourceVersionCode(
-                            getAppVersionCode(packageManager,
-                                    a11yServiceComponentName.getPackageName()));
+                    .setSourceVersionCode(getAppVersionCode(packageManager,
+                            a11yServiceComponentName.getPackageName()));
         } catch (PackageManager.NameNotFoundException e) {
             Slog.e(LOG_TAG, "Unknown package name", e);
+            return null;
+        }
+    }
+
+    static Set<AndroidAccessibilityCheckerResult> processResults(
+            AccessibilityNodeInfo nodeInfo,
+            List<AccessibilityHierarchyCheckResult> checkResults,
+            AndroidAccessibilityCheckerResult.Builder resultBuilder) {
+        String nodePath = AccessibilityNodePathBuilder.createNodePath(nodeInfo);
+        if (resultBuilder == null || nodePath == null) {
             return Set.of();
         }
-
         return checkResults.stream()
                 .filter(checkResult -> checkResult.getType()
                         == AccessibilityCheckResult.AccessibilityCheckResultType.ERROR
                         || checkResult.getType()
                         == AccessibilityCheckResult.AccessibilityCheckResultType.WARNING)
-                .map(checkResult -> new AndroidAccessibilityCheckerResult.Builder(
-                        commonBuilder).setResultCheckClass(
-                        getCheckClass(checkResult)).setResultType(
-                        getCheckResultType(checkResult)).setResultId(
-                        checkResult.getResultId()).build())
+                .map(checkResult -> new AndroidAccessibilityCheckerResult.Builder(resultBuilder)
+                        .setUiElementPath(nodePath)
+                        .setResultCheckClass(getCheckClass(checkResult))
+                        .setResultType(getCheckResultType(checkResult))
+                        .setResultId(checkResult.getResultId())
+                        .build())
                 .collect(Collectors.toUnmodifiableSet());
     }
 
diff --git a/services/accessibility/java/com/android/server/accessibility/magnification/MouseEventHandler.java b/services/accessibility/java/com/android/server/accessibility/magnification/MouseEventHandler.java
index 845249e..ab94e98 100644
--- a/services/accessibility/java/com/android/server/accessibility/magnification/MouseEventHandler.java
+++ b/services/accessibility/java/com/android/server/accessibility/magnification/MouseEventHandler.java
@@ -39,8 +39,14 @@
      * @param displayId The display that is being magnified
      */
     public void onEvent(MotionEvent event, int displayId) {
-        if (event.getAction() == ACTION_HOVER_MOVE
-                || (event.getAction() == ACTION_MOVE && event.getSource() == SOURCE_MOUSE)) {
+        // Ignore gesture events synthesized from the touchpad.
+        // TODO(b/354696546): Use synthesized pinch gestures to control scale.
+        boolean isSynthesizedFromTouchpad =
+                event.getClassification() != MotionEvent.CLASSIFICATION_NONE;
+
+        // Consume only move events from the mouse or hovers from any tool.
+        if (!isSynthesizedFromTouchpad && (event.getAction() == ACTION_HOVER_MOVE
+                || (event.getAction() == ACTION_MOVE && event.getSource() == SOURCE_MOUSE))) {
             final float eventX = event.getX();
             final float eventY = event.getY();
 
diff --git a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerService.java b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerService.java
index 954651d..a2d467c 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerService.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerService.java
@@ -16,8 +16,7 @@
 
 package com.android.server.appfunctions;
 
-import static android.app.appfunctions.flags.Flags.enableAppFunctionManager;
-
+import android.app.appfunctions.AppFunctionManagerConfiguration;
 import android.content.Context;
 
 import com.android.server.SystemService;
@@ -35,7 +34,7 @@
 
     @Override
     public void onStart() {
-        if (enableAppFunctionManager()) {
+        if (AppFunctionManagerConfiguration.isSupported(getContext())) {
             publishBinderService(Context.APP_FUNCTION_SERVICE, mServiceImpl);
         }
     }
diff --git a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
index 53885fc..32b8d6b 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/AppFunctionManagerServiceImpl.java
@@ -25,6 +25,7 @@
 import android.app.appfunctions.SafeOneTimeExecuteAppFunctionCallback;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Binder;
 import android.os.UserHandle;
 import android.text.TextUtils;
 import android.util.Slog;
@@ -94,36 +95,39 @@
             targetUser = mCallerValidator.verifyTargetUserHandle(
                     requestInternal.getUserHandle(), validatedCallingPackage);
         } catch (SecurityException exception) {
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse
-                    .Builder(ExecuteAppFunctionResponse.RESULT_DENIED,
-                    getExceptionMessage(exception)).build());
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse
+                    .newFailure(ExecuteAppFunctionResponse.RESULT_DENIED,
+                            exception.getMessage(),
+                            /*extras=*/  null));
             return;
         }
 
         // TODO(b/354956319): Add and honor the new enterprise policies.
         if (mCallerValidator.isUserOrganizationManaged(targetUser)) {
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse.Builder(
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse.newFailure(
                     ExecuteAppFunctionResponse.RESULT_INTERNAL_ERROR,
-                    "Cannot run on a device with a device owner or from the managed profile."
-            ).build());
+                    "Cannot run on a device with a device owner or from the managed profile.",
+                    /*extras=*/  null
+            ));
             return;
         }
 
         String targetPackageName = requestInternal.getClientRequest().getTargetPackageName();
         if (TextUtils.isEmpty(targetPackageName)) {
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse.Builder(
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse.newFailure(
                     ExecuteAppFunctionResponse.RESULT_INVALID_ARGUMENT,
-                    "Target package name cannot be empty."
-            ).build());
+                    "Target package name cannot be empty.",
+                    /*extras=*/  null
+            ));
             return;
         }
 
         if (!mCallerValidator.verifyCallerCanExecuteAppFunction(
                 validatedCallingPackage, targetPackageName)) {
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse
-                    .Builder(ExecuteAppFunctionResponse.RESULT_DENIED,
-                    "Caller does not have permission to execute the appfunction")
-                    .build());
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse
+                    .newFailure(ExecuteAppFunctionResponse.RESULT_DENIED,
+                            "Caller does not have permission to execute the appfunction",
+                            /*extras=*/  null));
             return;
         }
 
@@ -131,16 +135,23 @@
                 targetPackageName,
                 targetUser);
         if (serviceIntent == null) {
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse.Builder(
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse.newFailure(
                     ExecuteAppFunctionResponse.RESULT_INTERNAL_ERROR,
-                    "Cannot find the target service."
-            ).build());
+                    "Cannot find the target service.",
+                    /*extras=*/  null
+            ));
             return;
         }
-        bindAppFunctionServiceUnchecked(requestInternal, serviceIntent, targetUser,
+
+        final long token = Binder.clearCallingIdentity();
+        try {
+            bindAppFunctionServiceUnchecked(requestInternal, serviceIntent, targetUser,
                 safeExecuteAppFunctionCallback,
                 /*bindFlags=*/ Context.BIND_AUTO_CREATE,
                 /*timeoutInMillis=*/ mServiceConfig.getExecuteAppFunctionTimeoutMillis());
+        } finally {
+            Binder.restoreCallingIdentity(token);
+        }
     }
 
     private void bindAppFunctionServiceUnchecked(
@@ -171,9 +182,10 @@
                                     }
                             );
                         } catch (Exception e) {
-                            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse
-                                    .Builder(ExecuteAppFunctionResponse.RESULT_APP_UNKNOWN_ERROR,
-                                    getExceptionMessage(e)).build());
+                            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse
+                                    .newFailure(ExecuteAppFunctionResponse.RESULT_APP_UNKNOWN_ERROR,
+                                            e.getMessage(),
+                                            /*extras=*/  null));
                             serviceUsageCompleteListener.onCompleted();
                         }
                     }
@@ -181,33 +193,32 @@
                     @Override
                     public void onFailedToConnect() {
                         Slog.e(TAG, "Failed to connect to service");
-                        safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse
-                                .Builder(ExecuteAppFunctionResponse.RESULT_APP_UNKNOWN_ERROR,
-                                "Failed to connect to AppFunctionService").build());
+                        safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse
+                                .newFailure(ExecuteAppFunctionResponse.RESULT_APP_UNKNOWN_ERROR,
+                                        "Failed to connect to AppFunctionService",
+                                        /*extras=*/  null));
                     }
 
                     @Override
                     public void onTimedOut() {
                         Slog.e(TAG, "Timed out");
                         safeExecuteAppFunctionCallback.onResult(
-                                new ExecuteAppFunctionResponse.Builder(
+                                ExecuteAppFunctionResponse.newFailure(
                                         ExecuteAppFunctionResponse.RESULT_TIMED_OUT,
-                                        "Binding to AppFunctionService timed out."
-                                ).build());
+                                        "Binding to AppFunctionService timed out.",
+                                        /*extras=*/  null
+                                ));
                     }
                 }
         );
 
         if (!bindServiceResult) {
             Slog.e(TAG, "Failed to bind to the AppFunctionService");
-            safeExecuteAppFunctionCallback.onResult(new ExecuteAppFunctionResponse.Builder(
+            safeExecuteAppFunctionCallback.onResult(ExecuteAppFunctionResponse.newFailure(
                     ExecuteAppFunctionResponse.RESULT_TIMED_OUT,
-                    "Failed to bind the AppFunctionService."
-            ).build());
+                    "Failed to bind the AppFunctionService.",
+                    /*extras=*/  null
+            ));
         }
     }
-
-    private String getExceptionMessage(Exception exception) {
-        return exception.getMessage() == null ? "" : exception.getMessage();
-    }
 }
diff --git a/services/appfunctions/java/com/android/server/appfunctions/SyncAppSearchCallHelper.java b/services/appfunctions/java/com/android/server/appfunctions/SyncAppSearchCallHelper.java
new file mode 100644
index 0000000..5dd4c25
--- /dev/null
+++ b/services/appfunctions/java/com/android/server/appfunctions/SyncAppSearchCallHelper.java
@@ -0,0 +1,146 @@
+/*
+ * 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.appfunctions;
+
+import static android.app.appfunctions.flags.Flags.FLAG_ENABLE_APP_FUNCTION_MANAGER;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.WorkerThread;
+import android.app.appsearch.AppSearchManager;
+import android.app.appsearch.AppSearchManager.SearchContext;
+import android.app.appsearch.AppSearchResult;
+import android.app.appsearch.AppSearchSession;
+import android.app.appsearch.GetSchemaResponse;
+import android.app.appsearch.SetSchemaRequest;
+import android.app.appsearch.SetSchemaResponse;
+import android.util.Slog;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.Executor;
+
+/**
+ * Helper class for interacting with a system server local appsearch session asynchronously.
+ *
+ * <p>Converts the AppSearch Callback API to {@link AndroidFuture}.
+ */
+@FlaggedApi(FLAG_ENABLE_APP_FUNCTION_MANAGER)
+public class SyncAppSearchCallHelper implements Closeable {
+    private static final String TAG = SyncAppSearchCallHelper.class.getSimpleName();
+    private final Executor mExecutor;
+    private final AppSearchManager mAppSearchManager;
+    private final AndroidFuture<AppSearchResult<AppSearchSession>> mSettableSessionFuture;
+
+    public SyncAppSearchCallHelper(
+            @NonNull AppSearchManager appSearchManager,
+            @NonNull Executor executor,
+            @NonNull SearchContext appSearchContext) {
+        Objects.requireNonNull(appSearchManager);
+        Objects.requireNonNull(executor);
+        Objects.requireNonNull(appSearchContext);
+
+        mExecutor = executor;
+        mAppSearchManager = appSearchManager;
+        mSettableSessionFuture = new AndroidFuture<>();
+        mAppSearchManager.createSearchSession(
+                appSearchContext, mExecutor, mSettableSessionFuture::complete);
+    }
+
+    /** Converts a failed app search result codes into an exception. */
+    @NonNull
+    private static Exception failedResultToException(@NonNull AppSearchResult appSearchResult) {
+        return switch (appSearchResult.getResultCode()) {
+            case AppSearchResult.RESULT_INVALID_ARGUMENT ->
+                    new IllegalArgumentException(appSearchResult.getErrorMessage());
+            case AppSearchResult.RESULT_IO_ERROR ->
+                    new IOException(appSearchResult.getErrorMessage());
+            case AppSearchResult.RESULT_SECURITY_ERROR ->
+                    new SecurityException(appSearchResult.getErrorMessage());
+            default -> new IllegalStateException(appSearchResult.getErrorMessage());
+        };
+    }
+
+    private AndroidFuture<AppSearchSession> getSessionAsync() {
+        return mSettableSessionFuture.thenApply(
+                result -> {
+                    if (result.isSuccess()) {
+                        return result.getResultValue();
+                    } else {
+                        throw new RuntimeException(failedResultToException(result));
+                    }
+                });
+    }
+
+    /** Gets the schema for a given app search session. */
+    public AndroidFuture<GetSchemaResponse> getSchema() {
+        return getSessionAsync()
+                .thenComposeAsync(
+                        session -> {
+                            AndroidFuture<AppSearchResult<GetSchemaResponse>>
+                                    settableSchemaResponse = new AndroidFuture<>();
+                            session.getSchema(mExecutor, settableSchemaResponse::complete);
+                            return settableSchemaResponse.thenApply(
+                                    result -> {
+                                        if (result.isSuccess()) {
+                                            return result.getResultValue();
+                                        } else {
+                                            throw new RuntimeException(
+                                                    failedResultToException(result));
+                                        }
+                                    });
+                        },
+                        mExecutor);
+    }
+
+    /** Sets the schema for a given app search session. */
+    public AndroidFuture<SetSchemaResponse> setSchema(@NonNull SetSchemaRequest setSchemaRequest) {
+        return getSessionAsync()
+                .thenComposeAsync(
+                        session -> {
+                            AndroidFuture<AppSearchResult<SetSchemaResponse>>
+                                    settableSchemaResponse = new AndroidFuture<>();
+                            session.setSchema(
+                                    setSchemaRequest,
+                                    mExecutor,
+                                    mExecutor,
+                                    settableSchemaResponse::complete);
+                            return settableSchemaResponse.thenApply(
+                                    result -> {
+                                        if (result.isSuccess()) {
+                                            return result.getResultValue();
+                                        } else {
+                                            throw new RuntimeException(
+                                                    failedResultToException(result));
+                                        }
+                                    });
+                        },
+                        mExecutor);
+    }
+
+    @Override
+    public void close() throws IOException {
+        try {
+            getSessionAsync().get().close();
+        } catch (Exception ex) {
+            Slog.e(TAG, "Failed to close app search session", ex);
+        }
+    }
+}
diff --git a/services/autofill/java/com/android/server/autofill/ui/FillUi.java b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
index 14a3211..dc0b4b8 100644
--- a/services/autofill/java/com/android/server/autofill/ui/FillUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
@@ -48,6 +48,7 @@
 import android.view.WindowManager;
 import android.view.accessibility.AccessibilityManager;
 import android.view.autofill.AutofillId;
+import android.view.autofill.AutofillManager;
 import android.view.autofill.AutofillValue;
 import android.view.autofill.IAutofillWindowPresenter;
 import android.widget.BaseAdapter;
@@ -82,7 +83,6 @@
             com.android.internal.R.style.Theme_DeviceDefault_Light_Autofill;
     private static final int THEME_ID_DARK =
             com.android.internal.R.style.Theme_DeviceDefault_Autofill;
-    private static final int AUTOFILL_CREDMAN_MAX_VISIBLE_DATASETS = 5;
 
     private static final TypedValue sTempTypedValue = new TypedValue();
 
@@ -113,9 +113,11 @@
 
     private final @NonNull Callback mCallback;
 
+    private final @NonNull WindowManager mWindowManager;
+
     private final @Nullable View mHeader;
     private final @NonNull ListView mListView;
-    private final @Nullable View mFooter;
+    private @Nullable View mFooter;
 
     private final @Nullable ItemsAdapter mAdapter;
 
@@ -134,6 +136,8 @@
 
     private int mMaxInputLengthForAutofill;
 
+    private final boolean mIsCredmanAutofillSession;
+
     public static boolean isFullScreen(Context context) {
         if (sFullScreenMode != null) {
             if (sVerbose) Slog.v(TAG, "forcing full-screen mode to " + sFullScreenMode);
@@ -158,6 +162,9 @@
         mContext = new ContextThemeWrapper(context, mThemeId);
         mUserContext = Helper.getUserContext(mContext);
         mMaxInputLengthForAutofill = maxInputLengthForAutofill;
+        mIsCredmanAutofillSession = (Flags.autofillCredmanIntegration()
+            && ((response.getFlags() & FLAG_CREDENTIAL_MANAGER_RESPONSE) != 0));
+        mWindowManager = mContext.getSystemService(WindowManager.class);
 
         final LayoutInflater inflater = LayoutInflater.from(mContext);
 
@@ -167,7 +174,8 @@
         final ViewGroup decor;
         if (mFullScreen) {
             decor = (ViewGroup) inflater.inflate(R.layout.autofill_dataset_picker_fullscreen, null);
-        } else if (headerPresentation != null || footerPresentation != null) {
+        } else if (headerPresentation != null
+                || footerPresentation != null || mIsCredmanAutofillSession) {
             decor = (ViewGroup) inflater.inflate(R.layout.autofill_dataset_picker_header_footer,
                     null);
         } else {
@@ -219,11 +227,7 @@
             if (sVerbose) {
                 Slog.v(TAG, "overriding maximum visible datasets to " + mVisibleDatasetsMaxCount);
             }
-        } else if (Flags.autofillCredmanIntegration() && (
-                (response.getFlags() & FLAG_CREDENTIAL_MANAGER_RESPONSE) != 0)) {
-            mVisibleDatasetsMaxCount = AUTOFILL_CREDMAN_MAX_VISIBLE_DATASETS;
-        }
-        else {
+        } else {
             mVisibleDatasetsMaxCount = mContext.getResources()
                     .getInteger(com.android.internal.R.integer.autofill_max_visible_datasets);
         }
@@ -301,7 +305,7 @@
                 mHeader = null;
             }
 
-            if (footerPresentation != null) {
+            if (footerPresentation != null && !mIsCredmanAutofillSession) {
                 final LinearLayout footerContainer =
                         decor.findViewById(R.id.autofill_dataset_footer);
                 if (footerContainer != null) {
@@ -366,7 +370,22 @@
                     }
 
                     applyCancelAction(view, response.getCancelIds());
-                    items.add(new ViewItem(dataset, filterPattern, filterable, valueText, view));
+                    if (AutofillManager.PINNED_DATASET_ID.equals(dataset.getId())
+                            && mIsCredmanAutofillSession && !items.isEmpty()) {
+                        final LinearLayout footerContainer =
+                                decor.findViewById(R.id.autofill_dataset_footer);
+                        if (sVerbose) {
+                          Slog.v(TAG, "adding footer");
+                        }
+                        mFooter = view;
+                        footerContainer.addView(mFooter);
+                        footerContainer.setVisibility(View.VISIBLE);
+                        footerContainer.setClickable(true);
+                        footerContainer.setOnClickListener(v -> mCallback.onDatasetPicked(dataset));
+                    } else {
+                        items.add(
+                            new ViewItem(dataset, filterPattern, filterable, valueText, view));
+                    }
                 }
             }
 
@@ -459,12 +478,9 @@
                 if (updateContentSize()) {
                     requestShowFillUi();
                 }
-                if (mAdapter.getCount() > mVisibleDatasetsMaxCount) {
-                    mListView.setVerticalScrollBarEnabled(true);
-                    mListView.onVisibilityAggregated(true);
-                } else {
-                    mListView.setVerticalScrollBarEnabled(false);
-                }
+                mListView.setVerticalScrollBarEnabled(true);
+                mListView.onVisibilityAggregated(true);
+
                 if (mAdapter.getCount() != oldCount) {
                     mListView.requestLayout();
                 }
@@ -578,11 +594,18 @@
         return changed;
     }
 
+    private boolean heightLesserThanDisplayScreen(int height) {
+        // Don't update list height for credential options beyond 80% of display window even if we
+        // are still under the max visible number of datasets. This could happen when font or
+        // display size is set to large.
+        return height < (0.8 * mWindowManager.getCurrentWindowMetrics().getBounds().height());
+    }
+
     private boolean updateHeight(View view, Point maxSize) {
         boolean changed = false;
         final int clampedMeasuredHeight = Math.min(view.getMeasuredHeight(), maxSize.y);
         final int newContentHeight = mContentHeight + clampedMeasuredHeight;
-        if (newContentHeight != mContentHeight) {
+        if (newContentHeight != mContentHeight && heightLesserThanDisplayScreen(newContentHeight)) {
             mContentHeight = newContentHeight;
             changed = true;
         }
diff --git a/services/companion/Android.bp b/services/companion/Android.bp
index 2bfdd0a..77650eb 100644
--- a/services/companion/Android.bp
+++ b/services/companion/Android.bp
@@ -28,7 +28,6 @@
     ],
     static_libs: [
         "ukey2_jni",
-        "virtualdevice_flags_lib",
         "virtual_camera_service_aidl-java",
     ],
     lint: {
diff --git a/services/companion/java/com/android/server/companion/devicepresence/ObservableUuidStore.java b/services/companion/java/com/android/server/companion/devicepresence/ObservableUuidStore.java
index 4678a16..5fd282d 100644
--- a/services/companion/java/com/android/server/companion/devicepresence/ObservableUuidStore.java
+++ b/services/companion/java/com/android/server/companion/devicepresence/ObservableUuidStore.java
@@ -190,7 +190,7 @@
             }
             mCachedPerUser.set(userId, cachedObservableUuids);
         }
-        return cachedObservableUuids;
+        return cachedObservableUuids == null ? new ArrayList<>() : cachedObservableUuids;
     }
 
     /**
diff --git a/services/companion/java/com/android/server/companion/virtual/Android.bp b/services/companion/java/com/android/server/companion/virtual/Android.bp
deleted file mode 100644
index 66313e6..0000000
--- a/services/companion/java/com/android/server/companion/virtual/Android.bp
+++ /dev/null
@@ -1,17 +0,0 @@
-package {
-    default_team: "trendy_team_xr_framework",
-}
-
-java_aconfig_library {
-    name: "virtualdevice_flags_lib",
-    aconfig_declarations: "virtualdevice_flags",
-}
-
-aconfig_declarations {
-    name: "virtualdevice_flags",
-    package: "com.android.server.companion.virtual",
-    container: "system",
-    srcs: [
-        "flags.aconfig",
-    ],
-}
diff --git a/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java b/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
index 2119622..4b9065b 100644
--- a/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
+++ b/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
@@ -27,7 +27,6 @@
 import android.annotation.UserIdInt;
 import android.app.WindowConfiguration;
 import android.app.compat.CompatChanges;
-import android.companion.virtual.VirtualDeviceManager.ActivityListener;
 import android.compat.annotation.ChangeId;
 import android.compat.annotation.EnabledSince;
 import android.content.AttributionSource;
@@ -61,6 +60,9 @@
 
     private static final String TAG = "GenericWindowPolicyController";
 
+    private static final ComponentName BLOCKED_APP_STREAMING_COMPONENT =
+            new ComponentName("android", BlockedAppStreamingActivity.class.getName());
+
     /** Interface to listen running applications change on virtual display. */
     public interface RunningAppsChangedListener {
         /**
@@ -69,29 +71,25 @@
         void onRunningAppsChanged(ArraySet<Integer> runningUids);
     }
 
-    /**
-     * For communicating when activities are blocked from running on the display by this policy
-     * controller.
-     */
-    public interface ActivityBlockedCallback {
+    /** Interface to react to activity changes on the virtual display. */
+    public interface ActivityListener {
+
+        /** Called when the top activity changes. */
+        void onTopActivityChanged(int displayId, @NonNull ComponentName topActivity,
+                @UserIdInt int userId);
+
+        /** Called when the display becomes empty. */
+        void onDisplayEmpty(int displayId);
+
         /** Called when an activity is blocked.*/
-        void onActivityBlocked(int displayId, ActivityInfo activityInfo, IntentSender intentSender);
-    }
-    private static final ComponentName BLOCKED_APP_STREAMING_COMPONENT =
-            new ComponentName("android", BlockedAppStreamingActivity.class.getName());
+        void onActivityLaunchBlocked(int displayId, @NonNull ActivityInfo activityInfo,
+                @Nullable IntentSender intentSender);
 
-    /**
-     * For communicating when a secure window shows on the virtual display.
-     */
-    public interface SecureWindowCallback {
         /** Called when a secure window shows on the virtual display. */
-        void onSecureWindowShown(int displayId, int uid);
-    }
+        void onSecureWindowShown(int displayId, @NonNull ActivityInfo activityInfo);
 
-    /** Interface to listen for interception of intents. */
-    public interface IntentListenerCallback {
         /** Returns true when an intent should be intercepted */
-        boolean shouldInterceptIntent(Intent intent);
+        boolean shouldInterceptIntent(@NonNull Intent intent);
     }
 
     /**
@@ -118,7 +116,6 @@
     private final ArraySet<ComponentName> mCrossTaskNavigationExemptions;
     @NonNull
     private final Object mGenericWindowPolicyControllerLock = new Object();
-    @Nullable private final ActivityBlockedCallback mActivityBlockedCallback;
 
     // Do not access mDisplayId and mIsMirrorDisplay directly, instead use waitAndGetDisplayId()
     // and waitAndGetIsMirrorDisplay()
@@ -129,14 +126,12 @@
     @NonNull
     @GuardedBy("mGenericWindowPolicyControllerLock")
     private final ArraySet<Integer> mRunningUids = new ArraySet<>();
-    @Nullable private final ActivityListener mActivityListener;
-    @Nullable private final IntentListenerCallback mIntentListenerCallback;
+    @NonNull private final ActivityListener mActivityListener;
     private final Handler mHandler = new Handler(Looper.getMainLooper());
     @NonNull
     @GuardedBy("mGenericWindowPolicyControllerLock")
     private final ArraySet<RunningAppsChangedListener> mRunningAppsChangedListeners =
             new ArraySet<>();
-    @Nullable private final SecureWindowCallback mSecureWindowCallback;
     @NonNull private final Set<String> mDisplayCategories;
 
     @GuardedBy("mGenericWindowPolicyControllerLock")
@@ -162,12 +157,6 @@
      * @param crossTaskNavigationExemptions The set of components explicitly exempt from the default
      *   navigation policy.
      * @param activityListener Activity listener to listen for activity changes.
-     * @param activityBlockedCallback Callback that is called when an activity is blocked from
-     *   launching.
-     * @param secureWindowCallback Callback that is called when a secure window shows on the
-     *   virtual display.
-     * @param intentListenerCallback Callback that is called to intercept intents when matching
-     *   passed in filters.
      * @param showTasksInHostDeviceRecents whether to show activities in recents on the host device.
      * @param customHomeComponent The component acting as a home activity on the virtual display. If
      *   {@code null}, then the system-default secondary home activity will be used. This is only
@@ -184,10 +173,7 @@
             @NonNull Set<String> activityPolicyPackageExemptions,
             boolean crossTaskNavigationAllowedByDefault,
             @NonNull Set<ComponentName> crossTaskNavigationExemptions,
-            @Nullable ActivityListener activityListener,
-            @Nullable ActivityBlockedCallback activityBlockedCallback,
-            @Nullable SecureWindowCallback secureWindowCallback,
-            @Nullable IntentListenerCallback intentListenerCallback,
+            @NonNull ActivityListener activityListener,
             @NonNull Set<String> displayCategories,
             boolean showTasksInHostDeviceRecents,
             @Nullable ComponentName customHomeComponent) {
@@ -199,11 +185,8 @@
         mActivityPolicyPackageExemptions = new ArraySet<>(activityPolicyPackageExemptions);
         mCrossTaskNavigationAllowedByDefault = crossTaskNavigationAllowedByDefault;
         mCrossTaskNavigationExemptions = new ArraySet<>(crossTaskNavigationExemptions);
-        mActivityBlockedCallback = activityBlockedCallback;
         setInterestedWindowFlags(windowFlags, systemWindowFlags);
         mActivityListener = activityListener;
-        mSecureWindowCallback = secureWindowCallback;
-        mIntentListenerCallback = intentListenerCallback;
         mDisplayCategories = displayCategories;
         mShowTasksInHostDeviceRecents = showTasksInHostDeviceRecents;
         mCustomHomeComponent = customHomeComponent;
@@ -306,8 +289,7 @@
             @Nullable Intent intent, @WindowConfiguration.WindowingMode int windowingMode,
             int launchingFromDisplayId, boolean isNewTask, boolean isResultExpected,
             @Nullable Supplier<IntentSender> intentSender) {
-        if (mIntentListenerCallback != null && intent != null
-                && mIntentListenerCallback.shouldInterceptIntent(intent)) {
+        if (intent != null && mActivityListener.shouldInterceptIntent(intent)) {
             logActivityLaunchBlocked("Virtual device intercepting intent");
             return false;
         }
@@ -391,11 +373,9 @@
         int displayId = waitAndGetDisplayId();
         // The callback is fired only when windowFlags are changed. To let VirtualDevice owner
         // aware that the virtual display has a secure window on top.
-        if ((windowFlags & FLAG_SECURE) != 0 && mSecureWindowCallback != null
-                && displayId != INVALID_DISPLAY) {
+        if ((windowFlags & FLAG_SECURE) != 0 && displayId != INVALID_DISPLAY) {
             // Post callback on the main thread, so it doesn't block activity launching.
-            mHandler.post(() -> mSecureWindowCallback.onSecureWindowShown(displayId,
-                    activityInfo.applicationInfo.uid));
+            mHandler.post(() -> mActivityListener.onSecureWindowShown(displayId, activityInfo));
         }
 
         if (!CompatChanges.isChangeEnabled(ALLOW_SECURE_ACTIVITY_DISPLAY_ON_REMOTE_DEVICE,
@@ -418,7 +398,7 @@
         // Don't send onTopActivityChanged() callback when topActivity is null because it's defined
         // as @NonNull in ActivityListener interface. Sends onDisplayEmpty() callback instead when
         // there is no activity running on virtual display.
-        if (mActivityListener != null && topActivity != null && displayId != INVALID_DISPLAY) {
+        if (topActivity != null && displayId != INVALID_DISPLAY) {
             // Post callback on the main thread so it doesn't block activity launching
             mHandler.post(() ->
                     mActivityListener.onTopActivityChanged(displayId, topActivity, userId));
@@ -431,8 +411,7 @@
             mRunningUids.clear();
             mRunningUids.addAll(runningUids);
             int displayId = waitAndGetDisplayId();
-            if (mActivityListener != null && mRunningUids.isEmpty()
-                    && displayId != INVALID_DISPLAY) {
+            if (mRunningUids.isEmpty() && displayId != INVALID_DISPLAY) {
                 // Post callback on the main thread so it doesn't block activity launching
                 mHandler.post(() -> mActivityListener.onDisplayEmpty(displayId));
             }
@@ -482,9 +461,8 @@
         int displayId = waitAndGetDisplayId();
         // Don't trigger activity blocked callback for mirror displays, because we can't show
         // any activity or presentation on it anyway.
-        if (!waitAndGetIsMirrorDisplay() && mActivityBlockedCallback != null
-                && displayId != INVALID_DISPLAY) {
-            mActivityBlockedCallback.onActivityBlocked(displayId, activityInfo,
+        if (!waitAndGetIsMirrorDisplay() && displayId != INVALID_DISPLAY) {
+            mActivityListener.onActivityLaunchBlocked(displayId, activityInfo,
                     intentSender == null ? null : intentSender.get());
         }
         Counter.logIncrementWithUid(
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
index 4eb50a9..cd2dd3a 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceImpl.java
@@ -40,6 +40,7 @@
 import android.app.ActivityOptions;
 import android.app.PendingIntent;
 import android.app.admin.DevicePolicyManager;
+import android.app.compat.CompatChanges;
 import android.companion.AssociationInfo;
 import android.companion.virtual.ActivityPolicyExemption;
 import android.companion.virtual.IVirtualDevice;
@@ -48,7 +49,6 @@
 import android.companion.virtual.IVirtualDeviceSoundEffectListener;
 import android.companion.virtual.VirtualDevice;
 import android.companion.virtual.VirtualDeviceManager;
-import android.companion.virtual.VirtualDeviceManager.ActivityListener;
 import android.companion.virtual.VirtualDeviceParams;
 import android.companion.virtual.audio.IAudioConfigChangedCallback;
 import android.companion.virtual.audio.IAudioRoutingCallback;
@@ -56,6 +56,8 @@
 import android.companion.virtual.flags.Flags;
 import android.companion.virtual.sensor.VirtualSensor;
 import android.companion.virtual.sensor.VirtualSensorEvent;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledAfter;
 import android.content.AttributionSource;
 import android.content.ComponentName;
 import android.content.Context;
@@ -88,6 +90,7 @@
 import android.media.AudioManager;
 import android.media.audiopolicy.AudioMix;
 import android.os.Binder;
+import android.os.Build;
 import android.os.IBinder;
 import android.os.LocaleList;
 import android.os.Looper;
@@ -132,6 +135,16 @@
 
     private static final String TAG = "VirtualDeviceImpl";
 
+    /**
+     * Do not show a toast on the virtual display when a secure surface is shown after
+     * {@link android.os.Build.VERSION_CODES#VANILLA_ICE_CREAM}. VDM clients should use
+     * {@link VirtualDeviceManager.ActivityListener#onSecureWindowShown} instead to provide
+     * a custom notification if desired.
+     */
+    @ChangeId
+    @EnabledAfter(targetSdkVersion = Build.VERSION_CODES.VANILLA_ICE_CREAM)
+    public static final long DO_NOT_SHOW_TOAST_WHEN_SECURE_SURFACE_SHOWN = 311101667L;
+
     private static final int DEFAULT_VIRTUAL_DISPLAY_FLAGS =
             DisplayManager.VIRTUAL_DISPLAY_FLAG_TOUCH_FEEDBACK_DISABLED
                     | DisplayManager.VIRTUAL_DISPLAY_FLAG_DESTROY_CONTENT_ON_REMOVAL
@@ -182,7 +195,7 @@
     @GuardedBy("mVirtualDeviceLock")
     private final SparseArray<VirtualDisplayWrapper> mVirtualDisplays = new SparseArray<>();
     private IVirtualDeviceActivityListener mActivityListener;
-    private ActivityListener mActivityListenerAdapter = null;
+    private GenericWindowPolicyController.ActivityListener mActivityListenerAdapter = null;
     private IVirtualDeviceSoundEffectListener mSoundEffectListener;
     private final DisplayManagerGlobal mDisplayManager;
     private final DisplayManagerInternal mDisplayManagerInternal;
@@ -207,50 +220,122 @@
     @NonNull
     private final Set<String> mActivityPolicyPackageExemptions = new ArraySet<>();
 
-    private ActivityListener createListenerAdapter() {
-        return new ActivityListener() {
+    private class GwpcActivityListener implements GenericWindowPolicyController.ActivityListener {
 
-            @Override
-            public void onTopActivityChanged(int displayId, @NonNull ComponentName topActivity) {
-                try {
-                    mActivityListener.onTopActivityChanged(displayId, topActivity,
-                            UserHandle.USER_NULL);
-                } catch (RemoteException e) {
-                    Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
-                }
+        @Override
+        public void onTopActivityChanged(int displayId, @NonNull ComponentName topActivity,
+                @UserIdInt int userId) {
+            try {
+                mActivityListener.onTopActivityChanged(displayId, topActivity, userId);
+            } catch (RemoteException e) {
+                Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
+            }
+        }
+
+        @Override
+        public void onDisplayEmpty(int displayId) {
+            try {
+                mActivityListener.onDisplayEmpty(displayId);
+            } catch (RemoteException e) {
+                Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
+            }
+        }
+
+        @Override
+        public void onActivityLaunchBlocked(int displayId, @NonNull ActivityInfo activityInfo,
+                @Nullable IntentSender intentSender) {
+            Intent intent =
+                    BlockedAppStreamingActivity.createIntent(activityInfo, getDisplayName());
+            if (shouldShowBlockedActivityDialog(
+                    activityInfo.getComponentName(), intent.getComponent())) {
+                mContext.startActivityAsUser(
+                        intent.addFlags(
+                                Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK),
+                        ActivityOptions.makeBasic().setLaunchDisplayId(displayId).toBundle(),
+                        UserHandle.SYSTEM);
             }
 
-            @Override
-            public void onTopActivityChanged(int displayId, @NonNull ComponentName topActivity,
-                    @UserIdInt int userId) {
-                try {
-                    mActivityListener.onTopActivityChanged(displayId, topActivity, userId);
-                } catch (RemoteException e) {
-                    Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
-                }
-            }
-
-            @Override
-            public void onDisplayEmpty(int displayId) {
-                try {
-                    mActivityListener.onDisplayEmpty(displayId);
-                } catch (RemoteException e) {
-                    Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
-                }
-            }
-
-            @Override
-            public void onActivityLaunchBlocked(int displayId,
-                    @NonNull ComponentName componentName, @NonNull UserHandle user,
-                    @Nullable IntentSender intentSender) {
+            if (android.companion.virtualdevice.flags.Flags.activityControlApi()) {
                 try {
                     mActivityListener.onActivityLaunchBlocked(
-                            displayId, componentName, user, intentSender);
+                            displayId,
+                            activityInfo.getComponentName(),
+                            UserHandle.getUserHandleForUid(activityInfo.applicationInfo.uid),
+                            intentSender);
                 } catch (RemoteException e) {
                     Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
                 }
             }
-        };
+        }
+
+        @Override
+        public void onSecureWindowShown(int displayId, @NonNull ActivityInfo activityInfo) {
+            if (android.companion.virtualdevice.flags.Flags.activityControlApi()) {
+                try {
+                    mActivityListener.onSecureWindowShown(
+                            displayId,
+                            activityInfo.getComponentName(),
+                            UserHandle.getUserHandleForUid(activityInfo.applicationInfo.uid));
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Unable to call mActivityListener for display: " + displayId, e);
+                }
+
+                if (CompatChanges.isChangeEnabled(DO_NOT_SHOW_TOAST_WHEN_SECURE_SURFACE_SHOWN,
+                        mOwnerPackageName,  UserHandle.getUserHandleForUid(mOwnerUid))) {
+                    return;
+                }
+            }
+
+            // If a virtual display isn't secure, the screen can't be captured. Show a warning toast
+            // if the secure window is shown on a non-secure virtual display.
+            DisplayManager displayManager = mContext.getSystemService(DisplayManager.class);
+            Display display = displayManager.getDisplay(displayId);
+            if ((display.getFlags() & Display.FLAG_SECURE) == 0) {
+                showToastWhereUidIsRunning(activityInfo.applicationInfo.uid,
+                        com.android.internal.R.string.vdm_secure_window,
+                        Toast.LENGTH_LONG, mContext.getMainLooper());
+
+                Counter.logIncrementWithUid(
+                        "virtual_devices.value_secure_window_blocked_count",
+                        mAttributionSource.getUid());
+            }
+        }
+
+        /**
+         * Intercepts intent when matching any of the IntentFilter of any interceptor. Returns true
+         * if the intent matches any filter notifying the DisplayPolicyController to abort the
+         * activity launch to be replaced by the interception.
+         */
+        @Override
+        public boolean shouldInterceptIntent(@NonNull Intent intent) {
+            synchronized (mVirtualDeviceLock) {
+                boolean hasInterceptedIntent = false;
+                for (Map.Entry<IBinder, IntentFilter> interceptor
+                        : mIntentInterceptors.entrySet()) {
+                    IntentFilter intentFilter = interceptor.getValue();
+                    // Explicitly match the actions because the intent filter will match any intent
+                    // without an explicit action. If the intent has no action, then require that
+                    // there are no actions specified in the filter either.
+                    boolean explicitActionMatch =
+                            intent.getAction() != null || intentFilter.countActions() == 0;
+                    if (explicitActionMatch && intentFilter.match(
+                            intent.getAction(), intent.getType(), intent.getScheme(),
+                            intent.getData(), intent.getCategories(), TAG) >= 0) {
+                        try {
+                            // For privacy reasons, only returning the intents action and data.
+                            // Any other required field will require a review.
+                            IVirtualDeviceIntentInterceptor.Stub.asInterface(interceptor.getKey())
+                                    .onIntentIntercepted(
+                                            new Intent(intent.getAction(), intent.getData()));
+                            hasInterceptedIntent = true;
+                        } catch (RemoteException e) {
+                            Slog.w(TAG, "Unable to call mActivityListener", e);
+                        }
+                    }
+                }
+                return hasInterceptedIntent;
+            }
+        }
     }
 
     VirtualDeviceImpl(
@@ -1290,7 +1375,7 @@
                 Flags.vdmCustomHome() ? mParams.getHomeComponent() : null;
 
         if (mActivityListenerAdapter == null) {
-            mActivityListenerAdapter = createListenerAdapter();
+            mActivityListenerAdapter = new GwpcActivityListener();
         }
 
         final GenericWindowPolicyController gwpc = new GenericWindowPolicyController(
@@ -1306,9 +1391,6 @@
                         ? mParams.getBlockedCrossTaskNavigations()
                         : mParams.getAllowedCrossTaskNavigations(),
                 mActivityListenerAdapter,
-                this::onActivityBlocked,
-                this::onSecureWindowShown,
-                this::shouldInterceptIntent,
                 displayCategories,
                 showTasksInHostDeviceRecents,
                 homeComponent);
@@ -1378,28 +1460,6 @@
         }
     }
 
-    @RequiresPermission(android.Manifest.permission.INTERACT_ACROSS_USERS)
-    private void onActivityBlocked(int displayId, ActivityInfo activityInfo,
-            IntentSender intentSender) {
-        Intent intent = BlockedAppStreamingActivity.createIntent(activityInfo, getDisplayName());
-        if (shouldShowBlockedActivityDialog(
-                activityInfo.getComponentName(), intent.getComponent())) {
-            mContext.startActivityAsUser(
-                    intent.addFlags(
-                            Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK),
-                    ActivityOptions.makeBasic().setLaunchDisplayId(displayId).toBundle(),
-                    UserHandle.SYSTEM);
-        }
-
-        if (android.companion.virtualdevice.flags.Flags.activityControlApi()) {
-            mActivityListenerAdapter.onActivityLaunchBlocked(
-                    displayId,
-                    activityInfo.getComponentName(),
-                    UserHandle.getUserHandleForUid(activityInfo.applicationInfo.uid),
-                    intentSender);
-        }
-    }
-
     private boolean shouldShowBlockedActivityDialog(ComponentName blockedComponent,
             ComponentName blockedAppStreamingActivityComponent) {
         if (Objects.equals(blockedComponent, blockedAppStreamingActivityComponent)) {
@@ -1414,27 +1474,6 @@
         return getDevicePolicy(POLICY_TYPE_BLOCKED_ACTIVITY) == DEVICE_POLICY_DEFAULT;
     }
 
-    private void onSecureWindowShown(int displayId, int uid) {
-        synchronized (mVirtualDeviceLock) {
-            if (!mVirtualDisplays.contains(displayId)) {
-                return;
-            }
-        }
-
-        // If a virtual display isn't secure, the screen can't be captured. Show a warning toast
-        // if the secure window is shown on a non-secure virtual display.
-        DisplayManager displayManager = mContext.getSystemService(DisplayManager.class);
-        Display display = displayManager.getDisplay(displayId);
-        if ((display.getFlags() & Display.FLAG_SECURE) == 0) {
-            showToastWhereUidIsRunning(uid, com.android.internal.R.string.vdm_secure_window,
-                    Toast.LENGTH_LONG, mContext.getMainLooper());
-
-            Counter.logIncrementWithUid(
-                    "virtual_devices.value_secure_window_blocked_count",
-                    mAttributionSource.getUid());
-        }
-    }
-
     private ArraySet<UserHandle> getAllowedUserHandles() {
         ArraySet<UserHandle> result = new ArraySet<>();
         final long token = Binder.clearCallingIdentity();
@@ -1621,40 +1660,6 @@
         }
     }
 
-    /**
-     * Intercepts intent when matching any of the IntentFilter of any interceptor. Returns true if
-     * the intent matches any filter notifying the DisplayPolicyController to abort the
-     * activity launch to be replaced by the interception.
-     */
-    private boolean shouldInterceptIntent(Intent intent) {
-        synchronized (mVirtualDeviceLock) {
-            boolean hasInterceptedIntent = false;
-            for (Map.Entry<IBinder, IntentFilter> interceptor : mIntentInterceptors.entrySet()) {
-                IntentFilter intentFilter = interceptor.getValue();
-                // Explicitly match the actions because the intent filter will match any intent
-                // without an explicit action. If the intent has no action, then require that there
-                // are no actions specified in the filter either.
-                boolean explicitActionMatch =
-                        intent.getAction() != null || intentFilter.countActions() == 0;
-                if (explicitActionMatch && intentFilter.match(
-                        intent.getAction(), intent.getType(), intent.getScheme(), intent.getData(),
-                        intent.getCategories(), TAG) >= 0) {
-                    try {
-                        // For privacy reasons, only returning the intents action and data. Any
-                        // other required field will require a review.
-                        IVirtualDeviceIntentInterceptor.Stub.asInterface(interceptor.getKey())
-                            .onIntentIntercepted(new Intent(intent.getAction(), intent.getData()));
-                        hasInterceptedIntent = true;
-                    } catch (RemoteException e) {
-                        Slog.w(TAG, "Unable to call mVirtualDeviceIntentInterceptor", e);
-                    }
-                }
-            }
-
-            return hasInterceptedIntent;
-        }
-    }
-
     interface PendingTrampolineCallback {
         /**
          * Called when the callback should start waiting for the given pending trampoline.
diff --git a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceLog.java b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceLog.java
index b0bacfd..fed153f 100644
--- a/services/companion/java/com/android/server/companion/virtual/VirtualDeviceLog.java
+++ b/services/companion/java/com/android/server/companion/virtual/VirtualDeviceLog.java
@@ -48,9 +48,6 @@
     void logCreated(int deviceId, int ownerUid) {
         final long token = Binder.clearCallingIdentity();
         try {
-            if (!Flags.dumpHistory()) {
-                return;
-            }
             addEntry(new LogEntry(TYPE_CREATED, deviceId, System.currentTimeMillis(), ownerUid));
         } finally {
             Binder.restoreCallingIdentity(token);
@@ -60,9 +57,6 @@
     void logClosed(int deviceId, int ownerUid) {
         final long token = Binder.clearCallingIdentity();
         try {
-            if (!Flags.dumpHistory()) {
-                return;
-            }
             addEntry(new LogEntry(TYPE_CLOSED, deviceId, System.currentTimeMillis(), ownerUid));
         } finally {
             Binder.restoreCallingIdentity(token);
@@ -79,9 +73,6 @@
     void dump(PrintWriter pw) {
         final long token = Binder.clearCallingIdentity();
         try {
-            if (!Flags.dumpHistory()) {
-                return;
-            }
             pw.println("VirtualDevice Log:");
             UidToPackageNameCache packageNameCache = new UidToPackageNameCache(
                     mContext.getPackageManager());
diff --git a/services/companion/java/com/android/server/companion/virtual/flags.aconfig b/services/companion/java/com/android/server/companion/virtual/flags.aconfig
deleted file mode 100644
index 616f5d0..0000000
--- a/services/companion/java/com/android/server/companion/virtual/flags.aconfig
+++ /dev/null
@@ -1,11 +0,0 @@
-# OLD PACKAGE, DO NOT USE: Prefer `flags.aconfig` in core/java/android/companion/virtual
-# (or other custom files) to define your flags
-package: "com.android.server.companion.virtual"
-container: "system"
-
-flag {
-  name: "dump_history"
-  namespace: "virtual_devices"
-  description: "This flag controls if a history of virtual devices is shown in dumpsys virtualdevices"
-  bug: "293114719"
-}
\ No newline at end of file
diff --git a/services/core/java/com/android/server/BatteryService.java b/services/core/java/com/android/server/BatteryService.java
index 2de4482..1470e9a 100644
--- a/services/core/java/com/android/server/BatteryService.java
+++ b/services/core/java/com/android/server/BatteryService.java
@@ -23,6 +23,7 @@
 import static com.android.server.health.Utils.copyV1Battery;
 
 import android.annotation.Nullable;
+import android.annotation.SuppressLint;
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
 import android.app.AppOpsManager;
@@ -67,6 +68,7 @@
 
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.logging.MetricsLogger;
+import com.android.internal.os.SomeArgs;
 import com.android.internal.util.DumpUtils;
 import com.android.server.am.BatteryStatsService;
 import com.android.server.health.HealthServiceWrapper;
@@ -207,18 +209,18 @@
     private final CopyOnWriteArraySet<BatteryManagerInternal.ChargingPolicyChangeListener>
             mChargingPolicyChangeListeners = new CopyOnWriteArraySet<>();
 
-    private Bundle mBatteryChangedOptions = BroadcastOptions.makeBasic()
+    private static final Bundle BATTERY_CHANGED_OPTIONS = BroadcastOptions.makeBasic()
             .setDeliveryGroupPolicy(BroadcastOptions.DELIVERY_GROUP_POLICY_MOST_RECENT)
             .setDeferralPolicy(BroadcastOptions.DEFERRAL_POLICY_UNTIL_ACTIVE)
             .toBundle();
     /** Used for both connected/disconnected, so match using key */
-    private Bundle mPowerOptions = BroadcastOptions.makeBasic()
+    private static final Bundle POWER_OPTIONS = BroadcastOptions.makeBasic()
             .setDeliveryGroupPolicy(BroadcastOptions.DELIVERY_GROUP_POLICY_MOST_RECENT)
             .setDeliveryGroupMatchingKey("android", Intent.ACTION_POWER_CONNECTED)
             .setDeferralPolicy(BroadcastOptions.DEFERRAL_POLICY_UNTIL_ACTIVE)
             .toBundle();
     /** Used for both low/okay, so match using key */
-    private Bundle mBatteryOptions = BroadcastOptions.makeBasic()
+    private static final Bundle BATTERY_OPTIONS = BroadcastOptions.makeBasic()
             .setDeliveryGroupPolicy(BroadcastOptions.DELIVERY_GROUP_POLICY_MOST_RECENT)
             .setDeliveryGroupMatchingKey("android", Intent.ACTION_BATTERY_OKAY)
             .setDeferralPolicy(BroadcastOptions.DEFERRAL_POLICY_UNTIL_ACTIVE)
@@ -226,11 +228,60 @@
 
     private MetricsLogger mMetricsLogger;
 
+    private static final int MSG_BROADCAST_BATTERY_CHANGED = 1;
+    private static final int MSG_BROADCAST_POWER_CONNECTION_CHANGED = 2;
+    private static final int MSG_BROADCAST_BATTERY_LOW_OKAY = 3;
+
+    private final Handler.Callback mLocalCallback = msg -> {
+        switch (msg.what) {
+            case MSG_BROADCAST_BATTERY_CHANGED: {
+                final SomeArgs args = (SomeArgs) msg.obj;
+                final Context context;
+                final Intent intent;
+                try {
+                    context = (Context) args.arg1;
+                    intent = (Intent) args.arg2;
+                } finally {
+                    args.recycle();
+                }
+                broadcastBatteryChangedIntent(context, intent, BATTERY_CHANGED_OPTIONS);
+                return true;
+            }
+            case MSG_BROADCAST_POWER_CONNECTION_CHANGED: {
+                final SomeArgs args = (SomeArgs) msg.obj;
+                final Context context;
+                final Intent intent;
+                try {
+                    context = (Context) args.arg1;
+                    intent = (Intent) args.arg2;
+                } finally {
+                    args.recycle();
+                }
+                sendBroadcastToAllUsers(context, intent, POWER_OPTIONS);
+                return true;
+            }
+            case MSG_BROADCAST_BATTERY_LOW_OKAY: {
+                final SomeArgs args = (SomeArgs) msg.obj;
+                final Context context;
+                final Intent intent;
+                try {
+                    context = (Context) args.arg1;
+                    intent = (Intent) args.arg2;
+                } finally {
+                    args.recycle();
+                }
+                sendBroadcastToAllUsers(context, intent, BATTERY_OPTIONS);
+                return true;
+            }
+        }
+        return false;
+    };
+
     public BatteryService(Context context) {
         super(context);
 
         mContext = context;
-        mHandler = new Handler(true /*async*/);
+        mHandler = new Handler(mLocalCallback, true /*async*/);
         mLed = new Led(context, getLocalService(LightsManager.class));
         mBatteryStats = BatteryStatsService.getService();
         mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
@@ -660,25 +711,43 @@
                 final Intent statusIntent = new Intent(Intent.ACTION_POWER_CONNECTED);
                 statusIntent.setFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT);
                 statusIntent.putExtra(BatteryManager.EXTRA_SEQUENCE, mSequence);
-                mHandler.post(new Runnable() {
-                    @Override
-                    public void run() {
-                        mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
-                                mPowerOptions);
-                    }
-                });
+                if (com.android.server.flags.Flags.consolidateBatteryChangeEvents()) {
+                    mHandler.removeMessages(MSG_BROADCAST_POWER_CONNECTION_CHANGED);
+                    final SomeArgs args = SomeArgs.obtain();
+                    args.arg1 = mContext;
+                    args.arg2 = statusIntent;
+                    mHandler.obtainMessage(MSG_BROADCAST_POWER_CONNECTION_CHANGED, args)
+                            .sendToTarget();
+                } else {
+                    mHandler.post(new Runnable() {
+                        @Override
+                        public void run() {
+                            mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
+                                    POWER_OPTIONS);
+                        }
+                    });
+                }
             }
             else if (mPlugType == 0 && mLastPlugType != 0) {
                 final Intent statusIntent = new Intent(Intent.ACTION_POWER_DISCONNECTED);
                 statusIntent.setFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT);
                 statusIntent.putExtra(BatteryManager.EXTRA_SEQUENCE, mSequence);
-                mHandler.post(new Runnable() {
-                    @Override
-                    public void run() {
-                        mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
-                                mPowerOptions);
-                    }
-                });
+                if (com.android.server.flags.Flags.consolidateBatteryChangeEvents()) {
+                    mHandler.removeMessages(MSG_BROADCAST_POWER_CONNECTION_CHANGED);
+                    final SomeArgs args = SomeArgs.obtain();
+                    args.arg1 = mContext;
+                    args.arg2 = statusIntent;
+                    mHandler.obtainMessage(MSG_BROADCAST_POWER_CONNECTION_CHANGED, args)
+                            .sendToTarget();
+                } else {
+                    mHandler.post(new Runnable() {
+                        @Override
+                        public void run() {
+                            mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
+                                    POWER_OPTIONS);
+                        }
+                    });
+                }
             }
 
             if (shouldSendBatteryLowLocked()) {
@@ -686,26 +755,44 @@
                 final Intent statusIntent = new Intent(Intent.ACTION_BATTERY_LOW);
                 statusIntent.setFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT);
                 statusIntent.putExtra(BatteryManager.EXTRA_SEQUENCE, mSequence);
-                mHandler.post(new Runnable() {
-                    @Override
-                    public void run() {
-                        mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
-                                mBatteryOptions);
-                    }
-                });
+                if (com.android.server.flags.Flags.consolidateBatteryChangeEvents()) {
+                    mHandler.removeMessages(MSG_BROADCAST_BATTERY_LOW_OKAY);
+                    final SomeArgs args = SomeArgs.obtain();
+                    args.arg1 = mContext;
+                    args.arg2 = statusIntent;
+                    mHandler.obtainMessage(MSG_BROADCAST_BATTERY_LOW_OKAY, args)
+                            .sendToTarget();
+                } else {
+                    mHandler.post(new Runnable() {
+                        @Override
+                        public void run() {
+                            mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
+                                    BATTERY_OPTIONS);
+                        }
+                    });
+                }
             } else if (mSentLowBatteryBroadcast &&
                     mHealthInfo.batteryLevel >= mLowBatteryCloseWarningLevel) {
                 mSentLowBatteryBroadcast = false;
                 final Intent statusIntent = new Intent(Intent.ACTION_BATTERY_OKAY);
                 statusIntent.setFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT);
                 statusIntent.putExtra(BatteryManager.EXTRA_SEQUENCE, mSequence);
-                mHandler.post(new Runnable() {
-                    @Override
-                    public void run() {
-                        mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
-                                mBatteryOptions);
-                    }
-                });
+                if (com.android.server.flags.Flags.consolidateBatteryChangeEvents()) {
+                    mHandler.removeMessages(MSG_BROADCAST_BATTERY_LOW_OKAY);
+                    final SomeArgs args = SomeArgs.obtain();
+                    args.arg1 = mContext;
+                    args.arg2 = statusIntent;
+                    mHandler.obtainMessage(MSG_BROADCAST_BATTERY_LOW_OKAY, args)
+                            .sendToTarget();
+                } else {
+                    mHandler.post(new Runnable() {
+                        @Override
+                        public void run() {
+                            mContext.sendBroadcastAsUser(statusIntent, UserHandle.ALL, null,
+                                    BATTERY_OPTIONS);
+                        }
+                    });
+                }
             }
 
             // We are doing this after sending the above broadcasts, so anything processing
@@ -777,8 +864,16 @@
                     + ", info:" + mHealthInfo.toString());
         }
 
-        mHandler.post(() -> broadcastBatteryChangedIntent(mContext,
-                intent, mBatteryChangedOptions));
+        if (com.android.server.flags.Flags.consolidateBatteryChangeEvents()) {
+            mHandler.removeMessages(MSG_BROADCAST_BATTERY_CHANGED);
+            final SomeArgs args = SomeArgs.obtain();
+            args.arg1 = mContext;
+            args.arg2 = intent;
+            mHandler.obtainMessage(MSG_BROADCAST_BATTERY_CHANGED, args).sendToTarget();
+        } else {
+            mHandler.post(() -> broadcastBatteryChangedIntent(mContext,
+                    intent, BATTERY_CHANGED_OPTIONS));
+        }
     }
 
     private static void broadcastBatteryChangedIntent(Context context, Intent intent,
@@ -1307,6 +1402,12 @@
         Trace.traceEnd(Trace.TRACE_TAG_SYSTEM_SERVER);
     }
 
+    @SuppressLint("AndroidFrameworkRequiresPermission")
+    private static void sendBroadcastToAllUsers(Context context, Intent intent,
+            Bundle options) {
+        context.sendBroadcastAsUser(intent, UserHandle.ALL, null, options);
+    }
+
     private final class Led {
         // must match: config_notificationsBatteryLowBehavior in config.xml
         static final int LOW_BATTERY_BEHAVIOR_DEFAULT = 0;
diff --git a/services/core/java/com/android/server/EventLogTags.logtags b/services/core/java/com/android/server/EventLogTags.logtags
index 361b818..fd512a6 100644
--- a/services/core/java/com/android/server/EventLogTags.logtags
+++ b/services/core/java/com/android/server/EventLogTags.logtags
@@ -94,6 +94,8 @@
 275534 notification_unautogrouped (key|3)
 # when a notification is adjusted via assistant
 27535 notification_adjusted (key|3),(adjustment_type|3),(new_value|3)
+# when a notification cancellation is prevented by the system
+27536 notification_cancel_prevented (key|3)
 
 # ---------------------------
 # Watchdog.java
diff --git a/services/core/java/com/android/server/PackageWatchdog.java b/services/core/java/com/android/server/PackageWatchdog.java
index 47203fb..fbe593f 100644
--- a/services/core/java/com/android/server/PackageWatchdog.java
+++ b/services/core/java/com/android/server/PackageWatchdog.java
@@ -20,6 +20,8 @@
 import static android.content.Intent.ACTION_SHUTDOWN;
 import static android.service.watchdog.ExplicitHealthCheckService.PackageConfig;
 
+import static com.android.server.crashrecovery.CrashRecoveryUtils.dumpCrashRecoveryEvents;
+
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.annotation.IntDef;
@@ -44,6 +46,7 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.AtomicFile;
+import android.util.IndentingPrintWriter;
 import android.util.LongArrayQueue;
 import android.util.Slog;
 import android.util.Xml;
@@ -51,7 +54,6 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.BackgroundThread;
-import com.android.internal.util.IndentingPrintWriter;
 import com.android.internal.util.XmlUtils;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
@@ -72,6 +74,7 @@
 import java.io.InputStream;
 import java.io.ObjectInputStream;
 import java.io.ObjectOutputStream;
+import java.io.PrintWriter;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
@@ -1265,18 +1268,21 @@
 
 
     /** Dump status of every observer in mAllObservers. */
-    public void dump(IndentingPrintWriter pw) {
-        pw.println("Package Watchdog status");
-        pw.increaseIndent();
+    public void dump(PrintWriter pw) {
+        IndentingPrintWriter ipw = new IndentingPrintWriter(pw, "  ");
+        ipw.println("Package Watchdog status");
+        ipw.increaseIndent();
         synchronized (mLock) {
             for (String observerName : mAllObservers.keySet()) {
-                pw.println("Observer name: " + observerName);
-                pw.increaseIndent();
+                ipw.println("Observer name: " + observerName);
+                ipw.increaseIndent();
                 ObserverInternal observerInternal = mAllObservers.get(observerName);
-                observerInternal.dump(pw);
-                pw.decreaseIndent();
+                observerInternal.dump(ipw);
+                ipw.decreaseIndent();
             }
         }
+        ipw.decreaseIndent();
+        dumpCrashRecoveryEvents(ipw);
     }
 
     @VisibleForTesting
diff --git a/services/core/java/com/android/server/RescueParty.java b/services/core/java/com/android/server/RescueParty.java
index bba97fa..cadceb5 100644
--- a/services/core/java/com/android/server/RescueParty.java
+++ b/services/core/java/com/android/server/RescueParty.java
@@ -18,7 +18,7 @@
 
 import static android.provider.DeviceConfig.Properties;
 
-import static com.android.server.pm.PackageManagerServiceUtils.logCriticalInfo;
+import static com.android.server.crashrecovery.CrashRecoveryUtils.logCrashRecoveryEvent;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
@@ -291,13 +291,13 @@
                 Properties properties = new Properties.Builder(namespaceToReset).build();
                 try {
                     if (!DeviceConfig.setProperties(properties)) {
-                        logCriticalInfo(Log.ERROR, "Failed to clear properties under "
+                        logCrashRecoveryEvent(Log.ERROR, "Failed to clear properties under "
                             + namespaceToReset
                             + ". Running `device_config get_sync_disabled_for_tests` will confirm"
                             + " if config-bulk-update is enabled.");
                     }
                 } catch (DeviceConfig.BadConfigException exception) {
-                    logCriticalInfo(Log.WARN, "namespace " + namespaceToReset
+                    logCrashRecoveryEvent(Log.WARN, "namespace " + namespaceToReset
                             + " is already banned, skip reset.");
                 }
             }
@@ -528,7 +528,7 @@
             if (!TextUtils.isEmpty(failedPackage)) {
                 successMsg += " for package " + failedPackage;
             }
-            logCriticalInfo(Log.DEBUG, successMsg);
+            logCrashRecoveryEvent(Log.DEBUG, successMsg);
         } catch (Throwable t) {
             logRescueException(level, failedPackage, t);
         }
@@ -687,7 +687,7 @@
         if (!TextUtils.isEmpty(failedPackageName)) {
             failureMsg += " for package " + failedPackageName;
         }
-        logCriticalInfo(Log.ERROR, failureMsg + ": " + msg);
+        logCrashRecoveryEvent(Log.ERROR, failureMsg + ": " + msg);
     }
 
     private static int mapRescueLevelToUserImpact(int rescueLevel) {
diff --git a/services/core/java/com/android/server/StorageManagerService.java b/services/core/java/com/android/server/StorageManagerService.java
index 07e5f2e..d86bae1 100644
--- a/services/core/java/com/android/server/StorageManagerService.java
+++ b/services/core/java/com/android/server/StorageManagerService.java
@@ -224,6 +224,9 @@
     /** Extended timeout for the system server watchdog for vold#partition operation. */
     private static final int PARTITION_OPERATION_WATCHDOG_TIMEOUT_MS = 3 * 60 * 1000;
 
+    private static final Pattern OBB_FILE_PATH = Pattern.compile(
+            "(?i)(^/storage/[^/]+/(?:([0-9]+)/)?Android/obb/)([^/]+)/([^/]+\\.obb)");
+
     @GuardedBy("mLock")
     private final Set<Integer> mFuseMountedUser = new ArraySet<>();
 
@@ -3144,7 +3147,9 @@
         Objects.requireNonNull(rawPath, "rawPath cannot be null");
         Objects.requireNonNull(canonicalPath, "canonicalPath cannot be null");
         Objects.requireNonNull(token, "token cannot be null");
-        Objects.requireNonNull(obbInfo, "obbIfno cannot be null");
+        Objects.requireNonNull(obbInfo, "obbInfo cannot be null");
+
+        validateObbInfo(obbInfo, rawPath);
 
         final int callingUid = Binder.getCallingUid();
         final ObbState obbState = new ObbState(rawPath, canonicalPath,
@@ -3156,6 +3161,34 @@
             Slog.i(TAG, "Send to OBB handler: " + action.toString());
     }
 
+    private void validateObbInfo(ObbInfo obbInfo, String rawPath) {
+        String obbFilePath;
+        try {
+            obbFilePath = new File(rawPath).getCanonicalPath();
+        } catch (IOException ex) {
+            throw new RuntimeException("Failed to resolve path" + rawPath + " : " + ex);
+        }
+
+        Matcher matcher = OBB_FILE_PATH.matcher(obbFilePath);
+
+        if (matcher.matches()) {
+            int userId = UserHandle.getUserId(Binder.getCallingUid());
+            String pathUserId = matcher.group(2);
+            String pathPackageName = matcher.group(3);
+            if ((pathUserId != null && Integer.parseInt(pathUserId) != userId)
+                    || (pathUserId == null && userId != mCurrentUserId)) {
+                throw new SecurityException(
+                        "Path " + obbFilePath + "does not correspond to calling userId " + userId);
+            }
+            if (obbInfo != null && !obbInfo.packageName.equals(pathPackageName)) {
+                throw new SecurityException("Path " + obbFilePath
+                        + " does not contain package name " + pathPackageName);
+            }
+        } else {
+            throw new SecurityException("Invalid path to Obb file : " + obbFilePath);
+        }
+    }
+
     @Override
     public void unmountObb(String rawPath, boolean force, IObbActionListener token, int nonce) {
         Objects.requireNonNull(rawPath, "rawPath cannot be null");
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index d80b38e..e13b0a4 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -5426,7 +5426,9 @@
             for (int i=0; i<intents.length; i++) {
                 Intent intent = intents[i];
                 if (intent != null) {
-                    intent.prepareToEnterSystemServer();
+                    if (intent.hasFileDescriptors()) {
+                        throw new IllegalArgumentException("File descriptors passed in Intent");
+                    }
                     if (type == ActivityManager.INTENT_SENDER_BROADCAST &&
                             (intent.getFlags()&Intent.FLAG_RECEIVER_BOOT_UPGRADE) != 0) {
                         throw new IllegalArgumentException(
@@ -5459,6 +5461,7 @@
                         }
                     }
                     intents[i] = new Intent(intent);
+                    intents[i].removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
                 }
             }
             if (resolvedTypes != null && resolvedTypes.length != intents.length) {
@@ -13591,7 +13594,12 @@
         enforceNotIsolatedCaller("startService");
         enforceAllowedToStartOrBindServiceIfSdkSandbox(service);
         if (service != null) {
-            service.prepareToEnterSystemServer();
+            // Refuse possible leaked file descriptors
+            if (service.hasFileDescriptors()) {
+                throw new IllegalArgumentException("File descriptors passed in Intent");
+            }
+            // Remove existing mismatch flag so it can be properly updated later
+            service.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
         }
 
         if (callingPackage == null) {
@@ -13828,7 +13836,12 @@
         enforceAllowedToStartOrBindServiceIfSdkSandbox(service);
 
         if (service != null) {
-            service.prepareToEnterSystemServer();
+            // Refuse possible leaked file descriptors
+            if (service.hasFileDescriptors()) {
+                throw new IllegalArgumentException("File descriptors passed in Intent");
+            }
+            // Remove existing mismatch flag so it can be properly updated later
+            service.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
         }
 
         if (callingPackage == null) {
@@ -16248,43 +16261,55 @@
 
         boolean closeFd = true;
         try {
-            synchronized (mProcLock) {
-                if (fd == null) {
-                    throw new IllegalArgumentException("null fd");
-                }
-                mBinderTransactionTrackingEnabled = false;
+            Objects.requireNonNull(fd);
 
-                PrintWriter pw = new FastPrintWriter(new FileOutputStream(fd.getFileDescriptor()));
-                pw.println("Binder transaction traces for all processes.\n");
-                mProcessList.forEachLruProcessesLOSP(true, process -> {
+            record ProcessToDump(String processName, IApplicationThread thread) { }
+
+            PrintWriter pw = new FastPrintWriter(new FileOutputStream(fd.getFileDescriptor()));
+            pw.println("Binder transaction traces for all processes.\n");
+            final ArrayList<ProcessToDump> processes = new ArrayList<>();
+            synchronized (mProcLock) {
+                // Since dumping binder transactions is a long-running operation, we can't do it
+                // with mProcLock held. Do the initial verification here, and save the processes
+                // to dump later outside the lock.
+                final ArrayList<ProcessRecord> unverifiedProcesses =
+                        new ArrayList<>(mProcessList.getLruProcessesLOSP());
+                for (int i = 0, size = unverifiedProcesses.size(); i < size; i++) {
+                    ProcessRecord process = unverifiedProcesses.get(i);
                     final IApplicationThread thread = process.getThread();
                     if (!processSanityChecksLPr(process, thread)) {
-                        return;
+                        continue;
                     }
-
-                    pw.println("Traces for process: " + process.processName);
-                    pw.flush();
-                    try {
-                        TransferPipe tp = new TransferPipe();
-                        try {
-                            thread.stopBinderTrackingAndDump(tp.getWriteFd());
-                            tp.go(fd.getFileDescriptor());
-                        } finally {
-                            tp.kill();
-                        }
-                    } catch (IOException e) {
-                        pw.println("Failure while dumping IPC traces from " + process +
-                                ".  Exception: " + e);
-                        pw.flush();
-                    } catch (RemoteException e) {
-                        pw.println("Got a RemoteException while dumping IPC traces from " +
-                                process + ".  Exception: " + e);
-                        pw.flush();
-                    }
-                });
-                closeFd = false;
-                return true;
+                    processes.add(new ProcessToDump(process.processName, process.getThread()));
+                }
+                mBinderTransactionTrackingEnabled = false;
             }
+            for (int i = 0, size = processes.size(); i < size; i++) {
+                final String processName = processes.get(i).processName();
+                final IApplicationThread thread = processes.get(i).thread();
+
+                pw.println("Traces for process: " + processName);
+                pw.flush();
+                try {
+                    TransferPipe tp = new TransferPipe();
+                    try {
+                        thread.stopBinderTrackingAndDump(tp.getWriteFd());
+                        tp.go(fd.getFileDescriptor());
+                    } finally {
+                        tp.kill();
+                    }
+                } catch (IOException e) {
+                    pw.println("Failure while dumping IPC traces from " + processName +
+                            ".  Exception: " + e);
+                    pw.flush();
+                } catch (RemoteException e) {
+                    pw.println("Got a RemoteException while dumping IPC traces from " +
+                            processName + ".  Exception: " + e);
+                    pw.flush();
+                }
+            }
+            closeFd = false;
+            return true;
         } finally {
             if (fd != null && closeFd) {
                 try {
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 8e87342..955b75d 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -123,31 +123,16 @@
 import com.android.server.net.BaseNetworkObserver;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.power.optimization.Flags;
-import com.android.server.power.stats.AggregatedPowerStatsConfig;
-import com.android.server.power.stats.AmbientDisplayPowerStatsProcessor;
-import com.android.server.power.stats.AudioPowerStatsProcessor;
 import com.android.server.power.stats.BatteryExternalStatsWorker;
 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.BluetoothPowerStatsProcessor;
-import com.android.server.power.stats.CameraPowerStatsProcessor;
-import com.android.server.power.stats.CpuPowerStatsProcessor;
-import com.android.server.power.stats.CustomEnergyConsumerPowerStatsProcessor;
-import com.android.server.power.stats.FlashlightPowerStatsProcessor;
-import com.android.server.power.stats.GnssPowerStatsProcessor;
-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.PowerAttributor;
 import com.android.server.power.stats.PowerStatsScheduler;
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.PowerStatsUidResolver;
-import com.android.server.power.stats.ScreenPowerStatsProcessor;
-import com.android.server.power.stats.SensorPowerStatsProcessor;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
-import com.android.server.power.stats.VideoPowerStatsProcessor;
-import com.android.server.power.stats.WifiPowerStatsProcessor;
+import com.android.server.power.stats.processor.MultiStatePowerAttributor;
 import com.android.server.power.stats.wakeups.CpuWakeupStats;
 
 import java.io.File;
@@ -207,7 +192,7 @@
     private final AtomicFile mConfigFile;
     private final BatteryStats.BatteryStatsDumpHelper mDumpHelper;
     private final PowerStatsUidResolver mPowerStatsUidResolver = new PowerStatsUidResolver();
-    private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
+    private final PowerAttributor mPowerAttributor;
 
     private volatile boolean mMonitorEnabled = true;
 
@@ -445,14 +430,12 @@
             mStats.startTrackingSystemServerCpuTime();
         }
 
-        mAggregatedPowerStatsConfig = createAggregatedPowerStatsConfig();
-        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler, mAggregatedPowerStatsConfig);
+        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler);
+        mPowerAttributor = new MultiStatePowerAttributor(mContext, mPowerStatsStore, mPowerProfile,
+                mCpuScalingPolicies, mPowerStatsUidResolver);
         mPowerStatsScheduler = createPowerStatsScheduler(mContext);
-        PowerStatsExporter powerStatsExporter =
-                new PowerStatsExporter(mPowerStatsStore,
-                        new PowerStatsAggregator(mAggregatedPowerStatsConfig, mStats.getHistory()));
         mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context,
-                powerStatsExporter, mPowerProfile, mCpuScalingPolicies,
+                mPowerAttributor, mPowerProfile, mCpuScalingPolicies,
                 mPowerStatsStore, Clock.SYSTEM_CLOCK);
         mStats.saveBatteryUsageStatsOnReset(mBatteryUsageStatsProvider, mPowerStatsStore);
         mDumpHelper = new BatteryStatsDumpHelperImpl(mBatteryUsageStatsProvider);
@@ -472,154 +455,11 @@
                             onAlarmListener, aHandler);
                 };
         return new PowerStatsScheduler(mStats::schedulePowerStatsSampleCollection,
-                new PowerStatsAggregator(mAggregatedPowerStatsConfig,
-                        mStats.getHistory()), aggregatedPowerStatsSpanDuration,
+                mStats.getHistory(), mPowerAttributor, aggregatedPowerStatsSpanDuration,
                 powerStatsAggregationPeriod, mPowerStatsStore, alarmScheduler, Clock.SYSTEM_CLOCK,
                 mMonotonicClock, () -> mStats.getHistory().getStartTime(), mHandler);
     }
 
-    private AggregatedPowerStatsConfig createAggregatedPowerStatsConfig() {
-        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CPU)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .setProcessorSupplier(
-                        () -> new ScreenPowerStatsProcessor(mPowerProfile));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
-                        BatteryConsumer.POWER_COMPONENT_SCREEN)
-                .setProcessorSupplier(AmbientDisplayPowerStatsProcessor::new);
-
-        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)
-                .setProcessorSupplier(
-                        () -> new MobileRadioPowerStatsProcessor(mPowerProfile));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
-                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
-                .setProcessorSupplier(PhoneCallPowerStatsProcessor::new);
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new WifiPowerStatsProcessor(mPowerProfile));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_BLUETOOTH)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new BluetoothPowerStatsProcessor(mPowerProfile));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AUDIO)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new AudioPowerStatsProcessor(mPowerProfile, mPowerStatsUidResolver));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_VIDEO)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new VideoPowerStatsProcessor(mPowerProfile, mPowerStatsUidResolver));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new FlashlightPowerStatsProcessor(mPowerProfile,
-                                mPowerStatsUidResolver));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CAMERA)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new CameraPowerStatsProcessor(mPowerProfile, mPowerStatsUidResolver));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_GNSS)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(
-                        () -> new GnssPowerStatsProcessor(mPowerProfile, mPowerStatsUidResolver));
-
-        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SENSORS)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
-                .setProcessorSupplier(() -> new SensorPowerStatsProcessor(
-                        () -> mContext.getSystemService(SensorManager.class)));
-
-        config.trackCustomPowerComponents(CustomEnergyConsumerPowerStatsProcessor::new)
-                .trackDeviceStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN)
-                .trackUidStates(
-                        AggregatedPowerStatsConfig.STATE_POWER,
-                        AggregatedPowerStatsConfig.STATE_SCREEN,
-                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
-        return config;
-    }
-
     private void setPowerStatsThrottlePeriods(BatteryStatsImpl.BatteryStatsConfig.Builder builder,
             String configString) {
         if (configString == null) {
@@ -664,83 +504,84 @@
     }
 
     public void systemServicesReady() {
+        MultiStatePowerAttributor attributor = (MultiStatePowerAttributor) mPowerAttributor;
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_SCREEN,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_SCREEN,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_PHONE,
                 Flags.streamlinedConnectivityBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_WIFI,
                 Flags.streamlinedConnectivityBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_WIFI,
                 Flags.streamlinedConnectivityBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_BLUETOOTH,
                 Flags.streamlinedConnectivityBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_BLUETOOTH,
                 Flags.streamlinedConnectivityBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_AUDIO,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_AUDIO,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_VIDEO,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_VIDEO,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_FLASHLIGHT,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_GNSS,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_GNSS,
                 Flags.streamlinedMiscBatteryStats());
 
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_SENSORS,
                 Flags.streamlinedMiscBatteryStats());
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CAMERA,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_CAMERA,
                 Flags.streamlinedMiscBatteryStats());
 
         // By convention POWER_COMPONENT_ANY represents custom Energy Consumers
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_ANY,
                 Flags.streamlinedMiscBatteryStats());
-        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+        attributor.setPowerComponentSupported(
                 BatteryConsumer.POWER_COMPONENT_ANY,
                 Flags.streamlinedMiscBatteryStats());
 
diff --git a/services/core/java/com/android/server/am/BroadcastController.java b/services/core/java/com/android/server/am/BroadcastController.java
index 32026b2..f7085b4 100644
--- a/services/core/java/com/android/server/am/BroadcastController.java
+++ b/services/core/java/com/android/server/am/BroadcastController.java
@@ -1808,7 +1808,12 @@
 
     final Intent verifyBroadcastLocked(Intent intent) {
         if (intent != null) {
-            intent.prepareToEnterSystemServer();
+            // Refuse possible leaked file descriptors
+            if (intent.hasFileDescriptors()) {
+                throw new IllegalArgumentException("File descriptors passed in Intent");
+            }
+            // Remove existing mismatch flag so it can be properly updated later
+            intent.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
         }
 
         int flags = intent.getFlags();
diff --git a/services/core/java/com/android/server/am/OomAdjuster.java b/services/core/java/com/android/server/am/OomAdjuster.java
index ab63e24..0e266f5 100644
--- a/services/core/java/com/android/server/am/OomAdjuster.java
+++ b/services/core/java/com/android/server/am/OomAdjuster.java
@@ -1201,6 +1201,7 @@
                     >= UNKNOWN_ADJ) {
                     final ProcessServiceRecord psr = app.mServices;
                     switch (state.getCurProcState()) {
+                        case PROCESS_STATE_LAST_ACTIVITY:
                         case PROCESS_STATE_CACHED_ACTIVITY:
                         case ActivityManager.PROCESS_STATE_CACHED_ACTIVITY_CLIENT:
                         case ActivityManager.PROCESS_STATE_CACHED_RECENT:
@@ -2180,7 +2181,6 @@
                 procState = PROCESS_STATE_LAST_ACTIVITY;
                 schedGroup = SCHED_GROUP_BACKGROUND;
                 state.setAdjType("previous-expired");
-                adj = CACHED_APP_MIN_ADJ;
                 if (DEBUG_OOM_ADJ_REASON || logUid == appUid) {
                     reportOomAdjMessageLocked(TAG_OOM_ADJ, "Expire prev adj: " + app);
                 }
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index 154b52b..6ae6f3d 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -33,7 +33,6 @@
 import static android.app.AppOpsManager.MODE_ERRORED;
 import static android.app.AppOpsManager.MODE_FOREGROUND;
 import static android.app.AppOpsManager.MODE_IGNORED;
-import static android.app.AppOpsManager.OP_BLUETOOTH_CONNECT;
 import static android.app.AppOpsManager.OP_CAMERA;
 import static android.app.AppOpsManager.OP_CAMERA_SANDBOXED;
 import static android.app.AppOpsManager.OP_FLAGS_ALL;
@@ -3115,11 +3114,6 @@
                     packageName);
         }
         if (!isIncomingPackageValid(packageName, UserHandle.getUserId(uid))) {
-            // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-            if (code == OP_BLUETOOTH_CONNECT) {
-                Slog.e(TAG, "noting OP_BLUETOOTH_CONNECT returned MODE_ERRORED as incoming "
-                        + "package: " + packageName + " and uid: " + uid + " is invalid");
-            }
             return new SyncNotedAppOp(AppOpsManager.MODE_ERRORED, code, attributionTag,
                     packageName);
         }
@@ -3149,13 +3143,6 @@
             }
         } catch (SecurityException e) {
             logVerifyAndGetBypassFailure(uid, e, "noteOperation");
-            // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-            if (code == OP_BLUETOOTH_CONNECT) {
-                Slog.e(TAG, "noting OP_BLUETOOTH_CONNECT returned MODE_ERRORED as"
-                        + " verifyAndGetBypass returned a SecurityException for package: "
-                        + packageName + " and uid: " + uid + " and attributionTag: "
-                        + attributionTag, e);
-            }
             return new SyncNotedAppOp(AppOpsManager.MODE_ERRORED, code, attributionTag,
                     packageName);
         }
@@ -3173,17 +3160,6 @@
                 if (DEBUG) Slog.d(TAG, "noteOperation: no op for code " + code + " uid " + uid
                         + " package " + packageName + "flags: " +
                         AppOpsManager.flagsToString(flags));
-                // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-                if (code == OP_BLUETOOTH_CONNECT) {
-                    Slog.e(TAG, "noting OP_BLUETOOTH_CONNECT returned MODE_ERRORED as"
-                            + " #getOpsLocked returned null for"
-                            + " uid: " + uid
-                            + " packageName: " + packageName
-                            + " attributionTag: " + attributionTag
-                            + " pvr.isAttributionTagValid: " + pvr.isAttributionTagValid
-                            + " pvr.bypass: " + pvr.bypass);
-                    Slog.e(TAG, "mUidStates.get(" + uid + "): " + mUidStates.get(uid));
-                }
                 return new SyncNotedAppOp(AppOpsManager.MODE_ERRORED, code, attributionTag,
                         packageName);
             }
@@ -3228,11 +3204,6 @@
                     attributedOp.rejected(uidState.getState(), flags);
                     scheduleOpNotedIfNeededLocked(code, uid, packageName, attributionTag,
                             virtualDeviceId, flags, uidMode);
-                    // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-                    if (code == OP_BLUETOOTH_CONNECT && uidMode == MODE_ERRORED) {
-                        Slog.e(TAG, "noting OP_BLUETOOTH_CONNECT returned MODE_ERRORED as"
-                                + " uid mode is MODE_ERRORED");
-                    }
                     return new SyncNotedAppOp(uidMode, code, attributionTag, packageName);
                 }
             } else {
@@ -3252,11 +3223,6 @@
                     attributedOp.rejected(uidState.getState(), flags);
                     scheduleOpNotedIfNeededLocked(code, uid, packageName, attributionTag,
                             virtualDeviceId, flags, mode);
-                    // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-                    if (code == OP_BLUETOOTH_CONNECT && mode == MODE_ERRORED) {
-                        Slog.e(TAG, "noting OP_BLUETOOTH_CONNECT returned MODE_ERRORED as"
-                                + " package mode is MODE_ERRORED");
-                    }
                     return new SyncNotedAppOp(mode, code, attributionTag, packageName);
                 }
             }
diff --git a/services/core/java/com/android/server/appop/TEST_MAPPING b/services/core/java/com/android/server/appop/TEST_MAPPING
index 2a9dfa2..9317c1e 100644
--- a/services/core/java/com/android/server/appop/TEST_MAPPING
+++ b/services/core/java/com/android/server/appop/TEST_MAPPING
@@ -18,24 +18,7 @@
             "name": "FrameworksMockingServicesTests_android_server_appop"
         },
         {
-            "name": "CtsPermissionTestCases",
-            "options": [
-                {
-                    "exclude-annotation": "androidx.test.filters.FlakyTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.BackgroundPermissionsTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.SplitPermissionTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.PermissionFlagsTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.SharedUidPermissionsTest"
-                }
-            ]
+            "name": "CtsPermissionTestCases_Platform"
         },
         {
             "name": "CtsAppTestCases",
diff --git a/services/core/java/com/android/server/audio/AudioService.java b/services/core/java/com/android/server/audio/AudioService.java
index b738482..4e24cf3 100644
--- a/services/core/java/com/android/server/audio/AudioService.java
+++ b/services/core/java/com/android/server/audio/AudioService.java
@@ -790,8 +790,8 @@
     private final BroadcastReceiver mReceiver = new AudioServiceBroadcastReceiver();
 
     private final Executor mAudioServerLifecycleExecutor;
-    private final ConcurrentLinkedQueue<Future> mScheduledPermissionTasks =
-            new ConcurrentLinkedQueue();
+    private long mSysPropListenerNativeHandle;
+    private final List<Future> mScheduledPermissionTasks = new ArrayList();
 
     private IMediaProjectionManager mProjectionService; // to validate projection token
 
@@ -10600,46 +10600,50 @@
         // instanceof to simplify the construction requirements of AudioService for testing: no
         // delayed execution during unit tests.
         if (mAudioServerLifecycleExecutor instanceof ScheduledExecutorService exec) {
-            // We schedule and add from a this callback thread only (serially), so the task order on
-            // the serial executor matches the order on the task list.  This list should almost
-            // always have only two elements, except in cases of serious system contention.
-            Runnable task = () -> mScheduledPermissionTasks.add(exec.schedule(() -> {
-                    try {
-                        // Clean up completed tasks before us to bound the queue length.  Cancel any
-                        // pending permission refresh tasks, after our own, since we are about to
-                        // fulfill all of them.  We must be the first non-completed task in the
-                        // queue, since the execution order matches the queue order.  Note, this
-                        // task is the only writer on elements in the queue, and the task is
-                        // serialized, so
-                        //  => no in-flight cancellation
-                        //  => exists at least one non-completed task (ourselves)
-                        //  => the queue is non-empty (only completed tasks removed)
-                        final var iter = mScheduledPermissionTasks.iterator();
-                        while (iter.next().isDone()) {
-                            iter.remove();
-                        }
-                        // iter is on the first element which is not completed (us)
-                        while (iter.hasNext()) {
-                            if (!iter.next().cancel(false)) {
-                                throw new AssertionError(
-                                        "Cancel should be infallible since we" +
-                                        "cancel from the executor");
+            // The order on the task list is an embedding on the scheduling order of the executor,
+            // since we synchronously add the scheduled task to our local queue. This list should
+            // almost always have only two elements, except in cases of serious system contention.
+            Runnable task = () -> {
+                synchronized (mScheduledPermissionTasks) {
+                    mScheduledPermissionTasks.add(exec.schedule(() -> {
+                        try {
+                            // Our goal is to remove all tasks which don't correspond to ourselves
+                            // on this queue. Either they are already done (ahead of us), or we
+                            // should cancel them (behind us), since their work is redundant after
+                            // we fire.
+                            // We must be the first non-completed task in the queue, since the
+                            // execution order matches the queue order. Note, this task is the only
+                            // writer on elements in the queue, and the task is serialized, so
+                            //  => no in-flight cancellation
+                            //  => exists at least one non-completed task (ourselves)
+                            //  => the queue is non-empty (only completed tasks removed)
+                            synchronized (mScheduledPermissionTasks) {
+                                final var iter = mScheduledPermissionTasks.iterator();
+                                while (iter.next().isDone()) {
+                                    iter.remove();
+                                }
+                                // iter is on the first element which is not completed (us)
+                                while (iter.hasNext()) {
+                                    if (!iter.next().cancel(false)) {
+                                        throw new AssertionError(
+                                                "Cancel should be infallible since we" +
+                                                "cancel from the executor");
+                                    }
+                                    iter.remove();
+                                }
                             }
-                            iter.remove();
+                            mPermissionProvider.onPermissionStateChanged();
+                        } catch (Exception e) {
+                            // Handle executor routing exceptions to nowhere
+                            Thread.getDefaultUncaughtExceptionHandler()
+                                    .uncaughtException(Thread.currentThread(), e);
                         }
-                        mPermissionProvider.onPermissionStateChanged();
-                    } catch (Exception e) {
-                        // Handle executor routing exceptions to nowhere
-                        Thread.getDefaultUncaughtExceptionHandler()
-                                .uncaughtException(Thread.currentThread(), e);
-                    }
-                },
-                UPDATE_DELAY_MS,
-                TimeUnit.MILLISECONDS));
-            mAudioSystem.listenForSystemPropertyChange(
+                    }, UPDATE_DELAY_MS, TimeUnit.MILLISECONDS));
+                }
+            };
+            mSysPropListenerNativeHandle = mAudioSystem.listenForSystemPropertyChange(
                     PermissionManager.CACHE_KEY_PACKAGE_INFO,
                     task);
-            task.run();
         } else {
             mAudioSystem.listenForSystemPropertyChange(
                     PermissionManager.CACHE_KEY_PACKAGE_INFO,
@@ -14710,7 +14714,12 @@
     @Override
     /** @see AudioManager#permissionUpdateBarrier() */
     public void permissionUpdateBarrier() {
-        for (var x : List.copyOf(mScheduledPermissionTasks)) {
+        mAudioSystem.triggerSystemPropertyUpdate(mSysPropListenerNativeHandle);
+        List<Future> snapshot;
+        synchronized (mScheduledPermissionTasks) {
+            snapshot = List.copyOf(mScheduledPermissionTasks);
+        }
+        for (var x : snapshot) {
             try {
                 x.get();
             } catch (CancellationException e) {
diff --git a/services/core/java/com/android/server/audio/AudioSystemAdapter.java b/services/core/java/com/android/server/audio/AudioSystemAdapter.java
index d083c68..5cabdde 100644
--- a/services/core/java/com/android/server/audio/AudioSystemAdapter.java
+++ b/services/core/java/com/android/server/audio/AudioSystemAdapter.java
@@ -748,8 +748,12 @@
         return AudioSystem.setMasterMute(mute);
     }
 
-    public void listenForSystemPropertyChange(String systemPropertyName, Runnable callback) {
-        AudioSystem.listenForSystemPropertyChange(systemPropertyName, callback);
+    public long listenForSystemPropertyChange(String systemPropertyName, Runnable callback) {
+        return AudioSystem.listenForSystemPropertyChange(systemPropertyName, callback);
+    }
+
+    public void triggerSystemPropertyUpdate(long handle) {
+        AudioSystem.triggerSystemPropertyUpdate(handle);
     }
 
     /**
diff --git a/services/core/java/com/android/server/audio/MediaFocusControl.java b/services/core/java/com/android/server/audio/MediaFocusControl.java
index 70f3193..7e26356 100644
--- a/services/core/java/com/android/server/audio/MediaFocusControl.java
+++ b/services/core/java/com/android/server/audio/MediaFocusControl.java
@@ -1302,7 +1302,7 @@
         mEventLogger.enqueue((new EventLogger.StringEvent(
                 "abandonAudioFocus() from uid/pid " + Binder.getCallingUid()
                     + "/" + Binder.getCallingPid()
-                    + " clientId=" + clientId))
+                    + " clientId=" + clientId + " callingPack=" + callingPackageName))
                 .printLog(TAG));
         try {
             // this will take care of notifying the new focus owner if needed
diff --git a/services/core/java/com/android/server/audio/SoundDoseHelper.java b/services/core/java/com/android/server/audio/SoundDoseHelper.java
index dc79ab2..643f330 100644
--- a/services/core/java/com/android/server/audio/SoundDoseHelper.java
+++ b/services/core/java/com/android/server/audio/SoundDoseHelper.java
@@ -900,8 +900,10 @@
 
         try {
             if (!isAbsoluteVolume) {
-                mLogger.enqueue(
-                        SoundDoseEvent.getAbsVolumeAttenuationEvent(/*attenuation=*/0.f, device));
+                if (mSafeMediaVolumeDevices.indexOfKey(device) >= 0) {
+                    mLogger.enqueue(SoundDoseEvent.getAbsVolumeAttenuationEvent(/*attenuation=*/0.f,
+                            device));
+                }
                 // remove any possible previous attenuation
                 soundDose.updateAttenuation(/* attenuationDB= */0.f, device);
 
@@ -912,8 +914,12 @@
                     && safeDevicesContains(device)) {
                 float attenuationDb = -AudioSystem.getStreamVolumeDB(AudioSystem.STREAM_MUSIC,
                         (newIndex + 5) / 10, device);
-                mLogger.enqueue(
-                        SoundDoseEvent.getAbsVolumeAttenuationEvent(attenuationDb, device));
+
+                if (mSafeMediaVolumeDevices.indexOfKey(device) >= 0) {
+                    mLogger.enqueue(
+                            SoundDoseEvent.getAbsVolumeAttenuationEvent(attenuationDb, device));
+                }
+
                 soundDose.updateAttenuation(attenuationDb, device);
             }
         } catch (RemoteException e) {
diff --git a/services/core/java/com/android/server/biometrics/AuthSession.java b/services/core/java/com/android/server/biometrics/AuthSession.java
index 276ab03..abfbddc 100644
--- a/services/core/java/com/android/server/biometrics/AuthSession.java
+++ b/services/core/java/com/android/server/biometrics/AuthSession.java
@@ -221,7 +221,8 @@
         mFingerprintSensorProperties = fingerprintSensorProperties;
         mCancelled = false;
         mBiometricFrameworkStatsLogger = logger;
-        mOperationContext = new OperationContextExt(true /* isBP */);
+        mOperationContext = new OperationContextExt(true /* isBP */,
+                preAuthInfo.getIsMandatoryBiometricsAuthentication() /* isMandatoryBiometrics */);
         mBiometricManager = mContext.getSystemService(BiometricManager.class);
 
         mSfpsSensorIds = mFingerprintSensorProperties.stream().filter(
@@ -285,7 +286,8 @@
             sensor.goToStateWaitingForCookie(requireConfirmation, mToken, mOperationId,
                     mUserId, mSensorReceiver, mOpPackageName, mRequestId, cookie,
                     mPromptInfo.isAllowBackgroundAuthentication(),
-                    mPromptInfo.isForLegacyFingerprintManager());
+                    mPromptInfo.isForLegacyFingerprintManager(),
+                    mOperationContext.getIsMandatoryBiometrics());
         }
     }
 
diff --git a/services/core/java/com/android/server/biometrics/BiometricSensor.java b/services/core/java/com/android/server/biometrics/BiometricSensor.java
index 42dd36a..c7532d4 100644
--- a/services/core/java/com/android/server/biometrics/BiometricSensor.java
+++ b/services/core/java/com/android/server/biometrics/BiometricSensor.java
@@ -107,12 +107,13 @@
     void goToStateWaitingForCookie(boolean requireConfirmation, IBinder token, long sessionId,
             int userId, IBiometricSensorReceiver sensorReceiver, String opPackageName,
             long requestId, int cookie, boolean allowBackgroundAuthentication,
-            boolean isForLegacyFingerprintManager)
+            boolean isForLegacyFingerprintManager, boolean isMandatoryBiometrics)
             throws RemoteException {
         mCookie = cookie;
         impl.prepareForAuthentication(requireConfirmation, token,
                 sessionId, userId, sensorReceiver, opPackageName, requestId, mCookie,
-                allowBackgroundAuthentication, isForLegacyFingerprintManager);
+                allowBackgroundAuthentication, isForLegacyFingerprintManager,
+                isMandatoryBiometrics);
         mSensorState = STATE_WAITING_FOR_COOKIE;
     }
 
diff --git a/services/core/java/com/android/server/biometrics/PreAuthInfo.java b/services/core/java/com/android/server/biometrics/PreAuthInfo.java
index f0da67b..eaf4f81 100644
--- a/services/core/java/com/android/server/biometrics/PreAuthInfo.java
+++ b/services/core/java/com/android/server/biometrics/PreAuthInfo.java
@@ -447,6 +447,12 @@
                         getInternalStatus().second));
     }
 
+    /** Returns if mandatory biometrics authentication is in effect */
+    boolean getIsMandatoryBiometricsAuthentication() {
+        return mIsMandatoryBiometricsAuthentication;
+    }
+
+
     /**
      * For the given request, generate the appropriate reason why authentication cannot be started.
      * Note that for some errors, modality is intentionally cleared.
diff --git a/services/core/java/com/android/server/biometrics/log/BiometricFrameworkStatsLogger.java b/services/core/java/com/android/server/biometrics/log/BiometricFrameworkStatsLogger.java
index 2c52e3d..bc58501 100644
--- a/services/core/java/com/android/server/biometrics/log/BiometricFrameworkStatsLogger.java
+++ b/services/core/java/com/android/server/biometrics/log/BiometricFrameworkStatsLogger.java
@@ -98,7 +98,8 @@
                 foldType(operationContext.getFoldState()),
                 operationContext.getOrderAndIncrement(),
                 toProtoWakeReason(operationContext),
-                toProtoWakeReasonDetails(operationContext));
+                toProtoWakeReasonDetails(operationContext),
+                operationContext.getIsMandatoryBiometrics());
     }
 
     /** {@see FrameworkStatsLog.BIOMETRIC_AUTHENTICATED}. */
@@ -149,7 +150,8 @@
                 foldType(operationContext.getFoldState()),
                 operationContext.getOrderAndIncrement(),
                 toProtoWakeReason(operationContext),
-                toProtoWakeReasonDetails(operationContext));
+                toProtoWakeReasonDetails(operationContext),
+                operationContext.getIsMandatoryBiometrics());
     }
 
     @VisibleForTesting
diff --git a/services/core/java/com/android/server/biometrics/log/OperationContextExt.java b/services/core/java/com/android/server/biometrics/log/OperationContextExt.java
index da4e515..4df63e2 100644
--- a/services/core/java/com/android/server/biometrics/log/OperationContextExt.java
+++ b/services/core/java/com/android/server/biometrics/log/OperationContextExt.java
@@ -50,21 +50,33 @@
     @Surface.Rotation private int mOrientation = Surface.ROTATION_0;
     private int mFoldState = IBiometricContextListener.FoldState.UNKNOWN;
     private final boolean mIsBP;
+    private final boolean mIsMandatoryBiometrics;
 
     /** Create a context. */
     public OperationContextExt(boolean isBP) {
         this(new OperationContext(), isBP, BiometricAuthenticator.TYPE_NONE);
     }
 
-    public OperationContextExt(boolean isBP, @BiometricAuthenticator.Modality int modality) {
-        this(new OperationContext(), isBP, modality);
+    public OperationContextExt(boolean isBP, boolean isMandatoryBiometrics) {
+        this(new OperationContext(), isBP, BiometricAuthenticator.TYPE_NONE, isMandatoryBiometrics);
+    }
+
+    public OperationContextExt(boolean isBP, @BiometricAuthenticator.Modality int modality,
+            boolean isMandatoryBiometrics) {
+        this(new OperationContext(), isBP, modality, isMandatoryBiometrics);
     }
 
     /** Create a wrapped context. */
     public OperationContextExt(@NonNull OperationContext context, boolean isBP,
             @BiometricAuthenticator.Modality int modality) {
+        this(context, isBP, modality, false /* isMandatoryBiometrics */);
+    }
+
+    public OperationContextExt(@NonNull OperationContext context, boolean isBP,
+            @BiometricAuthenticator.Modality int modality, boolean isMandatoryBiometrics) {
         mAidlContext = context;
         mIsBP = isBP;
+        mIsMandatoryBiometrics = isMandatoryBiometrics;
 
         if (modality == BiometricAuthenticator.TYPE_FINGERPRINT) {
             mAidlContext.operationState = OperationState.fingerprintOperationState(
@@ -285,6 +297,11 @@
         return mAidlContext.operationState;
     }
 
+    /** If mandatory biometrics is active. */
+    public boolean getIsMandatoryBiometrics() {
+        return mIsMandatoryBiometrics;
+    }
+
     /** Update this object with the latest values from the given context. */
     OperationContextExt update(@NonNull BiometricContext biometricContext, boolean isCrypto) {
         mAidlContext.isAod = biometricContext.isAod();
diff --git a/services/core/java/com/android/server/biometrics/sensors/AcquisitionClient.java b/services/core/java/com/android/server/biometrics/sensors/AcquisitionClient.java
index fbd32a6..04522e3 100644
--- a/services/core/java/com/android/server/biometrics/sensors/AcquisitionClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/AcquisitionClient.java
@@ -59,9 +59,10 @@
     public AcquisitionClient(@NonNull Context context, @NonNull Supplier<T> lazyDaemon,
             @NonNull IBinder token, @NonNull ClientMonitorCallbackConverter listener, int userId,
             @NonNull String owner, int cookie, int sensorId, boolean shouldVibrate,
-            @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext) {
+            @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
+            boolean isMandatoryBiometrics) {
         super(context, lazyDaemon, token, listener, userId, owner, cookie, sensorId,
-                logger, biometricContext);
+                logger, biometricContext, isMandatoryBiometrics);
         mPowerManager = context.getSystemService(PowerManager.class);
         mShouldVibrate = shouldVibrate;
     }
diff --git a/services/core/java/com/android/server/biometrics/sensors/AuthenticationClient.java b/services/core/java/com/android/server/biometrics/sensors/AuthenticationClient.java
index daaafcb..09386ae28 100644
--- a/services/core/java/com/android/server/biometrics/sensors/AuthenticationClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/AuthenticationClient.java
@@ -99,7 +99,7 @@
             boolean shouldVibrate, int sensorStrength) {
         super(context, lazyDaemon, token, listener, options.getUserId(),
                 options.getOpPackageName(), cookie, options.getSensorId(), shouldVibrate,
-                biometricLogger, biometricContext);
+                biometricLogger, biometricContext, options.isMandatoryBiometrics());
         mIsStrongBiometric = isStrongBiometric;
         mOperationId = operationId;
         mRequireConfirmation = requireConfirmation;
diff --git a/services/core/java/com/android/server/biometrics/sensors/EnrollClient.java b/services/core/java/com/android/server/biometrics/sensors/EnrollClient.java
index 438367d..32c0bd3 100644
--- a/services/core/java/com/android/server/biometrics/sensors/EnrollClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/EnrollClient.java
@@ -60,7 +60,7 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             int enrollReason) {
         super(context, lazyDaemon, token, listener, userId, owner, 0 /* cookie */, sensorId,
-                shouldVibrate, logger, biometricContext);
+                shouldVibrate, logger, biometricContext, false /* isMandatoryBiometrics */);
         mBiometricUtils = utils;
         mHardwareAuthToken = Arrays.copyOf(hardwareAuthToken, hardwareAuthToken.length);
         mTimeoutSec = timeoutSec;
diff --git a/services/core/java/com/android/server/biometrics/sensors/GenerateChallengeClient.java b/services/core/java/com/android/server/biometrics/sensors/GenerateChallengeClient.java
index 2adf0cb..b573b56 100644
--- a/services/core/java/com/android/server/biometrics/sensors/GenerateChallengeClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/GenerateChallengeClient.java
@@ -37,7 +37,7 @@
             int userId, @NonNull String owner, int sensorId,
             @NonNull BiometricLogger biometricLogger, @NonNull BiometricContext biometricContext) {
         super(context, lazyDaemon, token, listener, userId, owner, 0 /* cookie */, sensorId,
-                biometricLogger, biometricContext);
+                biometricLogger, biometricContext, false /* isMandatoryBiometrics */);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/biometrics/sensors/HalClientMonitor.java b/services/core/java/com/android/server/biometrics/sensors/HalClientMonitor.java
index 0f01510..3bc51a9 100644
--- a/services/core/java/com/android/server/biometrics/sensors/HalClientMonitor.java
+++ b/services/core/java/com/android/server/biometrics/sensors/HalClientMonitor.java
@@ -41,26 +41,29 @@
     private final OperationContextExt mOperationContext;
 
     /**
-     * @param context    system_server context
-     * @param lazyDaemon pointer for lazy retrieval of the HAL
-     * @param token      a unique token for the client
-     * @param listener   recipient of related events (e.g. authentication)
-     * @param userId     target user id for operation
-     * @param owner      name of the client that owns this
-     * @param cookie     BiometricPrompt authentication cookie (to be moved into a subclass soon)
-     * @param sensorId   ID of the sensor that the operation should be requested of
-     * @param biometricLogger framework stats logger
+     * @param context          system_server context
+     * @param lazyDaemon       pointer for lazy retrieval of the HAL
+     * @param token            a unique token for the client
+     * @param listener         recipient of related events (e.g. authentication)
+     * @param userId           target user id for operation
+     * @param owner            name of the client that owns this
+     * @param cookie           BiometricPrompt authentication cookie (to be moved into a subclass
+     *                         soon)
+     * @param sensorId         ID of the sensor that the operation should be requested of
+     * @param biometricLogger  framework stats logger
      * @param biometricContext system context metadata
      */
     public HalClientMonitor(@NonNull Context context, @NonNull Supplier<T> lazyDaemon,
             @Nullable IBinder token, @Nullable ClientMonitorCallbackConverter listener, int userId,
             @NonNull String owner, int cookie, int sensorId,
-            @NonNull BiometricLogger biometricLogger, @NonNull BiometricContext biometricContext) {
+            @NonNull BiometricLogger biometricLogger, @NonNull BiometricContext biometricContext,
+            boolean isMandatoryBiometrics) {
         super(context, token, listener, userId, owner, cookie, sensorId,
                 biometricLogger, biometricContext);
         mLazyDaemon = lazyDaemon;
         int modality = listener != null ? listener.getModality() : BiometricAuthenticator.TYPE_NONE;
-        mOperationContext = new OperationContextExt(isBiometricPrompt(), modality);
+        mOperationContext = new OperationContextExt(isBiometricPrompt(), modality,
+                isMandatoryBiometrics);
     }
 
     @Nullable
diff --git a/services/core/java/com/android/server/biometrics/sensors/InternalCleanupClient.java b/services/core/java/com/android/server/biometrics/sensors/InternalCleanupClient.java
index 7bd905b..6c30559 100644
--- a/services/core/java/com/android/server/biometrics/sensors/InternalCleanupClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/InternalCleanupClient.java
@@ -143,7 +143,8 @@
             @NonNull BiometricUtils<S> utils,
             @NonNull Map<Integer, Long> authenticatorIds) {
         super(context, lazyDaemon, null /* token */, null /* ClientMonitorCallbackConverter */,
-                userId, owner, 0 /* cookie */, sensorId, logger, biometricContext);
+                userId, owner, 0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mBiometricUtils = utils;
         mAuthenticatorIds = authenticatorIds;
         mHasEnrollmentsBeforeStarting = !utils.getBiometricsForUser(context, userId).isEmpty();
diff --git a/services/core/java/com/android/server/biometrics/sensors/InternalEnumerateClient.java b/services/core/java/com/android/server/biometrics/sensors/InternalEnumerateClient.java
index 81ab26d..2c2c685 100644
--- a/services/core/java/com/android/server/biometrics/sensors/InternalEnumerateClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/InternalEnumerateClient.java
@@ -55,7 +55,8 @@
         // Internal enumerate does not need to send results to anyone. Cleanup (enumerate + remove)
         // is all done internally.
         super(context, lazyDaemon, token, null /* ClientMonitorCallbackConverter */, userId, owner,
-                0 /* cookie */, sensorId, logger, biometricContext);
+                0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mEnrolledList = enrolledList;
         mInitialEnrolledSize = mEnrolledList.size();
         mUtils = utils;
diff --git a/services/core/java/com/android/server/biometrics/sensors/InvalidationClient.java b/services/core/java/com/android/server/biometrics/sensors/InvalidationClient.java
index d5aa5e2..6c93366 100644
--- a/services/core/java/com/android/server/biometrics/sensors/InvalidationClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/InvalidationClient.java
@@ -49,7 +49,7 @@
             @NonNull IInvalidationCallback callback) {
         super(context, lazyDaemon, null /* token */, null /* listener */, userId,
                 context.getOpPackageName(), 0 /* cookie */, sensorId,
-                logger, biometricContext);
+                logger, biometricContext, false /* isMandatoryBiometrics */);
         mAuthenticatorIds = authenticatorIds;
         mInvalidationCallback = callback;
     }
diff --git a/services/core/java/com/android/server/biometrics/sensors/RemovalClient.java b/services/core/java/com/android/server/biometrics/sensors/RemovalClient.java
index d2ef278..ad5877a 100644
--- a/services/core/java/com/android/server/biometrics/sensors/RemovalClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/RemovalClient.java
@@ -49,7 +49,7 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             @NonNull Map<Integer, Long> authenticatorIds) {
         super(context, lazyDaemon, token, listener, userId, owner, 0 /* cookie */, sensorId,
-                logger, biometricContext);
+                logger, biometricContext, false /* isMandatoryBiometrics */);
         mBiometricUtils = utils;
         mAuthenticatorIds = authenticatorIds;
         mHasEnrollmentsBeforeStarting = !utils.getBiometricsForUser(context, userId).isEmpty();
diff --git a/services/core/java/com/android/server/biometrics/sensors/RevokeChallengeClient.java b/services/core/java/com/android/server/biometrics/sensors/RevokeChallengeClient.java
index 88f4da2..0c8a2dd 100644
--- a/services/core/java/com/android/server/biometrics/sensors/RevokeChallengeClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/RevokeChallengeClient.java
@@ -32,7 +32,8 @@
             @NonNull IBinder token, int userId, @NonNull String owner, int sensorId,
             @NonNull BiometricLogger biometricLogger, @NonNull BiometricContext biometricContext) {
         super(context, lazyDaemon, token, null /* listener */, userId, owner,
-                0 /* cookie */, sensorId, biometricLogger, biometricContext);
+                0 /* cookie */, sensorId, biometricLogger, biometricContext,
+                false /* isMandatoryBiometrics */);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/biometrics/sensors/StartUserClient.java b/services/core/java/com/android/server/biometrics/sensors/StartUserClient.java
index 21c9f64..ff694cd 100644
--- a/services/core/java/com/android/server/biometrics/sensors/StartUserClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/StartUserClient.java
@@ -51,7 +51,8 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             @NonNull UserStartedCallback<U> callback) {
         super(context, lazyDaemon, token, null /* listener */, userId, context.getOpPackageName(),
-                0 /* cookie */, sensorId, logger, biometricContext);
+                0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mUserStartedCallback = callback;
     }
 
diff --git a/services/core/java/com/android/server/biometrics/sensors/StopUserClient.java b/services/core/java/com/android/server/biometrics/sensors/StopUserClient.java
index e01c4ec..9119bc7 100644
--- a/services/core/java/com/android/server/biometrics/sensors/StopUserClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/StopUserClient.java
@@ -54,7 +54,8 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             @NonNull UserStoppedCallback callback) {
         super(context, lazyDaemon, token, null /* listener */, userId, context.getOpPackageName(),
-                0 /* cookie */, sensorId, logger, biometricContext);
+                0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mUserStoppedCallback = callback;
     }
 
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/FaceAuthenticator.java b/services/core/java/com/android/server/biometrics/sensors/face/FaceAuthenticator.java
index 2211003..67d7505 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/FaceAuthenticator.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/FaceAuthenticator.java
@@ -63,13 +63,14 @@
     public void prepareForAuthentication(boolean requireConfirmation, IBinder token,
             long operationId, int userId, IBiometricSensorReceiver sensorReceiver,
             String opPackageName, long requestId, int cookie, boolean allowBackgroundAuthentication,
-            boolean isForLegacyFingerprintManager)
+            boolean isForLegacyFingerprintManager, boolean isMandatoryBiometrics)
             throws RemoteException {
         mFaceService.prepareForAuthentication(requireConfirmation, token, operationId,
                 sensorReceiver, new FaceAuthenticateOptions.Builder()
                         .setUserId(userId)
                         .setSensorId(mSensorId)
                         .setOpPackageName(opPackageName)
+                        .setIsMandatoryBiometrics(isMandatoryBiometrics)
                         .build(),
                 requestId, cookie, allowBackgroundAuthentication);
     }
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceDetectClient.java b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceDetectClient.java
index 8b4da31..8eb62eb 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceDetectClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceDetectClient.java
@@ -83,7 +83,8 @@
             boolean isStrongBiometric, SensorPrivacyManager sensorPrivacyManager) {
         super(context, lazyDaemon, token, listener, options.getUserId(),
                 options.getOpPackageName(), 0 /* cookie */, options.getSensorId(),
-                false /* shouldVibrate */, logger, biometricContext);
+                false /* shouldVibrate */, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         setRequestId(requestId);
         mAuthenticationStateListeners = authenticationStateListeners;
         mIsStrongBiometric = isStrongBiometric;
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetAuthenticatorIdClient.java b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetAuthenticatorIdClient.java
index 1f4f612..dbd293c 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetAuthenticatorIdClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetAuthenticatorIdClient.java
@@ -42,7 +42,8 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             Map<Integer, Long> authenticatorIds) {
         super(context, lazyDaemon, null /* token */, null /* listener */, userId, opPackageName,
-                0 /* cookie */, sensorId, logger, biometricContext);
+                0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mAuthenticatorIds = authenticatorIds;
     }
 
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetFeatureClient.java b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetFeatureClient.java
index c41b706..8d1336f 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetFeatureClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceGetFeatureClient.java
@@ -55,7 +55,7 @@
             @NonNull String owner, int sensorId, @NonNull BiometricLogger logger,
             @NonNull BiometricContext biometricContext, int feature) {
         super(context, lazyDaemon, token, listener, userId, owner, 0 /* cookie */, sensorId,
-                logger, biometricContext);
+                logger, biometricContext, false /* isMandatoryBiometrics */);
         mUserId = userId;
         mFeature = feature;
     }
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceResetLockoutClient.java b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceResetLockoutClient.java
index d02eefa..93bc1f2 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceResetLockoutClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceResetLockoutClient.java
@@ -60,7 +60,8 @@
             @NonNull LockoutResetDispatcher lockoutResetDispatcher,
             @Authenticators.Types int biometricStrength) {
         super(context, lazyDaemon, null /* token */, null /* listener */, userId, owner,
-                0 /* cookie */, sensorId, logger, biometricContext);
+                0 /* cookie */, sensorId, logger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mHardwareAuthToken = HardwareAuthTokenUtils.toHardwareAuthToken(hardwareAuthToken);
         mLockoutTracker = lockoutTracker;
         mLockoutResetDispatcher = lockoutResetDispatcher;
diff --git a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceSetFeatureClient.java b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceSetFeatureClient.java
index f6da872..fc73e58 100644
--- a/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceSetFeatureClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/face/aidl/FaceSetFeatureClient.java
@@ -52,7 +52,7 @@
             @NonNull BiometricLogger logger, @NonNull BiometricContext biometricContext,
             int feature, boolean enabled, byte[] hardwareAuthToken) {
         super(context, lazyDaemon, token, listener, userId, owner, 0 /* cookie */, sensorId,
-                logger, biometricContext);
+                logger, biometricContext, false /* isMandatoryBiometrics */);
         mFeature = feature;
         mEnabled = enabled;
         mHardwareAuthToken = HardwareAuthTokenUtils.toHardwareAuthToken(hardwareAuthToken);
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/FingerprintAuthenticator.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/FingerprintAuthenticator.java
index b6fa0c1..d50ab8d 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/FingerprintAuthenticator.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/FingerprintAuthenticator.java
@@ -63,13 +63,14 @@
     public void prepareForAuthentication(boolean requireConfirmation, IBinder token,
             long operationId, int userId, IBiometricSensorReceiver sensorReceiver,
             String opPackageName, long requestId, int cookie, boolean allowBackgroundAuthentication,
-            boolean isForLegacyFingerprintManager)
+            boolean isForLegacyFingerprintManager, boolean isMandatoryBiometrics)
             throws RemoteException {
         mFingerprintService.prepareForAuthentication(token, operationId, sensorReceiver,
                 new FingerprintAuthenticateOptions.Builder()
                         .setSensorId(mSensorId)
                         .setUserId(userId)
                         .setOpPackageName(opPackageName)
+                        .setIsMandatoryBiometrics(isMandatoryBiometrics)
                         .build(),
                 requestId, cookie, allowBackgroundAuthentication, isForLegacyFingerprintManager);
     }
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintDetectClient.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintDetectClient.java
index fb48053..a81ab8e 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintDetectClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintDetectClient.java
@@ -71,7 +71,8 @@
             boolean isStrongBiometric) {
         super(context, lazyDaemon, token, listener, options.getUserId(),
                 options.getOpPackageName(), 0 /* cookie */, options.getSensorId(),
-                true /* shouldVibrate */, biometricLogger, biometricContext);
+                true /* shouldVibrate */, biometricLogger, biometricContext,
+                false /* isMandatoryBiometrics */);
         setRequestId(requestId);
         mAuthenticationStateListeners = authenticationStateListeners;
         mIsStrongBiometric = isStrongBiometric;
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintGetAuthenticatorIdClient.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintGetAuthenticatorIdClient.java
index 0353969..f77e116 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintGetAuthenticatorIdClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintGetAuthenticatorIdClient.java
@@ -42,7 +42,8 @@
             @NonNull BiometricLogger biometricLogger, @NonNull BiometricContext biometricContext,
             Map<Integer, Long> authenticatorIds) {
         super(context, lazyDaemon, null /* token */, null /* listener */, userId, owner,
-                0 /* cookie */, sensorId, biometricLogger, biometricContext);
+                0 /* cookie */, sensorId, biometricLogger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mAuthenticatorIds = authenticatorIds;
     }
 
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintResetLockoutClient.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintResetLockoutClient.java
index 387ae12..81a3d83 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintResetLockoutClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintResetLockoutClient.java
@@ -59,7 +59,8 @@
             @NonNull LockoutResetDispatcher lockoutResetDispatcher,
             @Authenticators.Types int biometricStrength) {
         super(context, lazyDaemon, null /* token */, null /* listener */, userId, owner,
-                0 /* cookie */, sensorId, biometricLogger, biometricContext);
+                0 /* cookie */, sensorId, biometricLogger, biometricContext,
+                false /* isMandatoryBiometrics */);
         mHardwareAuthToken = hardwareAuthToken == null ? null :
                 HardwareAuthTokenUtils.toHardwareAuthToken(hardwareAuthToken);
         mLockoutCache = lockoutTracker;
diff --git a/services/core/java/com/android/server/biometrics/sensors/iris/IrisAuthenticator.java b/services/core/java/com/android/server/biometrics/sensors/iris/IrisAuthenticator.java
index 01d1e378..76b5c4e 100644
--- a/services/core/java/com/android/server/biometrics/sensors/iris/IrisAuthenticator.java
+++ b/services/core/java/com/android/server/biometrics/sensors/iris/IrisAuthenticator.java
@@ -60,7 +60,7 @@
     public void prepareForAuthentication(boolean requireConfirmation, IBinder token,
             long sessionId, int userId, IBiometricSensorReceiver sensorReceiver,
             String opPackageName, long requestId, int cookie, boolean allowBackgroundAuthentication,
-            boolean isForLegacyFingerprintManager)
+            boolean isForLegacyFingerprintManager, boolean isMandatoryBiometrics)
             throws RemoteException {
     }
 
diff --git a/services/core/java/com/android/server/crashrecovery/CrashRecoveryUtils.java b/services/core/java/com/android/server/crashrecovery/CrashRecoveryUtils.java
new file mode 100644
index 0000000..3eb3380
--- /dev/null
+++ b/services/core/java/com/android/server/crashrecovery/CrashRecoveryUtils.java
@@ -0,0 +1,85 @@
+/*
+ * 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.crashrecovery;
+
+import android.os.Environment;
+import android.util.IndentingPrintWriter;
+import android.util.Slog;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.FileReader;
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.time.LocalDateTime;
+import java.time.ZoneId;
+
+/**
+ * Class containing helper methods for the CrashRecoveryModule.
+ *
+ * @hide
+ */
+public class CrashRecoveryUtils {
+    private static final String TAG = "CrashRecoveryUtils";
+    private static final long MAX_CRITICAL_INFO_DUMP_SIZE = 1000 * 1000; // ~1MB
+    private static final Object sFileLock = new Object();
+
+    /** Persist recovery related events in crashrecovery events file.**/
+    public static void logCrashRecoveryEvent(int priority, String msg) {
+        Slog.println(priority, TAG, msg);
+        try {
+            File fname = getCrashRecoveryEventsFile();
+            synchronized (sFileLock) {
+                FileOutputStream out = new FileOutputStream(fname, true);
+                PrintWriter pw = new PrintWriter(out);
+                String dateString = LocalDateTime.now(ZoneId.systemDefault()).toString();
+                pw.println(dateString + ": " + msg);
+                pw.close();
+            }
+        } catch (IOException e) {
+            Slog.e(TAG, "Unable to log CrashRecoveryEvents " + e.getMessage());
+        }
+    }
+
+    /** Dump recovery related events from crashrecovery events file.**/
+    public static void dumpCrashRecoveryEvents(IndentingPrintWriter pw) {
+        pw.println("CrashRecovery Events: ");
+        pw.increaseIndent();
+        final File file = getCrashRecoveryEventsFile();
+        final long skipSize = file.length() - MAX_CRITICAL_INFO_DUMP_SIZE;
+        synchronized (sFileLock) {
+            try (BufferedReader in = new BufferedReader(new FileReader(file))) {
+                if (skipSize > 0) {
+                    in.skip(skipSize);
+                }
+                String line;
+                while ((line = in.readLine()) != null) {
+                    pw.println(line);
+                }
+            } catch (IOException e) {
+                Slog.e(TAG, "Unable to dump CrashRecoveryEvents " + e.getMessage());
+            }
+        }
+        pw.decreaseIndent();
+    }
+
+    private static File getCrashRecoveryEventsFile() {
+        File systemDir = new File(Environment.getDataDirectory(), "system");
+        return new File(systemDir, "crashrecovery-events.txt");
+    }
+}
diff --git a/services/core/java/com/android/server/display/AutomaticBrightnessController.java b/services/core/java/com/android/server/display/AutomaticBrightnessController.java
index 240e91b..907e7c6 100644
--- a/services/core/java/com/android/server/display/AutomaticBrightnessController.java
+++ b/services/core/java/com/android/server/display/AutomaticBrightnessController.java
@@ -45,6 +45,7 @@
 import android.os.SystemClock;
 import android.os.Trace;
 import android.util.EventLog;
+import android.util.IndentingPrintWriter;
 import android.util.MathUtils;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -562,6 +563,7 @@
     public void resetShortTermModel() {
         mCurrentBrightnessMapper.clearUserDataPoints();
         mShortTermModel.reset();
+        Slog.i(TAG, "Resetting short term model");
     }
 
     public boolean setBrightnessConfiguration(BrightnessConfiguration configuration,
@@ -598,74 +600,79 @@
     }
 
     public void dump(PrintWriter pw) {
+        IndentingPrintWriter ipw = new IndentingPrintWriter(pw);
+        ipw.increaseIndent();
         pw.println();
         pw.println("Automatic Brightness Controller Configuration:");
-        pw.println("  mState=" + configStateToString(mState));
-        pw.println("  mScreenBrightnessRangeMinimum=" + mScreenBrightnessRangeMinimum);
-        pw.println("  mScreenBrightnessRangeMaximum=" + mScreenBrightnessRangeMaximum);
-        pw.println("  mDozeScaleFactor=" + mDozeScaleFactor);
-        pw.println("  mInitialLightSensorRate=" + mInitialLightSensorRate);
-        pw.println("  mNormalLightSensorRate=" + mNormalLightSensorRate);
-        pw.println("  mLightSensorWarmUpTimeConfig=" + mLightSensorWarmUpTimeConfig);
-        pw.println("  mBrighteningLightDebounceConfig=" + mBrighteningLightDebounceConfig);
-        pw.println("  mDarkeningLightDebounceConfig=" + mDarkeningLightDebounceConfig);
-        pw.println("  mBrighteningLightDebounceConfigIdle=" + mBrighteningLightDebounceConfigIdle);
-        pw.println("  mDarkeningLightDebounceConfigIdle=" + mDarkeningLightDebounceConfigIdle);
-        pw.println("  mResetAmbientLuxAfterWarmUpConfig=" + mResetAmbientLuxAfterWarmUpConfig);
-        pw.println("  mAmbientLightHorizonLong=" + mAmbientLightHorizonLong);
-        pw.println("  mAmbientLightHorizonShort=" + mAmbientLightHorizonShort);
-        pw.println("  mWeightingIntercept=" + mWeightingIntercept);
+        pw.println("----------------------------------------------");
+        ipw.println("mState=" + configStateToString(mState));
+        ipw.println("mScreenBrightnessRangeMinimum=" + mScreenBrightnessRangeMinimum);
+        ipw.println("mScreenBrightnessRangeMaximum=" + mScreenBrightnessRangeMaximum);
+        ipw.println("mDozeScaleFactor=" + mDozeScaleFactor);
+        ipw.println("mInitialLightSensorRate=" + mInitialLightSensorRate);
+        ipw.println("mNormalLightSensorRate=" + mNormalLightSensorRate);
+        ipw.println("mLightSensorWarmUpTimeConfig=" + mLightSensorWarmUpTimeConfig);
+        ipw.println("mBrighteningLightDebounceConfig=" + mBrighteningLightDebounceConfig);
+        ipw.println("mDarkeningLightDebounceConfig=" + mDarkeningLightDebounceConfig);
+        ipw.println("mBrighteningLightDebounceConfigIdle=" + mBrighteningLightDebounceConfigIdle);
+        ipw.println("mDarkeningLightDebounceConfigIdle=" + mDarkeningLightDebounceConfigIdle);
+        ipw.println("mResetAmbientLuxAfterWarmUpConfig=" + mResetAmbientLuxAfterWarmUpConfig);
+        ipw.println("mAmbientLightHorizonLong=" + mAmbientLightHorizonLong);
+        ipw.println("mAmbientLightHorizonShort=" + mAmbientLightHorizonShort);
+        ipw.println("mWeightingIntercept=" + mWeightingIntercept);
 
         pw.println();
         pw.println("Automatic Brightness Controller State:");
-        pw.println("  mLightSensor=" + mLightSensor);
-        pw.println("  mLightSensorEnabled=" + mLightSensorEnabled);
-        pw.println("  mLightSensorEnableTime=" + TimeUtils.formatUptime(mLightSensorEnableTime));
-        pw.println("  mCurrentLightSensorRate=" + mCurrentLightSensorRate);
-        pw.println("  mAmbientLux=" + mAmbientLux);
-        pw.println("  mAmbientLuxValid=" + mAmbientLuxValid);
-        pw.println("  mPreThresholdLux=" + mPreThresholdLux);
-        pw.println("  mPreThresholdBrightness=" + mPreThresholdBrightness);
-        pw.println("  mAmbientBrighteningThreshold=" + mAmbientBrighteningThreshold);
-        pw.println("  mAmbientDarkeningThreshold=" + mAmbientDarkeningThreshold);
-        pw.println("  mScreenBrighteningThreshold=" + mScreenBrighteningThreshold);
-        pw.println("  mScreenDarkeningThreshold=" + mScreenDarkeningThreshold);
-        pw.println("  mLastObservedLux=" + mLastObservedLux);
-        pw.println("  mLastObservedLuxTime=" + TimeUtils.formatUptime(mLastObservedLuxTime));
-        pw.println("  mRecentLightSamples=" + mRecentLightSamples);
-        pw.println("  mAmbientLightRingBuffer=" + mAmbientLightRingBuffer);
-        pw.println("  mScreenAutoBrightness=" + mScreenAutoBrightness);
-        pw.println("  mDisplayPolicy=" + DisplayPowerRequest.policyToString(mDisplayPolicy));
-        pw.println("  mShortTermModel=");
-        mShortTermModel.dump(pw);
-        pw.println("  mPausedShortTermModel=");
-        mPausedShortTermModel.dump(pw);
+        pw.println("--------------------------------------");
+        ipw.println("mLightSensor=" + mLightSensor);
+        ipw.println("mLightSensorEnabled=" + mLightSensorEnabled);
+        ipw.println("mLightSensorEnableTime=" + TimeUtils.formatUptime(mLightSensorEnableTime));
+        ipw.println("mCurrentLightSensorRate=" + mCurrentLightSensorRate);
+        ipw.println("mAmbientLux=" + mAmbientLux);
+        ipw.println("mAmbientLuxValid=" + mAmbientLuxValid);
+        ipw.println("mPreThresholdLux=" + mPreThresholdLux);
+        ipw.println("mPreThresholdBrightness=" + mPreThresholdBrightness);
+        ipw.println("mAmbientBrighteningThreshold=" + mAmbientBrighteningThreshold);
+        ipw.println("mAmbientDarkeningThreshold=" + mAmbientDarkeningThreshold);
+        ipw.println("mScreenBrighteningThreshold=" + mScreenBrighteningThreshold);
+        ipw.println("mScreenDarkeningThreshold=" + mScreenDarkeningThreshold);
+        ipw.println("mLastObservedLux=" + mLastObservedLux);
+        ipw.println("mLastObservedLuxTime=" + TimeUtils.formatUptime(mLastObservedLuxTime));
+        ipw.println("mRecentLightSamples=" + mRecentLightSamples);
+        ipw.println("mAmbientLightRingBuffer=" + mAmbientLightRingBuffer);
+        ipw.println("mScreenAutoBrightness=" + mScreenAutoBrightness);
+        ipw.println("mDisplayPolicy=" + DisplayPowerRequest.policyToString(mDisplayPolicy));
+        ipw.println("mShortTermModel=");
 
-        pw.println();
-        pw.println("  mBrightnessAdjustmentSamplePending=" + mBrightnessAdjustmentSamplePending);
-        pw.println("  mBrightnessAdjustmentSampleOldLux=" + mBrightnessAdjustmentSampleOldLux);
-        pw.println("  mBrightnessAdjustmentSampleOldBrightness="
+        mShortTermModel.dump(ipw);
+        ipw.println("mPausedShortTermModel=");
+        mPausedShortTermModel.dump(ipw);
+
+        ipw.println();
+        ipw.println("mBrightnessAdjustmentSamplePending=" + mBrightnessAdjustmentSamplePending);
+        ipw.println("mBrightnessAdjustmentSampleOldLux=" + mBrightnessAdjustmentSampleOldLux);
+        ipw.println("mBrightnessAdjustmentSampleOldBrightness="
                 + mBrightnessAdjustmentSampleOldBrightness);
-        pw.println("  mForegroundAppPackageName=" + mForegroundAppPackageName);
-        pw.println("  mPendingForegroundAppPackageName=" + mPendingForegroundAppPackageName);
-        pw.println("  mForegroundAppCategory=" + mForegroundAppCategory);
-        pw.println("  mPendingForegroundAppCategory=" + mPendingForegroundAppCategory);
-        pw.println("  Current mode="
+        ipw.println("mForegroundAppPackageName=" + mForegroundAppPackageName);
+        ipw.println("mPendingForegroundAppPackageName=" + mPendingForegroundAppPackageName);
+        ipw.println("mForegroundAppCategory=" + mForegroundAppCategory);
+        ipw.println("mPendingForegroundAppCategory=" + mPendingForegroundAppCategory);
+        ipw.println("Current mode="
                 + autoBrightnessModeToString(mCurrentBrightnessMapper.getMode()));
 
         for (int i = 0; i < mBrightnessMappingStrategyMap.size(); i++) {
-            pw.println();
-            pw.println("  Mapper for mode "
+            ipw.println();
+            ipw.println("Mapper for mode "
                     + autoBrightnessModeToString(mBrightnessMappingStrategyMap.keyAt(i)) + ":");
-            mBrightnessMappingStrategyMap.valueAt(i).dump(pw,
+            mBrightnessMappingStrategyMap.valueAt(i).dump(ipw,
                     mBrightnessRangeController.getNormalBrightnessMax());
         }
 
-        pw.println();
-        pw.println("  mAmbientBrightnessThresholds=" + mAmbientBrightnessThresholds);
-        pw.println("  mAmbientBrightnessThresholdsIdle=" + mAmbientBrightnessThresholdsIdle);
-        pw.println("  mScreenBrightnessThresholds=" + mScreenBrightnessThresholds);
-        pw.println("  mScreenBrightnessThresholdsIdle=" + mScreenBrightnessThresholdsIdle);
+        ipw.println();
+        ipw.println("mAmbientBrightnessThresholds=" + mAmbientBrightnessThresholds);
+        ipw.println("mAmbientBrightnessThresholdsIdle=" + mAmbientBrightnessThresholdsIdle);
+        ipw.println("mScreenBrightnessThresholds=" + mScreenBrightnessThresholds);
+        ipw.println("mScreenBrightnessThresholdsIdle=" + mScreenBrightnessThresholdsIdle);
     }
 
     public float[] getLastSensorValues() {
@@ -1339,8 +1346,10 @@
                     + "\n mIsValid: " + mIsValid;
         }
 
-        void dump(PrintWriter pw) {
-            pw.println(this);
+        void dump(IndentingPrintWriter ipw) {
+            ipw.increaseIndent();
+            ipw.println(this);
+            ipw.decreaseIndent();
         }
 
     }
diff --git a/services/core/java/com/android/server/display/BrightnessMappingStrategy.java b/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
index b0507fb..6a019f3 100644
--- a/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
+++ b/services/core/java/com/android/server/display/BrightnessMappingStrategy.java
@@ -1073,7 +1073,9 @@
             pw.println("  mBrightnessRangeAdjustmentApplied=" + mBrightnessRangeAdjustmentApplied);
             pw.println("  shortTermModelTimeout=" + getShortTermModelTimeout());
 
-            pw.println("  Previous short-term models (oldest to newest): ");
+            if (!mPreviousBrightnessSplines.isEmpty()) {
+                pw.println("  Previous short-term models (oldest to newest): ");
+            }
             for (int i = 0; i < mPreviousBrightnessSplines.size(); i++) {
                 pw.println("  Computed at "
                         + FORMAT.format(new Date(mBrightnessSplineChangeTimes.get(i))) + ": ");
diff --git a/services/core/java/com/android/server/display/BrightnessThrottler.java b/services/core/java/com/android/server/display/BrightnessThrottler.java
index 631e751..b56a234 100644
--- a/services/core/java/com/android/server/display/BrightnessThrottler.java
+++ b/services/core/java/com/android/server/display/BrightnessThrottler.java
@@ -265,6 +265,7 @@
 
     private void dumpLocal(PrintWriter pw) {
         pw.println("BrightnessThrottler:");
+        pw.println("--------------------");
         pw.println("  mThermalBrightnessThrottlingDataId=" + mThermalBrightnessThrottlingDataId);
         pw.println("  mThermalThrottlingData=" + mThermalThrottlingData);
         pw.println("  mUniqueDisplayId=" + mUniqueDisplayId);
diff --git a/services/core/java/com/android/server/display/BrightnessTracker.java b/services/core/java/com/android/server/display/BrightnessTracker.java
index ac5dd20..0f65360 100644
--- a/services/core/java/com/android/server/display/BrightnessTracker.java
+++ b/services/core/java/com/android/server/display/BrightnessTracker.java
@@ -782,6 +782,7 @@
 
     public void dump(final PrintWriter pw) {
         pw.println("BrightnessTracker state:");
+        pw.println("------------------------");
         synchronized (mDataCollectionLock) {
             pw.println("  mStarted=" + mStarted);
             pw.println("  mLightSensor=" + mLightSensor);
diff --git a/services/core/java/com/android/server/display/DeviceStateToLayoutMap.java b/services/core/java/com/android/server/display/DeviceStateToLayoutMap.java
index 146810f..4c133ff 100644
--- a/services/core/java/com/android/server/display/DeviceStateToLayoutMap.java
+++ b/services/core/java/com/android/server/display/DeviceStateToLayoutMap.java
@@ -95,6 +95,7 @@
 
     public void dumpLocked(IndentingPrintWriter ipw) {
         ipw.println("DeviceStateToLayoutMap:");
+        ipw.println("-----------------------");
         ipw.increaseIndent();
 
         ipw.println("mIsPortInDisplayLayoutEnabled=" + mIsPortInDisplayLayoutEnabled);
diff --git a/services/core/java/com/android/server/display/DisplayBrightnessState.java b/services/core/java/com/android/server/display/DisplayBrightnessState.java
index dc611fc..334dda0 100644
--- a/services/core/java/com/android/server/display/DisplayBrightnessState.java
+++ b/services/core/java/com/android/server/display/DisplayBrightnessState.java
@@ -16,6 +16,7 @@
 
 package com.android.server.display;
 
+import android.hardware.display.BrightnessInfo;
 import android.text.TextUtils;
 
 import com.android.server.display.brightness.BrightnessEvent;
@@ -50,6 +51,8 @@
 
     private final boolean mIsUserInitiatedChange;
 
+    private @BrightnessInfo.BrightnessMaxReason int mBrightnessMaxReason;
+
     private DisplayBrightnessState(Builder builder) {
         mBrightness = builder.getBrightness();
         mHdrBrightness = builder.getHdrBrightness();
@@ -64,6 +67,7 @@
         mBrightnessEvent = builder.getBrightnessEvent();
         mBrightnessAdjustmentFlag = builder.getBrightnessAdjustmentFlag();
         mIsUserInitiatedChange = builder.isUserInitiatedChange();
+        mBrightnessMaxReason = builder.getBrightnessMaxReason();
     }
 
     /**
@@ -159,6 +163,13 @@
         return mIsUserInitiatedChange;
     }
 
+    /**
+     * Gets reason for max brightness restriction
+     */
+    public @BrightnessInfo.BrightnessMaxReason int getBrightnessMaxReason() {
+        return mBrightnessMaxReason;
+    }
+
     @Override
     public String toString() {
         StringBuilder stringBuilder = new StringBuilder("DisplayBrightnessState:");
@@ -180,6 +191,8 @@
                 .append(Objects.toString(mBrightnessEvent, "null"));
         stringBuilder.append("\n    mBrightnessAdjustmentFlag:").append(mBrightnessAdjustmentFlag);
         stringBuilder.append("\n    mIsUserInitiatedChange:").append(mIsUserInitiatedChange);
+        stringBuilder.append("\n    mBrightnessMaxReason:")
+                .append(BrightnessInfo.briMaxReasonToString(mBrightnessMaxReason));
         return stringBuilder.toString();
     }
 
@@ -212,7 +225,8 @@
                     == otherState.shouldUpdateScreenBrightnessSetting()
                 && Objects.equals(mBrightnessEvent, otherState.getBrightnessEvent())
                 && mBrightnessAdjustmentFlag == otherState.getBrightnessAdjustmentFlag()
-                && mIsUserInitiatedChange == otherState.isUserInitiatedChange();
+                && mIsUserInitiatedChange == otherState.isUserInitiatedChange()
+                && mBrightnessMaxReason == otherState.getBrightnessMaxReason();
     }
 
     @Override
@@ -221,7 +235,7 @@
                 mShouldUseAutoBrightness, mIsSlowChange, mMaxBrightness, mMinBrightness,
                 mCustomAnimationRate,
                 mShouldUpdateScreenBrightnessSetting, mBrightnessEvent, mBrightnessAdjustmentFlag,
-                mIsUserInitiatedChange);
+                mIsUserInitiatedChange, mBrightnessMaxReason);
     }
 
     /**
@@ -245,12 +259,11 @@
         private float mMinBrightness;
         private float mCustomAnimationRate = CUSTOM_ANIMATION_RATE_NOT_SET;
         private boolean mShouldUpdateScreenBrightnessSetting;
-
         private BrightnessEvent mBrightnessEvent;
-
-        public int mBrightnessAdjustmentFlag = 0;
-
+        private int mBrightnessAdjustmentFlag = 0;
         private boolean mIsUserInitiatedChange;
+        private @BrightnessInfo.BrightnessMaxReason int mBrightnessMaxReason =
+                BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
 
         /**
          * Create a builder starting with the values from the specified {@link
@@ -274,6 +287,7 @@
             builder.setBrightnessEvent(state.getBrightnessEvent());
             builder.setBrightnessAdjustmentFlag(state.getBrightnessAdjustmentFlag());
             builder.setIsUserInitiatedChange(state.isUserInitiatedChange());
+            builder.setBrightnessMaxReason(state.getBrightnessMaxReason());
             return builder;
         }
 
@@ -506,5 +520,21 @@
             mIsUserInitiatedChange = isUserInitiatedChange;
             return this;
         }
+
+        /**
+         * Gets reason for max brightness restriction
+         */
+        public @BrightnessInfo.BrightnessMaxReason int getBrightnessMaxReason() {
+            return mBrightnessMaxReason;
+        }
+
+        /**
+         * Sets reason for max brightness restriction
+         */
+        public Builder setBrightnessMaxReason(
+                @BrightnessInfo.BrightnessMaxReason int brightnessMaxReason) {
+            mBrightnessMaxReason = brightnessMaxReason;
+            return this;
+        }
     }
 }
diff --git a/services/core/java/com/android/server/display/DisplayDeviceConfig.java b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
index cc115f1..d78fdfa 100644
--- a/services/core/java/com/android/server/display/DisplayDeviceConfig.java
+++ b/services/core/java/com/android/server/display/DisplayDeviceConfig.java
@@ -150,7 +150,9 @@
  *      <screenBrightnessDefault>0.65</screenBrightnessDefault>
  *      <powerThrottlingConfig>
  *        <brightnessLowestCapAllowed>0.1</brightnessLowestCapAllowed>
- *        <pollingWindowMillis>15</pollingWindowMillis>
+ *        <customAnimationRateSec>0.004</customAnimationRateSec>
+ *        <pollingWindowMaxMillis>30000</pollingWindowMaxMillis>
+ *        <pollingWindowMinMillis>10000</pollingWindowMinMillis>
  *          <powerThrottlingMap>
  *              <powerThrottlingPoint>
  *                  <thermalStatus>severe</thermalStatus>
@@ -2184,9 +2186,13 @@
             return;
         }
         float lowestBrightnessCap = powerThrottlingCfg.getBrightnessLowestCapAllowed().floatValue();
-        int pollingWindowMillis = powerThrottlingCfg.getPollingWindowMillis().intValue();
+        float customAnimationRateSec = powerThrottlingCfg.getCustomAnimationRateSec().floatValue();
+        int pollingWindowMaxMillis = powerThrottlingCfg.getPollingWindowMaxMillis().intValue();
+        int pollingWindowMinMillis = powerThrottlingCfg.getPollingWindowMinMillis().intValue();
         mPowerThrottlingConfigData = new PowerThrottlingConfigData(lowestBrightnessCap,
-                                                                   pollingWindowMillis);
+                                                                   customAnimationRateSec,
+                                                                   pollingWindowMaxMillis,
+                                                                   pollingWindowMinMillis);
     }
 
     private void loadRefreshRateSetting(DisplayConfiguration config) {
@@ -2980,12 +2986,19 @@
     public static class PowerThrottlingConfigData {
         /** Lowest brightness cap allowed for this device. */
         public final float brightnessLowestCapAllowed;
-        /** Time window for polling power in seconds. */
-        public final int pollingWindowMillis;
+        /** Time take to animate brightness in seconds. */
+        public final float customAnimationRateSec;
+        /** Time window for maximum polling power in milliseconds. */
+        public final int pollingWindowMaxMillis;
+        /** Time window for minimum polling power in milliseconds. */
+        public final int pollingWindowMinMillis;
         public PowerThrottlingConfigData(float brightnessLowestCapAllowed,
-                int pollingWindowMillis) {
+                float customAnimationRateSec, int pollingWindowMaxMillis,
+                int pollingWindowMinMillis) {
             this.brightnessLowestCapAllowed = brightnessLowestCapAllowed;
-            this.pollingWindowMillis = pollingWindowMillis;
+            this.customAnimationRateSec = customAnimationRateSec;
+            this.pollingWindowMaxMillis = pollingWindowMaxMillis;
+            this.pollingWindowMinMillis = pollingWindowMinMillis;
         }
 
         @Override
@@ -2993,7 +3006,9 @@
             return "PowerThrottlingConfigData{"
                     + "brightnessLowestCapAllowed: "
                     + brightnessLowestCapAllowed
-                    + ", pollingWindowMillis: " + pollingWindowMillis
+                    + ", customAnimationRateSec: " + customAnimationRateSec
+                    + ", pollingWindowMaxMillis: " + pollingWindowMaxMillis
+                    + ", pollingWindowMinMillis: " + pollingWindowMinMillis
                     + "} ";
         }
     }
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 187caba..0abd9bc 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -354,7 +354,7 @@
             new CopyOnWriteArrayList<>();
 
     /** All {@link DisplayPowerController}s indexed by {@link LogicalDisplay} ID. */
-    private final SparseArray<DisplayPowerControllerInterface> mDisplayPowerControllers =
+    private final SparseArray<DisplayPowerController> mDisplayPowerControllers =
             new SparseArray<>();
 
     /** {@link DisplayBlanker} used by all {@link DisplayPowerController}s. */
@@ -726,7 +726,7 @@
                 if (logicalDisplay.getDisplayInfoLocked().type != Display.TYPE_INTERNAL) {
                     return;
                 }
-                final DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(
+                final DisplayPowerController dpc = mDisplayPowerControllers.get(
                         logicalDisplay.getDisplayIdLocked());
                 if (dpc == null) {
                     return;
@@ -2058,7 +2058,7 @@
             configurePreferredDisplayModeLocked(display);
         }
 
-        DisplayPowerControllerInterface dpc = addDisplayPowerControllerLocked(display);
+        DisplayPowerController dpc = addDisplayPowerControllerLocked(display);
         if (dpc != null) {
             final int leadDisplayId = display.getLeadDisplayIdLocked();
             updateDisplayPowerControllerLeaderLocked(dpc, leadDisplayId);
@@ -2067,7 +2067,7 @@
             // that the follower display was added before the lead display.
             mLogicalDisplayMapper.forEachLocked(d -> {
                 if (d.getLeadDisplayIdLocked() == displayId) {
-                    DisplayPowerControllerInterface followerDpc =
+                    DisplayPowerController followerDpc =
                             mDisplayPowerControllers.get(d.getDisplayIdLocked());
                     if (followerDpc != null) {
                         updateDisplayPowerControllerLeaderLocked(followerDpc, displayId);
@@ -2151,7 +2151,7 @@
         scheduleTraversalLocked(false);
         mPersistentDataStore.saveIfNeeded();
 
-        DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(displayId);
+        DisplayPowerController dpc = mDisplayPowerControllers.get(displayId);
         if (dpc != null) {
             final int leadDisplayId = display.getLeadDisplayIdLocked();
             updateDisplayPowerControllerLeaderLocked(dpc, leadDisplayId);
@@ -2165,7 +2165,7 @@
     }
 
     private void updateDisplayPowerControllerLeaderLocked(
-            @NonNull DisplayPowerControllerInterface dpc, int leadDisplayId) {
+            @NonNull DisplayPowerController dpc, int leadDisplayId) {
         if (dpc.getLeadDisplayId() == leadDisplayId) {
             // Lead display hasn't changed, nothing to do.
             return;
@@ -2174,7 +2174,7 @@
         // If it has changed, then we need to unregister from the previous leader if there was one.
         final int prevLeaderId = dpc.getLeadDisplayId();
         if (prevLeaderId != Layout.NO_LEAD_DISPLAY) {
-            final DisplayPowerControllerInterface prevLeader =
+            final DisplayPowerController prevLeader =
                     mDisplayPowerControllers.get(prevLeaderId);
             if (prevLeader != null) {
                 prevLeader.removeDisplayBrightnessFollower(dpc);
@@ -2183,7 +2183,7 @@
 
         // And then, if it's following, register it with the new one.
         if (leadDisplayId != Layout.NO_LEAD_DISPLAY) {
-            final DisplayPowerControllerInterface newLeader =
+            final DisplayPowerController newLeader =
                     mDisplayPowerControllers.get(leadDisplayId);
             if (newLeader != null) {
                 newLeader.addDisplayBrightnessFollower(dpc);
@@ -2224,7 +2224,7 @@
     private void releaseDisplayAndEmitEvent(LogicalDisplay display, int event) {
         final int displayId = display.getDisplayIdLocked();
 
-        final DisplayPowerControllerInterface dpc =
+        final DisplayPowerController dpc =
                 mDisplayPowerControllers.removeReturnOld(displayId);
         if (dpc != null) {
             updateDisplayPowerControllerLeaderLocked(dpc, Layout.NO_LEAD_DISPLAY);
@@ -2271,7 +2271,7 @@
 
     private void handleLogicalDisplayDeviceStateTransitionLocked(@NonNull LogicalDisplay display) {
         final int displayId = display.getDisplayIdLocked();
-        final DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(displayId);
+        final DisplayPowerController dpc = mDisplayPowerControllers.get(displayId);
         if (dpc != null) {
             final int leadDisplayId = display.getLeadDisplayIdLocked();
             updateDisplayPowerControllerLeaderLocked(dpc, leadDisplayId);
@@ -2692,14 +2692,14 @@
             if (userId != mCurrentUserId) {
                 return;
             }
-            DisplayPowerControllerInterface dpc = getDpcFromUniqueIdLocked(uniqueId);
+            DisplayPowerController dpc = getDpcFromUniqueIdLocked(uniqueId);
             if (dpc != null) {
                 dpc.setBrightnessConfiguration(c, /* shouldResetShortTermModel= */ true);
             }
         }
     }
 
-    private DisplayPowerControllerInterface getDpcFromUniqueIdLocked(String uniqueId) {
+    private DisplayPowerController getDpcFromUniqueIdLocked(String uniqueId) {
         final DisplayDevice displayDevice = mDisplayDeviceRepo.getByUniqueIdLocked(uniqueId);
         final LogicalDisplay logicalDisplay = mLogicalDisplayMapper.getDisplayLocked(displayDevice);
         if (logicalDisplay != null) {
@@ -2740,7 +2740,7 @@
                 final BrightnessConfiguration config =
                         getBrightnessConfigForDisplayWithPdsFallbackLocked(uniqueId, userSerial);
                 if (config != null) {
-                    final DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(
+                    final DisplayPowerController dpc = mDisplayPowerControllers.get(
                             logicalDisplay.getDisplayIdLocked());
                     if (dpc != null) {
                         dpc.setBrightnessConfiguration(config,
@@ -2987,7 +2987,7 @@
 
     void setAutoBrightnessLoggingEnabled(boolean enabled) {
         synchronized (mSyncRoot) {
-            final DisplayPowerControllerInterface displayPowerController =
+            final DisplayPowerController displayPowerController =
                     mDisplayPowerControllers.get(Display.DEFAULT_DISPLAY);
             if (displayPowerController != null) {
                 displayPowerController.setAutoBrightnessLoggingEnabled(enabled);
@@ -2997,7 +2997,7 @@
 
     void setDisplayWhiteBalanceLoggingEnabled(boolean enabled) {
         synchronized (mSyncRoot) {
-            final DisplayPowerControllerInterface displayPowerController =
+            final DisplayPowerController displayPowerController =
                     mDisplayPowerControllers.get(Display.DEFAULT_DISPLAY);
             if (displayPowerController != null) {
                 displayPowerController.setDisplayWhiteBalanceLoggingEnabled(enabled);
@@ -3023,7 +3023,7 @@
 
     void setAmbientColorTemperatureOverride(float cct) {
         synchronized (mSyncRoot) {
-            final DisplayPowerControllerInterface displayPowerController =
+            final DisplayPowerController displayPowerController =
                     mDisplayPowerControllers.get(Display.DEFAULT_DISPLAY);
             if (displayPowerController != null) {
                 displayPowerController.setAmbientColorTemperatureOverride(cct);
@@ -3033,7 +3033,7 @@
 
     void setDockedAndIdleEnabled(boolean enabled, int displayId) {
         synchronized (mSyncRoot) {
-            final DisplayPowerControllerInterface displayPowerController =
+            final DisplayPowerController displayPowerController =
                     mDisplayPowerControllers.get(displayId);
             if (displayPowerController != null) {
                 displayPowerController.setAutomaticScreenBrightnessMode(enabled
@@ -3340,6 +3340,7 @@
             pw.println();
             final int displayStateCount = mDisplayStates.size();
             pw.println("Display States: size=" + displayStateCount);
+            pw.println("---------------------");
             for (int i = 0; i < displayStateCount; i++) {
                 final int displayId = mDisplayStates.keyAt(i);
                 final int displayState = mDisplayStates.valueAt(i);
@@ -3355,6 +3356,7 @@
 
             pw.println();
             pw.println("Display Adapters: size=" + mDisplayAdapters.size());
+            pw.println("------------------------");
             for (DisplayAdapter adapter : mDisplayAdapters) {
                 pw.println("  " + adapter.getName());
                 adapter.dumpLocked(ipw);
@@ -3362,6 +3364,7 @@
 
             pw.println();
             pw.println("Display Devices: size=" + mDisplayDeviceRepo.sizeLocked());
+            pw.println("-----------------------");
             mDisplayDeviceRepo.forEachLocked(device -> {
                 pw.println("  " + device.getDisplayDeviceInfoLocked());
                 device.dumpLocked(ipw);
@@ -3373,6 +3376,7 @@
             final int callbackCount = mCallbacks.size();
             pw.println();
             pw.println("Callbacks: size=" + callbackCount);
+            pw.println("-----------------");
             for (int i = 0; i < callbackCount; i++) {
                 CallbackRecord callback = mCallbacks.valueAt(i);
                 pw.println("  " + i + ": mPid=" + callback.mPid
@@ -3385,6 +3389,7 @@
             for (int i = 0; i < displayPowerControllerCount; i++) {
                 mDisplayPowerControllers.valueAt(i).dump(pw);
             }
+
             pw.println();
             mPersistentDataStore.dump(pw);
 
@@ -3403,8 +3408,10 @@
         }
         pw.println();
         mDisplayModeDirector.dump(pw);
+        pw.println();
         mBrightnessSynchronizer.dump(pw);
         if (mSmallAreaDetectionController != null) {
+            pw.println();
             mSmallAreaDetectionController.dump(pw);
         }
 
@@ -3564,7 +3571,7 @@
     }
 
     @RequiresPermission(Manifest.permission.READ_DEVICE_CONFIG)
-    private DisplayPowerControllerInterface addDisplayPowerControllerLocked(
+    private DisplayPowerController addDisplayPowerControllerLocked(
             LogicalDisplay display) {
         if (mPowerHandler == null) {
             // initPowerManagement has not yet been called.
@@ -3578,7 +3585,7 @@
         final int userSerial = getUserManager().getUserSerialNumber(mContext.getUserId());
         final BrightnessSetting brightnessSetting = new BrightnessSetting(userSerial,
                 mPersistentDataStore, display, mSyncRoot);
-        final DisplayPowerControllerInterface displayPowerController;
+        final DisplayPowerController displayPowerController;
 
         // If display is internal and has a HighBrightnessModeMetadata mapping, use that.
         // Or create a new one and use that.
@@ -4366,7 +4373,7 @@
                                     uniqueId, userSerial);
                     if (config == null) {
                         // Get default configuration
-                        DisplayPowerControllerInterface dpc = getDpcFromUniqueIdLocked(uniqueId);
+                        DisplayPowerController dpc = getDpcFromUniqueIdLocked(uniqueId);
                         if (dpc != null) {
                             config = dpc.getDefaultBrightnessConfiguration();
                         }
@@ -4420,7 +4427,7 @@
                     if (display == null || !display.isEnabledLocked()) {
                         return null;
                     }
-                    DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(displayId);
+                    DisplayPowerController dpc = mDisplayPowerControllers.get(displayId);
                     if (dpc != null) {
                         return dpc.getBrightnessInfo();
                     }
@@ -4465,7 +4472,7 @@
             final long token = Binder.clearCallingIdentity();
             try {
                 synchronized (mSyncRoot) {
-                    DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(displayId);
+                    DisplayPowerController dpc = mDisplayPowerControllers.get(displayId);
                     if (dpc != null) {
                         dpc.setBrightness(brightness);
                     }
@@ -4485,7 +4492,7 @@
             final long token = Binder.clearCallingIdentity();
             try {
                 synchronized (mSyncRoot) {
-                    DisplayPowerControllerInterface dpc = mDisplayPowerControllers.get(displayId);
+                    DisplayPowerController dpc = mDisplayPowerControllers.get(displayId);
                     if (dpc != null) {
                         brightness = dpc.getScreenBrightnessSetting();
                     }
@@ -4812,7 +4819,7 @@
                             id).getPrimaryDisplayDeviceLocked();
                     final int flags = displayDevice.getDisplayDeviceInfoLocked().flags;
                     if ((flags & DisplayDeviceInfo.FLAG_NEVER_BLANK) == 0) {
-                        final DisplayPowerControllerInterface displayPowerController =
+                        final DisplayPowerController displayPowerController =
                                 mDisplayPowerControllers.get(id);
                         if (displayPowerController != null) {
                             ready &= displayPowerController.requestPowerState(request,
@@ -5193,7 +5200,7 @@
                     return null;
                 }
 
-                DisplayPowerControllerInterface displayPowerController =
+                DisplayPowerController displayPowerController =
                         mDisplayPowerControllers.get(logicalDisplay.getDisplayIdLocked());
                 if (displayPowerController == null) {
                     Slog.w(TAG,
diff --git a/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java b/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
index a188e79..b05a96e 100644
--- a/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
+++ b/services/core/java/com/android/server/display/DisplayOffloadSessionImpl.java
@@ -39,12 +39,12 @@
 
     @Nullable
     private final DisplayManagerInternal.DisplayOffloader mDisplayOffloader;
-    private final DisplayPowerControllerInterface mDisplayPowerController;
+    private final DisplayPowerController mDisplayPowerController;
     private boolean mIsActive;
 
     public DisplayOffloadSessionImpl(
             @Nullable DisplayManagerInternal.DisplayOffloader displayOffloader,
-            DisplayPowerControllerInterface displayPowerController) {
+            DisplayPowerController displayPowerController) {
         mDisplayOffloader = displayOffloader;
         mDisplayPowerController = displayPowerController;
     }
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index a887f6d..bb2bed7 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -126,7 +126,7 @@
  * slower by changing the "animator duration scale" option in Development Settings.
  */
 final class DisplayPowerController implements AutomaticBrightnessController.Callbacks,
-        DisplayWhiteBalanceController.Callbacks, DisplayPowerControllerInterface {
+        DisplayWhiteBalanceController.Callbacks{
     private static final String SCREEN_ON_BLOCKED_TRACE_NAME = "Screen on blocked";
     private static final String SCREEN_OFF_BLOCKED_TRACE_NAME = "Screen off blocked";
 
@@ -253,10 +253,10 @@
     // The display blanker.
     private final DisplayBlanker mBlanker;
 
-    // The LogicalDisplay tied to this DisplayPowerController2.
+    // The LogicalDisplay tied to this DisplayPowerController.
     private final LogicalDisplay mLogicalDisplay;
 
-    // The ID of the LogicalDisplay tied to this DisplayPowerController2.
+    // The ID of the LogicalDisplay tied to this DisplayPowerController.
     private final int mDisplayId;
 
     // The ID of the display which this display follows for brightness purposes.
@@ -466,7 +466,7 @@
     private ObjectAnimator mColorFadeOffAnimator;
     private DualRampAnimator<DisplayPowerState> mScreenBrightnessRampAnimator;
 
-    // True if this DisplayPowerController2 has been stopped and should no longer be running.
+    // True if this DisplayPowerController has been stopped and should no longer be running.
     private boolean mStopped;
 
     private DisplayDeviceConfig mDisplayDeviceConfig;
@@ -481,7 +481,7 @@
     // DPCs following the brightness of this DPC. This is used in concurrent displays mode - there
     // is one lead display, the additional displays follow the brightness value of the lead display.
     @GuardedBy("mLock")
-    private final SparseArray<DisplayPowerControllerInterface> mDisplayBrightnessFollowers =
+    private final SparseArray<DisplayPowerController> mDisplayBrightnessFollowers =
             new SparseArray();
 
     private boolean mBootCompleted;
@@ -591,7 +591,8 @@
                         mThermalBrightnessThrottlingDataId,
                         logicalDisplay.getPowerThrottlingDataIdLocked(),
                         mDisplayDeviceConfig, displayDeviceInfo.width, displayDeviceInfo.height,
-                        displayToken, mDisplayId), mContext, flags, mSensorManager);
+                        displayToken, mDisplayId), mContext, flags, mSensorManager,
+                        mDisplayBrightnessController.getCurrentBrightness());
         // Seed the cached brightness
         saveBrightnessInfo(getScreenBrightnessSetting());
         mAutomaticBrightnessStrategy =
@@ -678,7 +679,6 @@
     /**
      * Returns true if the proximity sensor screen-off function is available.
      */
-    @Override
     public boolean isProximitySensorAvailable() {
         return mDisplayPowerProximityStateController.isProximitySensorAvailable();
     }
@@ -690,7 +690,6 @@
      * @param includePackage if false will null out the package name in events
      */
     @Nullable
-    @Override
     public ParceledListSlice<BrightnessChangeEvent> getBrightnessEvents(
             @UserIdInt int userId, boolean includePackage) {
         if (mBrightnessTracker == null) {
@@ -699,7 +698,6 @@
         return mBrightnessTracker.getEvents(userId, includePackage);
     }
 
-    @Override
     public void onSwitchUser(@UserIdInt int newUserId, int userSerial, float newBrightness) {
         Message msg = mHandler.obtainMessage(MSG_SWITCH_USER, newUserId, userSerial, newBrightness);
         mHandler.sendMessageAtTime(msg, mClock.uptimeMillis());
@@ -736,7 +734,6 @@
     }
 
     @Nullable
-    @Override
     public ParceledListSlice<AmbientBrightnessDayStats> getAmbientBrightnessStats(
             @UserIdInt int userId) {
         if (mBrightnessTracker == null) {
@@ -748,7 +745,6 @@
     /**
      * Persist the brightness slider events and ambient brightness stats to disk.
      */
-    @Override
     public void persistBrightnessTrackerState() {
         if (mBrightnessTracker != null) {
             mBrightnessTracker.persistBrightnessTrackerState();
@@ -805,7 +801,6 @@
         }
     }
 
-    @Override
     public void overrideDozeScreenState(int displayState, @Display.StateReason int reason) {
         Slog.i(TAG, "New offload doze override: " + Display.stateToString(displayState));
         if (mDisplayOffloadSession != null
@@ -832,7 +827,6 @@
         }
     }
 
-    @Override
     public void setDisplayOffloadSession(DisplayOffloadSession session) {
         if (session == mDisplayOffloadSession) {
             return;
@@ -841,7 +835,6 @@
         mDisplayOffloadSession = session;
     }
 
-    @Override
     public BrightnessConfiguration getDefaultBrightnessConfiguration() {
         if (mAutomaticBrightnessController == null) {
             return null;
@@ -856,12 +849,11 @@
      *
      * Make sure DisplayManagerService.mSyncRoot lock is held when this is called
      */
-    @Override
     public void onDisplayChanged(HighBrightnessModeMetadata hbmMetadata, int leadDisplayId) {
         mLeadDisplayId = leadDisplayId;
         final DisplayDevice device = mLogicalDisplay.getPrimaryDisplayDeviceLocked();
         if (device == null) {
-            Slog.wtf(mTag, "Display Device is null in DisplayPowerController2 for display: "
+            Slog.wtf(mTag, "Display Device is null in DisplayPowerController for display: "
                     + mLogicalDisplay.getDisplayIdLocked());
             return;
         }
@@ -935,10 +927,9 @@
     /**
      * Unregisters all listeners and interrupts all running threads; halting future work.
      *
-     * This method should be called when the DisplayPowerController2 is no longer in use; i.e. when
+     * This method should be called when the DisplayPowerController is no longer in use; i.e. when
      * the {@link #mDisplayId display} has been removed.
      */
-    @Override
     public void stop() {
         synchronized (mLock) {
             clearDisplayBrightnessFollowersLocked();
@@ -1215,7 +1206,6 @@
         }
     }
 
-    @Override
     public void setAutomaticScreenBrightnessMode(
             @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
         Message msg = mHandler.obtainMessage();
@@ -1313,7 +1303,7 @@
         boolean mustInitialize = false;
         mBrightnessReasonTemp.set(null);
         mTempBrightnessEvent.reset();
-        SparseArray<DisplayPowerControllerInterface> displayBrightnessFollowers;
+        SparseArray<DisplayPowerController> displayBrightnessFollowers;
         synchronized (mLock) {
             if (mStopped) {
                 return;
@@ -1546,7 +1536,7 @@
         float ambientLux = mAutomaticBrightnessController == null ? 0
                 : mAutomaticBrightnessController.getAmbientLux();
         for (int i = 0; i < displayBrightnessFollowers.size(); i++) {
-            DisplayPowerControllerInterface follower = displayBrightnessFollowers.valueAt(i);
+            DisplayPowerController follower = displayBrightnessFollowers.valueAt(i);
             follower.setBrightnessToFollow(rawBrightnessState,
                     mDisplayBrightnessController.convertToNits(rawBrightnessState),
                     ambientLux, slowChange);
@@ -1588,7 +1578,7 @@
         // brightness sources (such as an app override) are not saved to the setting, but should be
         // reflected in HBM calculations.
         mBrightnessRangeController.onBrightnessChanged(brightnessState, unthrottledBrightnessState,
-                mBrightnessClamperController.getBrightnessMaxReason());
+                clampedState.getBrightnessMaxReason());
 
         // Animate the screen brightness when the screen is on or dozing.
         // Skip the animation when the screen is off or suspended.
@@ -1804,7 +1794,7 @@
 
             if (userSetBrightnessChanged
                     || newEvent.getReason().getReason() != BrightnessReason.REASON_TEMPORARY) {
-                logBrightnessEvent(newEvent, unthrottledBrightnessState);
+                logBrightnessEvent(newEvent, unthrottledBrightnessState, clampedState);
             }
             if (mBrightnessEventRingBuffer != null) {
                 mBrightnessEventRingBuffer.append(newEvent);
@@ -1903,7 +1893,6 @@
         }
     }
 
-    @Override
     public void updateBrightness() {
         sendUpdatePowerState();
     }
@@ -1912,12 +1901,10 @@
      * Ignores the proximity sensor until the sensor state changes, but only if the sensor is
      * currently enabled and forcing the screen to be dark.
      */
-    @Override
     public void ignoreProximitySensorUntilChanged() {
         mDisplayPowerProximityStateController.ignoreProximitySensorUntilChanged();
     }
 
-    @Override
     public void setBrightnessConfiguration(BrightnessConfiguration c,
             boolean shouldResetShortTermModel) {
         Message msg = mHandler.obtainMessage(MSG_CONFIGURE_BRIGHTNESS,
@@ -1925,28 +1912,24 @@
         msg.sendToTarget();
     }
 
-    @Override
     public void setTemporaryBrightness(float brightness) {
         Message msg = mHandler.obtainMessage(MSG_SET_TEMPORARY_BRIGHTNESS,
                 Float.floatToIntBits(brightness), 0 /*unused*/);
         msg.sendToTarget();
     }
 
-    @Override
     public void setTemporaryAutoBrightnessAdjustment(float adjustment) {
         Message msg = mHandler.obtainMessage(MSG_SET_TEMPORARY_AUTO_BRIGHTNESS_ADJUSTMENT,
                 Float.floatToIntBits(adjustment), 0 /*unused*/);
         msg.sendToTarget();
     }
 
-    @Override
     public void setBrightnessFromOffload(float brightness) {
         Message msg = mHandler.obtainMessage(MSG_SET_BRIGHTNESS_FROM_OFFLOAD,
                 Float.floatToIntBits(brightness), 0 /*unused*/);
         mHandler.sendMessageAtTime(msg, mClock.uptimeMillis());
     }
 
-    @Override
     public float[] getAutoBrightnessLevels(
             @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
         int preset = Settings.System.getIntForUser(mContext.getContentResolver(),
@@ -1955,7 +1938,6 @@
         return mDisplayDeviceConfig.getAutoBrightnessBrighteningLevels(mode, preset);
     }
 
-    @Override
     public float[] getAutoBrightnessLuxLevels(
             @AutomaticBrightnessController.AutomaticBrightnessMode int mode) {
         int preset = Settings.System.getIntForUser(mContext.getContentResolver(),
@@ -1964,7 +1946,6 @@
         return mDisplayDeviceConfig.getAutoBrightnessBrighteningLevelsLux(mode, preset);
     }
 
-    @Override
     public BrightnessInfo getBrightnessInfo() {
         synchronized (mCachedBrightnessInfo) {
             return new BrightnessInfo(
@@ -1978,7 +1959,6 @@
         }
     }
 
-    @Override
     public void onBootCompleted() {
         Message msg = mHandler.obtainMessage(MSG_BOOT_COMPLETED);
         mHandler.sendMessageAtTime(msg, mClock.uptimeMillis());
@@ -1997,6 +1977,9 @@
         synchronized (mCachedBrightnessInfo) {
             float stateMax = state != null ? state.getMaxBrightness() : PowerManager.BRIGHTNESS_MAX;
             float stateMin = state != null ? state.getMinBrightness() : PowerManager.BRIGHTNESS_MAX;
+            @BrightnessInfo.BrightnessMaxReason int maxReason =
+                    state != null ? state.getBrightnessMaxReason()
+                            : BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
             final float minBrightness = Math.max(stateMin, Math.min(
                     mBrightnessRangeController.getCurrentBrightnessMin(), stateMax));
             final float maxBrightness = Math.min(
@@ -2023,7 +2006,7 @@
                             mBrightnessRangeController.getTransitionPoint());
             changed |=
                     mCachedBrightnessInfo.checkAndSetInt(mCachedBrightnessInfo.brightnessMaxReason,
-                            mBrightnessClamperController.getBrightnessMaxReason());
+                            maxReason);
             return changed;
         }
     }
@@ -2491,18 +2474,14 @@
         }
     }
 
-
-    @Override
     public float getScreenBrightnessSetting() {
         return mDisplayBrightnessController.getScreenBrightnessSetting();
     }
 
-    @Override
     public float getDozeBrightnessForOffload() {
         return mDisplayBrightnessController.getCurrentBrightness() * mDozeScaleFactor;
     }
 
-    @Override
     public void setBrightness(float brightness) {
         // After HBMController and NBMController migration to Clampers framework
         // currentBrightnessMax should be taken from clampers controller
@@ -2511,7 +2490,6 @@
                 mBrightnessRangeController.getCurrentBrightnessMax());
     }
 
-    @Override
     public void setBrightness(float brightness, int userSerial) {
         // After HBMController and NBMController migration to Clampers framework
         // currentBrightnessMax should be taken from clampers controller
@@ -2520,17 +2498,14 @@
                 mBrightnessRangeController.getCurrentBrightnessMax());
     }
 
-    @Override
     public int getDisplayId() {
         return mDisplayId;
     }
 
-    @Override
     public int getLeadDisplayId() {
         return mLeadDisplayId;
     }
 
-    @Override
     public void setBrightnessToFollow(float leadDisplayBrightness, float nits, float ambientLux,
             boolean slowChange) {
         mBrightnessRangeController.onAmbientLuxChange(ambientLux);
@@ -2591,16 +2566,14 @@
                 mAutomaticBrightnessController.getLastSensorTimestamps());
     }
 
-    @Override
-    public void addDisplayBrightnessFollower(DisplayPowerControllerInterface follower) {
+    public void addDisplayBrightnessFollower(DisplayPowerController follower) {
         synchronized (mLock) {
             mDisplayBrightnessFollowers.append(follower.getDisplayId(), follower);
             sendUpdatePowerStateLocked();
         }
     }
 
-    @Override
-    public void removeDisplayBrightnessFollower(DisplayPowerControllerInterface follower) {
+    public void removeDisplayBrightnessFollower(DisplayPowerController follower) {
         synchronized (mLock) {
             mDisplayBrightnessFollowers.remove(follower.getDisplayId());
             mHandler.postAtTime(() -> follower.setBrightnessToFollow(
@@ -2612,7 +2585,7 @@
     @GuardedBy("mLock")
     private void clearDisplayBrightnessFollowersLocked() {
         for (int i = 0; i < mDisplayBrightnessFollowers.size(); i++) {
-            DisplayPowerControllerInterface follower = mDisplayBrightnessFollowers.valueAt(i);
+            DisplayPowerController follower = mDisplayBrightnessFollowers.valueAt(i);
             mHandler.postAtTime(() -> follower.setBrightnessToFollow(
                     PowerManager.BRIGHTNESS_INVALID_FLOAT, BrightnessMappingStrategy.INVALID_NITS,
                     /* ambientLux= */ 0, /* slowChange= */ false), mClock.uptimeMillis());
@@ -2620,11 +2593,12 @@
         mDisplayBrightnessFollowers.clear();
     }
 
-    @Override
     public void dump(final PrintWriter pw) {
         synchronized (mLock) {
             pw.println();
             pw.println("Display Power Controller:");
+            pw.println("-------------------------");
+
             pw.println("  mDisplayId=" + mDisplayId);
             pw.println("  mLeadDisplayId=" + mLeadDisplayId);
             pw.println("  mLightSensor=" + mLightSensor);
@@ -2699,25 +2673,39 @@
                     + mColorFadeOffAnimator.isStarted());
         }
 
+        pw.println();
         if (mPowerState != null) {
             mPowerState.dump(pw);
         }
 
-        if (mAutomaticBrightnessController != null) {
-            mAutomaticBrightnessController.dump(pw);
-            dumpBrightnessEvents(pw);
+        pw.println();
+        if (mDisplayBrightnessController != null) {
+            mDisplayBrightnessController.dump(pw);
         }
 
-        dumpRbcEvents(pw);
-
-        if (mScreenOffBrightnessSensorController != null) {
-            mScreenOffBrightnessSensorController.dump(pw);
+        pw.println();
+        if (mBrightnessClamperController != null) {
+            mBrightnessClamperController.dump(ipw);
         }
 
+        pw.println();
         if (mBrightnessRangeController != null) {
             mBrightnessRangeController.dump(pw);
         }
 
+        pw.println();
+        if (mAutomaticBrightnessController != null) {
+            mAutomaticBrightnessController.dump(pw);
+            dumpBrightnessEvents(pw);
+        }
+        dumpRbcEvents(pw);
+
+        pw.println();
+        if (mScreenOffBrightnessSensorController != null) {
+            mScreenOffBrightnessSensorController.dump(pw);
+        }
+
+        pw.println();
         if (mBrightnessThrottler != null) {
             mBrightnessThrottler.dump(pw);
         }
@@ -2729,24 +2717,13 @@
         }
 
         pw.println();
-
         if (mWakelockController != null) {
             mWakelockController.dumpLocal(pw);
         }
 
         pw.println();
-        if (mDisplayBrightnessController != null) {
-            mDisplayBrightnessController.dump(pw);
-        }
-
-        pw.println();
         if (mDisplayStateController != null) {
-            mDisplayStateController.dumpsys(pw);
-        }
-
-        pw.println();
-        if (mBrightnessClamperController != null) {
-            mBrightnessClamperController.dump(ipw);
+            mDisplayStateController.dump(pw);
         }
     }
 
@@ -2926,7 +2903,8 @@
         return FrameworkStatsLog.DISPLAY_BRIGHTNESS_CHANGED__ENTIRE_REASON__REASON_UNKNOWN;
     }
 
-    private void logBrightnessEvent(BrightnessEvent event, float unmodifiedBrightness) {
+    private void logBrightnessEvent(BrightnessEvent event, float unmodifiedBrightness,
+            DisplayBrightnessState brightnessState) {
         int modifier = event.getReason().getModifier();
         int flags = event.getFlags();
         // It's easier to check if the brightness is at maximum level using the brightness
@@ -2963,7 +2941,7 @@
                     event.getHbmMode() == BrightnessInfo.HIGH_BRIGHTNESS_MODE_SUNLIGHT,
                     event.getHbmMode() == BrightnessInfo.HIGH_BRIGHTNESS_MODE_HDR,
                     (modifier & BrightnessReason.MODIFIER_LOW_POWER) > 0,
-                    mBrightnessClamperController.getBrightnessMaxReason(),
+                    brightnessState.getBrightnessMaxReason(),
                     // TODO: (flc) add brightnessMinReason here too.
                     (modifier & BrightnessReason.MODIFIER_DIMMED) > 0,
                     event.isRbcEnabled(),
@@ -3151,19 +3129,17 @@
         }
     }
 
-    @Override
     public void setAutoBrightnessLoggingEnabled(boolean enabled) {
         if (mAutomaticBrightnessController != null) {
             mAutomaticBrightnessController.setLoggingEnabled(enabled);
         }
     }
 
-    @Override // DisplayWhiteBalanceController.Callbacks
+    // DisplayWhiteBalanceController.Callbacks
     public void updateWhiteBalance() {
         sendUpdatePowerState();
     }
 
-    @Override
     public void setDisplayWhiteBalanceLoggingEnabled(boolean enabled) {
         Message msg = mHandler.obtainMessage();
         msg.what = MSG_SET_DWBC_LOGGING_ENABLED;
@@ -3171,7 +3147,6 @@
         msg.sendToTarget();
     }
 
-    @Override
     public void setAmbientColorTemperatureOverride(float cct) {
         Message msg = mHandler.obtainMessage();
         msg.what = MSG_SET_DWBC_COLOR_OVERRIDE;
@@ -3296,10 +3271,10 @@
         BrightnessClamperController getBrightnessClamperController(Handler handler,
                 BrightnessClamperController.ClamperChangeListener clamperChangeListener,
                 BrightnessClamperController.DisplayDeviceData data, Context context,
-                DisplayManagerFlags flags, SensorManager sensorManager) {
+                DisplayManagerFlags flags, SensorManager sensorManager, float currentBrightness) {
 
             return new BrightnessClamperController(handler, clamperChangeListener, data, context,
-                    flags, sensorManager);
+                    flags, sensorManager, currentBrightness);
         }
 
         DisplayWhiteBalanceController getDisplayWhiteBalanceController(Handler handler,
diff --git a/services/core/java/com/android/server/display/DisplayPowerControllerInterface.java b/services/core/java/com/android/server/display/DisplayPowerControllerInterface.java
deleted file mode 100644
index d28578a..0000000
--- a/services/core/java/com/android/server/display/DisplayPowerControllerInterface.java
+++ /dev/null
@@ -1,267 +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.display;
-
-import android.content.pm.ParceledListSlice;
-import android.hardware.display.AmbientBrightnessDayStats;
-import android.hardware.display.BrightnessChangeEvent;
-import android.hardware.display.BrightnessConfiguration;
-import android.hardware.display.BrightnessInfo;
-import android.hardware.display.DisplayManagerInternal;
-import android.os.PowerManager;
-import android.view.Display;
-
-import java.io.PrintWriter;
-
-/**
- * An interface to manage the display's power state and brightness
- */
-public interface DisplayPowerControllerInterface {
-    /**
-     * Notified when the display is changed.
-     *
-     * We use this to apply any changes that might be needed when displays get swapped on foldable
-     * devices, when layouts change, etc.
-     *
-     * Must be called while holding the SyncRoot lock.
-     *
-     * @param hbmInfo The high brightness mode metadata, like
-     *                remaining time and hbm events, for the corresponding
-     *                physical display, to make sure we stay within the safety margins.
-     * @param leadDisplayId The display who is considered our "leader" for things like brightness.
-     */
-    void onDisplayChanged(HighBrightnessModeMetadata hbmInfo, int leadDisplayId);
-
-    /**
-     * Unregisters all listeners and interrupts all running threads; halting future work.
-     *
-     * This method should be called when the DisplayPowerController is no longer in use; i.e. when
-     * the display has been removed.
-     */
-    void stop();
-
-    /**
-     * Used to update the display's BrightnessConfiguration
-     * @param config The new BrightnessConfiguration
-     */
-    void setBrightnessConfiguration(BrightnessConfiguration config,
-            boolean shouldResetShortTermModel);
-
-    /**
-     * Used to set the ambient color temperature of the Display
-     * @param ambientColorTemperature The target ambientColorTemperature
-     */
-    void setAmbientColorTemperatureOverride(float ambientColorTemperature);
-
-    /**
-     * Used to decide the associated AutomaticBrightnessController's BrightnessMode
-     * @param mode The auto-brightness mode
-     */
-    void setAutomaticScreenBrightnessMode(
-            @AutomaticBrightnessController.AutomaticBrightnessMode int mode);
-
-    /**
-     * Used to enable/disable the logging of the WhileBalance associated entities
-     * @param enabled Flag which represents if the logging is the be enabled
-     */
-    void setDisplayWhiteBalanceLoggingEnabled(boolean enabled);
-
-    /**
-     * Used to dump the state.
-     * @param writer The PrintWriter used to dump the state.
-     */
-    void dump(PrintWriter writer);
-
-    /**
-     * Used to get the ambient brightness stats
-     */
-    ParceledListSlice<AmbientBrightnessDayStats> getAmbientBrightnessStats(int userId);
-
-    /**
-     * Get the default brightness configuration
-     */
-    BrightnessConfiguration getDefaultBrightnessConfiguration();
-
-    /**
-     * Set the screen brightness of the associated display
-     * @param brightness The value to which the brightness is to be set
-     */
-    void setBrightness(float brightness);
-
-    /**
-     * Set the screen brightness of the associated display
-     * @param brightness The value to which the brightness is to be set
-     * @param userSerial The user for which the brightness value is to be set.
-     */
-    void setBrightness(float brightness, int userSerial);
-
-    /**
-     * Checks if the proximity sensor is available
-     */
-    boolean isProximitySensorAvailable();
-
-    /**
-     * Persist the brightness slider events and ambient brightness stats to disk.
-     */
-    void persistBrightnessTrackerState();
-
-    /**
-     * Ignores the proximity sensor until the sensor state changes, but only if the sensor is
-     * currently enabled and forcing the screen to be dark.
-     */
-    void ignoreProximitySensorUntilChanged();
-
-    /**
-     * Requests a new power state.
-     *
-     * @param request The requested power state.
-     * @param waitForNegativeProximity If true, issues a request to wait for
-     * negative proximity before turning the screen back on,
-     * assuming the screen was turned off by the proximity sensor.
-     * @return True if display is ready, false if there are important changes that must
-     * be made asynchronously.
-     */
-    boolean requestPowerState(DisplayManagerInternal.DisplayPowerRequest request,
-            boolean waitForNegativeProximity);
-
-    /**
-     * Overrides the current doze screen state.
-     *
-     * @param displayState the new doze display state.
-     * @param reason the reason behind the new doze display state.
-     */
-    void overrideDozeScreenState(int displayState, @Display.StateReason int reason);
-
-    void setDisplayOffloadSession(DisplayManagerInternal.DisplayOffloadSession session);
-
-    /**
-     * Sets up the temporary autobrightness adjustment when the user is yet to settle down to a
-     * value.
-     */
-    void setTemporaryAutoBrightnessAdjustment(float adjustment);
-
-    /**
-     * Sets temporary brightness from the offload chip until we get a brightness value from
-     * the light sensor.
-     * @param brightness The brightness value between {@link PowerManager.BRIGHTNESS_MIN} and
-     * {@link PowerManager.BRIGHTNESS_MAX}. Values outside of that range will be ignored.
-     */
-    void setBrightnessFromOffload(float brightness);
-
-    /**
-     * Gets the screen brightness setting
-     */
-    float getScreenBrightnessSetting();
-
-    /**
-     * Gets the brightness value used when the device is in doze
-     */
-    float getDozeBrightnessForOffload();
-
-    /**
-     * Sets up the temporary brightness for the associated display
-     */
-    void setTemporaryBrightness(float brightness);
-
-    /**
-     * Gets the associated {@link BrightnessInfo}
-     */
-    BrightnessInfo getBrightnessInfo();
-
-    /**
-     * Get the {@link BrightnessChangeEvent}s for the specified user.
-     */
-    ParceledListSlice<BrightnessChangeEvent> getBrightnessEvents(int userId, boolean hasUsageStats);
-
-    /**
-     * Sets up the logging for the associated {@link AutomaticBrightnessController}
-     * @param enabled Flag to represent if the logging is to be enabled
-     */
-    void setAutoBrightnessLoggingEnabled(boolean enabled);
-
-    /**
-     * Handles the changes to be done to update the brightness when the user is changed
-     * @param newUserId The new userId
-     * @param userSerial The serial number of the new user
-     * @param newBrightness The brightness for the new user
-     */
-    void onSwitchUser(int newUserId, int userSerial, float newBrightness);
-
-    /**
-     * Get the ID of the display associated with this DPC.
-     * @return The display ID
-     */
-    int getDisplayId();
-
-    /**
-     * Get the ID of the display that is the leader of this DPC.
-     *
-     * Note that this is different than the display associated with the DPC. The leader is another
-     * display which we follow for things like brightness.
-     *
-     * Must be called while holding the SyncRoot lock.
-     */
-    int getLeadDisplayId();
-
-    /**
-     * Set the brightness to follow if this is an additional display in a set of concurrent
-     * displays.
-     * @param leadDisplayBrightness The brightness of the lead display in the set of concurrent
-     *                              displays
-     * @param nits The brightness value in nits if the device supports nits. Set to a negative
-     *             number otherwise.
-     * @param ambientLux The lux value that will be passed to {@link HighBrightnessModeController}
-     * @param slowChange Indicates whether we should slowly animate to the given brightness value.
-     */
-    void setBrightnessToFollow(float leadDisplayBrightness, float nits, float ambientLux,
-            boolean slowChange);
-
-    /**
-     * Add an additional display that will copy the brightness value from this display. This is used
-     * when the device is in concurrent displays mode.
-     * @param follower The DPC that should copy the brightness value from this DPC
-     */
-    void addDisplayBrightnessFollower(DisplayPowerControllerInterface follower);
-
-    /**
-     * Removes the given display from the list of brightness followers.
-     * @param follower The DPC to remove from the followers list
-     */
-    void removeDisplayBrightnessFollower(DisplayPowerControllerInterface follower);
-
-    /**
-     * Indicate that boot has been completed and the screen is ready to update.
-     */
-    void onBootCompleted();
-
-    /**
-     * Get the brightness levels used to determine automatic brightness based on lux levels.
-     * @param mode The auto-brightness mode
-     * @return The brightness levels for the specified mode. The values are between
-     * {@link PowerManager.BRIGHTNESS_MIN} and {@link PowerManager.BRIGHTNESS_MAX}.
-     */
-    float[] getAutoBrightnessLevels(
-            @AutomaticBrightnessController.AutomaticBrightnessMode int mode);
-
-    /**
-     * Get the lux levels used to determine automatic brightness.
-     * @param mode The auto-brightness mode
-     * @return The lux levels for the specified mode
-     */
-    float[] getAutoBrightnessLuxLevels(
-            @AutomaticBrightnessController.AutomaticBrightnessMode int mode);
-}
diff --git a/services/core/java/com/android/server/display/DisplayPowerProximityStateController.java b/services/core/java/com/android/server/display/DisplayPowerProximityStateController.java
index 882c02f..215932c 100644
--- a/services/core/java/com/android/server/display/DisplayPowerProximityStateController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerProximityStateController.java
@@ -315,6 +315,7 @@
     public void dumpLocal(PrintWriter pw) {
         pw.println();
         pw.println("DisplayPowerProximityStateController:");
+        pw.println("-------------------------------------");
         synchronized (mLock) {
             pw.println("  mPendingWaitForNegativeProximityLocked="
                     + mPendingWaitForNegativeProximityLocked);
diff --git a/services/core/java/com/android/server/display/DisplayPowerState.java b/services/core/java/com/android/server/display/DisplayPowerState.java
index e5efebc..2fbb114 100644
--- a/services/core/java/com/android/server/display/DisplayPowerState.java
+++ b/services/core/java/com/android/server/display/DisplayPowerState.java
@@ -344,7 +344,6 @@
     }
 
     public void dump(PrintWriter pw) {
-        pw.println();
         pw.println("Display Power State:");
         pw.println("  mStopped=" + mStopped);
         pw.println("  mScreenState=" + Display.stateToString(mScreenState));
diff --git a/services/core/java/com/android/server/display/LocalDisplayAdapter.java b/services/core/java/com/android/server/display/LocalDisplayAdapter.java
index 6e0b9cf..0570b2a 100644
--- a/services/core/java/com/android/server/display/LocalDisplayAdapter.java
+++ b/services/core/java/com/android/server/display/LocalDisplayAdapter.java
@@ -1286,7 +1286,10 @@
                 pw.println("  " + mSupportedModes.valueAt(i));
             }
             pw.println("mSupportedColorModes=" + mSupportedColorModes);
-            pw.println("mDisplayDeviceConfig=" + mDisplayDeviceConfig);
+            pw.println("");
+            pw.println("DisplayDeviceConfig: ");
+            pw.println("---------------------");
+            pw.println(mDisplayDeviceConfig);
         }
 
         private int findSfDisplayModeIdLocked(int displayModeId, int modeGroup) {
diff --git a/services/core/java/com/android/server/display/LogicalDisplay.java b/services/core/java/com/android/server/display/LogicalDisplay.java
index 5d55d190..e8be8a4 100644
--- a/services/core/java/com/android/server/display/LogicalDisplay.java
+++ b/services/core/java/com/android/server/display/LogicalDisplay.java
@@ -1040,6 +1040,9 @@
 
     public void dumpLocked(PrintWriter pw) {
         pw.println("mDisplayId=" + mDisplayId);
+        pw.println("mPrimaryDisplayDevice=" + (mPrimaryDisplayDevice != null
+                ? mPrimaryDisplayDevice.getNameLocked() + "(" + mPrimaryDisplayDevice.getUniqueId()
+                        + ")" : "null"));
         pw.println("mIsEnabled=" + mIsEnabled);
         pw.println("mIsInTransition=" + mIsInTransition);
         pw.println("mLayerStack=" + mLayerStack);
@@ -1049,8 +1052,6 @@
         pw.println("mRequestedColorMode=" + mRequestedColorMode);
         pw.println("mDisplayOffset=(" + mDisplayOffsetX + ", " + mDisplayOffsetY + ")");
         pw.println("mDisplayScalingDisabled=" + mDisplayScalingDisabled);
-        pw.println("mPrimaryDisplayDevice=" + (mPrimaryDisplayDevice != null ?
-                mPrimaryDisplayDevice.getNameLocked() : "null"));
         pw.println("mBaseDisplayInfo=" + mBaseDisplayInfo);
         pw.println("mOverrideDisplayInfo=" + mOverrideDisplayInfo);
         pw.println("mRequestedMinimalPostProcessing=" + mRequestedMinimalPostProcessing);
diff --git a/services/core/java/com/android/server/display/LogicalDisplayMapper.java b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
index e9ecfc6..c3f6a52 100644
--- a/services/core/java/com/android/server/display/LogicalDisplayMapper.java
+++ b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
@@ -427,6 +427,7 @@
 
     public void dumpLocked(PrintWriter pw) {
         pw.println("LogicalDisplayMapper:");
+        pw.println("---------------------");
         IndentingPrintWriter ipw = new IndentingPrintWriter(pw, "  ");
         ipw.increaseIndent();
 
@@ -477,14 +478,14 @@
             // The boot animation might still be in progress, we do not want to switch states now
             // as the boot animation would end up with an incorrect size.
             if (DEBUG) {
-                Slog.d(TAG, "Postponing transition to state: " + mPendingDeviceState
-                        + " until boot is completed");
+                Slog.d(TAG, "Postponing transition to state: "
+                        + mPendingDeviceState.getIdentifier() + " until boot is completed");
             }
             mDeviceStateToBeAppliedAfterBoot = state;
             return;
         }
 
-        Slog.i(TAG, "Requesting Transition to state: " + state + ", from state="
+        Slog.i(TAG, "Requesting Transition to state: " + state.getIdentifier() + ", from state="
                 + mDeviceState.getIdentifier() + ", interactive=" + mInteractive
                 + ", mBootCompleted=" + mBootCompleted);
         // As part of a state transition, we may need to turn off some displays temporarily so that
diff --git a/services/core/java/com/android/server/display/PersistentDataStore.java b/services/core/java/com/android/server/display/PersistentDataStore.java
index 2d7792d..9cdc918 100644
--- a/services/core/java/com/android/server/display/PersistentDataStore.java
+++ b/services/core/java/com/android/server/display/PersistentDataStore.java
@@ -628,7 +628,9 @@
     }
 
     public void dump(PrintWriter pw) {
-        pw.println("PersistentDataStore");
+        pw.println("PersistentDataStore:");
+        pw.println("--------------------");
+
         pw.println("  mLoaded=" + mLoaded);
         pw.println("  mDirty=" + mDirty);
         pw.println("  RememberedWifiDisplays:");
diff --git a/services/core/java/com/android/server/display/ScreenOffBrightnessSensorController.java b/services/core/java/com/android/server/display/ScreenOffBrightnessSensorController.java
index 0a884c9..b63eba3 100644
--- a/services/core/java/com/android/server/display/ScreenOffBrightnessSensorController.java
+++ b/services/core/java/com/android/server/display/ScreenOffBrightnessSensorController.java
@@ -121,6 +121,7 @@
     /** Dump current state */
     public void dump(PrintWriter pw) {
         pw.println("Screen Off Brightness Sensor Controller:");
+        pw.println("----------------------------------------");
         IndentingPrintWriter idpw = new IndentingPrintWriter(pw);
         idpw.increaseIndent();
         idpw.println("registered=" + mRegistered);
diff --git a/services/core/java/com/android/server/display/SmallAreaDetectionController.java b/services/core/java/com/android/server/display/SmallAreaDetectionController.java
index bf384b0..3ed7e57 100644
--- a/services/core/java/com/android/server/display/SmallAreaDetectionController.java
+++ b/services/core/java/com/android/server/display/SmallAreaDetectionController.java
@@ -133,7 +133,8 @@
     }
 
     void dump(PrintWriter pw) {
-        pw.println("Small area detection allowlist");
+        pw.println("Small area detection allowlist:");
+        pw.println("-------------------------------");
         pw.println("  Packages:");
         synchronized (mLock) {
             for (String pkg : mAllowPkgMap.keySet()) {
diff --git a/services/core/java/com/android/server/display/WakelockController.java b/services/core/java/com/android/server/display/WakelockController.java
index 5b0229c..3520723 100644
--- a/services/core/java/com/android/server/display/WakelockController.java
+++ b/services/core/java/com/android/server/display/WakelockController.java
@@ -417,6 +417,7 @@
      */
     public void dumpLocal(PrintWriter pw) {
         pw.println("WakelockController State:");
+        pw.println("-------------------------");
         pw.println("  mDisplayId=" + mDisplayId);
         pw.println("  mUnfinishedBusiness=" + hasUnfinishedBusiness());
         pw.println("  mOnStateChangePending=" + isOnStateChangedPending());
diff --git a/services/core/java/com/android/server/display/brightness/DisplayBrightnessStrategySelector.java b/services/core/java/com/android/server/display/brightness/DisplayBrightnessStrategySelector.java
index 1b49bbc..06890e7 100644
--- a/services/core/java/com/android/server/display/brightness/DisplayBrightnessStrategySelector.java
+++ b/services/core/java/com/android/server/display/brightness/DisplayBrightnessStrategySelector.java
@@ -262,6 +262,7 @@
     public void dump(PrintWriter writer) {
         writer.println();
         writer.println("DisplayBrightnessStrategySelector:");
+        writer.println("----------------------------------");
         writer.println("  mDisplayId= " + mDisplayId);
         writer.println("  mOldBrightnessStrategyName= " + mOldBrightnessStrategyName);
         writer.println(
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
index 59fffe7..d3be33f 100644
--- a/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
+++ b/services/core/java/com/android/server/display/brightness/clamper/BrightnessClamperController.java
@@ -75,10 +75,13 @@
     private ModifiersAggregatedState mModifiersAggregatedState = new ModifiersAggregatedState();
 
     private final DeviceConfig.OnPropertiesChangedListener mOnPropertiesChangedListener;
+    private final DisplayManagerFlags mDisplayManagerFlags;
     private float mBrightnessCap = PowerManager.BRIGHTNESS_MAX;
 
     private float mCustomAnimationRate = DisplayBrightnessState.CUSTOM_ANIMATION_RATE_NOT_SET;
     @Nullable
+    private BrightnessPowerClamper mPowerClamper;
+    @Nullable
     private Type mClamperType = null;
 
     private boolean mClamperApplied = false;
@@ -93,16 +96,18 @@
 
     public BrightnessClamperController(Handler handler,
             ClamperChangeListener clamperChangeListener, DisplayDeviceData data, Context context,
-            DisplayManagerFlags flags, SensorManager sensorManager) {
-        this(new Injector(), handler, clamperChangeListener, data, context, flags, sensorManager);
+            DisplayManagerFlags flags, SensorManager sensorManager, float currentBrightness) {
+        this(new Injector(), handler, clamperChangeListener, data, context, flags, sensorManager,
+                currentBrightness);
     }
 
     @VisibleForTesting
     BrightnessClamperController(Injector injector, Handler handler,
             ClamperChangeListener clamperChangeListener, DisplayDeviceData data, Context context,
-            DisplayManagerFlags flags, SensorManager sensorManager) {
+            DisplayManagerFlags flags, SensorManager sensorManager, float currentBrightness) {
         mDeviceConfigParameterProvider = injector.getDeviceConfigParameterProvider();
         mHandler = handler;
+        mDisplayManagerFlags = flags;
         mLightSensorController = injector.getLightSensorController(sensorManager, context,
                 mLightSensorListener, mHandler);
 
@@ -117,7 +122,15 @@
         };
 
         mClampers = injector.getClampers(handler, clamperChangeListenerInternal, data, flags,
-                context);
+                context, currentBrightness);
+        if (mDisplayManagerFlags.isPowerThrottlingClamperEnabled()) {
+            for (BrightnessClamper clamper: mClampers) {
+                if (clamper.getType() == Type.POWER) {
+                    mPowerClamper = (BrightnessPowerClamper) clamper;
+                    break;
+                }
+            }
+        }
         mModifiers = injector.getModifiers(flags, context, handler, clamperChangeListener,
                 data);
 
@@ -161,6 +174,7 @@
         builder.setBrightness(cappedBrightness);
         builder.setMaxBrightness(mBrightnessCap);
         builder.setCustomAnimationRate(mCustomAnimationRate);
+        builder.setBrightnessMaxReason(getBrightnessMaxReason());
 
         if (mClamperType != null) {
             builder.getBrightnessReason().addModifier(BrightnessReason.MODIFIER_THROTTLED);
@@ -182,22 +196,17 @@
             mModifiers.get(i).apply(request, builder);
         }
 
+        if (mDisplayManagerFlags.isPowerThrottlingClamperEnabled()) {
+            if (mPowerClamper != null) {
+                mPowerClamper.updateCurrentBrightness(cappedBrightness);
+            }
+        }
+
         return builder.build();
     }
 
-    /**
-     * See BrightnessThrottler.getBrightnessMaxReason:
-     * used in:
-     * 1) DPC2.CachedBrightnessInfo to determine changes
-     * 2) DPC2.logBrightnessEvent
-     * 3) HBMController - for logging
-     * Method is called in mHandler thread (DisplayControllerHandler), in the same thread
-     * recalculateBrightnessCap and DPC2.updatePowerStateInternal are called.
-     * Should be moved to DisplayBrightnessState OR derived from DisplayBrightnessState
-     * TODO: b/263362199
-     */
     @BrightnessInfo.BrightnessMaxReason
-    public int getBrightnessMaxReason() {
+    private int getBrightnessMaxReason() {
         if (mClamperType == null) {
             return BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE;
         } else if (mClamperType == Type.THERMAL) {
@@ -224,6 +233,7 @@
      */
     public void dump(PrintWriter writer) {
         writer.println("BrightnessClamperController:");
+        writer.println("----------------------------");
         writer.println("  mBrightnessCap: " + mBrightnessCap);
         writer.println("  mClamperType: " + mClamperType);
         writer.println("  mClamperApplied: " + mClamperApplied);
@@ -321,13 +331,17 @@
 
         List<BrightnessClamper<? super DisplayDeviceData>> getClampers(Handler handler,
                 ClamperChangeListener clamperChangeListener, DisplayDeviceData data,
-                DisplayManagerFlags flags, Context context) {
+                DisplayManagerFlags flags, Context context, float currentBrightness) {
             List<BrightnessClamper<? super DisplayDeviceData>> clampers = new ArrayList<>();
             clampers.add(
                     new BrightnessThermalClamper(handler, clamperChangeListener, data));
             if (flags.isPowerThrottlingClamperEnabled()) {
-                clampers.add(new BrightnessPowerClamper(handler, clamperChangeListener,
-                        data));
+                // Check if power-throttling config is present.
+                PowerThrottlingConfigData configData = data.getPowerThrottlingConfigData();
+                if (configData != null) {
+                    clampers.add(new BrightnessPowerClamper(handler, clamperChangeListener,
+                            data, currentBrightness));
+                }
             }
             if (flags.isBrightnessWearBedtimeModeClamperEnabled()) {
                 clampers.add(new BrightnessWearBedtimeModeClamper(handler, context,
diff --git a/services/core/java/com/android/server/display/brightness/clamper/BrightnessPowerClamper.java b/services/core/java/com/android/server/display/brightness/clamper/BrightnessPowerClamper.java
index 790322d..85e81f9 100644
--- a/services/core/java/com/android/server/display/brightness/clamper/BrightnessPowerClamper.java
+++ b/services/core/java/com/android/server/display/brightness/clamper/BrightnessPowerClamper.java
@@ -21,16 +21,23 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
+import android.content.Context;
 import android.os.Handler;
+import android.os.IThermalEventListener;
+import android.os.IThermalService;
 import android.os.PowerManager;
+import android.os.RemoteException;
+import android.os.ServiceManager;
 import android.os.Temperature;
 import android.provider.DeviceConfigInterface;
 import android.util.Slog;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.DisplayDeviceConfig.PowerThrottlingConfigData;
 import com.android.server.display.DisplayDeviceConfig.PowerThrottlingData;
 import com.android.server.display.DisplayDeviceConfig.PowerThrottlingData.ThrottlingLevel;
+import com.android.server.display.brightness.BrightnessUtils;
 import com.android.server.display.feature.DeviceConfigParameterProvider;
 import com.android.server.display.utils.DeviceConfigParsingUtils;
 
@@ -65,14 +72,21 @@
     private PowerThrottlingData mPowerThrottlingDataActive = null;
     @Nullable
     private PowerThrottlingConfigData mPowerThrottlingConfigData = null;
-
+    @NonNull
+    private final ThermalLevelListener mThermalLevelListener;
+    @NonNull
+    private final PowerChangeListener mPowerChangeListener;
     private @Temperature.ThrottlingStatus int mCurrentThermalLevel = Temperature.THROTTLING_NONE;
+    private boolean mCurrentThermalLevelChanged = false;
     private float mCurrentAvgPowerConsumed = 0;
     @Nullable
     private String mUniqueDisplayId = null;
     @Nullable
     private String mDataId = null;
-
+    private float mCurrentBrightness = PowerManager.BRIGHTNESS_INVALID;
+    private float mCustomAnimationRateSec = DisplayBrightnessState.CUSTOM_ANIMATION_RATE_NOT_SET;
+    private float mCustomAnimationRateSecDeviceConfig =
+                        DisplayBrightnessState.CUSTOM_ANIMATION_RATE_NOT_SET;
     private final BiFunction<String, String, ThrottlingLevel> mDataPointMapper = (key, value) -> {
         try {
             int status = DeviceConfigParsingUtils.parseThermalStatus(key);
@@ -88,23 +102,41 @@
 
 
     BrightnessPowerClamper(Handler handler, ClamperChangeListener listener,
-            PowerData powerData) {
-        this(new Injector(), handler, listener, powerData);
+            PowerData powerData, float currentBrightness) {
+        this(new Injector(), handler, listener, powerData, currentBrightness);
     }
 
     @VisibleForTesting
     BrightnessPowerClamper(Injector injector, Handler handler, ClamperChangeListener listener,
-            PowerData powerData) {
+                           PowerData powerData, float currentBrightness) {
         super(handler, listener);
         mInjector = injector;
-        mConfigParameterProvider = injector.getDeviceConfigParameterProvider();
+        mCurrentBrightness = currentBrightness;
+        mPowerChangeListener = (powerConsumed, thermalStatus) -> {
+            recalculatePowerQuotaChange(powerConsumed, thermalStatus);
+        };
+        mPowerThrottlingConfigData = powerData.getPowerThrottlingConfigData();
+        if (mPowerThrottlingConfigData != null) {
+            mCustomAnimationRateSecDeviceConfig = mPowerThrottlingConfigData.customAnimationRateSec;
+        }
+        mThermalLevelListener = new ThermalLevelListener(handler);
+        mPmicMonitor =
+            mInjector.getPmicMonitor(mPowerChangeListener,
+                    mThermalLevelListener.getThermalService(),
+                    mPowerThrottlingConfigData.pollingWindowMaxMillis,
+                    mPowerThrottlingConfigData.pollingWindowMinMillis);
 
+        mConfigParameterProvider = injector.getDeviceConfigParameterProvider();
         mHandler.post(() -> {
             setDisplayData(powerData);
             loadOverrideData();
             start();
         });
+    }
 
+    @VisibleForTesting
+    PowerChangeListener getPowerChangeListener() {
+        return mPowerChangeListener;
     }
 
     @Override
@@ -114,6 +146,11 @@
     }
 
     @Override
+    float getCustomAnimationRate() {
+        return mCustomAnimationRateSec;
+    }
+
+    @Override
     void onDeviceConfigChanged() {
         mHandler.post(() -> {
             loadOverrideData();
@@ -134,6 +171,9 @@
         if (mPmicMonitor != null) {
             mPmicMonitor.shutdown();
         }
+        if (mThermalLevelListener != null) {
+            mThermalLevelListener.stop();
+        }
     }
 
     /**
@@ -144,11 +184,20 @@
         pw.println("  mCurrentAvgPowerConsumed=" + mCurrentAvgPowerConsumed);
         pw.println("  mUniqueDisplayId=" + mUniqueDisplayId);
         pw.println("  mCurrentThermalLevel=" + mCurrentThermalLevel);
+        pw.println("  mCurrentThermalLevelChanged=" + mCurrentThermalLevelChanged);
         pw.println("  mPowerThrottlingDataFromDDC=" + (mPowerThrottlingDataFromDDC == null ? "null"
                 : mPowerThrottlingDataFromDDC.toString()));
+        mThermalLevelListener.dump(pw);
         super.dump(pw);
     }
 
+    /**
+     * Updates current brightness, for power calculations.
+     */
+    public void updateCurrentBrightness(float currentBrightness) {
+        mCurrentBrightness = currentBrightness;
+    }
+
     private void recalculateActiveData() {
         if (mUniqueDisplayId == null || mDataId == null) {
             return;
@@ -156,17 +205,11 @@
         mPowerThrottlingDataActive = mPowerThrottlingDataOverride
                 .getOrDefault(mUniqueDisplayId, Map.of()).getOrDefault(mDataId,
                         mPowerThrottlingDataFromDDC);
-        if (mPowerThrottlingDataActive != null) {
-            if (mPmicMonitor != null) {
-                mPmicMonitor.stop();
-                mPmicMonitor.start();
-            }
-        } else {
+        if (mPowerThrottlingDataActive == null) {
             if (mPmicMonitor != null) {
                 mPmicMonitor.stop();
             }
         }
-        recalculateBrightnessCap();
     }
 
     private void loadOverrideData() {
@@ -198,21 +241,57 @@
         if (mPowerThrottlingDataActive == null) {
             return;
         }
-        if (powerQuota > 0 && mCurrentAvgPowerConsumed > powerQuota) {
-            isActive = true;
-            // calculate new brightness Cap.
-            // Brightness has a linear relation to power-consumed.
-            targetBrightnessCap =
-                    (powerQuota / mCurrentAvgPowerConsumed) * PowerManager.BRIGHTNESS_MAX;
-            // Cap to lowest allowed brightness on device.
+        if (powerQuota > 0) {
+            if (BrightnessUtils.isValidBrightnessValue(mCurrentBrightness)
+                    && (mCurrentAvgPowerConsumed > powerQuota)) {
+                isActive = true;
+                // calculate new brightness Cap.
+                // Brightness has a linear relation to power-consumed.
+                targetBrightnessCap =
+                    (powerQuota / mCurrentAvgPowerConsumed) * mCurrentBrightness;
+            } else if (mCurrentThermalLevelChanged) {
+                if (mCurrentThermalLevel == Temperature.THROTTLING_NONE) {
+                    // reset pmic and remove the power-throttling cap.
+                    isActive = true;
+                    targetBrightnessCap = PowerManager.BRIGHTNESS_MAX;
+                    mPmicMonitor.stop();
+                } else {
+                    isActive = true;
+                    // Since the thermal status has changed, we need to remove power-throttling cap.
+                    // Instead of recalculating and changing brightness again, adding flicker,
+                    // we will wait for the next pmic cycle to re-evaluate this value
+                    // make act on it, if needed.
+                    targetBrightnessCap = PowerManager.BRIGHTNESS_MAX;
+                    if (mPmicMonitor.isStopped()) {
+                        mPmicMonitor.start();
+                    }
+                }
+            } else { // Current power consumed is under the quota.
+                isActive = true;
+                targetBrightnessCap = PowerManager.BRIGHTNESS_MAX;
+            }
+        }
+
+        // Cap to lowest allowed brightness on device.
+        if (mPowerThrottlingConfigData != null) {
             targetBrightnessCap = Math.max(targetBrightnessCap,
-                    mPowerThrottlingConfigData.brightnessLowestCapAllowed);
+                                mPowerThrottlingConfigData.brightnessLowestCapAllowed);
         }
 
         if (mBrightnessCap != targetBrightnessCap || mIsActive != isActive) {
             mIsActive = isActive;
+            Slog.i(TAG, "Power clamper changing current brightness cap mBrightnessCap: "
+                    + mBrightnessCap + " to target brightness cap:" + targetBrightnessCap
+                    + " for current screen brightness: " + mCurrentBrightness);
             mBrightnessCap = targetBrightnessCap;
+            Slog.i(TAG, "Power clamper changed state: thermalStatus:" + mCurrentThermalLevel
+                    + " mCurrentThermalLevelChanged:" + mCurrentThermalLevelChanged
+                    + " mCurrentAvgPowerConsumed:" + mCurrentAvgPowerConsumed
+                    + " mCustomAnimationRateSec:" + mCustomAnimationRateSecDeviceConfig);
+            mCustomAnimationRateSec = mCustomAnimationRateSecDeviceConfig;
             mChangeListener.onChanged();
+        } else {
+            mCustomAnimationRateSec = DisplayBrightnessState.CUSTOM_ANIMATION_RATE_NOT_SET;
         }
     }
 
@@ -234,6 +313,11 @@
 
     private void recalculatePowerQuotaChange(float avgPowerConsumed, int thermalStatus) {
         mHandler.post(() -> {
+            if (mCurrentThermalLevel != thermalStatus) {
+                mCurrentThermalLevelChanged = true;
+            } else {
+                mCurrentThermalLevelChanged = false;
+            }
             mCurrentThermalLevel = thermalStatus;
             mCurrentAvgPowerConsumed = avgPowerConsumed;
             recalculateBrightnessCap();
@@ -244,14 +328,107 @@
         if (mPowerThrottlingConfigData == null) {
             return;
         }
-        PowerChangeListener listener = (powerConsumed, thermalStatus) -> {
-            recalculatePowerQuotaChange(powerConsumed, thermalStatus);
-        };
-        mPmicMonitor =
-            mInjector.getPmicMonitor(listener, mPowerThrottlingConfigData.pollingWindowMillis);
+        if (mPowerThrottlingConfigData.pollingWindowMaxMillis
+                <= mPowerThrottlingConfigData.pollingWindowMinMillis) {
+            Slog.e(TAG, "Brightness power max polling window:"
+                    + mPowerThrottlingConfigData.pollingWindowMaxMillis
+                    + " msec, should be greater than brightness min polling window:"
+                    + mPowerThrottlingConfigData.pollingWindowMinMillis + " msec.");
+            return;
+        }
+        if ((mPowerThrottlingConfigData.pollingWindowMaxMillis
+                % mPowerThrottlingConfigData.pollingWindowMinMillis) != 0) {
+            Slog.e(TAG, "Brightness power max polling window:"
+                    + mPowerThrottlingConfigData.pollingWindowMaxMillis
+                    + " msec, is not divisible by brightness min polling window:"
+                    + mPowerThrottlingConfigData.pollingWindowMinMillis + " msec.");
+            return;
+        }
+        mCustomAnimationRateSecDeviceConfig = mPowerThrottlingConfigData.customAnimationRateSec;
+        mThermalLevelListener.start();
+    }
+
+    private void activatePmicMonitor() {
+        if (!mPmicMonitor.isStopped()) {
+            return;
+        }
         mPmicMonitor.start();
     }
 
+    private void deactivatePmicMonitor(@Temperature.ThrottlingStatus int status) {
+        if (status != Temperature.THROTTLING_NONE) {
+            return;
+        }
+        if (mPmicMonitor.isStopped()) {
+            return;
+        }
+        mPmicMonitor.stop();
+    }
+
+    private final class ThermalLevelListener extends IThermalEventListener.Stub {
+        private final Handler mHandler;
+        private IThermalService mThermalService;
+        private boolean mStarted;
+
+        ThermalLevelListener(Handler handler) {
+            mHandler = handler;
+            mStarted = false;
+            mThermalService = IThermalService.Stub.asInterface(
+                    ServiceManager.getService(Context.THERMAL_SERVICE));
+        }
+
+        IThermalService getThermalService() {
+            return mThermalService;
+        }
+
+        void start() {
+            if (mStarted) {
+                return;
+            }
+            if (mThermalService == null) {
+                return;
+            }
+            try {
+                // TODO b/279114539 Try DISPLAY first and then fallback to SKIN.
+                mThermalService.registerThermalEventListenerWithType(this, Temperature.TYPE_SKIN);
+                mStarted = true;
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to register thermal status listener", e);
+            }
+        }
+
+        @Override
+        public void notifyThrottling(Temperature temp) {
+            @Temperature.ThrottlingStatus int status = temp.getStatus();
+            if (status >= Temperature.THROTTLING_LIGHT) {
+                Slog.d(TAG, "Activating pmic monitor due to thermal state:" + status);
+                mHandler.post(() -> activatePmicMonitor());
+            } else {
+                if (!mPmicMonitor.isStopped()) {
+                    mHandler.post(() -> deactivatePmicMonitor(status));
+                }
+            }
+        }
+
+        void stop() {
+            if (!mStarted) {
+                return;
+            }
+            try {
+                mThermalService.unregisterThermalEventListener(this);
+                mStarted = false;
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to unregister thermal status listener", e);
+            }
+            mThermalService = null;
+        }
+
+        void dump(PrintWriter writer) {
+            writer.println("  ThermalLevelObserver:");
+            writer.println("    mStarted: " + mStarted);
+        }
+    }
+
     public interface PowerData {
         @NonNull
         String getUniqueDisplayId();
@@ -279,8 +456,12 @@
 
     @VisibleForTesting
     static class Injector {
-        PmicMonitor getPmicMonitor(PowerChangeListener listener, int pollingTime) {
-            return new PmicMonitor(listener, pollingTime);
+        PmicMonitor getPmicMonitor(PowerChangeListener powerChangeListener,
+                                   IThermalService thermalService,
+                                   int pollingMaxTimeMillis,
+                                   int pollingMinTimeMillis) {
+            return new PmicMonitor(powerChangeListener, thermalService, pollingMaxTimeMillis,
+                                        pollingMinTimeMillis);
         }
 
         DeviceConfigParameterProvider getDeviceConfigParameterProvider() {
diff --git a/services/core/java/com/android/server/display/brightness/clamper/PmicMonitor.java b/services/core/java/com/android/server/display/brightness/clamper/PmicMonitor.java
index 26784f23..355f4fe 100644
--- a/services/core/java/com/android/server/display/brightness/clamper/PmicMonitor.java
+++ b/services/core/java/com/android/server/display/brightness/clamper/PmicMonitor.java
@@ -18,15 +18,12 @@
 
 import static com.android.server.display.brightness.clamper.BrightnessPowerClamper.PowerChangeListener;
 
-import android.annotation.NonNull;
 import android.annotation.Nullable;
-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.IThermalService;
 import android.os.RemoteException;
-import android.os.ServiceManager;
 import android.os.Temperature;
 import android.power.PowerStatsInternal;
 import android.util.IntArray;
@@ -51,25 +48,30 @@
 
     // The executor to periodically monitor the display power.
     private final ScheduledExecutorService mExecutor;
-    @NonNull
-    private final PowerChangeListener mPowerChangeListener;
-    private final long mPowerMonitorPeriodConfigSecs;
+    private final long mPowerMonitorPeriodConfigMillis;
     private final PowerStatsInternal mPowerStatsInternal;
     @VisibleForTesting final IThermalService mThermalService;
+    @VisibleForTesting PowerChangeListener mPowerChangeListener;
     private ScheduledFuture<?> mPmicMonitorFuture;
     private float mLastEnergyConsumed = 0;
-    private float mCurrentAvgPower = 0;
+    private float mCurrentTotalAvgPower = 0;
     private Temperature mCurrentTemperature;
     private long mCurrentTimestampMillis = 0;
+    private float[] mAvgPowerCircularList;
+    private int mPowerListStart = 0;
+    private int mPowerListEnd = 0;
 
-    PmicMonitor(PowerChangeListener listener, int powerMonitorPeriodConfigSecs) {
+    PmicMonitor(PowerChangeListener listener,
+                IThermalService thermalService,
+                int pollingMaxTimeMillis,
+                int pollingMinTimeMillis) {
         mPowerChangeListener = listener;
+        mThermalService = thermalService;
         mPowerStatsInternal = LocalServices.getService(PowerStatsInternal.class);
-        mThermalService = IThermalService.Stub.asInterface(
-                ServiceManager.getService(Context.THERMAL_SERVICE));
+        mAvgPowerCircularList = new float[pollingMaxTimeMillis / pollingMinTimeMillis];
         // start a periodic worker thread.
         mExecutor = Executors.newSingleThreadScheduledExecutor();
-        mPowerMonitorPeriodConfigSecs = (long) powerMonitorPeriodConfigSecs;
+        mPowerMonitorPeriodConfigMillis = pollingMinTimeMillis;
     }
 
     @Nullable
@@ -141,12 +143,28 @@
 
         // capture thermal state.
         Temperature displayTemperature = getDisplayTemperature();
-        mCurrentAvgPower = currentPower;
+        boolean isBufferFull = false;
+        mAvgPowerCircularList[mPowerListEnd] = currentPower;
+        mCurrentTotalAvgPower += currentPower;
+        mPowerListEnd =
+            (mPowerListEnd + 1) % mAvgPowerCircularList.length;
+        if (mPowerListStart == mPowerListEnd) {
+            isBufferFull = true;
+        }
+
         mCurrentTemperature = displayTemperature;
         mLastEnergyConsumed = displayResults[0].energyUWs;
         mCurrentTimestampMillis = displayResults[0].timestampMs;
-        if (mCurrentTemperature != null) {
-            mPowerChangeListener.onChanged(mCurrentAvgPower, mCurrentTemperature.getStatus());
+
+        if (mCurrentTemperature != null && isBufferFull) {
+            mPowerChangeListener.onChanged(mCurrentTotalAvgPower / mAvgPowerCircularList.length,
+                                            mCurrentTemperature.getStatus());
+        }
+
+        // average power long-term list is full, reset values for next cycle.
+        if (isBufferFull) {
+            mCurrentTotalAvgPower = mCurrentTotalAvgPower - mAvgPowerCircularList[mPowerListStart];
+            mPowerListStart = (mPowerListStart + 1) % mAvgPowerCircularList.length;
         }
     }
 
@@ -165,11 +183,11 @@
         if (mPmicMonitorFuture == null) {
             mPmicMonitorFuture = mExecutor.scheduleAtFixedRate(
                                     this::capturePeriodicDisplayPower,
-                                    mPowerMonitorPeriodConfigSecs,
-                                    mPowerMonitorPeriodConfigSecs,
-                                    TimeUnit.SECONDS);
+                                    mPowerMonitorPeriodConfigMillis,
+                                    mPowerMonitorPeriodConfigMillis,
+                                    TimeUnit.MILLISECONDS);
         } else {
-            Slog.e(TAG, "already scheduled, stop() called before start.");
+            Slog.e(TAG, "PMIC already scheduled, stop() called before start.");
         }
     }
 
@@ -184,6 +202,23 @@
     }
 
     /**
+     * Updates PMIC configuration.
+     */
+    public void updateConfiguration() {
+        if (mPmicMonitorFuture != null) {
+            mPmicMonitorFuture.cancel(true);
+            mPmicMonitorFuture = null;
+        }
+    }
+
+    /**
+     * Returns if PMIC monitor is stopped.
+     */
+    public boolean isStopped() {
+        return mPmicMonitorFuture == null;
+    }
+
+    /**
      * Shutdown power IC service and worker thread.
      */
     public void shutdown() {
diff --git a/services/core/java/com/android/server/display/brightness/strategy/AutoBrightnessFallbackStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/AutoBrightnessFallbackStrategy.java
index 1db9bbe..91c985830 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/AutoBrightnessFallbackStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/AutoBrightnessFallbackStrategy.java
@@ -100,6 +100,7 @@
         writer.println("AutoBrightnessFallbackStrategy:");
         writer.println("  mLeadDisplayId=" + mLeadDisplayId);
         writer.println("  mIsDisplayEnabled=" + mIsDisplayEnabled);
+        writer.println("");
         if (mScreenOffBrightnessSensorController != null) {
             IndentingPrintWriter ipw = new IndentingPrintWriter(writer, " ");
             mScreenOffBrightnessSensorController.dump(ipw);
diff --git a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
index 35be0f3..69b67c8 100644
--- a/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
+++ b/services/core/java/com/android/server/display/feature/DisplayManagerFlags.java
@@ -420,6 +420,7 @@
      */
     public void dump(PrintWriter pw) {
         pw.println("DisplayManagerFlags:");
+        pw.println("--------------------");
         pw.println(" " + mAdaptiveToneImprovements1);
         pw.println(" " + mAdaptiveToneImprovements2);
         pw.println(" " + mBackUpSmoothDisplayAndForcePeakRefreshRateFlagState);
diff --git a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
index d909004..18e0d6e 100644
--- a/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
+++ b/services/core/java/com/android/server/display/mode/DisplayModeDirector.java
@@ -578,7 +578,8 @@
      * @param pw The stream to dump information to.
      */
     public void dump(PrintWriter pw) {
-        pw.println("DisplayModeDirector");
+        pw.println("DisplayModeDirector:");
+        pw.println("--------------------");
         synchronized (mLock) {
             pw.println("  mSupportedModesByDisplay:");
             for (int i = 0; i < mSupportedModesByDisplay.size(); i++) {
diff --git a/services/core/java/com/android/server/display/state/DisplayStateController.java b/services/core/java/com/android/server/display/state/DisplayStateController.java
index 21bb208..dba6874 100644
--- a/services/core/java/com/android/server/display/state/DisplayStateController.java
+++ b/services/core/java/com/android/server/display/state/DisplayStateController.java
@@ -114,9 +114,9 @@
      *
      * @param pw The PrintWriter used to dump the state.
      */
-    public void dumpsys(PrintWriter pw) {
-        pw.println();
+    public void dump(PrintWriter pw) {
         pw.println("DisplayStateController:");
+        pw.println("-----------------------");
         pw.println("  mPerformScreenOffTransition:" + mPerformScreenOffTransition);
         pw.println("  mDozeStateOverride=" + mDozeStateOverride);
 
diff --git a/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceController.java b/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceController.java
index 7ea576d..49c45a7 100644
--- a/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceController.java
+++ b/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceController.java
@@ -427,7 +427,8 @@
      *      The writer used to dump the state.
      */
     public void dump(PrintWriter writer) {
-        writer.println("DisplayWhiteBalanceController");
+        writer.println("DisplayWhiteBalanceController:");
+        writer.println("------------------------------");
         writer.println("  mLoggingEnabled=" + mLoggingEnabled);
         writer.println("  mEnabled=" + mEnabled);
         writer.println("  mStrongModeEnabled=" + mStrongModeEnabled);
diff --git a/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceSettings.java b/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceSettings.java
index 0efb749..a5755ac 100644
--- a/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceSettings.java
+++ b/services/core/java/com/android/server/display/whitebalance/DisplayWhiteBalanceSettings.java
@@ -23,7 +23,6 @@
 import android.os.Message;
 import android.util.Slog;
 
-import com.android.internal.util.Preconditions;
 import com.android.server.LocalServices;
 import com.android.server.display.color.ColorDisplayService;
 import com.android.server.display.color.ColorDisplayService.ColorDisplayServiceInternal;
@@ -129,7 +128,8 @@
      *      The writer used to dump the state.
      */
     public void dump(PrintWriter writer) {
-        writer.println("DisplayWhiteBalanceSettings");
+        writer.println("DisplayWhiteBalanceSettings:");
+        writer.println("----------------------------");
         writer.println("  mLoggingEnabled=" + mLoggingEnabled);
         writer.println("  mContext=" + mContext);
         writer.println("  mHandler=" + mHandler);
diff --git a/services/core/java/com/android/server/flags/services.aconfig b/services/core/java/com/android/server/flags/services.aconfig
index c361aee..649678c 100644
--- a/services/core/java/com/android/server/flags/services.aconfig
+++ b/services/core/java/com/android/server/flags/services.aconfig
@@ -45,3 +45,14 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    namespace: "backstage_power"
+    name: "consolidate_battery_change_events"
+    description: "Optimize battery status updates by delivering only the most recent battery information"
+    bug: "361334584"
+    is_fixed_read_only: true
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/services/core/java/com/android/server/hdmi/HdmiControlService.java b/services/core/java/com/android/server/hdmi/HdmiControlService.java
index bbe7b2b..3c7b9d3 100644
--- a/services/core/java/com/android/server/hdmi/HdmiControlService.java
+++ b/services/core/java/com/android/server/hdmi/HdmiControlService.java
@@ -4328,6 +4328,7 @@
         HdmiCecLocalDevicePlayback playback = playback();
         HdmiCecLocalDeviceAudioSystem audioSystem = audioSystem();
         if (playback != null) {
+            playback.dismissUiOnActiveSourceStatusRecovered();
             playback.setActiveSource(playback.getDeviceInfo().getLogicalAddress(), physicalAddress,
                     caller);
             playback.wakeUpIfActiveSource();
diff --git a/services/core/java/com/android/server/input/debug/TouchpadDebugView.java b/services/core/java/com/android/server/input/debug/TouchpadDebugView.java
index 5fca771..7785ffb 100644
--- a/services/core/java/com/android/server/input/debug/TouchpadDebugView.java
+++ b/services/core/java/com/android/server/input/debug/TouchpadDebugView.java
@@ -16,23 +16,69 @@
 
 package com.android.server.input.debug;
 
+import android.annotation.NonNull;
 import android.content.Context;
+import android.content.res.Configuration;
 import android.graphics.Color;
+import android.graphics.PixelFormat;
+import android.graphics.Rect;
+import android.util.Slog;
 import android.view.Gravity;
+import android.view.MotionEvent;
+import android.view.ViewConfiguration;
+import android.view.WindowManager;
 import android.widget.LinearLayout;
 import android.widget.TextView;
 
-public class TouchpadDebugView extends LinearLayout {
+import java.util.Objects;
 
+public class TouchpadDebugView extends LinearLayout {
     /**
      * Input device ID for the touchpad that this debug view is displaying.
      */
     private final int mTouchpadId;
 
+    @NonNull
+    private final WindowManager mWindowManager;
+
+    @NonNull
+    private final WindowManager.LayoutParams mWindowLayoutParams;
+
+    private final int mTouchSlop;
+
+    private float mTouchDownX;
+    private float mTouchDownY;
+    private int mScreenWidth;
+    private int mScreenHeight;
+    private int mWindowLocationBeforeDragX;
+    private int mWindowLocationBeforeDragY;
+
     public TouchpadDebugView(Context context, int touchpadId) {
         super(context);
         mTouchpadId = touchpadId;
+        mWindowManager =
+                Objects.requireNonNull(getContext().getSystemService(WindowManager.class));
         init(context);
+        mTouchSlop = ViewConfiguration.get(context).getScaledTouchSlop();
+
+        // TODO(b/360137366): Use the hardware properties to initialise layout parameters.
+        mWindowLayoutParams = new WindowManager.LayoutParams();
+        mWindowLayoutParams.type = WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
+        mWindowLayoutParams.flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE
+                | WindowManager.LayoutParams.FLAG_LAYOUT_IN_SCREEN;
+        mWindowLayoutParams.privateFlags |=
+                WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS;
+        mWindowLayoutParams.setFitInsetsTypes(0);
+        mWindowLayoutParams.layoutInDisplayCutoutMode =
+                WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
+        mWindowLayoutParams.format = PixelFormat.TRANSLUCENT;
+        mWindowLayoutParams.setTitle("TouchpadDebugView - display " + mContext.getDisplayId());
+
+        mWindowLayoutParams.x = 40;
+        mWindowLayoutParams.y = 100;
+        mWindowLayoutParams.width = WindowManager.LayoutParams.WRAP_CONTENT;
+        mWindowLayoutParams.height = WindowManager.LayoutParams.WRAP_CONTENT;
+        mWindowLayoutParams.gravity = Gravity.TOP | Gravity.LEFT;
     }
 
     private void init(Context context) {
@@ -43,14 +89,12 @@
         setBackgroundColor(Color.TRANSPARENT);
 
         // TODO(b/286551975): Replace this content with the touchpad debug view.
-
         TextView textView1 = new TextView(context);
         textView1.setBackgroundColor(Color.parseColor("#FFFF0000"));
         textView1.setTextSize(20);
         textView1.setText("Touchpad Debug View 1");
         textView1.setGravity(Gravity.CENTER);
         textView1.setTextColor(Color.WHITE);
-
         textView1.setLayoutParams(new LayoutParams(1000, 200));
 
         TextView textView2 = new TextView(context);
@@ -63,9 +107,98 @@
 
         addView(textView1);
         addView(textView2);
+
+        updateScreenDimensions();
+    }
+
+    @Override
+    public boolean onTouchEvent(MotionEvent event) {
+        float deltaX;
+        float deltaY;
+        switch (event.getAction()) {
+            case MotionEvent.ACTION_DOWN:
+                mWindowLocationBeforeDragX = mWindowLayoutParams.x;
+                mWindowLocationBeforeDragY = mWindowLayoutParams.y;
+                mTouchDownX = event.getRawX() - mWindowLocationBeforeDragX;
+                mTouchDownY = event.getRawY() - mWindowLocationBeforeDragY;
+                return true;
+
+            case MotionEvent.ACTION_MOVE:
+                deltaX = event.getRawX() - mWindowLayoutParams.x - mTouchDownX;
+                deltaY = event.getRawY() - mWindowLayoutParams.y - mTouchDownY;
+                Slog.d("TouchpadDebugView", "Slop = " + mTouchSlop);
+                if (isSlopExceeded(deltaX, deltaY)) {
+                    Slog.d("TouchpadDebugView", "Slop exceeded");
+                    mWindowLayoutParams.x =
+                            Math.max(0, Math.min((int) (event.getRawX() - mTouchDownX),
+                                    mScreenWidth - this.getWidth()));
+                    mWindowLayoutParams.y =
+                            Math.max(0, Math.min((int) (event.getRawY() - mTouchDownY),
+                                    mScreenHeight - this.getHeight()));
+
+                    Slog.d("TouchpadDebugView", "New position X: "
+                            + mWindowLayoutParams.x + ", Y: " + mWindowLayoutParams.y);
+
+                    mWindowManager.updateViewLayout(this, mWindowLayoutParams);
+                }
+                return true;
+
+            case MotionEvent.ACTION_UP:
+                deltaX = event.getRawX() - mWindowLayoutParams.x - mTouchDownX;
+                deltaY = event.getRawY() - mWindowLayoutParams.y - mTouchDownY;
+                if (!isSlopExceeded(deltaX, deltaY)) {
+                    performClick();
+                }
+                return true;
+
+            case MotionEvent.ACTION_CANCEL:
+                // Move the window back to the original position
+                mWindowLayoutParams.x = mWindowLocationBeforeDragX;
+                mWindowLayoutParams.y = mWindowLocationBeforeDragY;
+                mWindowManager.updateViewLayout(this, mWindowLayoutParams);
+                return true;
+
+            default:
+                return super.onTouchEvent(event);
+        }
+    }
+
+    @Override
+    public boolean performClick() {
+        super.performClick();
+        Slog.d("TouchpadDebugView", "You clicked me!");
+        return true;
+    }
+
+    @Override
+    protected void onConfigurationChanged(Configuration newConfig) {
+        super.onConfigurationChanged(newConfig);
+        updateScreenDimensions();
+
+        // Adjust view position to stay within screen bounds after rotation
+        mWindowLayoutParams.x =
+                Math.max(0, Math.min(mWindowLayoutParams.x, mScreenWidth - getWidth()));
+        mWindowLayoutParams.y =
+                Math.max(0, Math.min(mWindowLayoutParams.y, mScreenHeight - getHeight()));
+        mWindowManager.updateViewLayout(this, mWindowLayoutParams);
+    }
+
+    private boolean isSlopExceeded(float deltaX, float deltaY) {
+        return deltaX * deltaX + deltaY * deltaY >= mTouchSlop * mTouchSlop;
+    }
+
+    private void updateScreenDimensions() {
+        Rect windowBounds =
+                mWindowManager.getCurrentWindowMetrics().getBounds();
+        mScreenWidth = windowBounds.width();
+        mScreenHeight = windowBounds.height();
     }
 
     public int getTouchpadId() {
         return mTouchpadId;
     }
+
+    public WindowManager.LayoutParams getWindowLayoutParams() {
+        return mWindowLayoutParams;
+    }
 }
diff --git a/services/core/java/com/android/server/input/debug/TouchpadDebugViewController.java b/services/core/java/com/android/server/input/debug/TouchpadDebugViewController.java
index c7760c6..1e59167 100644
--- a/services/core/java/com/android/server/input/debug/TouchpadDebugViewController.java
+++ b/services/core/java/com/android/server/input/debug/TouchpadDebugViewController.java
@@ -18,14 +18,12 @@
 
 import android.annotation.Nullable;
 import android.content.Context;
-import android.graphics.PixelFormat;
 import android.hardware.display.DisplayManager;
 import android.hardware.input.InputManager;
 import android.os.Handler;
 import android.os.Looper;
 import android.util.Slog;
 import android.view.Display;
-import android.view.Gravity;
 import android.view.InputDevice;
 import android.view.WindowManager;
 
@@ -36,12 +34,14 @@
 
 public class TouchpadDebugViewController {
 
-    private static final String TAG = "TouchpadDebugViewController";
+    private static final String TAG = "TouchpadDebugView";
 
     private final Context mContext;
     private final Handler mHandler;
+
     @Nullable
     private TouchpadDebugView mTouchpadDebugView;
+
     private final InputManagerService mInputManagerService;
 
     public TouchpadDebugViewController(Context context, Looper looper,
@@ -95,32 +95,15 @@
                 mContext.getSystemService(WindowManager.class));
 
         mTouchpadDebugView = new TouchpadDebugView(mContext, touchpadId);
+        final WindowManager.LayoutParams mWindowLayoutParams =
+                mTouchpadDebugView.getWindowLayoutParams();
 
-        final WindowManager.LayoutParams lp = new WindowManager.LayoutParams();
-        lp.type = WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY;
-        lp.flags = WindowManager.LayoutParams.FLAG_NOT_FOCUSABLE
-                | WindowManager.LayoutParams.FLAG_LAYOUT_IN_SCREEN;
-        lp.privateFlags |= WindowManager.LayoutParams.SYSTEM_FLAG_SHOW_FOR_ALL_USERS;
-        lp.setFitInsetsTypes(0);
-        lp.layoutInDisplayCutoutMode =
-                WindowManager.LayoutParams.LAYOUT_IN_DISPLAY_CUTOUT_MODE_ALWAYS;
-        lp.format = PixelFormat.TRANSLUCENT;
-        lp.setTitle("TouchpadDebugView - display " + mContext.getDisplayId());
-        lp.inputFeatures |= WindowManager.LayoutParams.INPUT_FEATURE_NO_INPUT_CHANNEL;
-
-        lp.x = 40;
-        lp.y = 100;
-        lp.width = WindowManager.LayoutParams.WRAP_CONTENT;
-        lp.height = WindowManager.LayoutParams.WRAP_CONTENT;
-        lp.gravity = Gravity.TOP | Gravity.LEFT;
-
-        wm.addView(mTouchpadDebugView, lp);
+        wm.addView(mTouchpadDebugView, mWindowLayoutParams);
         Slog.d(TAG, "Touchpad debug view created.");
 
         TouchpadHardwareProperties mTouchpadHardwareProperties =
                 mInputManagerService.getTouchpadHardwareProperties(
                         touchpadId);
-        // TODO(b/360137366): Use the hardware properties to initialise layout parameters.
         if (mTouchpadHardwareProperties != null) {
             Slog.d(TAG, mTouchpadHardwareProperties.toString());
         } else {
diff --git a/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
index 58e3452..e1f26d6 100644
--- a/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
+++ b/services/core/java/com/android/server/inputmethod/IInputMethodManagerImpl.java
@@ -116,11 +116,11 @@
         boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
                 @Nullable ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
                 @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
-                @SoftInputShowHideReason int reason);
+                @SoftInputShowHideReason int reason, boolean async);
 
         boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
                 @Nullable ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-                ResultReceiver resultReceiver, @SoftInputShowHideReason int reason);
+                ResultReceiver resultReceiver, @SoftInputShowHideReason int reason, boolean async);
 
         @PermissionVerified(Manifest.permission.TEST_INPUT_METHOD)
         void hideSoftInputFromServerForTest();
@@ -132,7 +132,8 @@
                 @Nullable EditorInfo editorInfo, IRemoteInputConnection inputConnection,
                 IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
                 int unverifiedTargetSdkVersion, @UserIdInt int userId,
-                @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq);
+                @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq,
+                boolean useAsyncShowHideMethod);
 
         InputBindResult startInputOrWindowGainedFocus(
                 @StartInputReason int startInputReason, IInputMethodClient client,
@@ -290,17 +291,17 @@
     public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
-            @SoftInputShowHideReason int reason) {
+            @SoftInputShowHideReason int reason, boolean async) {
         return mCallback.showSoftInput(client, windowToken, statsToken, flags, lastClickToolType,
-                resultReceiver, reason);
+                resultReceiver, reason, async);
     }
 
     @Override
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason, boolean async) {
         return mCallback.hideSoftInput(client, windowToken, statsToken, flags, resultReceiver,
-                reason);
+                reason, async);
     }
 
     @EnforcePermission(Manifest.permission.TEST_INPUT_METHOD)
@@ -336,11 +337,13 @@
             IRemoteInputConnection inputConnection,
             IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq) {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq,
+            boolean useAsyncShowHideMethod) {
         mCallback.startInputOrWindowGainedFocusAsync(
                 startInputReason, client, windowToken, startInputFlags, softInputMode,
                 windowFlags, editorInfo, inputConnection, remoteAccessibilityInputConnection,
-                unverifiedTargetSdkVersion, userId, imeDispatcher, startInputSeq);
+                unverifiedTargetSdkVersion, userId, imeDispatcher, startInputSeq,
+                useAsyncShowHideMethod);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 2ad0d2a..7cfd2cc 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -3086,7 +3086,7 @@
     public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             int lastClickToolType, ResultReceiver resultReceiver,
-            @SoftInputShowHideReason int reason) {
+            @SoftInputShowHideReason int reason, boolean async) {
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "IMMS.showSoftInput");
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#showSoftInput", mDumper);
@@ -3533,7 +3533,7 @@
     @Override
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @NonNull ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason, boolean async) {
         ImeTracing.getInstance().triggerManagerServiceDump(
                 "InputMethodManagerService#hideSoftInput", mDumper);
         synchronized (ImfLock.class) {
@@ -3675,7 +3675,8 @@
             IRemoteInputConnection inputConnection,
             IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq) {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq,
+            boolean useAsyncShowHideMethod) {
         // implemented by ZeroJankProxy
     }
 
diff --git a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
index c940a9c..214aa1d 100644
--- a/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
+++ b/services/core/java/com/android/server/inputmethod/ZeroJankProxy.java
@@ -77,6 +77,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A proxy that processes all {@link IInputMethodManager} calls asynchronously.
@@ -175,19 +176,45 @@
     public boolean showSoftInput(IInputMethodClient client, IBinder windowToken,
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.ShowFlags int flags,
             @MotionEvent.ToolType int lastClickToolType, ResultReceiver resultReceiver,
-            @SoftInputShowHideReason int reason) {
-        offload(() -> mInner.showSoftInput(
-                client, windowToken, statsToken, flags, lastClickToolType, resultReceiver, reason));
-        return true;
+            @SoftInputShowHideReason int reason, boolean async) {
+
+        if (async) {
+            offload(() -> mInner.showSoftInput(
+                    client, windowToken, statsToken, flags, lastClickToolType, resultReceiver,
+                    reason, async));
+            return true;
+        } else {
+            final var future = CompletableFuture.supplyAsync(
+                    () -> mInner.showSoftInput(
+                            client,
+                            windowToken,
+                            statsToken,
+                            flags,
+                            lastClickToolType,
+                            resultReceiver,
+                            reason,
+                            async),
+                    this::offload);
+            return future.completeOnTimeout(false, 1, TimeUnit.SECONDS).join();
+        }
     }
 
     @Override
     public boolean hideSoftInput(IInputMethodClient client, IBinder windowToken,
             @Nullable ImeTracker.Token statsToken, @InputMethodManager.HideFlags int flags,
-            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason) {
-        offload(() -> mInner.hideSoftInput(
-                client, windowToken, statsToken, flags, resultReceiver, reason));
-        return true;
+            ResultReceiver resultReceiver, @SoftInputShowHideReason int reason, boolean async) {
+
+        if (async) {
+            offload(() -> mInner.hideSoftInput(
+                    client, windowToken, statsToken, flags, resultReceiver, reason, async));
+            return true;
+        } else {
+            final var future = CompletableFuture.supplyAsync(
+                    () -> mInner.hideSoftInput(
+                            client, windowToken, statsToken, flags, resultReceiver, reason, async),
+                    this::offload);
+            return future.completeOnTimeout(false, 1, TimeUnit.SECONDS).join();
+        }
     }
 
     @Override
@@ -207,7 +234,8 @@
             IRemoteInputConnection inputConnection,
             IRemoteAccessibilityInputConnection remoteAccessibilityInputConnection,
             int unverifiedTargetSdkVersion, @UserIdInt int userId,
-            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq) {
+            @NonNull ImeOnBackInvokedDispatcher imeDispatcher, int startInputSeq,
+            boolean useAsyncShowHideMethod) {
         offload(() -> {
             InputBindResult result = mInner.startInputOrWindowGainedFocus(startInputReason, client,
                     windowToken, startInputFlags, softInputMode, windowFlags,
diff --git a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
index d0b8990..f44b852 100644
--- a/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
+++ b/services/core/java/com/android/server/locksettings/RebootEscrowManager.java
@@ -142,7 +142,6 @@
             ERROR_KEYSTORE_FAILURE,
             ERROR_NO_NETWORK,
             ERROR_TIMEOUT_EXHAUSTED,
-            ERROR_NO_REBOOT_ESCROW_DATA,
     })
     @Retention(RetentionPolicy.SOURCE)
     @interface RebootEscrowErrorCode {
@@ -158,7 +157,6 @@
     static final int ERROR_KEYSTORE_FAILURE = 7;
     static final int ERROR_NO_NETWORK = 8;
     static final int ERROR_TIMEOUT_EXHAUSTED = 9;
-    static final int ERROR_NO_REBOOT_ESCROW_DATA = 10;
 
     private @RebootEscrowErrorCode int mLoadEscrowDataErrorCode = ERROR_NONE;
 
@@ -507,9 +505,6 @@
         if (rebootEscrowUsers.isEmpty()) {
             Slog.i(TAG, "No reboot escrow data found for users,"
                     + " skipping loading escrow data");
-            setLoadEscrowDataErrorCode(ERROR_NO_REBOOT_ESCROW_DATA, retryHandler);
-            reportMetricOnRestoreComplete(
-                    /* success= */ false, /* attemptCount= */ 1, retryHandler);
             clearMetricsStorage();
             return;
         }
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 54bc6fb..54e9189 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -7600,16 +7600,14 @@
                     + " trying to post for invalid pkg " + pkg + " in user " + incomingUserId);
         }
 
-        if (android.app.Flags.secureAllowlistToken()) {
-            IBinder allowlistToken = notification.getAllowlistToken();
-            if (allowlistToken != null && allowlistToken != ALLOWLIST_TOKEN) {
-                throw new SecurityException(
-                        "Unexpected allowlist token received from " + callingUid);
-            }
-            // allowlistToken is populated by unparceling, so it can be null if the notification was
-            // posted from inside system_server. Ensure it's the expected value.
-            notification.overrideAllowlistToken(ALLOWLIST_TOKEN);
+        IBinder allowlistToken = notification.getAllowlistToken();
+        if (allowlistToken != null && allowlistToken != ALLOWLIST_TOKEN) {
+            throw new SecurityException(
+                    "Unexpected allowlist token received from " + callingUid);
         }
+        // allowlistToken is populated by unparceling, so it can be null if the notification was
+        // posted from inside system_server. Ensure it's the expected value.
+        notification.overrideAllowlistToken(ALLOWLIST_TOKEN);
 
         checkRestrictedCategories(notification);
 
@@ -8774,12 +8772,10 @@
          */
         private boolean enqueueNotification() {
             synchronized (mNotificationLock) {
-                if (android.app.Flags.secureAllowlistToken()) {
-                    // allowlistToken is populated by unparceling, so it will be absent if the
-                    // EnqueueNotificationRunnable is created directly by NMS (as we do for group
-                    // summaries) instead of via notify(). Fix that.
-                    r.getNotification().overrideAllowlistToken(ALLOWLIST_TOKEN);
-                }
+                // allowlistToken is populated by unparceling, so it will be absent if the
+                // EnqueueNotificationRunnable is created directly by NMS (as we do for group
+                // summaries) instead of via notify(). Fix that.
+                r.getNotification().overrideAllowlistToken(ALLOWLIST_TOKEN);
 
                 final long snoozeAt =
                         mSnoozeHelper.getSnoozeTimeForUnpostedNotification(
@@ -11777,6 +11773,8 @@
 
             mHandler.post(new EnqueueNotificationRunnable(record.getUser().getIdentifier(),
                     record, isAppForeground, /* isAppProvided= */ false, tracker));
+
+            EventLogTags.writeNotificationCancelPrevented(record.getKey());
         }
     }
 
diff --git a/services/core/java/com/android/server/notification/ZenModeHelper.java b/services/core/java/com/android/server/notification/ZenModeHelper.java
index 6ff0e04..2ada9ae4 100644
--- a/services/core/java/com/android/server/notification/ZenModeHelper.java
+++ b/services/core/java/com/android/server/notification/ZenModeHelper.java
@@ -226,10 +226,8 @@
         mDefaultConfig = Flags.modesUi()
                 ? ZenModeConfig.getDefaultConfig()
                 : readDefaultConfig(mContext.getResources());
-        updateDefaultConfigAutomaticRules();
-        if (Flags.modesApi()) {
-            updateDefaultAutomaticRulePolicies();
-        }
+        updateDefaultConfig(mContext, mDefaultConfig);
+
         mConfig = mDefaultConfig.copy();
         synchronized (mConfigsArrayLock) {
             mConfigs.put(UserHandle.USER_SYSTEM, mConfig);
@@ -1073,7 +1071,7 @@
     }
 
     void updateZenRulesOnLocaleChange() {
-        updateDefaultConfigAutomaticRules();
+        updateRuleStringsForCurrentLocale(mContext, mDefaultConfig);
         synchronized (mConfigLock) {
             if (mConfig == null) {
                 return;
@@ -2127,17 +2125,25 @@
     @GuardedBy("mConfigLock")
     private void applyCustomPolicy(ZenPolicy policy, ZenRule rule, boolean useManualConfig) {
         if (rule.zenMode == Global.ZEN_MODE_NO_INTERRUPTIONS) {
-            policy.apply(new ZenPolicy.Builder()
-                    .disallowAllSounds()
-                    .allowPriorityChannels(false)
-                    .build());
+            if (Flags.modesApi() && Flags.modesUi()) {
+                policy.apply(ZenPolicy.getBasePolicyInterruptionFilterNone());
+            } else {
+                policy.apply(new ZenPolicy.Builder()
+                        .disallowAllSounds()
+                        .allowPriorityChannels(false)
+                        .build());
+            }
         } else if (rule.zenMode == Global.ZEN_MODE_ALARMS) {
-            policy.apply(new ZenPolicy.Builder()
-                    .disallowAllSounds()
-                    .allowAlarms(true)
-                    .allowMedia(true)
-                    .allowPriorityChannels(false)
-                    .build());
+            if (Flags.modesApi() && Flags.modesUi()) {
+                policy.apply(ZenPolicy.getBasePolicyInterruptionFilterAlarms());
+            } else {
+                policy.apply(new ZenPolicy.Builder()
+                        .disallowAllSounds()
+                        .allowAlarms(true)
+                        .allowMedia(true)
+                        .allowPriorityChannels(false)
+                        .build());
+            }
         } else if (rule.zenPolicy != null) {
             policy.apply(rule.zenPolicy);
         } else {
@@ -2229,30 +2235,49 @@
         }
     }
 
-    private void updateDefaultConfigAutomaticRules() {
-        for (ZenRule rule : mDefaultConfig.automaticRules.values()) {
+    /**
+     * Apply changes to the <em>default</em> {@link ZenModeConfig} so that the rules included by
+     * default (Events / Sleeping) support the latest Zen features and are ready for new users.
+     *
+     * <p>This includes: setting a fully populated ZenPolicy, setting correct type and
+     * allowManualInvocation=true, and ensuring default names and trigger descriptions correspond
+     * to the current locale.
+     */
+    private static void updateDefaultConfig(Context context, ZenModeConfig defaultConfig) {
+        if (Flags.modesApi()) {
+            updateDefaultAutomaticRulePolicies(defaultConfig);
+        }
+        if (Flags.modesApi() && Flags.modesUi()) {
+            SystemZenRules.maybeUpgradeRules(context, defaultConfig);
+        }
+        updateRuleStringsForCurrentLocale(context, defaultConfig);
+    }
+
+    private static void updateRuleStringsForCurrentLocale(Context context,
+            ZenModeConfig defaultConfig) {
+        for (ZenRule rule : defaultConfig.automaticRules.values()) {
             if (ZenModeConfig.EVENTS_DEFAULT_RULE_ID.equals(rule.id)) {
-                rule.name = mContext.getResources()
+                rule.name = context.getResources()
                         .getString(R.string.zen_mode_default_events_name);
             } else if (ZenModeConfig.EVERY_NIGHT_DEFAULT_RULE_ID.equals(rule.id)) {
-                rule.name = mContext.getResources()
+                rule.name = context.getResources()
                         .getString(R.string.zen_mode_default_every_night_name);
             }
             if (Flags.modesApi() && Flags.modesUi()) {
-                SystemZenRules.updateTriggerDescription(mContext, rule);
+                SystemZenRules.updateTriggerDescription(context, rule);
             }
         }
     }
 
     // Updates the policies in the default automatic rules (provided via default XML config) to
     // be fully filled in default values.
-    private void updateDefaultAutomaticRulePolicies() {
+    private static void updateDefaultAutomaticRulePolicies(ZenModeConfig defaultConfig) {
         if (!Flags.modesApi()) {
             // Should be checked before calling, but just in case.
             return;
         }
-        ZenPolicy defaultPolicy = mDefaultConfig.getZenPolicy();
-        for (ZenRule rule : mDefaultConfig.automaticRules.values()) {
+        ZenPolicy defaultPolicy = defaultConfig.getZenPolicy();
+        for (ZenRule rule : defaultConfig.automaticRules.values()) {
             if (ZenModeConfig.DEFAULT_RULE_IDS.contains(rule.id) && rule.zenPolicy == null) {
                 rule.zenPolicy = defaultPolicy.copy();
             }
diff --git a/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java b/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
index e3061a7..4135161 100644
--- a/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
+++ b/services/core/java/com/android/server/pm/BackgroundUserSoundNotifier.java
@@ -69,6 +69,7 @@
         mUserManager = mSystemUserContext.getSystemService(UserManager.class);
         NotificationChannel channel = new NotificationChannel(BUSN_CHANNEL_ID, BUSN_CHANNEL_NAME,
                 NotificationManager.IMPORTANCE_HIGH);
+        channel.setSound(null, null);
         mNotificationManager.createNotificationChannel(channel);
         setupFocusControlAudioPolicy();
     }
diff --git a/services/core/java/com/android/server/pm/InstallPackageHelper.java b/services/core/java/com/android/server/pm/InstallPackageHelper.java
index 1317866..662c792 100644
--- a/services/core/java/com/android/server/pm/InstallPackageHelper.java
+++ b/services/core/java/com/android/server/pm/InstallPackageHelper.java
@@ -209,6 +209,7 @@
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 
+
 final class InstallPackageHelper {
     private final PackageManagerService mPm;
     private final AppDataHelper mAppDataHelper;
@@ -989,15 +990,6 @@
         return false;
     }
 
-    void installPackagesTraced(List<InstallRequest> requests) {
-        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
-            Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "installPackages");
-            installPackagesLI(requests);
-        } finally {
-            Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
-        }
-    }
-
     /**
      * Installs one or more packages atomically. This operation is broken up into four phases:
      * <ul>
@@ -1017,15 +1009,31 @@
      *
      * Failure at any phase will result in a full failure to install all packages.
      */
-    @GuardedBy("mPm.mInstallLock")
-    private void installPackagesLI(List<InstallRequest> requests) {
-        final Set<String> scannedPackages = new ArraySet<>(requests.size());
-        final Map<String, Settings.VersionInfo> versionInfos = new ArrayMap<>(requests.size());
-        final Map<String, Boolean> createdAppId = new ArrayMap<>(requests.size());
-        CriticalEventLog.getInstance().logInstallPackagesStarted();
+    void installPackagesTraced(List<InstallRequest> requests) {
+        Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "installPackages");
         boolean success = false;
+        final Map<String, Boolean> createdAppId = new ArrayMap<>(requests.size());
+        final Map<String, Settings.VersionInfo> versionInfos = new ArrayMap<>(requests.size());
         try {
-            Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "installPackagesLI");
+            CriticalEventLog.getInstance().logInstallPackagesStarted();
+
+            if (prepareInstallPackages(requests)
+                    && scanInstallPackages(requests, createdAppId, versionInfos)) {
+                List<ReconciledPackage> reconciledPackages =
+                        reconcileInstallPackages(requests, versionInfos);
+                if (reconciledPackages != null && commitInstallPackages(reconciledPackages)) {
+                    success = true;
+                }
+            }
+        } finally {
+            completeInstallProcess(requests, createdAppId, success);
+            Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
+        }
+    }
+
+    private boolean prepareInstallPackages(List<InstallRequest> requests) {
+        // TODO: will remove the locking after doRename is moved out of prepare
+        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
             for (InstallRequest request : requests) {
                 try {
                     Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "preparePackage");
@@ -1036,17 +1044,27 @@
                             prepareFailure.getMessage());
                     request.setOriginPackage(prepareFailure.mConflictingPackage);
                     request.setOriginPermission(prepareFailure.mConflictingPermission);
-                    return;
+                    return false;
                 } finally {
                     request.onPrepareFinished();
                     Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
                 }
+            }
+        }
+        return true;
+    }
 
+    private boolean scanInstallPackages(List<InstallRequest> requests,
+            Map<String, Boolean> createdAppId, Map<String, Settings.VersionInfo> versionInfos) {
+        // TODO(b/362840929): remove locker
+        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
+            final Set<String> scannedPackages = new ArraySet<>(requests.size());
+            for (InstallRequest request : requests) {
                 final ParsedPackage packageToScan = request.getParsedPackage();
                 if (packageToScan == null) {
                     request.setError(INSTALL_FAILED_SESSION_INVALID,
                             "Failed to obtain package to scan");
-                    return;
+                    return false;
                 }
                 request.setReturnCode(PackageManager.INSTALL_SUCCEEDED);
                 final String packageName = packageToScan.getPackageName();
@@ -1064,7 +1082,7 @@
                                 "Duplicate package "
                                         + packageName
                                         + " in multi-package install request.");
-                        return;
+                        return false;
                     }
                     if (!checkNoAppStorageIsConsistent(
                             request.getScanRequestOldPackage(), packageToScan)) {
@@ -1074,7 +1092,7 @@
                                 INSTALL_FAILED_UPDATE_INCOMPATIBLE,
                                 "Update attempted to change value of "
                                         + PackageManager.PROPERTY_NO_APP_DATA_STORAGE);
-                        return;
+                        return false;
                     }
                     final boolean isApex = (request.getScanFlags() & SCAN_AS_APEX) != 0;
                     final boolean isSdkLibrary = packageToScan.isSdkLibrary();
@@ -1087,7 +1105,7 @@
                             mPm.getSettingsVersionForPackage(packageToScan));
                 } catch (PackageManagerException e) {
                     request.setError("Scanning Failed.", e);
-                    return;
+                    return false;
                 }
                 if (request.isArchived()) {
                     final SparseArray<String> responsibleInstallerTitles =
@@ -1099,17 +1117,22 @@
                         request.setError(PackageManagerException.ofInternalError(
                                 "Failed to obtain the responsible installer info",
                                 INTERNAL_ERROR_ARCHIVE_NO_INSTALLER_TITLE));
-                        return;
+                        return false;
                     }
                     request.setResponsibleInstallerTitles(responsibleInstallerTitles);
                 }
             }
+        }
+        return true;
+    }
 
-            List<ReconciledPackage> reconciledPackages;
+    private List<ReconciledPackage> reconcileInstallPackages(List<InstallRequest> requests,
+            Map<String, Settings.VersionInfo> versionInfos) {
+        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
             synchronized (mPm.mLock) {
                 try {
                     Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "reconcilePackages");
-                    reconciledPackages = ReconcilePackageUtils.reconcilePackages(
+                    return ReconcilePackageUtils.reconcilePackages(
                             requests, Collections.unmodifiableMap(mPm.mPackages),
                             versionInfos, mSharedLibraries, mPm.mSettings.getKeySetManagerService(),
                             mPm.mSettings, mPm.mInjector.getSystemConfig());
@@ -1117,70 +1140,80 @@
                     for (InstallRequest request : requests) {
                         request.setError("Reconciliation failed...", e);
                     }
-                    return;
+                    return null;
                 } finally {
                     Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
                 }
-                if (Flags.improveInstallFreeze()) {
-                    // Postpone freezer until after reconcile
-                    for (ReconciledPackage reconciledPkg : reconciledPackages) {
-                        InstallRequest installRequest = reconciledPkg.mInstallRequest;
-                        String packageName = installRequest.getParsedPackage().getPackageName();
-                        PackageFreezer freezer = freezePackageForInstall(packageName,
-                                UserHandle.USER_ALL, installRequest.getInstallFlags(),
-                                "installPackageLI", ApplicationExitInfo.REASON_PACKAGE_UPDATED,
-                                installRequest);
-                        installRequest.setFreezer(freezer);
-                    }
+            }
+        }
+    }
+
+
+    private boolean commitInstallPackages(List<ReconciledPackage> reconciledPackages) {
+        try (PackageManagerTracedLock installLock = mPm.mInstallLock.acquireLock()) {
+            if (Flags.improveInstallFreeze()) {
+                // Postpone freezer until after reconcile
+                for (ReconciledPackage reconciledPkg : reconciledPackages) {
+                    InstallRequest installRequest = reconciledPkg.mInstallRequest;
+                    String packageName = installRequest.getParsedPackage().getPackageName();
+                    PackageFreezer freezer = freezePackageForInstall(packageName,
+                            UserHandle.USER_ALL, installRequest.getInstallFlags(),
+                            "installPackageLI", ApplicationExitInfo.REASON_PACKAGE_UPDATED,
+                            installRequest);
+                    installRequest.setFreezer(freezer);
                 }
+            }
+            synchronized (mPm.mLock) {
                 try {
                     Trace.traceBegin(TRACE_TAG_PACKAGE_MANAGER, "commitPackages");
                     commitPackagesLocked(reconciledPackages, mPm.mUserManager.getUserIds());
-                    success = true;
                 } finally {
                     Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
                 }
             }
             executePostCommitStepsLIF(reconciledPackages);
-        } finally {
-            if (success) {
-                for (InstallRequest request : requests) {
-                    if (request.getDataLoaderType() != DataLoaderType.INCREMENTAL) {
-                        continue;
-                    }
-                    if (request.getSignatureSchemeVersion() != SIGNING_BLOCK_V4) {
-                        continue;
-                    }
-                    // For incremental installs, we bypass the verifier prior to install. Now
-                    // that we know the package is valid, send a notice to the verifier with
-                    // the root hash of the base.apk.
-                    final String baseCodePath = request.getPkg().getBaseApkPath();
-                    final String[] splitCodePaths = request.getPkg().getSplitCodePaths();
-                    final Uri originUri = request.getOriginUri();
-                    final int verificationId = mPm.mPendingVerificationToken++;
-                    final String rootHashString = PackageManagerServiceUtils
-                            .buildVerificationRootHashString(baseCodePath, splitCodePaths);
-                    VerificationUtils.broadcastPackageVerified(verificationId, originUri,
-                            PackageManager.VERIFICATION_ALLOW, rootHashString,
-                            request.getDataLoaderType(), request.getUser(), mContext);
+        }
+        return true;
+    }
+
+    private void completeInstallProcess(List<InstallRequest> requests,
+            Map<String, Boolean> createdAppId, boolean success) {
+        if (success) {
+            for (InstallRequest request : requests) {
+                if (request.getDataLoaderType() != DataLoaderType.INCREMENTAL) {
+                    continue;
                 }
-            } else {
-                for (InstallRequest installRequest : requests) {
-                    if (installRequest.getParsedPackage() != null && createdAppId.getOrDefault(
-                            installRequest.getParsedPackage().getPackageName(), false)) {
-                        cleanUpAppIdCreation(installRequest);
-                    }
+                if (request.getSignatureSchemeVersion() != SIGNING_BLOCK_V4) {
+                    continue;
                 }
-                // TODO(b/194319951): create a more descriptive reason than unknown
-                // mark all non-failure installs as UNKNOWN so we do not treat them as success
-                for (InstallRequest request : requests) {
-                    request.closeFreezer();
-                    if (request.getReturnCode() == PackageManager.INSTALL_SUCCEEDED) {
-                        request.setReturnCode(PackageManager.INSTALL_UNKNOWN);
-                    }
+                // For incremental installs, we bypass the verifier prior to install. Now
+                // that we know the package is valid, send a notice to the verifier with
+                // the root hash of the base.apk.
+                final String baseCodePath = request.getPkg().getBaseApkPath();
+                final String[] splitCodePaths = request.getPkg().getSplitCodePaths();
+                final Uri originUri = request.getOriginUri();
+                final int verificationId = mPm.mPendingVerificationToken++;
+                final String rootHashString = PackageManagerServiceUtils
+                        .buildVerificationRootHashString(baseCodePath, splitCodePaths);
+                VerificationUtils.broadcastPackageVerified(verificationId, originUri,
+                        PackageManager.VERIFICATION_ALLOW, rootHashString,
+                        request.getDataLoaderType(), request.getUser(), mContext);
+            }
+        } else {
+            for (InstallRequest installRequest : requests) {
+                if (installRequest.getParsedPackage() != null && createdAppId.getOrDefault(
+                        installRequest.getParsedPackage().getPackageName(), false)) {
+                    cleanUpAppIdCreation(installRequest);
                 }
             }
-            Trace.traceEnd(TRACE_TAG_PACKAGE_MANAGER);
+            // TODO(b/194319951): create a more descriptive reason than unknown
+            // mark all non-failure installs as UNKNOWN so we do not treat them as success
+            for (InstallRequest request : requests) {
+                request.closeFreezer();
+                if (request.getReturnCode() == PackageManager.INSTALL_SUCCEEDED) {
+                    request.setReturnCode(PackageManager.INSTALL_UNKNOWN);
+                }
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index 2124ff6..ff9c3e5 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -60,6 +60,7 @@
 import android.app.ApplicationPackageManager;
 import android.app.BroadcastOptions;
 import android.app.IActivityManager;
+import android.app.admin.DevicePolicyManagerInternal;
 import android.app.admin.IDevicePolicyManager;
 import android.app.admin.SecurityLog;
 import android.app.backup.IBackupManager;
@@ -3372,8 +3373,10 @@
     // TODO(b/261957226): centralise this logic in DPM
     boolean isPackageDeviceAdmin(String packageName, int userId) {
         final IDevicePolicyManager dpm = getDevicePolicyManager();
+        final DevicePolicyManagerInternal dpmi =
+                mInjector.getLocalService(DevicePolicyManagerInternal.class);
         try {
-            if (dpm != null) {
+            if (dpm != null && dpmi != null) {
                 final ComponentName deviceOwnerComponentName = dpm.getDeviceOwnerComponent(
                         /* callingUserOnly =*/ false);
                 final String deviceOwnerPackageName = deviceOwnerComponentName == null ? null
@@ -3396,7 +3399,8 @@
                     if (dpm.packageHasActiveAdmins(packageName, users[i])) {
                         return true;
                     }
-                    if (isDeviceManagementRoleHolder(packageName, users[i])) {
+                    if (isDeviceManagementRoleHolder(packageName, users[i])
+                            && dpmi.isUserOrganizationManaged(users[i])) {
                         return true;
                     }
                 }
diff --git a/services/core/java/com/android/server/pm/Settings.java b/services/core/java/com/android/server/pm/Settings.java
index 4c9be21..1f672a0 100644
--- a/services/core/java/com/android/server/pm/Settings.java
+++ b/services/core/java/com/android/server/pm/Settings.java
@@ -5575,18 +5575,18 @@
                     }
                     if (!paths.getOverlayPaths().isEmpty()) {
                         pw.print(prefix);
-                        pw.println("    ");
+                        pw.print("    ");
                         pw.print(libOverlayPaths.getKey());
                         pw.println(" overlay paths:");
                         for (String path : paths.getOverlayPaths()) {
                             pw.print(prefix);
-                            pw.print("        ");
+                            pw.print("      ");
                             pw.println(path);
                         }
                     }
                     if (!paths.getResourceDirs().isEmpty()) {
                         pw.print(prefix);
-                        pw.println("      ");
+                        pw.print("    ");
                         pw.print(libOverlayPaths.getKey());
                         pw.println(" legacy overlay paths:");
                         for (String path : paths.getResourceDirs()) {
diff --git a/services/core/java/com/android/server/pm/ShortcutService.java b/services/core/java/com/android/server/pm/ShortcutService.java
index 25468fa..a3ff195 100644
--- a/services/core/java/com/android/server/pm/ShortcutService.java
+++ b/services/core/java/com/android/server/pm/ShortcutService.java
@@ -588,7 +588,7 @@
 
     void handleOnDefaultLauncherChanged(int userId) {
         if (DEBUG) {
-            Slog.v(TAG, "Default launcher changed for user: " + userId);
+            Slog.v(TAG, "Default launcher changed for userId=" + userId);
         }
 
         // Default launcher is removed or changed, revoke all URI permissions.
@@ -712,7 +712,7 @@
     /** lifecycle event */
     void handleUnlockUser(int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "handleUnlockUser: user=" + userId);
+            Slog.d(TAG, "handleUnlockUser: userId=" + userId);
         }
         synchronized (mUnlockedUsers) {
             mUnlockedUsers.put(userId, true);
@@ -739,7 +739,7 @@
     /** lifecycle event */
     void handleStopUser(int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "handleStopUser: user=" + userId);
+            Slog.d(TAG, "handleStopUser: userId=" + userId);
         }
         Trace.traceBegin(Trace.TRACE_TAG_SYSTEM_SERVER, "shortcutHandleStopUser");
         synchronized (mServiceLock) {
@@ -755,7 +755,7 @@
     @GuardedBy("mServiceLock")
     private void unloadUserLocked(int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "unloadUserLocked: user=" + userId);
+            Slog.d(TAG, "unloadUserLocked: userId=" + userId);
         }
         // Cancel any ongoing background tasks.
         getUserShortcutsLocked(userId).cancelAllInFlightTasks();
@@ -1221,7 +1221,7 @@
 
     private void scheduleSaveInner(@UserIdInt int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "Scheduling to save for " + userId);
+            Slog.d(TAG, "Scheduling to save for userId=" + userId);
         }
         synchronized (mServiceLock) {
             if (!mDirtyUserIds.contains(userId)) {
@@ -1333,7 +1333,7 @@
     // Requires mServiceLock held, but "Locked" prefix would look weird so we just say "L".
     void throwIfUserLockedL(@UserIdInt int userId) {
         if (!isUserUnlockedL(userId)) {
-            throw new IllegalStateException("User " + userId + " is locked or not running");
+            throw new IllegalStateException("User (with userId=" + userId + ") is locked or not running");
         }
     }
 
@@ -1720,7 +1720,7 @@
 
         // Otherwise, make sure the arguments are valid.
         if (UserHandle.getUserId(callingUid) != userId) {
-            throw new SecurityException("Invalid user-ID");
+            throw new SecurityException("Invalid userId");
         }
     }
 
@@ -1735,7 +1735,7 @@
 
         // Otherwise, make sure the arguments are valid.
         if (UserHandle.getUserId(callingUid) != userId) {
-            throw new SecurityException("Invalid user-ID");
+            throw new SecurityException("Invalid userId");
         }
         if (injectGetPackageUid(packageName, userId) != callingUid) {
             throw new SecurityException("Calling package name mismatch");
@@ -1755,7 +1755,7 @@
         }
         final int callingUid = injectBinderCallingUid();
         if (UserHandle.getUserId(callingUid) != si.getUserId()) {
-            throw new SecurityException("User-ID in shortcut doesn't match the caller");
+            throw new SecurityException("UserId in shortcut doesn't match the caller");
         }
     }
 
@@ -1822,7 +1822,7 @@
         final int userId = sp.getPackageUserId();
         if (DEBUG) {
             Slog.d(TAG, String.format(
-                    "Shortcut changes: package=%s, user=%d", packageName, userId));
+                    "Shortcut changes: package=%s, userId=%d", packageName, userId));
         }
         injectPostToHandlerDebounced(sp, notifyListenerRunnable(packageName, userId));
         notifyShortcutChangeCallbacks(packageName, userId, changedShortcuts, removedShortcuts);
@@ -1832,7 +1832,7 @@
     private void notifyListeners(@NonNull final String packageName, @UserIdInt final int userId) {
         if (DEBUG) {
             Slog.d(TAG, String.format(
-                    "Shortcut changes: package=%s, user=%d", packageName, userId));
+                    "Shortcut changes: package=%s, userId=%d", packageName, userId));
         }
         injectPostToHandler(notifyListenerRunnable(packageName, userId));
     }
@@ -2736,14 +2736,14 @@
     void resetThrottlingInner(@UserIdInt int userId) {
         synchronized (mServiceLock) {
             if (!isUserUnlockedL(userId)) {
-                Log.w(TAG, "User " + userId + " is locked or not running");
+                Log.w(TAG, "User (with userId=" + userId + ") is locked or not running");
                 return;
             }
 
             getUserShortcutsLocked(userId).resetThrottling();
         }
         scheduleSaveUser(userId);
-        Slog.i(TAG, "ShortcutManager: throttling counter reset for user " + userId);
+        Slog.i(TAG, "ShortcutManager: throttling counter reset for userId=" + userId);
     }
 
     void resetAllThrottlingInner() {
@@ -2755,7 +2755,7 @@
     @Override
     public void onApplicationActive(String packageName, int userId) {
         if (DEBUG) {
-            Slog.d(TAG, "onApplicationActive: package=" + packageName + "  userid=" + userId);
+            Slog.d(TAG, "onApplicationActive: package=" + packageName + "  userId=" + userId);
         }
         enforceResetThrottlingPermission();
         synchronized (mServiceLock) {
@@ -2822,7 +2822,7 @@
 
             if (defaultLauncher != null) {
                 if (DEBUG) {
-                    Slog.v(TAG, "Detected launcher: " + defaultLauncher + " user: " + userId);
+                    Slog.v(TAG, "Detected launcher: " + defaultLauncher + " userId=" + userId);
                 }
                 return defaultLauncher.equals(packageName);
             } else {
@@ -2875,11 +2875,11 @@
                 if (defaultLauncher != null) {
                     if (DEBUG) {
                         Slog.v(TAG, "Default launcher from RoleManager: " + defaultLauncher
-                                + " user: " + userId);
+                                + " userId=" + userId);
                     }
                     user.setCachedLauncher(defaultLauncher);
                 } else {
-                    Slog.e(TAG, "Default launcher not found." + " user: " + userId);
+                    Slog.e(TAG, "Default launcher not found." + " userId=" + userId);
                 }
 
                 return defaultLauncher;
@@ -2974,7 +2974,7 @@
                 int queryFlags, int userId, int callingPid, int callingUid) {
             if (DEBUG_REBOOT) {
                 Slog.d(TAG, "Getting shortcuts for launcher= " + callingPackage
-                        + "user=" + userId + " pkg=" + packageName);
+                        + "userId=" + userId + " pkg=" + packageName);
             }
             final ArrayList<ShortcutInfo> ret = new ArrayList<>();
 
@@ -3793,7 +3793,7 @@
                     if (!isUserUnlockedL(userId)) {
                         if (DEBUG) {
                             Slog.d(TAG, "Ignoring package broadcast " + action
-                                    + " for locked/stopped user " + userId);
+                                    + " for locked/stopped userId=" + userId);
                         }
                         return;
                     }
@@ -3814,10 +3814,10 @@
                 switch (action) {
                     case Intent.ACTION_PACKAGE_ADDED:
                         if (replacing) {
-                            Slog.d(TAG, "replacing package: " + packageName + " userId" + userId);
+                            Slog.d(TAG, "replacing package: " + packageName + " userId=" + userId);
                             handlePackageUpdateFinished(packageName, userId);
                         } else {
-                            Slog.d(TAG, "adding package: " + packageName + " userId" + userId);
+                            Slog.d(TAG, "adding package: " + packageName + " userId=" + userId);
                             handlePackageAdded(packageName, userId);
                         }
                         break;
@@ -3825,21 +3825,21 @@
                         if (!replacing || (replacing && archival)) {
                             if (!replacing) {
                                 Slog.d(TAG, "removing package: "
-                                        + packageName + " userId" + userId);
+                                        + packageName + " userId=" + userId);
                             } else if (archival) {
                                 Slog.d(TAG, "archiving package: "
-                                        + packageName + " userId" + userId);
+                                        + packageName + " userId=" + userId);
                             }
                             handlePackageRemoved(packageName, userId);
                         }
                         break;
                     case Intent.ACTION_PACKAGE_CHANGED:
-                        Slog.d(TAG, "changing package: " + packageName + " userId" + userId);
+                        Slog.d(TAG, "changing package: " + packageName + " userId=" + userId);
                         handlePackageChanged(packageName, userId);
                         break;
                     case Intent.ACTION_PACKAGE_DATA_CLEARED:
                         Slog.d(TAG, "clearing data for package: "
-                                + packageName + " userId" + userId);
+                                + packageName + " userId=" + userId);
                         handlePackageDataCleared(packageName, userId);
                         break;
                 }
@@ -3902,7 +3902,7 @@
                     if (!isPackageInstalled(spi.getPackageName(), spi.getPackageUserId())) {
                         if (DEBUG) {
                             Slog.d(TAG, "Uninstalled: " + spi.getPackageName()
-                                    + " user " + spi.getPackageUserId());
+                                    + " userId=" + spi.getPackageUserId());
                         }
                         gonePackages.add(
                                 UserPackage.of(spi.getPackageUserId(), spi.getPackageName()));
@@ -3927,7 +3927,7 @@
     @GuardedBy("mServiceLock")
     private void rescanUpdatedPackagesLocked(@UserIdInt int userId, long lastScanTime) {
         if (DEBUG_REBOOT) {
-            Slog.d(TAG, "rescan updated package user=" + userId + " last scanned=" + lastScanTime);
+            Slog.d(TAG, "rescan updated package userId=" + userId + " last scanned=" + lastScanTime);
         }
         final ShortcutUser user = getUserShortcutsLocked(userId);
 
@@ -3953,7 +3953,7 @@
 
     private void handlePackageAdded(String packageName, @UserIdInt int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, String.format("handlePackageAdded: %s user=%d", packageName, userId));
+            Slog.d(TAG, String.format("handlePackageAdded: %s userId=%d", packageName, userId));
         }
         synchronized (mServiceLock) {
             final ShortcutUser user = getUserShortcutsLocked(userId);
@@ -3965,7 +3965,7 @@
 
     private void handlePackageUpdateFinished(String packageName, @UserIdInt int userId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, String.format("handlePackageUpdateFinished: %s user=%d",
+            Slog.d(TAG, String.format("handlePackageUpdateFinished: %s userId=%d",
                     packageName, userId));
         }
         synchronized (mServiceLock) {
@@ -3981,7 +3981,7 @@
 
     private void handlePackageRemoved(String packageName, @UserIdInt int packageUserId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, String.format("handlePackageRemoved: %s user=%d", packageName,
+            Slog.d(TAG, String.format("handlePackageRemoved: %s userId=%d", packageName,
                     packageUserId));
         }
         cleanUpPackageForAllLoadedUsers(packageName, packageUserId, /* appStillExists = */ false);
@@ -3991,7 +3991,7 @@
 
     private void handlePackageDataCleared(String packageName, int packageUserId) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, String.format("handlePackageDataCleared: %s user=%d", packageName,
+            Slog.d(TAG, String.format("handlePackageDataCleared: %s userId=%d", packageName,
                     packageUserId));
         }
         cleanUpPackageForAllLoadedUsers(packageName, packageUserId, /* appStillExists = */ true);
@@ -4006,7 +4006,7 @@
             return;
         }
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, String.format("handlePackageChanged: %s user=%d", packageName,
+            Slog.d(TAG, String.format("handlePackageChanged: %s userId=%d", packageName,
                     packageUserId));
         }
 
@@ -4193,7 +4193,7 @@
     private void forUpdatedPackages(@UserIdInt int userId, long lastScanTime, boolean afterOta,
             Consumer<ApplicationInfo> callback) {
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "forUpdatedPackages for user " + userId + ", lastScanTime=" + lastScanTime
+            Slog.d(TAG, "forUpdatedPackages for userId=" + userId + ", lastScanTime=" + lastScanTime
                     + " afterOta=" + afterOta);
         }
         final List<PackageInfo> list = getInstalledPackages(userId);
@@ -4302,7 +4302,7 @@
             return mContext.createContextAsUser(UserHandle.of(userId), /* flags */ 0)
                     .getPackageManager().getResourcesForApplication(packageName);
         } catch (NameNotFoundException e) {
-            Slog.e(TAG, "Resources of package " + packageName + " for user " + userId
+            Slog.e(TAG, "Resources of package " + packageName + " for userId=" + userId
                     + " not found");
             return null;
         } finally {
@@ -4512,17 +4512,17 @@
     public byte[] getBackupPayload(@UserIdInt int userId) {
         enforceSystem();
         if (DEBUG) {
-            Slog.d(TAG, "Backing up user " + userId);
+            Slog.d(TAG, "Backing up user with userId=" + userId);
         }
         synchronized (mServiceLock) {
             if (!isUserUnlockedL(userId)) {
-                wtf("Can't backup: user " + userId + " is locked or not running");
+                wtf("Can't backup: userId=" + userId + " is locked or not running");
                 return null;
             }
 
             final ShortcutUser user = getUserShortcutsLocked(userId);
             if (user == null) {
-                wtf("Can't backup: user not found: id=" + userId);
+                wtf("Can't backup: user not found: userId=" + userId);
                 return null;
             }
 
@@ -4562,11 +4562,11 @@
     public void applyRestore(byte[] payload, @UserIdInt int userId) {
         enforceSystem();
         if (DEBUG || DEBUG_REBOOT) {
-            Slog.d(TAG, "Restoring user " + userId);
+            Slog.d(TAG, "Restoring user with userId=" + userId);
         }
         synchronized (mServiceLock) {
             if (!isUserUnlockedL(userId)) {
-                wtf("Can't restore: user " + userId + " is locked or not running");
+                wtf("Can't restore: user (with userId=" + userId + ") is locked or not running");
                 return;
             }
             // Note we print the file timestamps in dumpsys too, but also printing the timestamp
@@ -4989,7 +4989,7 @@
                             mUserId = UserHandle.parseUserArg(getNextArgRequired());
                             if (!isUserUnlockedL(mUserId)) {
                                 throw new CommandException(
-                                        "User " + mUserId + " is not running or locked");
+                                        "User (with userId=" + mUserId + ") is not running or locked");
                             }
                             break;
                         }
@@ -5094,7 +5094,7 @@
             synchronized (mServiceLock) {
                 parseOptionsLocked(/* takeUser =*/ true);
 
-                Slog.i(TAG, "cmd: handleResetThrottling: user=" + mUserId);
+                Slog.i(TAG, "cmd: handleResetThrottling: userId=" + mUserId);
 
                 resetThrottlingInner(mUserId);
             }
@@ -5136,7 +5136,7 @@
                 final String defaultLauncher = getDefaultLauncher(mUserId);
                 if (defaultLauncher == null) {
                     throw new CommandException(
-                            "Failed to get the default launcher for user " + mUserId);
+                            "Failed to get the default launcher for userId=" + mUserId);
                 }
 
                 // Get the class name of the component from PM to keep the old behaviour.
@@ -5157,7 +5157,7 @@
             synchronized (mServiceLock) {
                 parseOptionsLocked(/* takeUser =*/ true);
 
-                Slog.i(TAG, "cmd: handleUnloadUser: user=" + mUserId);
+                Slog.i(TAG, "cmd: handleUnloadUser: userId=" + mUserId);
 
                 ShortcutService.this.handleStopUser(mUserId);
             }
@@ -5168,7 +5168,7 @@
                 parseOptionsLocked(/* takeUser =*/ true);
                 final String packageName = getNextArgRequired();
 
-                Slog.i(TAG, "cmd: handleClearShortcuts: user" + mUserId + ", " + packageName);
+                Slog.i(TAG, "cmd: handleClearShortcuts: userId=" + mUserId + ", " + packageName);
 
                 ShortcutService.this.cleanUpPackageForAllLoadedUsers(packageName, mUserId,
                         /* appStillExists = */ true);
@@ -5180,7 +5180,7 @@
                 parseOptionsLocked(/* takeUser =*/ true);
                 final String packageName = getNextArgRequired();
 
-                Slog.i(TAG, "cmd: handleGetShortcuts: user=" + mUserId + ", flags="
+                Slog.i(TAG, "cmd: handleGetShortcuts: userId=" + mUserId + ", flags="
                         + mShortcutMatchFlags + ", package=" + packageName);
 
                 final ShortcutUser user = ShortcutService.this.getUserShortcutsLocked(mUserId);
diff --git a/services/core/java/com/android/server/pm/UserManagerService.java b/services/core/java/com/android/server/pm/UserManagerService.java
index a683a8c..13901c1 100644
--- a/services/core/java/com/android/server/pm/UserManagerService.java
+++ b/services/core/java/com/android/server/pm/UserManagerService.java
@@ -1194,11 +1194,11 @@
             // Avoid marking pre-created users for removal.
             return;
         }
-        if (ui.lastLoggedInTime == 0 && ui.isGuest() && Resources.getSystem().getBoolean(
-                com.android.internal.R.bool.config_guestUserAutoCreated)) {
-            // Avoid marking auto-created but not-yet-logged-in guest user for removal. Because a
-            // new one will be created anyway, and this one doesn't have any personal data in it yet
-            // due to not being logged in.
+        if (ui.lastLoggedInTime == 0) {
+            // Avoid marking a not-yet-logged-in ephemeral user for removal, since it doesn't have
+            // any personal data in it yet due to not being logged in.
+            // This will also avoid marking an auto-created not-yet-logged-in ephemeral guest user
+            // for removal, which would be recreated again later in the boot anyway.
             return;
         }
         // Mark the user for removal.
diff --git a/services/core/java/com/android/server/pm/UserRestrictionsUtils.java b/services/core/java/com/android/server/pm/UserRestrictionsUtils.java
index 95e5b84..2bc6d53 100644
--- a/services/core/java/com/android/server/pm/UserRestrictionsUtils.java
+++ b/services/core/java/com/android/server/pm/UserRestrictionsUtils.java
@@ -503,26 +503,21 @@
             String restriction,
             boolean isMainUser,
             boolean isProfileOwnerOnOrgOwnedDevice) {
-        if (android.app.admin.flags.Flags.esimManagementEnabled()) {
-            if (IMMUTABLE_BY_OWNERS.contains(restriction)) {
-                return false;
-            }
-            if (DEVICE_OWNER_ONLY_RESTRICTIONS.contains(restriction)) {
-                return false;
-            }
-            if (!isMainUser && MAIN_USER_ONLY_RESTRICTIONS.contains(restriction)) {
-                return false;
-            }
-            if (!isProfileOwnerOnOrgOwnedDevice
-                    && PROFILE_OWNER_ORGANIZATION_OWNED_PROFILE_RESTRICTIONS.contains(
-                            restriction)) {
-                return false;
-            }
-            return true;
+        if (IMMUTABLE_BY_OWNERS.contains(restriction)) {
+            return false;
         }
-        return !IMMUTABLE_BY_OWNERS.contains(restriction)
-                && !DEVICE_OWNER_ONLY_RESTRICTIONS.contains(restriction)
-                && !(!isMainUser && MAIN_USER_ONLY_RESTRICTIONS.contains(restriction));
+        if (DEVICE_OWNER_ONLY_RESTRICTIONS.contains(restriction)) {
+            return false;
+        }
+        if (!isMainUser && MAIN_USER_ONLY_RESTRICTIONS.contains(restriction)) {
+            return false;
+        }
+        if (!isProfileOwnerOnOrgOwnedDevice
+                && PROFILE_OWNER_ORGANIZATION_OWNED_PROFILE_RESTRICTIONS.contains(
+                restriction)) {
+            return false;
+        }
+        return true;
     }
 
     /**
diff --git a/services/core/java/com/android/server/pm/permission/AccessCheckDelegate.java b/services/core/java/com/android/server/pm/permission/AccessCheckDelegate.java
index f518769..e9cb279 100644
--- a/services/core/java/com/android/server/pm/permission/AccessCheckDelegate.java
+++ b/services/core/java/com/android/server/pm/permission/AccessCheckDelegate.java
@@ -375,24 +375,34 @@
                 @Nullable String message, boolean shouldCollectMessage, boolean skiProxyOperation,
                 @NonNull HexFunction<Integer, AttributionSource, Boolean, String, Boolean,
                         Boolean, SyncNotedAppOp> superImpl) {
-            if (attributionSource.getUid() == mDelegateAndOwnerUid && isDelegateOp(code)) {
-                final int shellUid = UserHandle.getUid(
-                        UserHandle.getUserId(attributionSource.getUid()), Process.SHELL_UID);
-                final long identity = Binder.clearCallingIdentity();
-                try {
-                    return superImpl.apply(code,
-                            new AttributionSource(shellUid, Process.INVALID_PID, SHELL_PKG,
-                                    attributionSource.getAttributionTag(),
-                                    attributionSource.getToken(), /*renouncedPermissions*/ null,
-                                    attributionSource.getDeviceId(), attributionSource.getNext()),
-                            shouldCollectAsyncNotedOp, message, shouldCollectMessage,
-                            skiProxyOperation);
-                } finally {
-                    Binder.restoreCallingIdentity(identity);
-                }
+            if (!isDelegateOp(code)) {
+                return superImpl.apply(code, attributionSource, shouldCollectAsyncNotedOp,
+                        message, shouldCollectMessage, skiProxyOperation);
             }
-            return superImpl.apply(code, attributionSource, shouldCollectAsyncNotedOp,
-                    message, shouldCollectMessage, skiProxyOperation);
+
+            final int shellUid = UserHandle.getUid(
+                    UserHandle.getUserId(attributionSource.getUid()), Process.SHELL_UID);
+            AttributionSource next = attributionSource.getNext();
+            if (next != null && next.getUid() == mDelegateAndOwnerUid) {
+                next = new AttributionSource(shellUid, Process.INVALID_PID, SHELL_PKG,
+                        next.getAttributionTag(), next.getToken(), /*renouncedPermissions*/ null,
+                        next.getDeviceId(), next.getNext());
+                attributionSource = new AttributionSource(attributionSource, next);
+            }
+            if (attributionSource.getUid() == mDelegateAndOwnerUid) {
+                attributionSource = new AttributionSource(shellUid, Process.INVALID_PID, SHELL_PKG,
+                        attributionSource.getAttributionTag(),
+                        attributionSource.getToken(), /*renouncedPermissions*/ null,
+                        attributionSource.getDeviceId(), attributionSource.getNext());
+            }
+            final long identity = Binder.clearCallingIdentity();
+            try {
+                return superImpl.apply(code, attributionSource,
+                        shouldCollectAsyncNotedOp, message, shouldCollectMessage,
+                        skiProxyOperation);
+            } finally {
+                Binder.restoreCallingIdentity(identity);
+            }
         }
 
         @Override
diff --git a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
index 46e6546..df9f7fb 100644
--- a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
+++ b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
@@ -26,7 +26,6 @@
 import static android.app.AppOpsManager.MODE_ALLOWED;
 import static android.app.AppOpsManager.MODE_ERRORED;
 import static android.app.AppOpsManager.MODE_IGNORED;
-import static android.app.AppOpsManager.OP_BLUETOOTH_CONNECT;
 import static android.content.pm.ApplicationInfo.AUTO_REVOKE_DISALLOWED;
 import static android.content.pm.ApplicationInfo.AUTO_REVOKE_DISCOURAGED;
 import static android.permission.flags.Flags.serverSideAttributionRegistration;
@@ -1640,22 +1639,7 @@
                         throw new SecurityException(msg + ":" + e.getMessage());
                     }
                 }
-                int result = Math.max(checkedOpResult, notedOpResult);
-                // TODO(b/302609140): Remove extra logging after this issue is diagnosed.
-                if (op == OP_BLUETOOTH_CONNECT && result == MODE_ERRORED) {
-                    if (result == checkedOpResult) {
-                        Slog.e(LOG_TAG, "BLUETOOTH_CONNECT permission hard denied as"
-                                + " checkOp for resolvedAttributionSource "
-                                + resolvedAttributionSource + " and op " + op
-                                + " returned MODE_ERRORED");
-                    } else {
-                        Slog.e(LOG_TAG, "BLUETOOTH_CONNECT permission hard denied as"
-                                + " noteOp for resolvedAttributionSource "
-                                + resolvedAttributionSource + " and op " + notedOp
-                                + " returned MODE_ERRORED");
-                    }
-                }
-                return result;
+                return Math.max(checkedOpResult, notedOpResult);
             }
         }
 
diff --git a/services/core/java/com/android/server/pm/permission/TEST_MAPPING b/services/core/java/com/android/server/pm/permission/TEST_MAPPING
index 24323c8..8a3c74b 100644
--- a/services/core/java/com/android/server/pm/permission/TEST_MAPPING
+++ b/services/core/java/com/android/server/pm/permission/TEST_MAPPING
@@ -1,24 +1,7 @@
 {
     "presubmit": [
         {
-            "name": "CtsPermissionTestCases",
-            "options": [
-                {
-                    "exclude-annotation": "androidx.test.filters.FlakyTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.BackgroundPermissionsTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.SplitPermissionTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.PermissionFlagsTest"
-                },
-                {
-                    "include-filter": "android.permission.cts.SharedUidPermissionsTest"
-                }
-            ]
+            "name": "CtsPermissionTestCases_Platform"
         },
         {
             "name": "CtsAppSecurityHostTestCases",
diff --git a/services/core/java/com/android/server/policy/ModifierShortcutManager.java b/services/core/java/com/android/server/policy/ModifierShortcutManager.java
index 7ed8972..027e69c 100644
--- a/services/core/java/com/android/server/policy/ModifierShortcutManager.java
+++ b/services/core/java/com/android/server/policy/ModifierShortcutManager.java
@@ -17,7 +17,9 @@
 package com.android.server.policy;
 
 import static com.android.server.flags.Flags.modifierShortcutManagerMultiuser;
+import static com.android.hardware.input.Flags.modifierShortcutManagerRefactor;
 
+import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SuppressLint;
 import android.app.role.RoleManager;
@@ -37,6 +39,7 @@
 import android.util.IndentingPrintWriter;
 import android.util.Log;
 import android.util.LongSparseArray;
+import android.util.Pair;
 import android.util.Slog;
 import android.util.SparseArray;
 import android.view.KeyCharacterMap;
@@ -81,8 +84,8 @@
     private static final String ATTRIBUTE_SHIFT = "shift";
     private static final String ATTRIBUTE_ROLE = "role";
 
-    private final SparseArray<Intent> mIntentShortcuts = new SparseArray<>();
-    private final SparseArray<Intent> mShiftShortcuts = new SparseArray<>();
+    private final SparseArray<Intent> mCategoryShortcuts = new SparseArray<>();
+    private final SparseArray<Intent> mShiftCategoryShortcuts = new SparseArray<>();
     private final SparseArray<String> mRoleShortcuts = new SparseArray<String>();
     private final SparseArray<String> mShiftRoleShortcuts = new SparseArray<String>();
     private final Map<String, Intent> mRoleIntents = new HashMap<String, Intent>();
@@ -127,6 +130,7 @@
     private boolean mSearchKeyShortcutPending = false;
     private boolean mConsumeSearchKeyUp = true;
     private UserHandle mCurrentUser;
+    private final Map<Pair<Character, Boolean>, Bookmark> mBookmarks = new HashMap<>();
 
     ModifierShortcutManager(Context context, Handler handler, UserHandle currentUser) {
         mContext = context;
@@ -134,7 +138,14 @@
         RoleManager rm = mContext.getSystemService(RoleManager.class);
         rm.addOnRoleHoldersChangedListenerAsUser(mContext.getMainExecutor(),
                 (String roleName, UserHandle user) -> {
-                    mRoleIntents.remove(roleName);
+                    if (modifierShortcutManagerRefactor()) {
+                        mBookmarks.values().stream().filter(b ->
+                                b instanceof RoleBookmark
+                                && ((RoleBookmark) b).getRole().equals(roleName))
+                                .forEach(Bookmark::clearIntent);
+                    } else {
+                        mRoleIntents.remove(roleName);
+                    }
                 }, UserHandle.ALL);
         mCurrentUser = currentUser;
         mInputManagerInternal = LocalServices.getService(InputManagerInternal.class);
@@ -146,8 +157,26 @@
 
         // Role based shortcuts may resolve to different apps for different users
         // so clear the cache.
-        mRoleIntents.clear();
-        mComponentIntents.clear();
+        clearRoleIntents();
+        clearComponentIntents();
+    }
+
+    void clearRoleIntents() {
+        if (modifierShortcutManagerRefactor()) {
+            mBookmarks.values().stream().filter(b ->
+                    b instanceof RoleBookmark).forEach(Bookmark::clearIntent);
+        } else {
+            mRoleIntents.clear();
+        }
+    }
+
+    void clearComponentIntents() {
+        if (modifierShortcutManagerRefactor()) {
+            mBookmarks.values().stream().filter(b ->
+                    b instanceof ComponentBookmark).forEach(Bookmark::clearIntent);
+        } else {
+            mComponentIntents.clear();
+        }
     }
 
     /**
@@ -176,77 +205,111 @@
 
         Intent shortcutIntent = null;
 
-        // If the Shift key is pressed, then search for the shift shortcuts.
-        SparseArray<Intent> shortcutMap = isShiftOn ? mShiftShortcuts : mIntentShortcuts;
-
         // First try the exact keycode (with modifiers).
         int shortcutChar = kcm.get(keyCode, metaState);
         if (shortcutChar == 0) {
             return null;
         }
-        shortcutIntent = shortcutMap.get(shortcutChar);
 
-        if (shortcutIntent == null) {
-            // Next try the primary character on that key.
-            shortcutChar = Character.toLowerCase(kcm.getDisplayLabel(keyCode));
-            if (shortcutChar == 0) {
-                return null;
+        if (modifierShortcutManagerRefactor()) {
+            Bookmark bookmark = mBookmarks.get(new Pair<>((char) shortcutChar, isShiftOn));
+            if (bookmark == null) {
+                // Next try the primary character on that key.
+                shortcutChar = Character.toLowerCase(kcm.getDisplayLabel(keyCode));
+                if (shortcutChar == 0) {
+                    return null;
+                }
+                bookmark = mBookmarks.get(new Pair<>((char) shortcutChar, isShiftOn));
             }
+
+            if (bookmark != null) {
+                Context context = modifierShortcutManagerMultiuser()
+                        ? mContext.createContextAsUser(mCurrentUser, 0) : mContext;
+                shortcutIntent = bookmark.getIntent(context);
+            } else {
+                Log.d(TAG, "No bookmark found for "
+                        + (isShiftOn ? "SHIFT+" : "") + (char) shortcutChar);
+            }
+        } else {
+            // If the Shift key is pressed, then search for the shift shortcuts.
+            SparseArray<Intent> shortcutMap = isShiftOn
+                    ? mShiftCategoryShortcuts : mCategoryShortcuts;
             shortcutIntent = shortcutMap.get(shortcutChar);
-        }
 
-        if (shortcutIntent == null) {
-            // Next check for role based shortcut with primary character.
-            String role = isShiftOn ? mShiftRoleShortcuts.get(shortcutChar)
-                    : mRoleShortcuts.get(shortcutChar);
-            if (role != null) {
-                shortcutIntent = getRoleLaunchIntent(role);
-            }
-        }
-
-        if (modifierShortcutManagerMultiuser()) {
             if (shortcutIntent == null) {
-                // Next check component based shortcuts with primary character.
-                ComponentName component = isShiftOn
-                        ? mShiftComponentShortcuts.get(shortcutChar)
-                        : mComponentShortcuts.get(shortcutChar);
-                if (component != null) {
-                    shortcutIntent = resolveComponentNameIntent(component);
+                // Next try the primary character on that key.
+                shortcutChar = Character.toLowerCase(kcm.getDisplayLabel(keyCode));
+                if (shortcutChar == 0) {
+                    return null;
+                }
+                shortcutIntent = shortcutMap.get(shortcutChar);
+            }
+
+            if (shortcutIntent == null) {
+                // Next check for role based shortcut with primary character.
+                String role = isShiftOn ? mShiftRoleShortcuts.get(shortcutChar)
+                        : mRoleShortcuts.get(shortcutChar);
+                if (role != null) {
+                    shortcutIntent = getRoleLaunchIntent(role);
+                }
+            }
+
+            if (modifierShortcutManagerMultiuser()) {
+                if (shortcutIntent == null) {
+                    // Next check component based shortcuts with primary character.
+                    ComponentName component = isShiftOn
+                            ? mShiftComponentShortcuts.get(shortcutChar)
+                            : mComponentShortcuts.get(shortcutChar);
+                    if (component != null) {
+                        shortcutIntent = resolveComponentNameIntent(component);
+                    }
                 }
             }
         }
         return shortcutIntent;
     }
 
+    @Nullable
+    private static Intent getRoleLaunchIntent(Context context, String role) {
+        Intent intent = null;
+        RoleManager rm = context.getSystemService(RoleManager.class);
+        PackageManager pm = context.getPackageManager();
+        if (rm.isRoleAvailable(role)) {
+            String rolePackage = rm.getDefaultApplication(role);
+            if (rolePackage != null) {
+                intent = pm.getLaunchIntentForPackage(rolePackage);
+                if (intent != null) {
+                    intent.putExtra(EXTRA_ROLE, role);
+
+                } else {
+                    Log.w(TAG, "No launch intent for role " + role);
+                }
+            } else {
+                Log.w(TAG, "No default application for role "
+                        + role + " user=" + context.getUser());
+            }
+        } else {
+            Log.w(TAG, "Role " + role + " is not available.");
+        }
+        return intent;
+    }
+
+    @Nullable
     private Intent getRoleLaunchIntent(String role) {
         Intent intent = mRoleIntents.get(role);
         if (intent == null) {
             Context context = modifierShortcutManagerMultiuser()
                     ? mContext.createContextAsUser(mCurrentUser, 0) : mContext;
-            RoleManager rm = context.getSystemService(RoleManager.class);
-            PackageManager pm = context.getPackageManager();
-            if (rm.isRoleAvailable(role)) {
-                String rolePackage = rm.getDefaultApplication(role);
-                if (rolePackage != null) {
-                    intent = pm.getLaunchIntentForPackage(rolePackage);
-                    if (intent != null) {
-                        intent.putExtra(EXTRA_ROLE, role);
-                        mRoleIntents.put(role, intent);
-                    } else {
-                        Log.w(TAG, "No launch intent for role " + role);
-                    }
-                } else {
-                    Log.w(TAG, "No default application for role " + role);
-                }
-            } else {
-                Log.w(TAG, "Role " + role + " is not available.");
+            intent = getRoleLaunchIntent(context, role);
+            if (intent != null) {
+                mRoleIntents.put(role, intent);
             }
         }
+
         return intent;
     }
 
     private void loadShortcuts() {
-
         try {
             XmlResourceParser parser = mContext.getResources().getXml(R.xml.bookmarks);
             XmlUtils.beginDocument(parser, TAG_BOOKMARKS);
@@ -276,57 +339,84 @@
                     continue;
                 }
 
-                final int shortcutChar = shortcutName.charAt(0);
                 final boolean isShiftShortcut = (shiftName != null && shiftName.equals("true"));
-                final Intent intent;
-                if (packageName != null && className != null) {
-                    if (roleName != null || categoryName != null) {
-                        Log.w(TAG, "Cannot specify role or category when package and class"
-                                + " are present for bookmark packageName=" + packageName
-                                + " className=" + className + " shortcutChar=" + shortcutChar);
-                        continue;
+
+                if (modifierShortcutManagerRefactor()) {
+                    final char shortcutChar = shortcutName.charAt(0);
+                    Bookmark bookmark = null;
+                    if (packageName != null && className != null) {
+                        bookmark = new ComponentBookmark(
+                                shortcutChar, isShiftShortcut, packageName, className);
+                    } else if (categoryName != null) {
+                        bookmark = new CategoryBookmark(
+                                shortcutChar, isShiftShortcut, categoryName);
+                    } else if (roleName != null) {
+                        bookmark = new RoleBookmark(shortcutChar, isShiftShortcut, roleName);
                     }
-                    if (modifierShortcutManagerMultiuser()) {
-                        ComponentName componentName = new ComponentName(packageName, className);
-                        if (isShiftShortcut) {
-                            mShiftComponentShortcuts.put(shortcutChar, componentName);
+                    if (bookmark != null) {
+                        Log.d(TAG, "adding shortcut " + bookmark + "shift="
+                                + isShiftShortcut + " char=" + shortcutChar);
+                        mBookmarks.put(new Pair<>(shortcutChar, isShiftShortcut), bookmark);
+                    }
+                } else {
+                    final int shortcutChar = shortcutName.charAt(0);
+                    if (packageName != null && className != null) {
+                        if (roleName != null || categoryName != null) {
+                            Log.w(TAG, "Cannot specify role or category when package and class"
+                                    + " are present for bookmark packageName=" + packageName
+                                    + " className=" + className + " shortcutChar=" + shortcutChar);
+                            continue;
+                        }
+                        if (modifierShortcutManagerMultiuser()) {
+                            ComponentName componentName =
+                                    new ComponentName(packageName, className);
+                            if (isShiftShortcut) {
+                                mShiftComponentShortcuts.put(shortcutChar, componentName);
+                            } else {
+                                mComponentShortcuts.put(shortcutChar, componentName);
+                            }
                         } else {
-                            mComponentShortcuts.put(shortcutChar, componentName);
+                            Intent intent = resolveComponentNameIntent(packageName, className);
+                            if (isShiftShortcut) {
+                                mShiftCategoryShortcuts.put(shortcutChar, intent);
+                            } else {
+                                mCategoryShortcuts.put(shortcutChar, intent);
+                            }
+                        }
+                        continue;
+                    } else if (categoryName != null) {
+                        if (roleName != null) {
+                            Log.w(TAG, "Cannot specify role bookmark when category is present for"
+                                    + " bookmark shortcutChar=" + shortcutChar
+                                    + " category= " + categoryName);
+                            continue;
+                        }
+                        Intent intent = Intent.makeMainSelectorActivity(
+                                Intent.ACTION_MAIN, categoryName);
+                        if (intent == null) {
+                            Log.w(TAG, "Null selector intent for " + categoryName);
+                        } else {
+                            if (isShiftShortcut) {
+                                mShiftCategoryShortcuts.put(shortcutChar, intent);
+                            } else {
+                                mCategoryShortcuts.put(shortcutChar, intent);
+                            }
+                        }
+                        continue;
+                    } else if (roleName != null) {
+                        // We can't resolve the role at the time of this file being parsed as the
+                        // device hasn't finished booting, so we will look it up lazily.
+                        if (isShiftShortcut) {
+                            mShiftRoleShortcuts.put(shortcutChar, roleName);
+                        } else {
+                            mRoleShortcuts.put(shortcutChar, roleName);
                         }
                         continue;
                     } else {
-                        intent = resolveComponentNameIntent(packageName, className);
-                    }
-                } else if (categoryName != null) {
-                    if (roleName != null) {
-                        Log.w(TAG, "Cannot specify role bookmark when category is present for"
-                                + " bookmark shortcutChar=" + shortcutChar
-                                + " category= " + categoryName);
+                        Log.w(TAG, "Unable to add bookmark for shortcut " + shortcutName
+                                + ": missing package/class, category or role attributes");
                         continue;
                     }
-                    intent = Intent.makeMainSelectorActivity(Intent.ACTION_MAIN, categoryName);
-                    if (intent == null) {
-                        Log.w(TAG, "Null selector intent for " + categoryName);
-                    }
-                } else if (roleName != null) {
-                    // We can't resolve the role at the time of this file being parsed as the
-                    // device hasn't finished booting, so we will look it up lazily.
-                    if (isShiftShortcut) {
-                        mShiftRoleShortcuts.put(shortcutChar, roleName);
-                    } else {
-                        mRoleShortcuts.put(shortcutChar, roleName);
-                    }
-                    continue;
-                } else {
-                    Log.w(TAG, "Unable to add bookmark for shortcut " + shortcutName
-                            + ": missing package/class, category or role attributes");
-                    continue;
-                }
-
-                if (isShiftShortcut) {
-                    mShiftShortcuts.put(shortcutChar, intent);
-                } else {
-                    mIntentShortcuts.put(shortcutChar, intent);
                 }
             }
         } catch (XmlPullParserException | IOException e) {
@@ -336,21 +426,35 @@
 
     @Nullable
     private Intent resolveComponentNameIntent(ComponentName componentName) {
-        Intent intent = mComponentIntents.get(componentName);
-        if (intent == null) {
-            intent = resolveComponentNameIntent(
-                    componentName.getPackageName(), componentName.getClassName());
-            if (intent != null) {
-                mComponentIntents.put(componentName, intent);
+        if (modifierShortcutManagerRefactor()) {
+            return null;
+        } else {
+            Intent intent = mComponentIntents.get(componentName);
+            if (intent == null) {
+                intent = resolveComponentNameIntent(
+                        componentName.getPackageName(), componentName.getClassName());
+                if (intent != null) {
+                    mComponentIntents.put(componentName, intent);
+                }
             }
+            return intent;
         }
-        return intent;
     }
 
     @Nullable
     private Intent resolveComponentNameIntent(String packageName, String className) {
-        Context context = modifierShortcutManagerMultiuser()
-                ? mContext.createContextAsUser(mCurrentUser, 0) : mContext;
+        if (modifierShortcutManagerRefactor()) {
+            return null;
+        } else {
+            Context context = modifierShortcutManagerMultiuser()
+                        ? mContext.createContextAsUser(mCurrentUser, 0) : mContext;
+            return resolveComponentNameIntent(context, packageName, className);
+        }
+    }
+
+    @Nullable
+    private static Intent resolveComponentNameIntent(
+            Context context, String packageName, String className) {
         PackageManager pm = context.getPackageManager();
         int flags = PackageManager.MATCH_DIRECT_BOOT_UNAWARE;
         if (!modifierShortcutManagerMultiuser()) {
@@ -562,64 +666,81 @@
      */
     public KeyboardShortcutGroup getApplicationLaunchKeyboardShortcuts(int deviceId) {
         List<KeyboardShortcutInfo> shortcuts = new ArrayList();
-        for (int i = 0; i <  mIntentShortcuts.size(); i++) {
-            KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                    (char) (mIntentShortcuts.keyAt(i)), mIntentShortcuts.valueAt(i), false);
-            if (info != null) {
-                shortcuts.add(info);
-            }
-        }
-
-        for (int i = 0; i <  mShiftShortcuts.size(); i++) {
-            KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                    (char) (mShiftShortcuts.keyAt(i)), mShiftShortcuts.valueAt(i), true);
-            if (info != null) {
-                shortcuts.add(info);
-            }
-        }
-
-        for (int i = 0; i <  mRoleShortcuts.size(); i++) {
-            String role = mRoleShortcuts.valueAt(i);
-            KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                    (char) (mRoleShortcuts.keyAt(i)), getRoleLaunchIntent(role), false);
-            if (info != null) {
-                shortcuts.add(info);
-            }
-        }
-
-        for (int i = 0; i <  mShiftRoleShortcuts.size(); i++) {
-            String role = mShiftRoleShortcuts.valueAt(i);
-            KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                    (char) (mShiftRoleShortcuts.keyAt(i)), getRoleLaunchIntent(role), true);
-            if (info != null) {
-                shortcuts.add(info);
-            }
-        }
-
-        if (modifierShortcutManagerMultiuser()) {
-            for (int i = 0; i < mComponentShortcuts.size(); i++) {
-                ComponentName component = mComponentShortcuts.valueAt(i);
+        if (modifierShortcutManagerRefactor()) {
+            for (Bookmark b : mBookmarks.values()) {
                 KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                        (char) (mComponentShortcuts.keyAt(i)),
-                        resolveComponentNameIntent(component),
+                        b.getShortcutChar(), b.getIntent(mContext), b.isShift());
+                if (info != null) {
+                    shortcuts.add(info);
+                }
+            }
+        } else {
+            for (int i = 0; i <  mCategoryShortcuts.size(); i++) {
+                KeyboardShortcutInfo info = shortcutInfoFromIntent(
+                        (char) (mCategoryShortcuts.keyAt(i)),
+                        mCategoryShortcuts.valueAt(i),
                         false);
                 if (info != null) {
                     shortcuts.add(info);
                 }
             }
 
-            for (int i = 0; i < mShiftComponentShortcuts.size(); i++) {
-                ComponentName component = mShiftComponentShortcuts.valueAt(i);
+            for (int i = 0; i <  mShiftCategoryShortcuts.size(); i++) {
                 KeyboardShortcutInfo info = shortcutInfoFromIntent(
-                        (char) (mShiftComponentShortcuts.keyAt(i)),
-                        resolveComponentNameIntent(component),
+                        (char) (mShiftCategoryShortcuts.keyAt(i)),
+                        mShiftCategoryShortcuts.valueAt(i),
                         true);
                 if (info != null) {
                     shortcuts.add(info);
                 }
             }
-        }
 
+            for (int i = 0; i <  mRoleShortcuts.size(); i++) {
+                String role = mRoleShortcuts.valueAt(i);
+                KeyboardShortcutInfo info = shortcutInfoFromIntent(
+                        (char) (mRoleShortcuts.keyAt(i)),
+                        getRoleLaunchIntent(role),
+                        false);
+                if (info != null) {
+                    shortcuts.add(info);
+                }
+            }
+
+            for (int i = 0; i <  mShiftRoleShortcuts.size(); i++) {
+                String role = mShiftRoleShortcuts.valueAt(i);
+                KeyboardShortcutInfo info = shortcutInfoFromIntent(
+                        (char) (mShiftRoleShortcuts.keyAt(i)),
+                        getRoleLaunchIntent(role),
+                        true);
+                if (info != null) {
+                    shortcuts.add(info);
+                }
+            }
+
+            if (modifierShortcutManagerMultiuser()) {
+                for (int i = 0; i < mComponentShortcuts.size(); i++) {
+                    ComponentName component = mComponentShortcuts.valueAt(i);
+                    KeyboardShortcutInfo info = shortcutInfoFromIntent(
+                            (char) (mComponentShortcuts.keyAt(i)),
+                            resolveComponentNameIntent(component),
+                            false);
+                    if (info != null) {
+                        shortcuts.add(info);
+                    }
+                }
+
+                for (int i = 0; i < mShiftComponentShortcuts.size(); i++) {
+                    ComponentName component = mShiftComponentShortcuts.valueAt(i);
+                    KeyboardShortcutInfo info = shortcutInfoFromIntent(
+                            (char) (mShiftComponentShortcuts.keyAt(i)),
+                            resolveComponentNameIntent(component),
+                            true);
+                    if (info != null) {
+                        shortcuts.add(info);
+                    }
+                }
+            }
+        }
         return new KeyboardShortcutGroup(
                 mContext.getString(R.string.keyboard_shortcut_group_applications),
                 shortcuts);
@@ -800,57 +921,171 @@
     void dump(String prefix, PrintWriter pw) {
         IndentingPrintWriter ipw = new IndentingPrintWriter(pw,  "  ", prefix);
         ipw.println("ModifierShortcutManager shortcuts:");
-
-        ipw.increaseIndent();
-        ipw.println("Roles");
-        ipw.increaseIndent();
-        for (int i = 0; i <  mRoleShortcuts.size(); i++) {
-            String role = mRoleShortcuts.valueAt(i);
-            char shortcutChar = (char) mRoleShortcuts.keyAt(i);
-            Intent intent = getRoleLaunchIntent(role);
-            ipw.println(shortcutChar + " " + role + " " + intent);
-        }
-
-        for (int i = 0; i <  mShiftRoleShortcuts.size(); i++) {
-            String role = mShiftRoleShortcuts.valueAt(i);
-            char shortcutChar = (char) mShiftRoleShortcuts.keyAt(i);
-            Intent intent = getRoleLaunchIntent(role);
-            ipw.println("SHIFT+" + shortcutChar + " " + role + " " + intent);
-        }
-
-        ipw.decreaseIndent();
-        ipw.println("Selectors");
-        ipw.increaseIndent();
-        for (int i = 0; i <  mIntentShortcuts.size(); i++) {
-            char shortcutChar = (char) mIntentShortcuts.keyAt(i);
-            Intent intent = mIntentShortcuts.valueAt(i);
-            ipw.println(shortcutChar + " " + intent);
-        }
-
-        for (int i = 0; i <  mShiftShortcuts.size(); i++) {
-            char shortcutChar = (char) mShiftShortcuts.keyAt(i);
-            Intent intent = mShiftShortcuts.valueAt(i);
-            ipw.println("SHIFT+" + shortcutChar + " " + intent);
-
-        }
-
-        if (modifierShortcutManagerMultiuser()) {
-            ipw.decreaseIndent();
-            ipw.println("ComponentNames");
+        if (modifierShortcutManagerRefactor()) {
             ipw.increaseIndent();
-            for (int i = 0; i < mComponentShortcuts.size(); i++) {
-                char shortcutChar = (char) mComponentShortcuts.keyAt(i);
-                ComponentName component = mComponentShortcuts.valueAt(i);
-                Intent intent = resolveComponentNameIntent(component);
-                ipw.println(shortcutChar + " " + component + " " + intent);
+            for (Bookmark b : mBookmarks.values()) {
+                boolean isShift = b.isShift();
+                char shortcutChar = b.getShortcutChar();
+                Context context = modifierShortcutManagerMultiuser()
+                        ? mContext.createContextAsUser(mCurrentUser, 0) : mContext;
+
+                Intent intent = b.getIntent(context);
+                ipw.print(isShift ? "SHIFT+" : "");
+                ipw.println(shortcutChar + " " + intent);
+                ipw.increaseIndent();
+                ipw.increaseIndent();
+                KeyboardShortcutInfo info = shortcutInfoFromIntent(shortcutChar, intent, isShift);
+                if (info != null) {
+                    ipw.println("Resolves to: " + info.getLabel());
+                } else {
+                    ipw.println("<No KeyboardShortcutInfo available for this shortcut>");
+                }
+                ipw.decreaseIndent();
+                ipw.decreaseIndent();
+            }
+        } else {
+            ipw.increaseIndent();
+            ipw.println("Roles");
+            ipw.increaseIndent();
+            for (int i = 0; i < mRoleShortcuts.size(); i++) {
+                String role = mRoleShortcuts.valueAt(i);
+                char shortcutChar = (char) mRoleShortcuts.keyAt(i);
+                Intent intent = getRoleLaunchIntent(role);
+                ipw.println(shortcutChar + " " + role + " " + intent);
             }
 
-            for (int i = 0; i < mShiftComponentShortcuts.size(); i++) {
-                char shortcutChar = (char) mShiftComponentShortcuts.keyAt(i);
-                ComponentName component = mShiftComponentShortcuts.valueAt(i);
-                Intent intent = resolveComponentNameIntent(component);
-                ipw.println("SHIFT+" + shortcutChar + " " + component + " " + intent);
+            for (int i = 0; i < mShiftRoleShortcuts.size(); i++) {
+                String role = mShiftRoleShortcuts.valueAt(i);
+                char shortcutChar = (char) mShiftRoleShortcuts.keyAt(i);
+                Intent intent = getRoleLaunchIntent(role);
+                ipw.println("SHIFT+" + shortcutChar + " " + role + " " + intent);
             }
+
+            ipw.decreaseIndent();
+            ipw.println("Selectors");
+            ipw.increaseIndent();
+            for (int i = 0; i < mCategoryShortcuts.size(); i++) {
+                char shortcutChar = (char) mCategoryShortcuts.keyAt(i);
+                Intent intent = mCategoryShortcuts.valueAt(i);
+                ipw.println(shortcutChar + " " + intent);
+            }
+
+            for (int i = 0; i < mShiftCategoryShortcuts.size(); i++) {
+                char shortcutChar = (char) mShiftCategoryShortcuts.keyAt(i);
+                Intent intent = mShiftCategoryShortcuts.valueAt(i);
+                ipw.println("SHIFT+" + shortcutChar + " " + intent);
+
+            }
+
+            if (modifierShortcutManagerMultiuser()) {
+                ipw.decreaseIndent();
+                ipw.println("ComponentNames");
+                ipw.increaseIndent();
+                for (int i = 0; i < mComponentShortcuts.size(); i++) {
+                    char shortcutChar = (char) mComponentShortcuts.keyAt(i);
+                    ComponentName component = mComponentShortcuts.valueAt(i);
+                    Intent intent = resolveComponentNameIntent(component);
+                    ipw.println(shortcutChar + " " + component + " " + intent);
+                }
+
+                for (int i = 0; i < mShiftComponentShortcuts.size(); i++) {
+                    char shortcutChar = (char) mShiftComponentShortcuts.keyAt(i);
+                    ComponentName component = mShiftComponentShortcuts.valueAt(i);
+                    Intent intent = resolveComponentNameIntent(component);
+                    ipw.println("SHIFT+" + shortcutChar + " " + component + " " + intent);
+                }
+            }
+        }
+    }
+
+    private abstract static  class Bookmark {
+        private final char mShortcutChar;
+        private final boolean mShift;
+        protected Intent mIntent;
+
+        Bookmark(char shortcutChar, boolean shift) {
+            mShortcutChar = shortcutChar;
+            mShift = shift;
+        }
+
+        public char getShortcutChar() {
+            return mShortcutChar;
+        }
+
+        public boolean isShift() {
+            return mShift;
+        }
+
+        public abstract Intent getIntent(Context context);
+
+        public void clearIntent() {
+            mIntent = null;
+        }
+
+    }
+
+    private static final class RoleBookmark extends Bookmark {
+        private final String mRole;
+
+        RoleBookmark(char shortcutChar, boolean shift, String role) {
+            super(shortcutChar, shift);
+            mRole = role;
+        }
+
+        public String getRole() {
+            return mRole;
+        }
+
+        @Nullable
+        @Override
+        public Intent getIntent(Context context) {
+            if (mIntent != null) {
+                return mIntent;
+            }
+            mIntent = getRoleLaunchIntent(context, mRole);
+            return mIntent;
+        }
+    }
+
+    private static final class CategoryBookmark extends Bookmark {
+        private final String mCategory;
+
+        CategoryBookmark(char shortcutChar, boolean shift, String category) {
+            super(shortcutChar, shift);
+            mCategory = category;
+        }
+
+        @NonNull
+        @Override
+        public Intent getIntent(Context context) {
+            if (mIntent != null) {
+                return mIntent;
+            }
+
+            mIntent = Intent.makeMainSelectorActivity(Intent.ACTION_MAIN, mCategory);
+            return mIntent;
+        }
+    }
+
+    private static final class ComponentBookmark extends Bookmark {
+        private final String mPackageName;
+        private final String mClassName;
+
+        ComponentBookmark(
+                char shortcutChar, boolean shift, String packageName, String className) {
+            super(shortcutChar, shift);
+            mPackageName = packageName;
+            mClassName = className;
+        }
+
+        @Nullable
+        @Override
+        public Intent getIntent(Context context) {
+            if (mIntent != null) {
+                return mIntent;
+            }
+            mIntent = resolveComponentNameIntent(context, mPackageName, mClassName);
+            return mIntent;
         }
     }
 }
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 04d2de9..749025b 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -3600,18 +3600,6 @@
                 if (down) {
                     int direction = keyCode == KeyEvent.KEYCODE_BRIGHTNESS_UP ? 1 : -1;
 
-                    // Disable autobrightness if it's on
-                    int auto = Settings.System.getIntForUser(
-                            mContext.getContentResolver(),
-                            Settings.System.SCREEN_BRIGHTNESS_MODE,
-                            Settings.System.SCREEN_BRIGHTNESS_MODE_MANUAL,
-                            UserHandle.USER_CURRENT_OR_SELF);
-                    if (auto != 0) {
-                        Settings.System.putIntForUser(mContext.getContentResolver(),
-                                Settings.System.SCREEN_BRIGHTNESS_MODE,
-                                Settings.System.SCREEN_BRIGHTNESS_MODE_MANUAL,
-                                UserHandle.USER_CURRENT_OR_SELF);
-                    }
                     int screenDisplayId = displayId < 0 ? DEFAULT_DISPLAY : displayId;
 
                     float minLinearBrightness = mPowerManager.getBrightnessConstraint(
diff --git a/services/core/java/com/android/server/policy/TEST_MAPPING b/services/core/java/com/android/server/policy/TEST_MAPPING
index 338b479..bdb174d 100644
--- a/services/core/java/com/android/server/policy/TEST_MAPPING
+++ b/services/core/java/com/android/server/policy/TEST_MAPPING
@@ -46,18 +46,7 @@
       ]
     },
     {
-      "name": "CtsPermissionTestCases",
-      "options": [
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        },
-        {
-          "include-filter": "android.permission.cts.SplitPermissionTest"
-        },
-        {
-          "include-filter": "android.permission.cts.BackgroundPermissionsTest"
-        }
-      ]
+      "name": "CtsPermissionTestCases_Platform"
     },
     {
       "name": "CtsBackupTestCases",
diff --git a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
index d0b70c3..da8b01a 100644
--- a/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
+++ b/services/core/java/com/android/server/policy/keyguard/KeyguardServiceDelegate.java
@@ -176,8 +176,9 @@
 
         final DreamManagerInternal dreamManager =
                 LocalServices.getService(DreamManagerInternal.class);
-
-        dreamManager.registerDreamManagerStateListener(mDreamManagerStateListener);
+        if(dreamManager != null){
+            dreamManager.registerDreamManagerStateListener(mDreamManagerStateListener);
+        }
     }
 
     private final ServiceConnection mKeyguardConnection = new ServiceConnection() {
diff --git a/services/core/java/com/android/server/power/Notifier.java b/services/core/java/com/android/server/power/Notifier.java
index 1a2a196..303828f 100644
--- a/services/core/java/com/android/server/power/Notifier.java
+++ b/services/core/java/com/android/server/power/Notifier.java
@@ -1064,9 +1064,9 @@
     private void notifyWakeLockListener(IWakeLockCallback callback, String tag, boolean isEnabled,
             int ownerUid, int ownerPid, int flags, WorkSource workSource, String packageName,
             String historyTag) {
+        long currentTime = mInjector.currentTimeMillis();
         mHandler.post(() -> {
             if (mFlags.improveWakelockLatency()) {
-                long currentTime = mInjector.currentTimeMillis();
                 if (isEnabled) {
                     notifyWakelockAcquisition(tag, ownerUid, ownerPid, flags,
                             workSource, packageName, historyTag, currentTime);
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index a27360d..12e7fd0 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -1379,8 +1379,10 @@
                     new DisplayGroupPowerChangeListener();
             mDisplayManagerInternal.registerDisplayGroupListener(displayGroupPowerChangeListener);
 
-            // This DreamManager method does not acquire a lock, so it should be safe to call.
-            mDreamManager.registerDreamManagerStateListener(new DreamManagerStateListener());
+            if(mDreamManager != null){
+                // This DreamManager method does not acquire a lock, so it should be safe to call.
+                mDreamManager.registerDreamManagerStateListener(new DreamManagerStateListener());
+            }
 
             mWirelessChargerDetector = mInjector.createWirelessChargerDetector(sensorManager,
                     mInjector.createSuspendBlocker(
@@ -3543,7 +3545,7 @@
         }
 
         // Stop dream.
-        if (isDreaming) {
+        if (isDreaming && mDreamManager != null) {
             mDreamManager.stopDream(/* immediate= */ false, "power manager request" /*reason*/);
         }
     }
diff --git a/services/core/java/com/android/server/power/ThermalManagerService.java b/services/core/java/com/android/server/power/ThermalManagerService.java
index 822ec2e..6847a5c 100644
--- a/services/core/java/com/android/server/power/ThermalManagerService.java
+++ b/services/core/java/com/android/server/power/ThermalManagerService.java
@@ -54,6 +54,7 @@
 import android.os.ShellCommand;
 import android.os.SystemClock;
 import android.os.Temperature;
+import android.os.Trace;
 import android.util.ArrayMap;
 import android.util.EventLog;
 import android.util.Slog;
@@ -247,6 +248,7 @@
 
     private void setStatusLocked(int newStatus) {
         if (newStatus != mStatus) {
+            Trace.traceCounter(Trace.TRACE_TAG_POWER, "ThermalManagerService.status", newStatus);
             mStatus = newStatus;
             notifyStatusListenersLocked();
         }
diff --git a/services/core/java/com/android/server/power/WakeLockLog.java b/services/core/java/com/android/server/power/WakeLockLog.java
index 968ff59..eda222e 100644
--- a/services/core/java/com/android/server/power/WakeLockLog.java
+++ b/services/core/java/com/android/server/power/WakeLockLog.java
@@ -19,6 +19,7 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.PowerManager;
+import android.os.Process;
 import android.text.TextUtils;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -122,6 +123,9 @@
 
     private static final SimpleDateFormat DATE_FORMAT = new SimpleDateFormat("MM-dd HH:mm:ss.SSS");
 
+    @VisibleForTesting
+    static final String SYSTEM_PACKAGE_NAME = "System";
+
     /**
      * Lock protects WakeLockLog.dump (binder thread) from conflicting with changes to the log
      * happening on the background thread.
@@ -516,21 +520,26 @@
                 return;
             }
 
-            String[] packages;
-            if (uidToPackagesCache.contains(tag.ownerUid)) {
-                packages = uidToPackagesCache.get(tag.ownerUid);
-            } else {
-                packages = packageManager.getPackagesForUid(tag.ownerUid);
-                uidToPackagesCache.put(tag.ownerUid, packages);
+            if (tag.ownerUid == Process.SYSTEM_UID) {
+                packageName = SYSTEM_PACKAGE_NAME;
             }
+            else {
+                String[] packages;
+                if (uidToPackagesCache.contains(tag.ownerUid)) {
+                    packages = uidToPackagesCache.get(tag.ownerUid);
+                } else {
+                    packages = packageManager.getPackagesForUid(tag.ownerUid);
+                    uidToPackagesCache.put(tag.ownerUid, packages);
+                }
 
-            if (packages != null && packages.length > 0) {
-                packageName = packages[0];
-                if (packages.length > 1) {
-                    StringBuilder sb = new StringBuilder();
-                    sb.append(packageName)
-                            .append(",...");
-                    packageName = sb.toString();
+                if (packages != null && packages.length > 0) {
+                    packageName = packages[0];
+                    if (packages.length > 1) {
+                        StringBuilder sb = new StringBuilder();
+                        sb.append(packageName)
+                                .append(",...");
+                        packageName = sb.toString();
+                    }
                 }
             }
         }
diff --git a/services/core/java/com/android/server/power/batterysaver/BatterySaverStateMachine.java b/services/core/java/com/android/server/power/batterysaver/BatterySaverStateMachine.java
index 9a4c60d..68760aa 100644
--- a/services/core/java/com/android/server/power/batterysaver/BatterySaverStateMachine.java
+++ b/services/core/java/com/android/server/power/batterysaver/BatterySaverStateMachine.java
@@ -864,7 +864,8 @@
                     buildNotification(DYNAMIC_MODE_NOTIF_CHANNEL_ID,
                             R.string.dynamic_mode_notification_title,
                             R.string.dynamic_mode_notification_summary,
-                            Settings.ACTION_BATTERY_SAVER_SETTINGS, 0L),
+                            Settings.ACTION_BATTERY_SAVER_SETTINGS, 0L,
+                            R.drawable.ic_settings),
                     UserHandle.ALL);
         });
     }
@@ -889,7 +890,8 @@
                             R.string.dynamic_mode_notification_summary_v2,
                             Settings.ACTION_BATTERY_SAVER_SETTINGS,
                             0L /* timeoutMs */,
-                            highlightBundle),
+                            highlightBundle,
+                            R.drawable.ic_qs_battery_saver),
                     UserHandle.ALL);
         });
     }
@@ -911,7 +913,8 @@
                             R.string.battery_saver_off_notification_title,
                             R.string.battery_saver_charged_notification_summary,
                             Settings.ACTION_BATTERY_SAVER_SETTINGS,
-                            STICKY_DISABLED_NOTIFY_TIMEOUT_MS),
+                            STICKY_DISABLED_NOTIFY_TIMEOUT_MS,
+                            R.drawable.ic_settings),
                     UserHandle.ALL);
         });
     }
@@ -926,7 +929,7 @@
     }
 
     private Notification buildNotification(@NonNull String channelId, @StringRes int titleId,
-            @StringRes int summaryId, @NonNull String intentAction, long timeoutMs) {
+            @StringRes int summaryId, @NonNull String intentAction, long timeoutMs, int iconResId) {
         Resources res = mContext.getResources();
         Intent intent = new Intent(intentAction);
         intent.setFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK);
@@ -937,7 +940,7 @@
         final String summary = res.getString(summaryId);
 
         return new Notification.Builder(mContext, channelId)
-                .setSmallIcon(R.drawable.ic_battery)
+                .setSmallIcon(iconResId)
                 .setContentTitle(title)
                 .setContentText(summary)
                 .setContentIntent(batterySaverIntent)
@@ -950,7 +953,7 @@
 
     private Notification buildNotificationV2(@NonNull String channelId, @StringRes int titleId,
             @StringRes int summaryId, @NonNull String intentAction, long timeoutMs,
-            @NonNull Bundle highlightBundle) {
+            @NonNull Bundle highlightBundle, int iconResId) {
         Resources res = mContext.getResources();
         Intent intent = new Intent(intentAction)
                 .setFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK)
@@ -963,7 +966,7 @@
         final String summary = res.getString(summaryId);
 
         return new Notification.Builder(mContext, channelId)
-                .setSmallIcon(R.drawable.ic_battery)
+                .setSmallIcon(iconResId)
                 .setContentTitle(title)
                 .setContentText(summary)
                 .setContentIntent(batterySaverIntent)
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java b/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java
deleted file mode 100644
index 7ba4330..0000000
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsSection.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2023 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.util.IndentingPrintWriter;
-
-import com.android.modules.utils.TypedXmlSerializer;
-
-import java.io.IOException;
-
-class AggregatedPowerStatsSection extends PowerStatsSpan.Section {
-    public static final String TYPE = "aggregated-power-stats";
-
-    private final AggregatedPowerStats mAggregatedPowerStats;
-
-    AggregatedPowerStatsSection(AggregatedPowerStats aggregatedPowerStats) {
-        super(TYPE);
-        mAggregatedPowerStats = aggregatedPowerStats;
-    }
-
-    public AggregatedPowerStats getAggregatedPowerStats() {
-        return mAggregatedPowerStats;
-    }
-
-    @Override
-    void write(TypedXmlSerializer serializer) throws IOException {
-        mAggregatedPowerStats.writeXml(serializer);
-    }
-
-    @Override
-    public void dump(IndentingPrintWriter ipw) {
-        mAggregatedPowerStats.dump(ipw);
-    }
-}
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 b45651d..680b1ac 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -23,8 +23,6 @@
 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;
@@ -149,6 +147,7 @@
 import com.android.server.LocalServices;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
 
 import libcore.util.EmptyArray;
 
@@ -180,7 +179,6 @@
 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;
 
@@ -2028,7 +2026,7 @@
         void setContext(Context context) {
             mPackageManager = context.getPackageManager();
             mConsumedEnergyRetriever = new PowerStatsCollector.ConsumedEnergyRetrieverImpl(
-                    LocalServices.getService(PowerStatsInternal.class));
+                    LocalServices.getService(PowerStatsInternal.class), () -> mBatteryVoltageMv);
             mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
             mTelephonyManager =
                     (TelephonyManager) context.getSystemService(Context.TELEPHONY_SERVICE);
@@ -2083,11 +2081,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> mBatteryVoltageMv;
-        }
-
-        @Override
         public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
             return mScreenUsageTimeRetriever;
         }
@@ -13191,7 +13184,8 @@
                 final int freq = deltaInfo.getSpecificInfoFrequencyRange(index);
 
                 // Map RadioAccessNetworkType to course grain RadioAccessTechnology.
-                final int ratBucket = mapRadioAccessNetworkTypeToRadioAccessTechnology(rat);
+                final int ratBucket = MobileRadioPowerStatsLayout
+                        .mapRadioAccessNetworkTypeToRadioAccessTechnology(rat);
                 final RadioAccessTechnologyBatteryStats ratStats = getRatBatteryStatsLocked(
                         ratBucket);
 
@@ -15300,15 +15294,6 @@
                 mHistory.writeHistoryItem(elapsedRealtimeMs, uptimeMs);
             }
         }
-        if (!onBattery &&
-                (status == BatteryManager.BATTERY_STATUS_FULL ||
-                        status == BatteryManager.BATTERY_STATUS_UNKNOWN)) {
-            // We don't record history while we are plugged in and fully charged
-            // (or when battery is not present).  The next time we are
-            // unplugged, history will be cleared.
-            mHistory.setHistoryRecordingEnabled(DEBUG);
-        }
-
         mLastLearnedBatteryCapacityUah = chargeFullUah;
         if (mMinLearnedBatteryCapacityUah == -1) {
             mMinLearnedBatteryCapacityUah = chargeFullUah;
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 d51cfea..87a3e5e 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -16,6 +16,7 @@
 
 package com.android.server.power.stats;
 
+import android.annotation.NonNull;
 import android.content.Context;
 import android.hardware.SensorManager;
 import android.os.BatteryConsumer;
@@ -27,7 +28,6 @@
 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;
@@ -44,8 +44,7 @@
 public class BatteryUsageStatsProvider {
     private static final String TAG = "BatteryUsageStatsProv";
     private final Context mContext;
-    private final SparseBooleanArray mPowerStatsExporterEnabled = new SparseBooleanArray();
-    private final PowerStatsExporter mPowerStatsExporter;
+    private final PowerAttributor mPowerAttributor;
     private final PowerStatsStore mPowerStatsStore;
     private final PowerProfile mPowerProfile;
     private final CpuScalingPolicies mCpuScalingPolicies;
@@ -53,16 +52,18 @@
     private final Object mLock = new Object();
     private List<PowerCalculator> mPowerCalculators;
 
-    public BatteryUsageStatsProvider(Context context,
-            PowerStatsExporter powerStatsExporter,
-            PowerProfile powerProfile, CpuScalingPolicies cpuScalingPolicies,
-            PowerStatsStore powerStatsStore, Clock clock) {
+    public BatteryUsageStatsProvider(@NonNull Context context,
+            @NonNull PowerAttributor powerAttributor,
+            @NonNull PowerProfile powerProfile, @NonNull CpuScalingPolicies cpuScalingPolicies,
+            @NonNull PowerStatsStore powerStatsStore, @NonNull Clock clock) {
         mContext = context;
-        mPowerStatsExporter = powerStatsExporter;
+        mPowerAttributor = powerAttributor;
         mPowerStatsStore = powerStatsStore;
         mPowerProfile = powerProfile;
         mCpuScalingPolicies = cpuScalingPolicies;
         mClock = clock;
+
+        mPowerStatsStore.addSectionReader(new BatteryUsageStatsSection.Reader());
     }
 
     private List<PowerCalculator> getPowerCalculators() {
@@ -72,55 +73,67 @@
 
                 // Power calculators are applied in the order of registration
                 mPowerCalculators.add(new BatteryChargeCalculator());
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_CPU)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_CPU)) {
                     mPowerCalculators.add(
                             new CpuPowerCalculator(mCpuScalingPolicies, mPowerProfile));
                 }
                 mPowerCalculators.add(new MemoryPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new WakelockPowerCalculator(mPowerProfile));
                 if (!BatteryStats.checkWifiOnly(mContext)) {
-                    if (!mPowerStatsExporterEnabled.get(
+                    if (!mPowerAttributor.isPowerComponentSupported(
                             BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)) {
                         mPowerCalculators.add(new MobileRadioPowerCalculator(mPowerProfile));
                     }
-                    if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_PHONE)) {
+                    if (!mPowerAttributor.isPowerComponentSupported(
+                            BatteryConsumer.POWER_COMPONENT_PHONE)) {
                         mPowerCalculators.add(new PhonePowerCalculator(mPowerProfile));
                     }
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_WIFI)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_WIFI)) {
                     mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_BLUETOOTH)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_BLUETOOTH)) {
                     mPowerCalculators.add(new BluetoothPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_SENSORS)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_SENSORS)) {
                     mPowerCalculators.add(new SensorPowerCalculator(
                             mContext.getSystemService(SensorManager.class)));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_GNSS)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_GNSS)) {
                     mPowerCalculators.add(new GnssPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_CAMERA)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_CAMERA)) {
                     mPowerCalculators.add(new CameraPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_FLASHLIGHT)) {
                     mPowerCalculators.add(new FlashlightPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_AUDIO)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_AUDIO)) {
                     mPowerCalculators.add(new AudioPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_VIDEO)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_VIDEO)) {
                     mPowerCalculators.add(new VideoPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)) {
                     mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
                 }
-                if (!mPowerStatsExporterEnabled.get(
+                if (!mPowerAttributor.isPowerComponentSupported(
                         BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY)) {
                     mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
                 }
                 mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
-                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_ANY)) {
+                if (!mPowerAttributor.isPowerComponentSupported(
+                        BatteryConsumer.POWER_COMPONENT_ANY)) {
                     mPowerCalculators.add(new CustomEnergyConsumerPowerCalculator(mPowerProfile));
                 }
                 mPowerCalculators.add(new UserPowerCalculator());
@@ -196,7 +209,7 @@
         final boolean includeProcessStateData = ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_PROCESS_STATE_DATA) != 0)
                 && stats.isProcessStateDataAvailable();
-        final boolean includeVirtualUids =  ((query.getFlags()
+        final boolean includeVirtualUids = ((query.getFlags()
                 & BatteryUsageStatsQuery.FLAG_BATTERY_USAGE_STATS_INCLUDE_VIRTUAL_UIDS) != 0);
         final double minConsumedPowerThreshold = query.getMinConsumedPowerThreshold();
 
@@ -258,10 +271,8 @@
             }
         }
 
-        if (mPowerStatsExporterEnabled.indexOfValue(true) >= 0) {
-            mPowerStatsExporter.exportAggregatedPowerStats(batteryUsageStatsBuilder,
-                    monotonicStartTime, monotonicEndTime);
-        }
+        mPowerAttributor.estimatePowerConsumption(batteryUsageStatsBuilder, stats.getHistory(),
+                monotonicStartTime, monotonicEndTime);
 
         BatteryUsageStats batteryUsageStats = batteryUsageStatsBuilder.build();
         if (includeProcessStateData) {
@@ -272,6 +283,7 @@
 
     // STOPSHIP(b/229906525): remove verification before shipping
     private static boolean sErrorReported;
+
     private void verify(BatteryUsageStats stats) {
         if (sErrorReported) {
             return;
@@ -390,7 +402,7 @@
             // while the "to" timestamp is *inclusive*.
             boolean isInRange =
                     (query.getFromTimestamp() == 0 || minTime > query.getFromTimestamp())
-                    && (query.getToTimestamp() == 0 || maxTime <= query.getToTimestamp());
+                            && (query.getToTimestamp() == 0 || maxTime <= query.getToTimestamp());
             if (!isInRange) {
                 continue;
             }
@@ -422,12 +434,4 @@
         }
         return builder.build();
     }
-
-    /**
-     * Specify whether PowerStats based attribution is supported for the specified component.
-     */
-    public void setPowerStatsExporterEnabled(int powerComponentId, boolean enabled) {
-        mPowerStatsExporterEnabled.put(powerComponentId, enabled);
-        mPowerStatsExporter.setPowerComponentEnabled(powerComponentId, enabled);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
index b95faac..af36524 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsSection.java
@@ -19,8 +19,11 @@
 import android.os.BatteryUsageStats;
 import android.util.IndentingPrintWriter;
 
+import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
 
+import org.xmlpull.v1.XmlPullParserException;
+
 import java.io.IOException;
 
 class BatteryUsageStatsSection extends PowerStatsSpan.Section {
@@ -38,7 +41,7 @@
     }
 
     @Override
-    void write(TypedXmlSerializer serializer) throws IOException {
+    public void write(TypedXmlSerializer serializer) throws IOException {
         mBatteryUsageStats.writeXml(serializer);
     }
 
@@ -46,4 +49,17 @@
     public void dump(IndentingPrintWriter ipw) {
         mBatteryUsageStats.dump(ipw, "");
     }
+
+    static class Reader implements PowerStatsSpan.SectionReader {
+        @Override
+        public String getType() {
+            return TYPE;
+        }
+
+        @Override
+        public PowerStatsSpan.Section read(String sectionType, TypedXmlPullParser parser)
+                throws IOException, XmlPullParserException {
+            return new BatteryUsageStatsSection(BatteryUsageStats.createFromXml(parser));
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
index 8a5085b..539c415 100644
--- a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/BluetoothPowerStatsCollector.java
@@ -28,13 +28,12 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.BluetoothPowerStatsLayout;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
 import java.util.concurrent.TimeUnit;
-import java.util.function.IntSupplier;
 
 public class BluetoothPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "BluetoothPowerStatsCollector";
@@ -43,7 +42,7 @@
 
     private static final long ENERGY_UNSPECIFIED = -1;
 
-    interface BluetoothStatsRetriever {
+    public interface BluetoothStatsRetriever {
         interface Callback {
             void onBluetoothScanTime(int uid, long scanTimeMs);
         }
@@ -54,29 +53,25 @@
                 BluetoothAdapter.OnBluetoothActivityEnergyInfoCallback callback);
     }
 
-    interface Injector {
+    public interface Injector {
         Handler getHandler();
         Clock getClock();
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         BluetoothStatsRetriever getBluetoothStatsRetriever();
     }
 
     private final Injector mInjector;
 
-    private BluetoothPowerStatsLayout mLayout;
+    private com.android.server.power.stats.format.BluetoothPowerStatsLayout mLayout;
     private boolean mIsInitialized;
     private PowerStats mPowerStats;
     private long[] mDeviceStats;
     private BluetoothStatsRetriever mBluetoothStatsRetriever;
     private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
 
     private long mLastRxTime;
     private long mLastTxTime;
@@ -93,7 +88,7 @@
 
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
 
-    BluetoothPowerStatsCollector(Injector injector) {
+    public BluetoothPowerStatsCollector(Injector injector) {
         super(injector.getHandler(),  injector.getPowerStatsCollectionThrottlePeriod(
                         BatteryConsumer.powerComponentIdToString(
                                 BatteryConsumer.POWER_COMPONENT_BLUETOOTH)),
@@ -122,21 +117,11 @@
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mBluetoothStatsRetriever = mInjector.getBluetoothStatsRetriever();
-        mEnergyConsumerIds =
-                mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
-
-        mLayout = new BluetoothPowerStatsLayout();
-        mLayout.addDeviceBluetoothControllerActivity();
-        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
-        mLayout.addDeviceSectionUsageDuration();
-        mLayout.addDeviceSectionPowerEstimate();
-        mLayout.addUidTrafficStats();
-        mLayout.addUidSectionPowerEstimate();
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mConsumedEnergyRetriever,
+                EnergyConsumerType.BLUETOOTH);
+        mLayout = new BluetoothPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount());
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -152,7 +137,7 @@
     }
 
     @Override
-    protected PowerStats collectStats() {
+    public PowerStats collectStats() {
         if (!ensureInitialized()) {
             return null;
         }
@@ -162,9 +147,7 @@
         collectBluetoothActivityInfo();
         collectBluetoothScanStats();
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         return mPowerStats;
     }
@@ -296,34 +279,6 @@
         mLayout.setDeviceScanTime(mDeviceStats, totalScanTime);
     }
 
-    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];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         super.onUidRemoved(uid);
diff --git a/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
index 8705bd5..0f70064 100644
--- a/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CameraPowerStatsCollector.java
@@ -19,12 +19,13 @@
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 
+import com.android.server.power.stats.format.BinaryStatePowerStatsLayout;
+
 public class CameraPowerStatsCollector extends EnergyConsumerPowerStatsCollector {
 
-    CameraPowerStatsCollector(Injector injector) {
+    public CameraPowerStatsCollector(Injector injector) {
         super(injector, BatteryConsumer.POWER_COMPONENT_CAMERA,
                 BatteryConsumer.powerComponentIdToString(BatteryConsumer.POWER_COMPONENT_CAMERA),
-                EnergyConsumerType.CAMERA, /* energy consumer name */ null,
-                new BinaryStatePowerStatsLayout());
+                EnergyConsumerType.CAMERA, new BinaryStatePowerStatsLayout());
     }
 }
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 b5ef67b..128f14a 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CpuPowerStatsCollector.java
@@ -31,11 +31,10 @@
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.CpuPowerStatsLayout;
 
 import java.io.PrintWriter;
-import java.util.Arrays;
 import java.util.Locale;
-import java.util.function.IntSupplier;
 
 /**
  * Collects snapshots of power-related system statistics.
@@ -46,7 +45,6 @@
 public class CpuPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "CpuPowerStatsCollector";
     private static final long NANOS_PER_MILLIS = 1000000;
-    private static final long ENERGY_UNSPECIFIED = -1;
     private static final int DEFAULT_CPU_POWER_BRACKETS = 3;
     private static final int DEFAULT_CPU_POWER_BRACKETS_PER_ENERGY_CONSUMER = 2;
 
@@ -58,7 +56,6 @@
         PowerProfile getPowerProfile();
         KernelCpuStatsReader getKernelCpuStatsReader();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
 
         default int getDefaultCpuPowerBrackets() {
@@ -76,8 +73,7 @@
     private CpuScalingPolicies mCpuScalingPolicies;
     private PowerProfile mPowerProfile;
     private KernelCpuStatsReader mKernelCpuStatsReader;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private int mDefaultCpuPowerBrackets;
     private int mDefaultCpuPowerBracketsPerEnergyConsumer;
     private long[] mCpuTimeByScalingStep;
@@ -85,15 +81,12 @@
     private long[] mTempUidStats;
     private final SparseArray<UidStats> mUidStats = new SparseArray<>();
     private boolean mIsPerUidTimeInStateSupported;
-    private int[] mCpuEnergyConsumerIds = new int[0];
     private PowerStats.Descriptor mPowerStatsDescriptor;
     // Reusable instance
     private PowerStats mCpuPowerStats;
     private CpuPowerStatsLayout mLayout;
     private long mLastUpdateTimestampNanos;
     private long mLastUpdateUptimeMillis;
-    private int mLastVoltageMv;
-    private long[] mLastConsumedEnergyUws;
 
     CpuPowerStatsCollector(Injector injector) {
         super(injector.getHandler(), injector.getPowerStatsCollectionThrottlePeriod(
@@ -115,31 +108,22 @@
         mCpuScalingPolicies = mInjector.getCpuScalingPolicies();
         mPowerProfile = mInjector.getPowerProfile();
         mKernelCpuStatsReader = mInjector.getKernelCpuStatsReader();
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mDefaultCpuPowerBrackets = mInjector.getDefaultCpuPowerBrackets();
         mDefaultCpuPowerBracketsPerEnergyConsumer =
                 mInjector.getDefaultCpuPowerBracketsPerEnergyConsumer();
 
         mIsPerUidTimeInStateSupported = mKernelCpuStatsReader.isSupportedFeature();
-        mCpuEnergyConsumerIds =
-                mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.CPU_CLUSTER);
-        mLastConsumedEnergyUws = new long[mCpuEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.CPU_CLUSTER);
 
         int cpuScalingStepCount = mCpuScalingPolicies.getScalingStepCount();
         mCpuTimeByScalingStep = new long[cpuScalingStepCount];
         mTempCpuTimeByScalingStep = new long[cpuScalingStepCount];
         int[] scalingStepToPowerBracketMap = initPowerBrackets();
 
-        mLayout = new CpuPowerStatsLayout();
-        mLayout.addDeviceSectionCpuTimeByScalingStep(cpuScalingStepCount);
-        mLayout.addDeviceSectionCpuTimeByCluster(mCpuScalingPolicies.getPolicies().length);
-        mLayout.addDeviceSectionUsageDuration();
-        mLayout.addDeviceSectionEnergyConsumers(mCpuEnergyConsumerIds.length);
-        mLayout.addDeviceSectionPowerEstimate();
-        mLayout.addUidSectionCpuTimeByPowerBracket(scalingStepToPowerBracketMap);
-        mLayout.addUidSectionPowerEstimate();
+        mLayout = new CpuPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
+                mCpuScalingPolicies.getPolicies().length, scalingStepToPowerBracketMap);
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -158,16 +142,18 @@
     private int[] initPowerBrackets() {
         if (mPowerProfile.getCpuPowerBracketCount() != PowerProfile.POWER_BRACKETS_UNSPECIFIED) {
             return initPowerBracketsFromPowerProfile();
-        } else if (mCpuEnergyConsumerIds.length == 0 || mCpuEnergyConsumerIds.length == 1) {
+        } else if (mConsumedEnergyHelper.getEnergyConsumerCount() == 0
+                || mConsumedEnergyHelper.getEnergyConsumerCount() == 1) {
             return initDefaultPowerBrackets(mDefaultCpuPowerBrackets);
-        } else if (mCpuScalingPolicies.getPolicies().length == mCpuEnergyConsumerIds.length) {
+        } else if (mCpuScalingPolicies.getPolicies().length
+                == mConsumedEnergyHelper.getEnergyConsumerCount()) {
             return initPowerBracketsByCluster(mDefaultCpuPowerBracketsPerEnergyConsumer);
         } else {
             Slog.i(TAG, "Assigning a single power brackets to each CPU_CLUSTER energy consumer."
                         + " Number of CPU clusters ("
                         + mCpuScalingPolicies.getPolicies().length
                         + ") does not match the number of energy consumers ("
-                        + mCpuEnergyConsumerIds.length + "). "
+                        + mConsumedEnergyHelper.getEnergyConsumerCount() + "). "
                         + " Using default power bucket assignment.");
             return initDefaultPowerBrackets(mDefaultCpuPowerBrackets);
         }
@@ -368,41 +354,11 @@
         }
         mLayout.setUsageDuration(mCpuPowerStats.stats, uptimeDelta);
 
-        if (mCpuEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mCpuPowerStats, mLayout);
 
         return mCpuPowerStats;
     }
 
-    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(mCpuEnergyConsumerIds);
-        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(mCpuPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
-            mLastConsumedEnergyUws[i] = energyUws[i];
-        }
-    }
-
     @VisibleForNative
     interface KernelCpuStatsCallback {
         @Keep // Called from native
diff --git a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
index 4bfe442..c1c7e12 100644
--- a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsCollector.java
@@ -19,6 +19,8 @@
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 
+import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
+
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.List;
@@ -29,7 +31,8 @@
     private final EnergyConsumerPowerStatsCollector.Injector mInjector;
     private List<EnergyConsumerPowerStatsCollector> mCollectors;
 
-    CustomEnergyConsumerPowerStatsCollector(EnergyConsumerPowerStatsCollector.Injector injector) {
+    public CustomEnergyConsumerPowerStatsCollector(
+            EnergyConsumerPowerStatsCollector.Injector injector) {
         super(injector.getHandler(), 0, injector.getUidResolver(), injector.getClock());
         mInjector = injector;
     }
@@ -46,7 +49,8 @@
         for (int i = 0; i < energyConsumerIds.length; i++) {
             String name = retriever.getEnergyConsumerName(energyConsumerIds[i]);
             EnergyConsumerPowerStatsCollector collector = new EnergyConsumerPowerStatsCollector(
-                    mInjector, powerComponentId++, name, energyConsumerIds[i], sLayout);
+                    mInjector, powerComponentId++, name, EnergyConsumerType.OTHER,
+                    energyConsumerIds[i], sLayout);
             collector.setEnabled(true);
             collector.addConsumer(this::deliverStats);
             mCollectors.add(collector);
diff --git a/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
index 79fbe8e..1d2e388 100644
--- a/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsCollector.java
@@ -16,78 +16,57 @@
 
 package com.android.server.power.stats;
 
-import android.hardware.power.stats.EnergyConsumerAttribution;
-import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.Handler;
 import android.os.PersistableBundle;
-import android.util.Slog;
-import android.util.SparseLongArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
-
-import java.util.function.IntSupplier;
+import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
 
 public class EnergyConsumerPowerStatsCollector extends PowerStatsCollector {
-    private static final String TAG = "EnergyConsumerPowerStatsCollector";
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
-    interface Injector {
+    public interface Injector {
         Handler getHandler();
         Clock getClock();
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
     }
 
+    private static final int UNSPECIFIED = -1;
+
     private final Injector mInjector;
     private final int mPowerComponentId;
     private final String mPowerComponentName;
+    private final int mEnergyConsumerId;
     private final int mEnergyConsumerType;
-    private final String mEnergyConsumerName;
 
     private final EnergyConsumerPowerStatsLayout mLayout;
     private boolean mIsInitialized;
 
     private PowerStats mPowerStats;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds;
-    private long mLastConsumedEnergyUws = ENERGY_UNSPECIFIED;
-    private SparseLongArray mLastConsumerEnergyPerUid = new SparseLongArray();
-    private int mLastVoltageMv;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private long mLastUpdateTimestamp;
-    private boolean mFirstCollection = true;
 
     EnergyConsumerPowerStatsCollector(Injector injector, int powerComponentId,
             String powerComponentName, @EnergyConsumerType int energyConsumerType,
-            String energyConsumerName, EnergyConsumerPowerStatsLayout statsLayout) {
-        super(injector.getHandler(),
-                injector.getPowerStatsCollectionThrottlePeriod(powerComponentName),
-                injector.getUidResolver(), injector.getClock());
-        mInjector = injector;
-        mPowerComponentId = powerComponentId;
-        mPowerComponentName = powerComponentName;
-        mEnergyConsumerType = energyConsumerType;
-        mEnergyConsumerName = energyConsumerName;
-        mLayout = statsLayout;
+            EnergyConsumerPowerStatsLayout statsLayout) {
+        this(injector, powerComponentId, powerComponentName, energyConsumerType, UNSPECIFIED,
+                statsLayout);
     }
 
     EnergyConsumerPowerStatsCollector(Injector injector, int powerComponentId,
-            String powerComponentName, int energyConsumerId,
-            EnergyConsumerPowerStatsLayout statsLayout) {
+            String powerComponentName, @EnergyConsumerType int energyConsumerType,
+            int energyConsumerId, EnergyConsumerPowerStatsLayout statsLayout) {
         super(injector.getHandler(),
                 injector.getPowerStatsCollectionThrottlePeriod(powerComponentName),
                 injector.getUidResolver(), injector.getClock());
         mInjector = injector;
         mPowerComponentId = powerComponentId;
         mPowerComponentName = powerComponentName;
-        mEnergyConsumerIds = new int[]{energyConsumerId};
-        mEnergyConsumerType = EnergyConsumerType.OTHER;
-        mEnergyConsumerName = null;
+        mEnergyConsumerId = energyConsumerId;
+        mEnergyConsumerType = energyConsumerType;
         mLayout = statsLayout;
     }
 
@@ -100,12 +79,14 @@
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
-        if (mEnergyConsumerIds == null) {
-            mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(mEnergyConsumerType,
-                    mEnergyConsumerName);
+        if (mEnergyConsumerId != UNSPECIFIED) {
+            mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                    mEnergyConsumerId, true /* perUidAttribution */);
+        } else {
+            mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                    mEnergyConsumerType);
         }
+
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
         PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
@@ -124,85 +105,15 @@
             return null;
         }
 
-        if (mEnergyConsumerIds.length == 0) {
-            return null;
-        }
-
-        int voltageMv = mVoltageSupplier.getAsInt();
-        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
-        if (averageVoltage <= 0) {
-            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
-                    + " mV) when querying energy consumers");
-            return null;
-        }
-
-        mLastVoltageMv = voltageMv;
-
-        EnergyConsumerResult[] energy =
-                    mConsumedEnergyRetriever.getConsumedEnergy(mEnergyConsumerIds);
-        long consumedEnergy = 0;
-        if (energy != null) {
-            for (int i = energy.length - 1; i >= 0; i--) {
-                if (energy[i].energyUWs != ENERGY_UNSPECIFIED) {
-                    consumedEnergy += energy[i].energyUWs;
-                }
-            }
-        }
-
-        long energyDelta = mLastConsumedEnergyUws != ENERGY_UNSPECIFIED
-                ? consumedEnergy - mLastConsumedEnergyUws : 0;
-        mLastConsumedEnergyUws = consumedEnergy;
-        if (energyDelta < 0) {
-            // Likely, restart of powerstats HAL
-            energyDelta = 0;
-        }
-
-        if (energyDelta == 0 && !mFirstCollection) {
-            return null;
-        }
-
-        mLayout.setConsumedEnergy(mPowerStats.stats, 0, uJtoUc(energyDelta, averageVoltage));
-
         mPowerStats.uidStats.clear();
-        if (energy != null) {
-            for (int i = energy.length - 1; i >= 0; i--) {
-                EnergyConsumerAttribution[] perUid = energy[i].attribution;
-                if (perUid == null) {
-                    continue;
-                }
 
-                for (EnergyConsumerAttribution attribution : perUid) {
-                    int uid = mUidResolver.mapUid(attribution.uid);
-                    long lastEnergy = mLastConsumerEnergyPerUid.get(uid, ENERGY_UNSPECIFIED);
-                    mLastConsumerEnergyPerUid.put(uid, attribution.energyUWs);
-                    if (lastEnergy == ENERGY_UNSPECIFIED) {
-                        continue;
-                    }
-                    long deltaEnergy = attribution.energyUWs - lastEnergy;
-                    if (deltaEnergy <= 0) {
-                        continue;
-                    }
-                    long[] uidStats = mPowerStats.uidStats.get(uid);
-                    if (uidStats == null) {
-                        uidStats = new long[mLayout.getUidStatsArrayLength()];
-                        mPowerStats.uidStats.put(uid, uidStats);
-                    }
-
-                    mLayout.setUidConsumedEnergy(uidStats, 0,
-                            mLayout.getUidConsumedEnergy(uidStats, 0)
-                                    + uJtoUc(deltaEnergy, averageVoltage));
-                }
-            }
+        if (!mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout)) {
+            return null;
         }
+
         long timestamp = mClock.elapsedRealtime();
         mPowerStats.durationMs = timestamp - mLastUpdateTimestamp;
         mLastUpdateTimestamp = timestamp;
-        mFirstCollection = false;
         return mPowerStats;
     }
-
-    @Override
-    protected void onUidRemoved(int uid) {
-        mLastConsumerEnergyPerUid.delete(uid);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
index 168a874..c1b4c31 100644
--- a/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/GnssPowerStatsCollector.java
@@ -19,12 +19,13 @@
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 
+import com.android.server.power.stats.format.GnssPowerStatsLayout;
+
 public class GnssPowerStatsCollector extends EnergyConsumerPowerStatsCollector {
 
-    GnssPowerStatsCollector(Injector injector) {
+    public GnssPowerStatsCollector(Injector injector) {
         super(injector, BatteryConsumer.POWER_COMPONENT_GNSS,
                 BatteryConsumer.powerComponentIdToString(BatteryConsumer.POWER_COMPONENT_GNSS),
-                EnergyConsumerType.GNSS, /* energy consumer name */ null,
-                new GnssPowerStatsLayout());
+                EnergyConsumerType.GNSS, new GnssPowerStatsLayout());
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
index c88e1b0..cbd6fab 100644
--- a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsCollector.java
@@ -35,12 +35,12 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
 
 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;
 
@@ -56,8 +56,6 @@
 
     private static final long MODEM_ACTIVITY_REQUEST_TIMEOUT = 20000;
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     @VisibleForTesting
     @AccessNetworkConstants.RadioAccessNetworkType
     static final int[] NETWORK_TYPES = {
@@ -77,14 +75,13 @@
                 long elapsedRealtimeMs, long uptimeMs);
     }
 
-    interface Injector {
+    public interface Injector {
         Handler getHandler();
         Clock getClock();
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         Supplier<NetworkStats> getMobileNetworkStatsSupplier();
         TelephonyManager getTelephonyManager();
         LongSupplier getCallDurationSupplier();
@@ -103,18 +100,14 @@
     private LongSupplier mCallDurationSupplier;
     private LongSupplier mScanDurationSupplier;
     private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private long mLastUpdateTimestampMillis;
     private ModemActivityInfo mLastModemActivityInfo;
     private NetworkStats mLastNetworkStats;
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
     private long mLastCallDuration;
     private long mLastScanDuration;
 
-    MobileRadioPowerStatsCollector(Injector injector, Observer observer) {
+    public MobileRadioPowerStatsCollector(Injector injector, Observer observer) {
         super(injector.getHandler(), injector.getPowerStatsCollectionThrottlePeriod(
                         BatteryConsumer.powerComponentIdToString(
                                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)),
@@ -144,34 +137,22 @@
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
-
         mTelephonyManager = mInjector.getTelephonyManager();
         mNetworkStatsSupplier = mInjector.getMobileNetworkStatsSupplier();
         mCallDurationSupplier = mInjector.getCallDurationSupplier();
         mScanDurationSupplier = mInjector.getPhoneSignalScanDurationSupplier();
 
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
                 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();
+        mLayout = new MobileRadioPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount());
 
         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);
+                int stateKey = MobileRadioPowerStatsLayout.makeStateKey(rat, freq);
                 StringBuilder sb = new StringBuilder();
                 if (rat != BatteryStats.RADIO_ACCESS_TECHNOLOGY_OTHER) {
                     sb.append(BatteryStats.RADIO_ACCESS_TECHNOLOGY_NAMES[rat]);
@@ -200,7 +181,7 @@
     }
 
     @Override
-    protected PowerStats collectStats() {
+    public PowerStats collectStats() {
         if (!ensureInitialized()) {
             return null;
         }
@@ -210,9 +191,8 @@
 
         ModemActivityInfo modemActivityDelta = collectModemActivityInfo();
         List<BatteryStatsImpl.NetworkStatsDelta> networkStatsDeltas = collectNetworkStats();
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         if (mPowerStats.durationMs == 0) {
             setTimestamp(mClock.elapsedRealtime());
@@ -346,65 +326,8 @@
         return delta;
     }
 
-    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/PowerAttributor.java b/services/core/java/com/android/server/power/stats/PowerAttributor.java
new file mode 100644
index 0000000..d1f8564
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/PowerAttributor.java
@@ -0,0 +1,59 @@
+/*
+ * 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.BatteryUsageStats;
+import android.util.IndentingPrintWriter;
+
+import com.android.internal.os.BatteryStatsHistory;
+
+public interface PowerAttributor {
+
+    /**
+     * Returns true if the specified power component can be handled by this PowerAttributor
+     */
+    boolean isPowerComponentSupported(@BatteryConsumer.PowerComponentId int powerComponentId);
+
+    /**
+     * Performs the power attribution calculations and returns the results by populating the
+     * supplied BatteryUsageStats.Builder
+     */
+    void estimatePowerConsumption(BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            BatteryStatsHistory batteryHistory, long monotonicStartTime, long monotonicEndTime);
+
+    /**
+     * Computes estimated power consumption attribution for the specified time range and stores
+     * it in PowerStatsStore for potential accumulation.
+     *
+     * Returns the monotonic timestamp of the last processed history item.
+     */
+    long storeEstimatedPowerConsumption(BatteryStatsHistory batteryStatsHistory, long startTime,
+            long endTimeMs);
+
+    /**
+     * Returns the monotonic timestamp of the last processed history item, stored in
+     * PowerStatsStore.
+     */
+    long getLastSavedEstimatesPowerConsumptionTimestamp();
+
+    /**
+     * Performs the power attribution calculation and prints the results.
+     */
+    void dumpEstimatedPowerConsumption(IndentingPrintWriter ipw,
+            BatteryStatsHistory batteryStatsHistory, long startTime, long endTime);
+}
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 f5b0005..291f0e3 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsCollector.java
@@ -19,6 +19,7 @@
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.hardware.power.stats.EnergyConsumer;
+import android.hardware.power.stats.EnergyConsumerAttribution;
 import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.ConditionVariable;
@@ -26,13 +27,16 @@
 import android.power.PowerStatsInternal;
 import android.util.IndentingPrintWriter;
 import android.util.Slog;
+import android.util.SparseLongArray;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.PowerStatsLayout;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
@@ -41,6 +45,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Consumer;
+import java.util.function.IntSupplier;
 
 /**
  * Collects snapshots of power-related system statistics.
@@ -53,6 +58,8 @@
     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 static final long ENERGY_UNSPECIFIED = -1;
+
     private final Handler mHandler;
     protected final PowerStatsUidResolver mUidResolver;
     protected final Clock mClock;
@@ -235,46 +242,31 @@
         return (deltaEnergyUj * MILLIVOLTS_PER_VOLT + (avgVoltageMv / 2)) / avgVoltageMv;
     }
 
-    interface ConsumedEnergyRetriever {
+    public interface ConsumedEnergyRetriever {
+
         @NonNull
-        int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType, String name);
+        int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType);
 
         String getEnergyConsumerName(int energyConsumerId);
 
         @Nullable
         EnergyConsumerResult[] getConsumedEnergy(int[] energyConsumerIds);
 
-        @Nullable
-        default long[] getConsumedEnergyUws(int[] energyConsumerIds) {
-            EnergyConsumerResult[] results = getConsumedEnergy(energyConsumerIds);
-            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;
-        }
-
-        default int[] getEnergyConsumerIds(@EnergyConsumerType int energyConsumerType) {
-            return getEnergyConsumerIds(energyConsumerType, null);
-        }
+        /**
+         * Returns the last known battery/charger voltage in milli-volts.
+         */
+        int getVoltageMv();
     }
 
     static class ConsumedEnergyRetrieverImpl implements ConsumedEnergyRetriever {
         private final PowerStatsInternal mPowerStatsInternal;
+        private final IntSupplier mVoltageSupplier;
         private EnergyConsumer[] mEnergyConsumers;
 
-        ConsumedEnergyRetrieverImpl(PowerStatsInternal powerStatsInternal) {
+        ConsumedEnergyRetrieverImpl(PowerStatsInternal powerStatsInternal,
+                IntSupplier voltageSupplier) {
             mPowerStatsInternal = powerStatsInternal;
+            mVoltageSupplier = voltageSupplier;
         }
 
         private void ensureEnergyConsumers() {
@@ -293,8 +285,9 @@
             }
         }
 
+        @NonNull
         @Override
-        public int[] getEnergyConsumerIds(int energyConsumerType, String name) {
+        public int[] getEnergyConsumerIds(int energyConsumerType) {
             ensureEnergyConsumers();
 
             if (mEnergyConsumers.length == 0) {
@@ -303,8 +296,7 @@
 
             List<EnergyConsumer> energyConsumers = new ArrayList<>();
             for (EnergyConsumer energyConsumer : mEnergyConsumers) {
-                if (energyConsumer.type == energyConsumerType
-                        && (name == null || name.equals(energyConsumer.name))) {
+                if (energyConsumer.type == energyConsumerType) {
                     energyConsumers.add(energyConsumer);
                 }
             }
@@ -335,6 +327,11 @@
         }
 
         @Override
+        public int getVoltageMv() {
+            return mVoltageSupplier.getAsInt();
+        }
+
+        @Override
         public String getEnergyConsumerName(int energyConsumerId) {
             ensureEnergyConsumers();
 
@@ -368,4 +365,152 @@
             return sb.toString();
         }
     }
+
+    class ConsumedEnergyHelper implements PowerStatsUidResolver.Listener {
+        private final ConsumedEnergyRetriever mConsumedEnergyRetriever;
+        private final @EnergyConsumerType int mEnergyConsumerType;
+        private final boolean mPerUidAttributionSupported;
+
+        private boolean mIsInitialized;
+        private boolean mFirstCollection = true;
+        private int[] mEnergyConsumerIds;
+        private long[] mLastConsumedEnergyUws;
+        private final SparseLongArray mLastConsumerEnergyPerUid;
+        private int mLastVoltageMv;
+
+        ConsumedEnergyHelper(ConsumedEnergyRetriever consumedEnergyRetriever,
+                @EnergyConsumerType int energyConsumerType) {
+            mConsumedEnergyRetriever = consumedEnergyRetriever;
+            mEnergyConsumerType = energyConsumerType;
+            mPerUidAttributionSupported = false;
+            mLastConsumerEnergyPerUid = null;
+        }
+
+        ConsumedEnergyHelper(ConsumedEnergyRetriever consumedEnergyRetriever,
+                int energyConsumerId, boolean perUidAttributionSupported) {
+            mConsumedEnergyRetriever = consumedEnergyRetriever;
+            mEnergyConsumerType = EnergyConsumerType.OTHER;
+            mEnergyConsumerIds = new int[]{energyConsumerId};
+            mPerUidAttributionSupported = perUidAttributionSupported;
+            mLastConsumerEnergyPerUid = mPerUidAttributionSupported ? new SparseLongArray() : null;
+        }
+
+        private void ensureInitialized() {
+            if (!mIsInitialized) {
+                if (mEnergyConsumerIds == null) {
+                    mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
+                            mEnergyConsumerType);
+                }
+                mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
+                Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+                mUidResolver.addListener(this);
+                mIsInitialized = true;
+            }
+        }
+
+        int getEnergyConsumerCount() {
+            ensureInitialized();
+            return mEnergyConsumerIds.length;
+        }
+
+        boolean collectConsumedEnergy(PowerStats powerStats, PowerStatsLayout layout) {
+            ensureInitialized();
+
+            if (mEnergyConsumerIds.length == 0) {
+                return false;
+            }
+
+            int voltageMv = mConsumedEnergyRetriever.getVoltageMv();
+            int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
+            if (averageVoltage <= 0) {
+                Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
+                        + " mV) when querying energy consumers");
+                return false;
+            }
+
+            mLastVoltageMv = voltageMv;
+
+            EnergyConsumerResult[] energy =
+                    mConsumedEnergyRetriever.getConsumedEnergy(mEnergyConsumerIds);
+            System.out.println("mEnergyConsumerIds = " + Arrays.toString(mEnergyConsumerIds) + " "
+                    + "energy = "
+                    + Arrays.toString(energy));
+            if (energy == null) {
+                return false;
+            }
+
+            for (int i = 0; i < mEnergyConsumerIds.length; i++) {
+                populatePowerStats(powerStats, layout, energy, i, averageVoltage);
+            }
+            mFirstCollection = false;
+            return true;
+        }
+
+        private void populatePowerStats(PowerStats powerStats, PowerStatsLayout layout,
+                @NonNull EnergyConsumerResult[] energy, int energyConsumerIndex,
+                int averageVoltage) {
+            long consumedEnergy = energy[energyConsumerIndex].energyUWs;
+            long energyDelta = mLastConsumedEnergyUws[energyConsumerIndex] != ENERGY_UNSPECIFIED
+                    ? consumedEnergy - mLastConsumedEnergyUws[energyConsumerIndex] : 0;
+            mLastConsumedEnergyUws[energyConsumerIndex] = consumedEnergy;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
+            }
+
+            if (energyDelta == 0 && !mFirstCollection) {
+                return;
+            }
+
+            layout.setConsumedEnergy(powerStats.stats, energyConsumerIndex,
+                    uJtoUc(energyDelta, averageVoltage));
+
+            if (!mPerUidAttributionSupported) {
+                return;
+            }
+
+            EnergyConsumerAttribution[] perUid = energy[energyConsumerIndex].attribution;
+            if (perUid == null) {
+                return;
+            }
+
+            for (EnergyConsumerAttribution attribution : perUid) {
+                int uid = mUidResolver.mapUid(attribution.uid);
+                long lastEnergy = mLastConsumerEnergyPerUid.get(uid, ENERGY_UNSPECIFIED);
+                mLastConsumerEnergyPerUid.put(uid, attribution.energyUWs);
+                if (lastEnergy == ENERGY_UNSPECIFIED) {
+                    continue;
+                }
+                long deltaEnergy = attribution.energyUWs - lastEnergy;
+                if (deltaEnergy <= 0) {
+                    continue;
+                }
+
+                long[] uidStats = powerStats.uidStats.get(uid);
+                if (uidStats == null) {
+                    uidStats = new long[layout.getUidStatsArrayLength()];
+                    powerStats.uidStats.put(uid, uidStats);
+                }
+
+                layout.setUidConsumedEnergy(uidStats, energyConsumerIndex,
+                        layout.getUidConsumedEnergy(uidStats, energyConsumerIndex)
+                                + uJtoUc(deltaEnergy, averageVoltage));
+            }
+        }
+
+        @Override
+        public void onAfterIsolatedUidRemoved(int isolatedUid, int parentUid) {
+            if (mLastConsumerEnergyPerUid != null) {
+                mHandler.post(() -> mLastConsumerEnergyPerUid.delete(isolatedUid));
+            }
+        }
+
+        @Override
+        public void onIsolatedUidAdded(int isolatedUid, int parentUid) {
+        }
+
+        @Override
+        public void onBeforeIsolatedUidRemoved(int isolatedUid, int parentUid) {
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java b/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
index abe4c0c..38ca087 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsScheduler.java
@@ -23,6 +23,7 @@
 import android.util.IndentingPrintWriter;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.Clock;
 import com.android.internal.os.MonotonicClock;
 
@@ -51,7 +52,8 @@
     private final Handler mHandler;
     private final Runnable mPowerStatsCollector;
     private final Supplier<Long> mEarliestAvailableBatteryHistoryTimeMs;
-    private final PowerStatsAggregator mPowerStatsAggregator;
+    private final BatteryStatsHistory mBatteryStatsHistory;
+    private final PowerAttributor mPowerAttributor;
     private long mLastSavedSpanEndMonotonicTime;
 
     /**
@@ -66,12 +68,13 @@
     }
 
     public PowerStatsScheduler(Runnable powerStatsCollector,
-            PowerStatsAggregator powerStatsAggregator,
+            BatteryStatsHistory batteryStatsHistory, PowerAttributor powerAttributor,
             @DurationMillisLong long aggregatedPowerStatsSpanDuration,
             @DurationMillisLong long powerStatsAggregationPeriod, PowerStatsStore powerStatsStore,
             AlarmScheduler alarmScheduler, Clock clock, MonotonicClock monotonicClock,
             Supplier<Long> earliestAvailableBatteryHistoryTimeMs, Handler handler) {
-        mPowerStatsAggregator = powerStatsAggregator;
+        mBatteryStatsHistory = batteryStatsHistory;
+        mPowerAttributor = powerAttributor;
         mAggregatedPowerStatsSpanDuration = aggregatedPowerStatsSpanDuration;
         mPowerStatsAggregationPeriod = powerStatsAggregationPeriod;
         mPowerStatsStore = powerStatsStore;
@@ -123,12 +126,8 @@
         long endTimeMs = alignToWallClock(startTime + mAggregatedPowerStatsSpanDuration,
                 mAggregatedPowerStatsSpanDuration, currentMonotonicTime, currentTimeMillis);
         while (endTimeMs <= currentMonotonicTime) {
-            mPowerStatsAggregator.aggregatePowerStats(startTime, endTimeMs,
-                    stats -> {
-                        storeAggregatedPowerStats(stats);
-                        mLastSavedSpanEndMonotonicTime = stats.getStartTime() + stats.getDuration();
-                    });
-
+            mLastSavedSpanEndMonotonicTime = mPowerAttributor.storeEstimatedPowerConsumption(
+                    mBatteryStatsHistory, startTime, endTimeMs);
             startTime = endTimeMs;
             endTimeMs += mAggregatedPowerStatsSpanDuration;
         }
@@ -153,15 +152,8 @@
             mPowerStatsStore.dump(ipw);
             // Aggregate the remainder of power stats and dump the results without storing them yet.
             long powerStoreEndMonotonicTime = getLastSavedSpanEndMonotonicTime();
-            mPowerStatsAggregator.aggregatePowerStats(powerStoreEndMonotonicTime,
-                    MonotonicClock.UNDEFINED,
-                    stats -> {
-                        // Create a PowerStatsSpan for consistency of the textual output
-                        PowerStatsSpan span = PowerStatsStore.createPowerStatsSpan(stats);
-                        if (span != null) {
-                            span.dump(ipw);
-                        }
-                    });
+            mPowerAttributor.dumpEstimatedPowerConsumption(ipw, mBatteryStatsHistory,
+                    powerStoreEndMonotonicTime, MonotonicClock.UNDEFINED);
         });
 
         awaitCompletion();
@@ -223,28 +215,13 @@
     }
 
     private long getLastSavedSpanEndMonotonicTime() {
-        if (mLastSavedSpanEndMonotonicTime != 0) {
-            return mLastSavedSpanEndMonotonicTime;
-        }
-
-        mLastSavedSpanEndMonotonicTime = -1;
-        for (PowerStatsSpan.Metadata metadata : mPowerStatsStore.getTableOfContents()) {
-            if (metadata.getSections().contains(AggregatedPowerStatsSection.TYPE)) {
-                for (PowerStatsSpan.TimeFrame timeFrame : metadata.getTimeFrames()) {
-                    long endMonotonicTime = timeFrame.startMonotonicTime + timeFrame.duration;
-                    if (endMonotonicTime > mLastSavedSpanEndMonotonicTime) {
-                        mLastSavedSpanEndMonotonicTime = endMonotonicTime;
-                    }
-                }
-            }
+        if (mLastSavedSpanEndMonotonicTime == 0) {
+            mLastSavedSpanEndMonotonicTime =
+                    mPowerAttributor.getLastSavedEstimatesPowerConsumptionTimestamp();
         }
         return mLastSavedSpanEndMonotonicTime;
     }
 
-    private void storeAggregatedPowerStats(AggregatedPowerStats stats) {
-        mPowerStatsStore.storeAggregatedPowerStats(stats);
-    }
-
     private void awaitCompletion() {
         ConditionVariable done = new ConditionVariable();
         mHandler.post(done::open);
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
index 4df919d..fc0611f 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsSpan.java
@@ -44,6 +44,7 @@
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
@@ -72,7 +73,7 @@
     private static final DateTimeFormatter DATE_FORMAT =
             DateTimeFormatter.ofPattern("MM-dd HH:mm:ss.SSS").withZone(ZoneId.systemDefault());
 
-    static class TimeFrame {
+    public static class TimeFrame {
         public final long startMonotonicTime;
         @CurrentTimeMillisLong
         public final long startTime;
@@ -119,7 +120,7 @@
         }
     }
 
-    static class Metadata {
+    public static class Metadata {
         static final Comparator<Metadata> COMPARATOR = Comparator.comparing(Metadata::getId);
 
         private final long mId;
@@ -262,7 +263,7 @@
     public abstract static class Section {
         private final String mType;
 
-        Section(String type) {
+        protected Section(String type) {
             mType = type;
         }
 
@@ -274,7 +275,10 @@
             return mType;
         }
 
-        abstract void write(TypedXmlSerializer serializer) throws IOException;
+        /**
+         * Adds the contents of this section to the XML doc.
+         */
+        public abstract void write(TypedXmlSerializer serializer) throws IOException;
 
         /**
          * Prints the section type.
@@ -290,6 +294,11 @@
      */
     public interface SectionReader {
         /**
+         * Returns the unique type of content handled by this reader.
+         */
+        String getType();
+
+        /**
          * Reads the contents of the section using the parser. The type of the object
          * read and the corresponding XML format are determined by the section type.
          */
@@ -316,12 +325,18 @@
         return mMetadata.mId;
     }
 
-    void addTimeFrame(long monotonicTime, @CurrentTimeMillisLong long wallClockTime,
+    /**
+     * Adds a time frame covered by this PowerStats span
+     */
+    public void addTimeFrame(long monotonicTime, @CurrentTimeMillisLong long wallClockTime,
             @DurationMillisLong long duration) {
         mMetadata.mTimeFrames.add(new TimeFrame(monotonicTime, wallClockTime, duration));
     }
 
-    void addSection(Section section) {
+    /**
+     * Adds the supplied section to the span.
+     */
+    public void addSection(Section section) {
         mMetadata.addSection(section.getType());
         mSections.add(section);
     }
@@ -354,7 +369,7 @@
 
     @Nullable
     static PowerStatsSpan read(InputStream in, TypedXmlPullParser parser,
-            SectionReader sectionReader, String... sectionTypes)
+            Map<String, SectionReader> sectionReaders, String... sectionTypes)
             throws IOException, XmlPullParserException {
         Set<String> neededSections = Sets.newArraySet(sectionTypes);
         boolean selectSections = !neededSections.isEmpty();
@@ -386,7 +401,11 @@
                 if (tag.equals(XML_TAG_SECTION)) {
                     String sectionType = parser.getAttributeValue(null, XML_ATTR_SECTION_TYPE);
                     if (!selectSections || neededSections.contains(sectionType)) {
-                        Section section = sectionReader.read(sectionType, parser);
+                        Section section = null;
+                        SectionReader sectionReader = sectionReaders.get(sectionType);
+                        if (sectionReader != null) {
+                            section = sectionReader.read(sectionType, parser);
+                        }
                         if (section == null) {
                             if (selectSections) {
                                 throw new XmlPullParserException(
@@ -396,11 +415,11 @@
                                     @Override
                                     public void dump(IndentingPrintWriter ipw) {
                                         ipw.println("Unsupported PowerStatsStore section type: "
-                                                    + sectionType);
+                                                + sectionType);
                                     }
 
                                     @Override
-                                    void write(TypedXmlSerializer serializer) {
+                                    public void write(TypedXmlSerializer serializer) {
                                     }
                                 };
                             }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsStore.java b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
index 7bcdc71..a875c30 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsStore.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsStore.java
@@ -42,6 +42,7 @@
 import java.nio.file.StandardOpenOption;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
@@ -66,28 +67,31 @@
     private FileLock mJvmLock;
     private final long mMaxStorageBytes;
     private final Handler mHandler;
-    private final PowerStatsSpan.SectionReader mSectionReader;
+    private final Map<String, PowerStatsSpan.SectionReader> mSectionReaders = new HashMap<>();
     private volatile List<PowerStatsSpan.Metadata> mTableOfContents;
 
-    public PowerStatsStore(@NonNull File systemDir, Handler handler,
-            AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
-        this(systemDir, MAX_POWER_STATS_SPAN_STORAGE_BYTES, handler,
-                new DefaultSectionReader(aggregatedPowerStatsConfig));
+    public PowerStatsStore(@NonNull File systemDir, Handler handler) {
+        this(systemDir, MAX_POWER_STATS_SPAN_STORAGE_BYTES, handler);
     }
 
     @VisibleForTesting
-    public PowerStatsStore(@NonNull File systemDir, long maxStorageBytes, Handler handler,
-            @NonNull PowerStatsSpan.SectionReader sectionReader) {
+    public PowerStatsStore(@NonNull File systemDir, long maxStorageBytes, Handler handler) {
         mSystemDir = systemDir;
         mStoreDir = new File(systemDir, POWER_STATS_DIR);
         mLockFile = new File(mStoreDir, DIR_LOCK_FILENAME);
         mHandler = handler;
         mMaxStorageBytes = maxStorageBytes;
-        mSectionReader = sectionReader;
         mHandler.post(this::maybeClearLegacyStore);
     }
 
     /**
+     * Registers a Reader for a section type, which is determined by `sectionReader.getType()`
+     */
+    public void addSectionReader(PowerStatsSpan.SectionReader sectionReader) {
+        mSectionReaders.put(sectionReader.getType(), sectionReader);
+    }
+
+    /**
      * Returns the metadata for all {@link PowerStatsSpan}'s contained in the store.
      */
     @NonNull
@@ -169,7 +173,7 @@
         try {
             File file = makePowerStatsSpanFilename(id);
             try (InputStream inputStream = new BufferedInputStream(new FileInputStream(file))) {
-                return PowerStatsSpan.read(inputStream, parser, mSectionReader, sectionTypes);
+                return PowerStatsSpan.read(inputStream, parser, mSectionReaders, sectionTypes);
             } catch (IOException | XmlPullParserException e) {
                 Slog.wtf(TAG, "Cannot read PowerStatsSpan file: " + file, e);
             }
@@ -179,41 +183,6 @@
         return null;
     }
 
-    void storeAggregatedPowerStats(AggregatedPowerStats stats) {
-        PowerStatsSpan span = createPowerStatsSpan(stats);
-        if (span == null) {
-            return;
-        }
-        storePowerStatsSpan(span);
-    }
-
-    static PowerStatsSpan createPowerStatsSpan(AggregatedPowerStats stats) {
-        List<AggregatedPowerStats.ClockUpdate> clockUpdates = stats.getClockUpdates();
-        if (clockUpdates.isEmpty()) {
-            Slog.w(TAG, "No clock updates in aggregated power stats " + stats);
-            return null;
-        }
-
-        long monotonicTime = clockUpdates.get(0).monotonicTime;
-        long durationSum = 0;
-        PowerStatsSpan span = new PowerStatsSpan(monotonicTime);
-        for (int i = 0; i < clockUpdates.size(); i++) {
-            AggregatedPowerStats.ClockUpdate clockUpdate = clockUpdates.get(i);
-            long duration;
-            if (i == clockUpdates.size() - 1) {
-                duration = stats.getDuration() - durationSum;
-            } else {
-                duration = clockUpdate.monotonicTime - monotonicTime;
-            }
-            span.addTimeFrame(clockUpdate.monotonicTime, clockUpdate.currentTime, duration);
-            monotonicTime = clockUpdate.monotonicTime;
-            durationSum += duration;
-        }
-
-        span.addSection(new AggregatedPowerStatsSection(stats));
-        return span;
-    }
-
     /**
      * Stores a {@link PowerStatsSpan} containing a single section for the supplied
      * battery usage stats.
@@ -344,28 +313,4 @@
         }
         ipw.decreaseIndent();
     }
-
-    private static class DefaultSectionReader implements PowerStatsSpan.SectionReader {
-        private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
-
-        DefaultSectionReader(AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
-            mAggregatedPowerStatsConfig = aggregatedPowerStatsConfig;
-        }
-
-        @Override
-        public PowerStatsSpan.Section read(String sectionType, TypedXmlPullParser parser)
-                throws IOException, XmlPullParserException {
-            switch (sectionType) {
-                case AggregatedPowerStatsSection.TYPE:
-                    return new AggregatedPowerStatsSection(
-                            AggregatedPowerStats.createFromXml(parser,
-                                    mAggregatedPowerStatsConfig));
-                case BatteryUsageStatsSection.TYPE:
-                    return new BatteryUsageStatsSection(
-                            BatteryUsageStats.createFromXml(parser));
-                default:
-                    return null;
-            }
-        }
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
index 291f289..8371e66 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsCollector.java
@@ -21,19 +21,16 @@
 import android.os.BatteryStats;
 import android.os.Handler;
 import android.os.PersistableBundle;
-import android.util.Slog;
 import android.util.SparseLongArray;
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
-
-import java.util.Arrays;
-import java.util.function.IntSupplier;
+import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
 public class ScreenPowerStatsCollector extends PowerStatsCollector {
     private static final String TAG = "ScreenPowerStatsCollector";
 
-    interface ScreenUsageTimeRetriever {
+    public interface ScreenUsageTimeRetriever {
         interface Callback {
             void onUidTopActivityTime(int uid, long topActivityTimeMs);
         }
@@ -45,30 +42,23 @@
         long getScreenDozeTimeMs(int display);
     }
 
-    interface Injector {
+    public interface Injector {
         Handler getHandler();
         Clock getClock();
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         ScreenUsageTimeRetriever getScreenUsageTimeRetriever();
         int getDisplayCount();
     }
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     private final Injector mInjector;
     private boolean mIsInitialized;
     private ScreenPowerStatsLayout mLayout;
     private int mDisplayCount;
     private PowerStats mPowerStats;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
-    private int[] mEnergyConsumerIds = new int[0];
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
     private boolean mFirstSample = true;
     private long[] mLastScreenOnTime;
     private long[][] mLastBrightnessLevelTime;
@@ -76,7 +66,7 @@
     private final SparseLongArray mLastTopActivityTime = new SparseLongArray();
     private long mLastCollectionTime;
 
-    ScreenPowerStatsCollector(Injector injector) {
+    public ScreenPowerStatsCollector(Injector injector) {
         super(injector.getHandler(),
                 injector.getPowerStatsCollectionThrottlePeriod(
                         BatteryConsumer.powerComponentIdToString(
@@ -95,21 +85,12 @@
         }
 
         mDisplayCount = mInjector.getDisplayCount();
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mScreenUsageTimeRetriever = mInjector.getScreenUsageTimeRetriever();
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(
-                EnergyConsumerType.DISPLAY);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
 
-        mLayout = new ScreenPowerStatsLayout();
-        mLayout.addDeviceScreenUsageDurationSection(mInjector.getDisplayCount());
-        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
-        mLayout.addDeviceSectionUsageDuration();
-        mLayout.addDeviceSectionPowerEstimate();
-        mLayout.addUidTopActivitiyDuration();
-        mLayout.addUidSectionPowerEstimate();
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.DISPLAY);
+        mLayout = new ScreenPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
+                mInjector.getDisplayCount());
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -129,14 +110,12 @@
     }
 
     @Override
-    protected PowerStats collectStats() {
+    public PowerStats collectStats() {
         if (!ensureInitialized()) {
             return null;
         }
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         for (int display = 0; display < mDisplayCount; display++) {
             long screenOnTimeMs = mScreenUsageTimeRetriever.getScreenOnTimeMs(display);
@@ -192,34 +171,6 @@
         return mPowerStats;
     }
 
-    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];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         mLastTopActivityTime.delete(uid);
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
index 90981ada..7a84b05 100644
--- a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
@@ -28,12 +28,11 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.WifiPowerStatsLayout;
 
-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.Supplier;
 
 public class WifiPowerStatsCollector extends PowerStatsCollector {
@@ -41,15 +40,13 @@
 
     private static final long WIFI_ACTIVITY_REQUEST_TIMEOUT = 20000;
 
-    private static final long ENERGY_UNSPECIFIED = -1;
-
     interface Observer {
         void onWifiPowerStatsRetrieved(WifiActivityEnergyInfo info,
                 List<BatteryStatsImpl.NetworkStatsDelta> delta, long elapsedRealtimeMs,
                 long uptimeMs);
     }
 
-    interface WifiStatsRetriever {
+    public interface WifiStatsRetriever {
         interface Callback {
             void onWifiScanTime(int uid, long scanTimeMs, long batchScanTimeMs);
         }
@@ -58,14 +55,13 @@
         long getWifiActiveDuration();
     }
 
-    interface Injector {
+    public interface Injector {
         Handler getHandler();
         Clock getClock();
         PowerStatsUidResolver getUidResolver();
         long getPowerStatsCollectionThrottlePeriod(String powerComponentName);
         PackageManager getPackageManager();
         ConsumedEnergyRetriever getConsumedEnergyRetriever();
-        IntSupplier getVoltageSupplier();
         Supplier<NetworkStats> getWifiNetworkStatsSupplier();
         WifiManager getWifiManager();
         WifiStatsRetriever getWifiStatsRetriever();
@@ -83,13 +79,9 @@
     private volatile WifiManager mWifiManager;
     private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
     private volatile WifiStatsRetriever mWifiStatsRetriever;
-    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
-    private IntSupplier mVoltageSupplier;
-    private int[] mEnergyConsumerIds = new int[0];
+    private ConsumedEnergyHelper mConsumedEnergyHelper;
     private WifiActivityEnergyInfo mLastWifiActivityInfo;
     private NetworkStats mLastNetworkStats;
-    private long[] mLastConsumedEnergyUws;
-    private int mLastVoltageMv;
 
     private static class WifiScanTimes {
         public long basicScanTimeMs;
@@ -99,7 +91,7 @@
     private final SparseArray<WifiScanTimes> mLastScanTimes = new SparseArray<>();
     private long mLastWifiActiveDuration;
 
-    WifiPowerStatsCollector(Injector injector, Observer observer) {
+    public WifiPowerStatsCollector(Injector injector, Observer observer) {
         super(injector.getHandler(), injector.getPowerStatsCollectionThrottlePeriod(
                         BatteryConsumer.powerComponentIdToString(
                                 BatteryConsumer.POWER_COMPONENT_WIFI)),
@@ -128,25 +120,17 @@
             return false;
         }
 
-        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
-        mVoltageSupplier = mInjector.getVoltageSupplier();
         mWifiManager = mInjector.getWifiManager();
         mNetworkStatsSupplier = mInjector.getWifiNetworkStatsSupplier();
         mWifiStatsRetriever = mInjector.getWifiStatsRetriever();
         mPowerReportingSupported =
                 mWifiManager != null && mWifiManager.isEnhancedPowerReportingSupported();
 
-        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI);
-        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
-        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+        mConsumedEnergyHelper = new ConsumedEnergyHelper(mInjector.getConsumedEnergyRetriever(),
+                EnergyConsumerType.WIFI);
 
-        mLayout = new WifiPowerStatsLayout();
-        mLayout.addDeviceWifiActivity(mPowerReportingSupported);
-        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
-        mLayout.addUidNetworkStats();
-        mLayout.addDeviceSectionUsageDuration();
-        mLayout.addDeviceSectionPowerEstimate();
-        mLayout.addUidSectionPowerEstimate();
+        mLayout = new WifiPowerStatsLayout(mConsumedEnergyHelper.getEnergyConsumerCount(),
+                mPowerReportingSupported);
 
         PersistableBundle extras = new PersistableBundle();
         mLayout.toExtras(extras);
@@ -162,7 +146,7 @@
     }
 
     @Override
-    protected PowerStats collectStats() {
+    public PowerStats collectStats() {
         if (!ensureInitialized()) {
             return null;
         }
@@ -176,9 +160,7 @@
         List<BatteryStatsImpl.NetworkStatsDelta> networkStatsDeltas = collectNetworkStats();
         collectWifiScanTime();
 
-        if (mEnergyConsumerIds.length != 0) {
-            collectEnergyConsumers();
-        }
+        mConsumedEnergyHelper.collectConsumedEnergy(mPowerStats, mLayout);
 
         if (mObserver != null) {
             mObserver.onWifiPowerStatsRetrieved(activityInfo, networkStatsDeltas,
@@ -318,34 +300,6 @@
         mLayout.setDeviceBatchedScanTime(mDeviceStats, mScanTimes.batchedScanTimeMs);
     }
 
-    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];
-        }
-    }
-
     @Override
     protected void onUidRemoved(int uid) {
         super.onUidRemoved(uid);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl b/services/core/java/com/android/server/power/stats/format/AmbientDisplayPowerStatsLayout.java
similarity index 73%
copy from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
copy to services/core/java/com/android/server/power/stats/format/AmbientDisplayPowerStatsLayout.java
index 3c5beeb..1b99b0d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
+++ b/services/core/java/com/android/server/power/stats/format/AmbientDisplayPowerStatsLayout.java
@@ -13,7 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+package com.android.server.power.stats.format;
 
-package com.android.wm.shell.common.bubbles;
-
-parcelable BubbleBarLocation;
\ No newline at end of file
+public class AmbientDisplayPowerStatsLayout extends PowerStatsLayout {
+    public AmbientDisplayPowerStatsLayout() {
+        addDeviceSectionPowerEstimate();
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/BinaryStatePowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
similarity index 68%
rename from services/core/java/com/android/server/power/stats/BinaryStatePowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
index 502337c..4a26d83 100644
--- a/services/core/java/com/android/server/power/stats/BinaryStatePowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/BinaryStatePowerStatsLayout.java
@@ -14,11 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
-class BinaryStatePowerStatsLayout extends EnergyConsumerPowerStatsLayout {
-    BinaryStatePowerStatsLayout() {
+import com.android.internal.os.PowerStats;
+
+public class BinaryStatePowerStatsLayout extends EnergyConsumerPowerStatsLayout {
+    public BinaryStatePowerStatsLayout() {
         addDeviceSectionUsageDuration();
         addUidSectionUsageDuration();
     }
+
+    public BinaryStatePowerStatsLayout(PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/BluetoothPowerStatsLayout.java
similarity index 88%
rename from services/core/java/com/android/server/power/stats/BluetoothPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/BluetoothPowerStatsLayout.java
index 9358b5e..534a9f7 100644
--- a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/BluetoothPowerStatsLayout.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.annotation.NonNull;
 import android.os.PersistableBundle;
@@ -37,21 +37,49 @@
     private int mUidTxBytesPosition;
     private int mUidScanTimePosition;
 
-    BluetoothPowerStatsLayout() {
+    public BluetoothPowerStatsLayout(int energyConsumerCount) {
+        addDeviceBluetoothControllerActivity();
+        addDeviceSectionEnergyConsumers(energyConsumerCount);
+        addDeviceSectionUsageDuration();
+        addDeviceSectionPowerEstimate();
+        addUidTrafficStats();
+        addUidSectionPowerEstimate();
     }
 
-    BluetoothPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+    public BluetoothPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
         super(descriptor);
+        PersistableBundle extras = descriptor.extras;
+        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
+        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
+        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
+        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
+        mUidRxBytesPosition = extras.getInt(EXTRA_UID_RX_BYTES_POSITION);
+        mUidTxBytesPosition = extras.getInt(EXTRA_UID_TX_BYTES_POSITION);
+        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
     }
 
-    void addDeviceBluetoothControllerActivity() {
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
+        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
+        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
+        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
+        extras.putInt(EXTRA_UID_RX_BYTES_POSITION, mUidRxBytesPosition);
+        extras.putInt(EXTRA_UID_TX_BYTES_POSITION, mUidTxBytesPosition);
+        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
+    }
+
+    private void addDeviceBluetoothControllerActivity() {
         mDeviceRxTimePosition = addDeviceSection(1, "rx");
         mDeviceTxTimePosition = addDeviceSection(1, "tx");
         mDeviceIdleTimePosition = addDeviceSection(1, "idle");
         mDeviceScanTimePosition = addDeviceSection(1, "scan", FLAG_OPTIONAL);
     }
 
-    void addUidTrafficStats() {
+    private void addUidTrafficStats() {
         mUidRxBytesPosition = addUidSection(1, "rx-B");
         mUidTxBytesPosition = addUidSection(1, "tx-B");
         mUidScanTimePosition = addUidSection(1, "scan", FLAG_OPTIONAL);
@@ -112,32 +140,4 @@
     public long getUidScanTime(long[] stats) {
         return stats[mUidScanTimePosition];
     }
-
-    /**
-     * Copies the elements of the stats array layout into <code>extras</code>
-     */
-    public void toExtras(PersistableBundle extras) {
-        super.toExtras(extras);
-        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
-        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
-        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
-        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
-        extras.putInt(EXTRA_UID_RX_BYTES_POSITION, mUidRxBytesPosition);
-        extras.putInt(EXTRA_UID_TX_BYTES_POSITION, mUidTxBytesPosition);
-        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
-    }
-
-    /**
-     * Retrieves elements of the stats array layout from <code>extras</code>
-     */
-    public void fromExtras(PersistableBundle extras) {
-        super.fromExtras(extras);
-        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
-        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
-        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
-        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
-        mUidRxBytesPosition = extras.getInt(EXTRA_UID_RX_BYTES_POSITION);
-        mUidTxBytesPosition = extras.getInt(EXTRA_UID_TX_BYTES_POSITION);
-        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/CpuPowerStatsLayout.java
similarity index 85%
rename from services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/CpuPowerStatsLayout.java
index 2a02bd0..3186d7d 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/CpuPowerStatsLayout.java
@@ -14,10 +14,13 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
+import android.annotation.NonNull;
 import android.os.PersistableBundle;
 
+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.
@@ -40,10 +43,59 @@
 
     private int[] mScalingStepToPowerBracketMap;
 
+    public CpuPowerStatsLayout(int energyConsumerCount, int cpuScalingPolicyCount,
+            int[] scalingStepToPowerBracketMap) {
+        addDeviceSectionCpuTimeByScalingStep(scalingStepToPowerBracketMap.length);
+        addDeviceSectionCpuTimeByCluster(cpuScalingPolicyCount);
+        addDeviceSectionUsageDuration();
+        addDeviceSectionEnergyConsumers(energyConsumerCount);
+        addDeviceSectionPowerEstimate();
+        addUidSectionCpuTimeByPowerBracket(scalingStepToPowerBracketMap);
+        addUidSectionPowerEstimate();
+    }
+
+    public CpuPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+        PersistableBundle extras = descriptor.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();
+    }
+
+    /**
+     * 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);
+    }
+
     /**
      * Declare that the stats array has a section capturing CPU time per scaling step
      */
-    public void addDeviceSectionCpuTimeByScalingStep(int scalingStepCount) {
+    private void addDeviceSectionCpuTimeByScalingStep(int scalingStepCount) {
         mDeviceCpuTimeByScalingStepPosition = addDeviceSection(scalingStepCount, "steps");
         mDeviceCpuTimeByScalingStepCount = scalingStepCount;
     }
@@ -71,7 +123,7 @@
     /**
      * Declare that the stats array has a section capturing CPU time in each cluster
      */
-    public void addDeviceSectionCpuTimeByCluster(int clusterCount) {
+    private void addDeviceSectionCpuTimeByCluster(int clusterCount) {
         mDeviceCpuTimeByClusterPosition = addDeviceSection(clusterCount, "clusters");
         mDeviceCpuTimeByClusterCount = clusterCount;
     }
@@ -99,7 +151,7 @@
     /**
      * Declare that the UID stats array has a section capturing CPU time per power bracket.
      */
-    public void addUidSectionCpuTimeByPowerBracket(int[] scalingStepToPowerBracketMap) {
+    private void addUidSectionCpuTimeByPowerBracket(int[] scalingStepToPowerBracketMap) {
         mScalingStepToPowerBracketMap = scalingStepToPowerBracketMap;
         updatePowerBracketCount();
         mUidPowerBracketsPosition = addUidSection(mUidPowerBracketCount, "time");
@@ -135,44 +187,4 @@
     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/EnergyConsumerPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/EnergyConsumerPowerStatsLayout.java
similarity index 80%
rename from services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/EnergyConsumerPowerStatsLayout.java
index 8430f56..e7a4822 100644
--- a/services/core/java/com/android/server/power/stats/EnergyConsumerPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/EnergyConsumerPowerStatsLayout.java
@@ -14,10 +14,12 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
-class EnergyConsumerPowerStatsLayout extends PowerStatsLayout {
-    EnergyConsumerPowerStatsLayout() {
+import com.android.internal.os.PowerStats;
+
+public class EnergyConsumerPowerStatsLayout extends PowerStatsLayout {
+    public EnergyConsumerPowerStatsLayout() {
         // Add a section for consumed energy, even if the specific device does not
         // have support EnergyConsumers.  This is done to guarantee format compatibility between
         // PowerStats created by a PowerStatsCollector and those produced by a PowerStatsProcessor.
@@ -30,4 +32,8 @@
         addUidSectionEnergyConsumers(1);
         addUidSectionPowerEstimate();
     }
+
+    public EnergyConsumerPowerStatsLayout(PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
 }
diff --git a/services/core/java/com/android/server/power/stats/GnssPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/GnssPowerStatsLayout.java
similarity index 81%
rename from services/core/java/com/android/server/power/stats/GnssPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/GnssPowerStatsLayout.java
index 9a1317d..b70b173 100644
--- a/services/core/java/com/android/server/power/stats/GnssPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/GnssPowerStatsLayout.java
@@ -14,28 +14,31 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
+import android.annotation.NonNull;
 import android.location.GnssSignalQuality;
 import android.os.PersistableBundle;
 
-class GnssPowerStatsLayout extends BinaryStatePowerStatsLayout {
+import com.android.internal.os.PowerStats;
+
+public class GnssPowerStatsLayout extends BinaryStatePowerStatsLayout {
     private static final String EXTRA_DEVICE_TIME_SIGNAL_LEVEL_POSITION = "dt-sig";
     private static final String EXTRA_UID_TIME_SIGNAL_LEVEL_POSITION = "ut-sig";
 
-    private int mDeviceSignalLevelTimePosition;
-    private int mUidSignalLevelTimePosition;
+    private final int mDeviceSignalLevelTimePosition;
+    private final int mUidSignalLevelTimePosition;
 
-    GnssPowerStatsLayout() {
+    public GnssPowerStatsLayout() {
         mDeviceSignalLevelTimePosition = addDeviceSection(
                 GnssSignalQuality.NUM_GNSS_SIGNAL_QUALITY_LEVELS, "level");
         mUidSignalLevelTimePosition = addUidSection(
                 GnssSignalQuality.NUM_GNSS_SIGNAL_QUALITY_LEVELS, "level");
     }
 
-    @Override
-    public void fromExtras(PersistableBundle extras) {
-        super.fromExtras(extras);
+    public GnssPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+        PersistableBundle extras = descriptor.extras;
         mDeviceSignalLevelTimePosition = extras.getInt(EXTRA_DEVICE_TIME_SIGNAL_LEVEL_POSITION);
         mUidSignalLevelTimePosition = extras.getInt(EXTRA_UID_TIME_SIGNAL_LEVEL_POSITION);
     }
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/MobileRadioPowerStatsLayout.java
similarity index 79%
rename from services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/MobileRadioPowerStatsLayout.java
index 07d78f8..da6fc41 100644
--- a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/MobileRadioPowerStatsLayout.java
@@ -14,10 +14,12 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.annotation.NonNull;
+import android.os.BatteryStats;
 import android.os.PersistableBundle;
+import android.telephony.AccessNetworkConstants;
 import android.telephony.ModemActivityInfo;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -28,7 +30,7 @@
  * Captures the positions and lengths of sections of the stats array, such as time-in-state,
  * power usage estimates etc.
  */
-class MobileRadioPowerStatsLayout extends PowerStatsLayout {
+public 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";
@@ -56,27 +58,95 @@
     private int mUidRxPacketsPosition;
     private int mUidTxPacketsPosition;
 
-    MobileRadioPowerStatsLayout() {
+    public MobileRadioPowerStatsLayout(int energyConsumerCount) {
+        addDeviceMobileActivity();
+        addDeviceSectionEnergyConsumers(energyConsumerCount);
+        addStateStats();
+        addUidNetworkStats();
+        addDeviceSectionUsageDuration();
+        addDeviceSectionPowerEstimate();
+        addUidSectionPowerEstimate();
     }
 
-    MobileRadioPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+    public MobileRadioPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
         super(descriptor);
+        PersistableBundle extras = descriptor.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);
     }
 
-    void addDeviceMobileActivity() {
+    /**
+     * 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);
+    }
+
+    public static int makeStateKey(int rat, int freqRange) {
+        if (rat == BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR) {
+            return rat | (freqRange << 8);
+        } else {
+            return rat;
+        }
+    }
+
+    @BatteryStats.RadioAccessTechnology
+    public 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;
+        }
+    }
+
+    private void addDeviceMobileActivity() {
         mDeviceSleepTimePosition = addDeviceSection(1, "sleep");
         mDeviceIdleTimePosition = addDeviceSection(1, "idle");
         mDeviceScanTimePosition = addDeviceSection(1, "scan");
         mDeviceCallTimePosition = addDeviceSection(1, "call", FLAG_OPTIONAL);
     }
 
-    void addStateStats() {
+    private void addStateStats() {
         mStateRxTimePosition = addStateSection(1, "rx");
         mStateTxTimesCount = ModemActivityInfo.getNumTxPowerLevels();
         mStateTxTimesPosition = addStateSection(mStateTxTimesCount, "tx");
     }
 
-    void addUidNetworkStats() {
+    private void addUidNetworkStats() {
         mUidRxPacketsPosition = addUidSection(1, "rx-pkts");
         mUidRxBytesPosition = addUidSection(1, "rx-B");
         mUidTxPacketsPosition = addUidSection(1, "tx-pkts");
@@ -84,7 +154,7 @@
     }
 
     @Override
-    public void addDeviceSectionPowerEstimate() {
+    protected void addDeviceSectionPowerEstimate() {
         super.addDeviceSectionPowerEstimate();
         // Printed as part of the PhoneCallPowerStatsProcessor
         mDeviceCallPowerPosition = addDeviceSection(1, "call-power", FLAG_HIDDEN);
@@ -178,44 +248,6 @@
         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) {
@@ -239,9 +271,8 @@
             return;
         }
 
-        int rat = MobileRadioPowerStatsCollector.mapRadioAccessNetworkTypeToRadioAccessTechnology(
-                networkType);
-        int stateKey = MobileRadioPowerStatsCollector.makeStateKey(rat, freqRange);
+        int rat = mapRadioAccessNetworkTypeToRadioAccessTechnology(networkType);
+        int stateKey = makeStateKey(rat, freqRange);
         long[] stats = stateStats.get(stateKey);
         if (stats == null) {
             stats = new long[getStateStatsArrayLength()];
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl b/services/core/java/com/android/server/power/stats/format/PhoneCallPowerStatsLayout.java
similarity index 74%
copy from libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
copy to services/core/java/com/android/server/power/stats/format/PhoneCallPowerStatsLayout.java
index 3c5beeb..5a34148 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/bubbles/BubbleBarLocation.aidl
+++ b/services/core/java/com/android/server/power/stats/format/PhoneCallPowerStatsLayout.java
@@ -13,7 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+package com.android.server.power.stats.format;
 
-package com.android.wm.shell.common.bubbles;
-
-parcelable BubbleBarLocation;
\ No newline at end of file
+public class PhoneCallPowerStatsLayout extends PowerStatsLayout {
+    public PhoneCallPowerStatsLayout() {
+        addDeviceSectionPowerEstimate();
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/PowerStatsLayout.java
similarity index 91%
rename from services/core/java/com/android/server/power/stats/PowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/PowerStatsLayout.java
index 62abfc6..d070919 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/PowerStatsLayout.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.os.PersistableBundle;
 import android.util.Slog;
@@ -36,7 +36,7 @@
     private static final String EXTRA_UID_ENERGY_CONSUMERS_COUNT = "uec";
     private static final String EXTRA_UID_POWER_POSITION = "up";
 
-    protected static final int UNSUPPORTED = -1;
+    public static final int UNSUPPORTED = -1;
     protected static final double MILLI_TO_NANO_MULTIPLIER = 1000000.0;
     protected static final int FLAG_OPTIONAL = 1;
     protected static final int FLAG_HIDDEN = 2;
@@ -46,9 +46,9 @@
     private int mStateStatsArrayLength;
     private int mUidStatsArrayLength;
 
-    private StringBuilder mDeviceFormat = new StringBuilder();
-    private StringBuilder mStateFormat = new StringBuilder();
-    private StringBuilder mUidFormat = new StringBuilder();
+    private final StringBuilder mDeviceFormat = new StringBuilder();
+    private final StringBuilder mStateFormat = new StringBuilder();
+    private final StringBuilder mUidFormat = new StringBuilder();
 
     protected int mDeviceDurationPosition = UNSUPPORTED;
     private int mDeviceEnergyConsumerPosition;
@@ -63,7 +63,32 @@
     }
 
     public PowerStatsLayout(PowerStats.Descriptor descriptor) {
-        fromExtras(descriptor.extras);
+        PersistableBundle extras = descriptor.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);
+        mUidDurationPosition = extras.getInt(EXTRA_UID_DURATION_POSITION);
+        mUidEnergyConsumerPosition = extras.getInt(EXTRA_UID_ENERGY_CONSUMERS_POSITION);
+        mUidEnergyConsumerCount = extras.getInt(EXTRA_UID_ENERGY_CONSUMERS_COUNT);
+        mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
+    }
+
+    /**
+     * 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_DURATION_POSITION, mUidDurationPosition);
+        extras.putInt(EXTRA_UID_ENERGY_CONSUMERS_POSITION, mUidEnergyConsumerPosition);
+        extras.putInt(EXTRA_UID_ENERGY_CONSUMERS_COUNT, mUidEnergyConsumerCount);
+        extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
+        extras.putString(PowerStats.Descriptor.EXTRA_DEVICE_STATS_FORMAT, mDeviceFormat.toString());
+        extras.putString(PowerStats.Descriptor.EXTRA_STATE_STATS_FORMAT, mStateFormat.toString());
+        extras.putString(PowerStats.Descriptor.EXTRA_UID_STATS_FORMAT, mUidFormat.toString());
     }
 
     public int getDeviceStatsArrayLength() {
@@ -141,7 +166,7 @@
     /**
      * Declare that the stats array has a section capturing usage duration
      */
-    public void addDeviceSectionUsageDuration() {
+    protected void addDeviceSectionUsageDuration() {
         mDeviceDurationPosition = addDeviceSection(1, "usage", FLAG_OPTIONAL);
     }
 
@@ -163,7 +188,7 @@
      * Declares that the stats array has a section capturing EnergyConsumer data from
      * PowerStatsService.
      */
-    public void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
+    protected void addDeviceSectionEnergyConsumers(int energyConsumerCount) {
         mDeviceEnergyConsumerPosition = addDeviceSection(energyConsumerCount, "energy",
                 FLAG_OPTIONAL);
         mDeviceEnergyConsumerCount = energyConsumerCount;
@@ -192,7 +217,7 @@
     /**
      * Declare that the stats array has a section capturing a power estimate
      */
-    public void addDeviceSectionPowerEstimate() {
+    protected void addDeviceSectionPowerEstimate() {
         mDevicePowerEstimatePosition = addDeviceSection(1, "power",
                 FLAG_FORMAT_AS_POWER | FLAG_OPTIONAL);
     }
@@ -215,14 +240,14 @@
     /**
      * Declare that the UID stats array has a section capturing usage duration
      */
-    public void addUidSectionUsageDuration() {
+    protected void addUidSectionUsageDuration() {
         mUidDurationPosition = addUidSection(1, "time");
     }
 
     /**
      * Declare that the UID stats array has a section capturing a power estimate
      */
-    public void addUidSectionPowerEstimate() {
+    protected void addUidSectionPowerEstimate() {
         mUidPowerEstimatePosition = addUidSection(1, "power", FLAG_FORMAT_AS_POWER | FLAG_OPTIONAL);
     }
 
@@ -251,7 +276,7 @@
      * Declares that the UID stats array has a section capturing EnergyConsumer data from
      * PowerStatsService.
      */
-    public void addUidSectionEnergyConsumers(int energyConsumerCount) {
+    protected void addUidSectionEnergyConsumers(int energyConsumerCount) {
         mUidEnergyConsumerPosition = addUidSection(energyConsumerCount, "energy",
                 FLAG_OPTIONAL);
         mUidEnergyConsumerCount = energyConsumerCount;
@@ -292,39 +317,6 @@
         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_DURATION_POSITION, mUidDurationPosition);
-        extras.putInt(EXTRA_UID_ENERGY_CONSUMERS_POSITION, mUidEnergyConsumerPosition);
-        extras.putInt(EXTRA_UID_ENERGY_CONSUMERS_COUNT, mUidEnergyConsumerCount);
-        extras.putInt(EXTRA_UID_POWER_POSITION, mUidPowerEstimatePosition);
-        extras.putString(PowerStats.Descriptor.EXTRA_DEVICE_STATS_FORMAT, mDeviceFormat.toString());
-        extras.putString(PowerStats.Descriptor.EXTRA_STATE_STATS_FORMAT, mStateFormat.toString());
-        extras.putString(PowerStats.Descriptor.EXTRA_UID_STATS_FORMAT, mUidFormat.toString());
-    }
-
-    /**
-     * 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);
-        mUidDurationPosition = extras.getInt(EXTRA_UID_DURATION_POSITION);
-        mUidEnergyConsumerPosition = extras.getInt(EXTRA_UID_ENERGY_CONSUMERS_POSITION);
-        mUidEnergyConsumerCount = extras.getInt(EXTRA_UID_ENERGY_CONSUMERS_COUNT);
-        mUidPowerEstimatePosition = extras.getInt(EXTRA_UID_POWER_POSITION);
-    }
-
     protected void putIntArray(PersistableBundle extras, String key, int[] array) {
         if (array == null) {
             return;
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/ScreenPowerStatsLayout.java
similarity index 89%
rename from services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/ScreenPowerStatsLayout.java
index f134aa8..6f6a7ff 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/ScreenPowerStatsLayout.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.annotation.NonNull;
 import android.os.BatteryStats;
@@ -41,14 +41,40 @@
     private int mDeviceScreenDozePowerPosition;
     private int mUidTopActivityTimePosition;
 
-    ScreenPowerStatsLayout() {
+    public ScreenPowerStatsLayout(int energyConsumerCount, int displayCount) {
+        addDeviceScreenUsageDurationSection(displayCount);
+        addDeviceSectionEnergyConsumers(energyConsumerCount);
+        addDeviceSectionUsageDuration();
+        addDeviceSectionPowerEstimate();
+        addUidTopActivitiyDuration();
+        addUidSectionPowerEstimate();
     }
 
-    ScreenPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+    public ScreenPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
         super(descriptor);
+        PersistableBundle extras = descriptor.extras;
+        mDisplayCount = extras.getInt(EXTRA_DEVICE_SCREEN_COUNT, 1);
+        mDeviceScreenOnDurationPosition = extras.getInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION);
+        mDeviceBrightnessDurationPositions = extras.getIntArray(
+                EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS);
+        mDeviceScreenDozeDurationPosition = extras.getInt(EXTRA_DEVICE_DOZE_DURATION_POSITION);
+        mDeviceScreenDozePowerPosition = extras.getInt(EXTRA_DEVICE_DOZE_POWER_POSITION);
+        mUidTopActivityTimePosition = extras.getInt(EXTRA_UID_FOREGROUND_DURATION);
     }
 
-    void addDeviceScreenUsageDurationSection(int displayCount) {
+    @Override
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putInt(EXTRA_DEVICE_SCREEN_COUNT, mDisplayCount);
+        extras.putInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION, mDeviceScreenOnDurationPosition);
+        extras.putIntArray(EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS,
+                mDeviceBrightnessDurationPositions);
+        extras.putInt(EXTRA_DEVICE_DOZE_DURATION_POSITION, mDeviceScreenDozeDurationPosition);
+        extras.putInt(EXTRA_DEVICE_DOZE_POWER_POSITION, mDeviceScreenDozePowerPosition);
+        extras.putInt(EXTRA_UID_FOREGROUND_DURATION, mUidTopActivityTimePosition);
+    }
+
+    private void addDeviceScreenUsageDurationSection(int displayCount) {
         mDisplayCount = displayCount;
         mDeviceScreenOnDurationPosition = addDeviceSection(displayCount, "on");
         mDeviceBrightnessDurationPositions = new int[BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS];
@@ -60,7 +86,7 @@
     }
 
     @Override
-    public void addDeviceSectionPowerEstimate() {
+    protected void addDeviceSectionPowerEstimate() {
         super.addDeviceSectionPowerEstimate();
         // Used by AmbientDisplayPowerStatsProcessor
         mDeviceScreenDozePowerPosition = addDeviceSection(1, "doze-power", FLAG_HIDDEN);
@@ -127,7 +153,7 @@
         return stats[mDeviceScreenDozePowerPosition] / MILLI_TO_NANO_MULTIPLIER;
     }
 
-    void addUidTopActivitiyDuration() {
+    private void addUidTopActivitiyDuration() {
         mUidTopActivityTimePosition = addUidSection(1, "top");
     }
 
@@ -144,28 +170,4 @@
     public long getUidTopActivityDuration(long[] stats) {
         return stats[mUidTopActivityTimePosition];
     }
-
-    @Override
-    public void toExtras(PersistableBundle extras) {
-        super.toExtras(extras);
-        extras.putInt(EXTRA_DEVICE_SCREEN_COUNT, mDisplayCount);
-        extras.putInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION, mDeviceScreenOnDurationPosition);
-        extras.putIntArray(EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS,
-                mDeviceBrightnessDurationPositions);
-        extras.putInt(EXTRA_DEVICE_DOZE_DURATION_POSITION, mDeviceScreenDozeDurationPosition);
-        extras.putInt(EXTRA_DEVICE_DOZE_POWER_POSITION, mDeviceScreenDozePowerPosition);
-        extras.putInt(EXTRA_UID_FOREGROUND_DURATION, mUidTopActivityTimePosition);
-    }
-
-    @Override
-    public void fromExtras(PersistableBundle extras) {
-        super.fromExtras(extras);
-        mDisplayCount = extras.getInt(EXTRA_DEVICE_SCREEN_COUNT, 1);
-        mDeviceScreenOnDurationPosition = extras.getInt(EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION);
-        mDeviceBrightnessDurationPositions = extras.getIntArray(
-                EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS);
-        mDeviceScreenDozeDurationPosition = extras.getInt(EXTRA_DEVICE_DOZE_DURATION_POSITION);
-        mDeviceScreenDozePowerPosition = extras.getInt(EXTRA_DEVICE_DOZE_POWER_POSITION);
-        mUidTopActivityTimePosition = extras.getInt(EXTRA_UID_FOREGROUND_DURATION);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/SensorPowerStatsLayout.java
similarity index 72%
rename from services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/SensorPowerStatsLayout.java
index e66cd39..e8df3dd 100644
--- a/services/core/java/com/android/server/power/stats/SensorPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/SensorPowerStatsLayout.java
@@ -14,12 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.os.PersistableBundle;
 import android.util.Slog;
 import android.util.SparseIntArray;
 
+import com.android.internal.os.PowerStats;
+
+import java.util.Arrays;
+import java.util.Map;
+
 public class SensorPowerStatsLayout extends PowerStatsLayout {
     private static final String TAG = "SensorPowerStatsLayout";
     private static final String EXTRA_DEVICE_SENSOR_HANDLES = "dsh";
@@ -27,7 +32,48 @@
 
     private final SparseIntArray mSensorPositions = new SparseIntArray();
 
-    void addUidSensorSection(int handle, String label) {
+    public SensorPowerStatsLayout(Map<Integer, String> idToLabelMap) {
+        Integer[] keys = new Integer[idToLabelMap.size()];
+        idToLabelMap.keySet().toArray(keys);
+        Arrays.sort(keys);
+        for (int i = 0; i < keys.length; i++) {
+            addUidSensorSection(keys[i], idToLabelMap.get(keys[i]));
+        }
+        addUidSectionPowerEstimate();
+        addDeviceSectionPowerEstimate();
+    }
+
+    public SensorPowerStatsLayout(PowerStats.Descriptor descriptor) {
+        super(descriptor);
+
+        PersistableBundle extras = descriptor.extras;
+        int[] handlers = extras.getIntArray(EXTRA_DEVICE_SENSOR_HANDLES);
+        int[] uidDurationPositions = extras.getIntArray(EXTRA_UID_SENSOR_POSITIONS);
+
+        if (handlers != null && uidDurationPositions != null) {
+            for (int i = 0; i < handlers.length; i++) {
+                mSensorPositions.put(handlers[i], uidDurationPositions[i]);
+            }
+        }
+    }
+
+    @Override
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+
+        int[] handlers = new int[mSensorPositions.size()];
+        int[] uidDurationPositions = new int[mSensorPositions.size()];
+
+        for (int i = 0; i < mSensorPositions.size(); i++) {
+            handlers[i] = mSensorPositions.keyAt(i);
+            uidDurationPositions[i] = mSensorPositions.valueAt(i);
+        }
+
+        extras.putIntArray(EXTRA_DEVICE_SENSOR_HANDLES, handlers);
+        extras.putIntArray(EXTRA_UID_SENSOR_POSITIONS, uidDurationPositions);
+    }
+
+    private void addUidSensorSection(int handle, String label) {
         mSensorPositions.put(handle, addUidSection(1, label, FLAG_OPTIONAL));
     }
 
@@ -50,32 +96,4 @@
         }
         stats[position] += durationMs;
     }
-
-    @Override
-    public void toExtras(PersistableBundle extras) {
-        super.toExtras(extras);
-
-        int[] handlers = new int[mSensorPositions.size()];
-        int[] uidDurationPositions = new int[mSensorPositions.size()];
-
-        for (int i = 0; i < mSensorPositions.size(); i++) {
-            handlers[i] = mSensorPositions.keyAt(i);
-            uidDurationPositions[i] = mSensorPositions.valueAt(i);
-        }
-
-        extras.putIntArray(EXTRA_DEVICE_SENSOR_HANDLES, handlers);
-        extras.putIntArray(EXTRA_UID_SENSOR_POSITIONS, uidDurationPositions);
-    }
-
-    @Override
-    public void fromExtras(PersistableBundle extras) {
-        super.fromExtras(extras);
-
-        int[] handlers = extras.getIntArray(EXTRA_DEVICE_SENSOR_HANDLES);
-        int[] uidDurationPositions = extras.getIntArray(EXTRA_UID_SENSOR_POSITIONS);
-
-        for (int i = 0; i < handlers.length; i++) {
-            mSensorPositions.put(handlers[i], uidDurationPositions[i]);
-        }
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/format/WifiPowerStatsLayout.java
similarity index 93%
rename from services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
rename to services/core/java/com/android/server/power/stats/format/WifiPowerStatsLayout.java
index e2e8226..ce7ef12 100644
--- a/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/format/WifiPowerStatsLayout.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.format;
 
 import android.annotation.NonNull;
 import android.os.PersistableBundle;
@@ -22,7 +22,6 @@
 import com.android.internal.os.PowerStats;
 
 public class WifiPowerStatsLayout extends PowerStatsLayout {
-    private static final String TAG = "WifiPowerStatsLayout";
     private static final int UNSPECIFIED = -1;
     private static final String EXTRA_POWER_REPORTING_SUPPORTED = "prs";
     private static final String EXTRA_DEVICE_RX_TIME_POSITION = "dt-rx";
@@ -54,14 +53,56 @@
     private int mUidScanTimePosition;
     private int mUidBatchScanTimePosition;
 
-    WifiPowerStatsLayout() {
+    public WifiPowerStatsLayout(int energyConsumerCount, boolean powerReportingSupported) {
+        addDeviceWifiActivity(powerReportingSupported);
+        addDeviceSectionEnergyConsumers(energyConsumerCount);
+        addUidNetworkStats();
+        addDeviceSectionUsageDuration();
+        addDeviceSectionPowerEstimate();
+        addUidSectionPowerEstimate();
     }
 
-    WifiPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+    public WifiPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
         super(descriptor);
+        PersistableBundle extras = descriptor.extras;
+        mPowerReportingSupported = extras.getBoolean(EXTRA_POWER_REPORTING_SUPPORTED);
+        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
+        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
+        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
+        mDeviceBasicScanTimePosition = extras.getInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION);
+        mDeviceBatchedScanTimePosition = extras.getInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION);
+        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
+        mDeviceActiveTimePosition = extras.getInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION);
+        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);
+        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
+        mUidBatchScanTimePosition = extras.getInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION);
     }
 
-    void addDeviceWifiActivity(boolean powerReportingSupported) {
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putBoolean(EXTRA_POWER_REPORTING_SUPPORTED, mPowerReportingSupported);
+        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
+        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
+        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION, mDeviceBasicScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION, mDeviceBatchedScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
+        extras.putInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION, mDeviceActiveTimePosition);
+        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);
+        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
+        extras.putInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION, mUidBatchScanTimePosition);
+    }
+
+    private void addDeviceWifiActivity(boolean powerReportingSupported) {
         mPowerReportingSupported = powerReportingSupported;
         if (mPowerReportingSupported) {
             mDeviceActiveTimePosition = UNSPECIFIED;
@@ -80,7 +121,7 @@
         mDeviceBatchedScanTimePosition = addDeviceSection(1, "batched-scan", FLAG_OPTIONAL);
     }
 
-    void addUidNetworkStats() {
+    private void addUidNetworkStats() {
         mUidRxPacketsPosition = addUidSection(1, "rx-pkts");
         mUidRxBytesPosition = addUidSection(1, "rx-B");
         mUidTxPacketsPosition = addUidSection(1, "tx-pkts");
@@ -196,46 +237,4 @@
     public long getUidBatchedScanTime(long[] stats) {
         return stats[mUidBatchScanTimePosition];
     }
-
-    /**
-     * Copies the elements of the stats array layout into <code>extras</code>
-     */
-    public void toExtras(PersistableBundle extras) {
-        super.toExtras(extras);
-        extras.putBoolean(EXTRA_POWER_REPORTING_SUPPORTED, mPowerReportingSupported);
-        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
-        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
-        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
-        extras.putInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION, mDeviceBasicScanTimePosition);
-        extras.putInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION, mDeviceBatchedScanTimePosition);
-        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
-        extras.putInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION, mDeviceActiveTimePosition);
-        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);
-        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
-        extras.putInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION, mUidBatchScanTimePosition);
-    }
-
-    /**
-     * Retrieves elements of the stats array layout from <code>extras</code>
-     */
-    public void fromExtras(PersistableBundle extras) {
-        super.fromExtras(extras);
-        mPowerReportingSupported = extras.getBoolean(EXTRA_POWER_REPORTING_SUPPORTED);
-        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
-        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
-        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
-        mDeviceBasicScanTimePosition = extras.getInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION);
-        mDeviceBatchedScanTimePosition = extras.getInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION);
-        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
-        mDeviceActiveTimePosition = extras.getInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION);
-        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);
-        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
-        mUidBatchScanTimePosition = extras.getInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION);
-    }
 }
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStats.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
rename to services/core/java/com/android/server/power/stats/processor/AggregatedPowerStats.java
index 674b4bc..a4758dd 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStats.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.annotation.CurrentTimeMillisLong;
 import android.annotation.DurationMillisLong;
@@ -33,7 +33,7 @@
 import com.android.internal.os.PowerStats;
 import com.android.modules.utils.TypedXmlPullParser;
 import com.android.modules.utils.TypedXmlSerializer;
-import com.android.server.power.stats.AggregatedPowerStatsConfig.PowerComponent;
+import com.android.server.power.stats.processor.AggregatedPowerStatsConfig.PowerComponent;
 
 import org.xmlpull.v1.XmlPullParser;
 import org.xmlpull.v1.XmlPullParserException;
@@ -143,7 +143,8 @@
         }
     }
 
-    List<ClockUpdate> getClockUpdates() {
+    // TODO - DO NOT SUBMIT public
+    public List<ClockUpdate> getClockUpdates() {
         return mClockUpdates;
     }
 
@@ -274,7 +275,8 @@
                         int powerComponentId = parser.getAttributeInt(null,
                                 PowerComponentAggregatedPowerStats.XML_ATTR_ID);
 
-                        PowerComponentAggregatedPowerStats powerComponentStats =
+                        PowerComponentAggregatedPowerStats
+                                powerComponentStats =
                                 stats.getPowerComponentStats(powerComponentId);
                         if (powerComponentStats == null) {
                             PowerComponent powerComponent =
diff --git a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsConfig.java
similarity index 95%
rename from services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
rename to services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsConfig.java
index ec12228..eaeda43 100644
--- a/services/core/java/com/android/server/power/stats/AggregatedPowerStatsConfig.java
+++ b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsConfig.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.annotation.IntDef;
 import android.annotation.NonNull;
@@ -32,7 +32,7 @@
  * WiFi, etc).  Also, it determines which states are tracked globally and which ones on a per-UID
  * basis.
  */
-public class AggregatedPowerStatsConfig {
+class AggregatedPowerStatsConfig {
     public static final int STATE_POWER = 0;
     public static final int STATE_SCREEN = 1;
     public static final int STATE_PROCESS_STATE = 2;
@@ -70,7 +70,7 @@
     /**
      * Configuration for a give power component (CPU, WiFi, etc)
      */
-    public static class PowerComponent {
+    static class PowerComponent {
         private final int mPowerComponentId;
         private @TrackedState int[] mTrackedDeviceStates;
         private @TrackedState int[] mTrackedUidStates;
@@ -174,7 +174,7 @@
      * standard power component IDs, e.g. {@link BatteryConsumer#POWER_COMPONENT_CPU}, or
      * a custom power component.
      */
-    public PowerComponent trackPowerComponent(int powerComponentId) {
+    PowerComponent trackPowerComponent(int powerComponentId) {
         PowerComponent builder = new PowerComponent(powerComponentId);
         mPowerComponents.add(builder);
         return builder;
@@ -185,7 +185,7 @@
      * of a different power component.  The tracked states will be the same as the parent
      * component's.
      */
-    public PowerComponent trackPowerComponent(int powerComponentId,
+    PowerComponent trackPowerComponent(int powerComponentId,
             int parentPowerComponentId) {
         PowerComponent parent = null;
         for (int i = 0; i < mPowerComponents.size(); i++) {
@@ -211,7 +211,7 @@
      * Creates a configuration for custom power components, which are yet to be discovered
      * dynamically through the integration with PowerStatsService.
      */
-    public PowerComponent trackCustomPowerComponents(
+    PowerComponent trackCustomPowerComponents(
             Supplier<PowerStatsProcessor> processorFactory) {
         mCustomPowerStatsProcessorFactory = processorFactory;
         mCustomPowerComponent = new PowerComponent(BatteryConsumer.POWER_COMPONENT_ANY);
@@ -221,7 +221,7 @@
     /**
      * Returns configurations for all registered or dynamically discovered power components.
      */
-    public List<PowerComponent> getPowerComponentsAggregatedStatsConfigs() {
+    List<PowerComponent> getPowerComponentsAggregatedStatsConfigs() {
         return mPowerComponents;
     }
 
@@ -230,7 +230,7 @@
      * integration with PowerStatsService.
      */
     @Nullable
-    public PowerComponent createPowerComponent(int powerComponentId) {
+    PowerComponent createPowerComponent(int powerComponentId) {
         if (mCustomPowerComponent == null) {
             return null;
         }
diff --git a/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsSection.java b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsSection.java
new file mode 100644
index 0000000..4a9730c
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/processor/AggregatedPowerStatsSection.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright (C) 2023 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.processor;
+
+import android.util.IndentingPrintWriter;
+
+import com.android.modules.utils.TypedXmlPullParser;
+import com.android.modules.utils.TypedXmlSerializer;
+import com.android.server.power.stats.PowerStatsSpan;
+
+import org.xmlpull.v1.XmlPullParserException;
+
+import java.io.IOException;
+
+class AggregatedPowerStatsSection extends PowerStatsSpan.Section {
+    public static final String TYPE = "aggregated-power-stats";
+
+    private final AggregatedPowerStats mAggregatedPowerStats;
+
+    AggregatedPowerStatsSection(AggregatedPowerStats aggregatedPowerStats) {
+        super(TYPE);
+        mAggregatedPowerStats = aggregatedPowerStats;
+    }
+
+    public AggregatedPowerStats getAggregatedPowerStats() {
+        return mAggregatedPowerStats;
+    }
+
+    @Override
+    public void write(TypedXmlSerializer serializer) throws IOException {
+        mAggregatedPowerStats.writeXml(serializer);
+    }
+
+    @Override
+    public void dump(IndentingPrintWriter ipw) {
+        mAggregatedPowerStats.dump(ipw);
+    }
+
+    static class Reader implements PowerStatsSpan.SectionReader {
+        private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
+
+        Reader(AggregatedPowerStatsConfig config) {
+            mAggregatedPowerStatsConfig = config;
+        }
+
+        @Override
+        public String getType() {
+            return TYPE;
+        }
+
+        @Override
+        public PowerStatsSpan.Section read(String sectionType, TypedXmlPullParser parser)
+                throws IOException, XmlPullParserException {
+            return new AggregatedPowerStatsSection(
+                    AggregatedPowerStats.createFromXml(parser, mAggregatedPowerStatsConfig));
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessor.java
similarity index 86%
rename from services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessor.java
index a42929f..32dfdf9 100644
--- a/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessor.java
@@ -13,24 +13,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.PersistableBundle;
 
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.AmbientDisplayPowerStatsLayout;
+import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
-public class AmbientDisplayPowerStatsProcessor extends PowerStatsProcessor {
-    private final PowerStatsLayout mStatsLayout;
+class AmbientDisplayPowerStatsProcessor extends PowerStatsProcessor {
+    private final AmbientDisplayPowerStatsLayout mStatsLayout;
     private final PowerStats.Descriptor mDescriptor;
     private final long[] mTmpDeviceStats;
     private PowerStats.Descriptor mScreenPowerStatsDescriptor;
     private ScreenPowerStatsLayout mScreenPowerStatsLayout;
     private long[] mTmpScreenStats;
 
-    public AmbientDisplayPowerStatsProcessor() {
-        mStatsLayout = new PowerStatsLayout();
-        mStatsLayout.addDeviceSectionPowerEstimate();
+    AmbientDisplayPowerStatsProcessor() {
+        mStatsLayout = new AmbientDisplayPowerStatsLayout();
         PersistableBundle extras = new PersistableBundle();
         mStatsLayout.toExtras(extras);
         mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
diff --git a/services/core/java/com/android/server/power/stats/AudioPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/AudioPowerStatsProcessor.java
similarity index 82%
rename from services/core/java/com/android/server/power/stats/AudioPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/AudioPowerStatsProcessor.java
index a48f162..ad1b4a7 100644
--- a/services/core/java/com/android/server/power/stats/AudioPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/AudioPowerStatsProcessor.java
@@ -14,15 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.PowerStatsUidResolver;
 
-public class AudioPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
-    public AudioPowerStatsProcessor(PowerProfile powerProfile,
+class AudioPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
+    AudioPowerStatsProcessor(PowerProfile powerProfile,
             PowerStatsUidResolver uidResolver) {
         super(BatteryConsumer.POWER_COMPONENT_AUDIO, uidResolver,
                 powerProfile.getAveragePower(PowerProfile.POWER_AUDIO));
diff --git a/services/core/java/com/android/server/power/stats/BinaryStatePowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessor.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/BinaryStatePowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessor.java
index 03df46a..e45a707 100644
--- a/services/core/java/com/android/server/power/stats/BinaryStatePowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessor.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.annotation.IntDef;
 import android.os.BatteryStats;
@@ -22,6 +22,9 @@
 import android.os.Process;
 
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.BinaryStatePowerStatsLayout;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
diff --git a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/BluetoothPowerStatsProcessor.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/BluetoothPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/BluetoothPowerStatsProcessor.java
index 077b057..4c1a0db 100644
--- a/services/core/java/com/android/server/power/stats/BluetoothPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/BluetoothPowerStatsProcessor.java
@@ -14,17 +14,17 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.BluetoothPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class BluetoothPowerStatsProcessor extends PowerStatsProcessor {
-    private static final String TAG = "BluetoothPowerStatsProcessor";
-
+class BluetoothPowerStatsProcessor extends PowerStatsProcessor {
     private final UsageBasedPowerEstimator mRxPowerEstimator;
     private final UsageBasedPowerEstimator mTxPowerEstimator;
     private final UsageBasedPowerEstimator mIdlePowerEstimator;
@@ -37,7 +37,7 @@
     private long[] mTmpDeviceStatsArray;
     private long[] mTmpUidStatsArray;
 
-    public BluetoothPowerStatsProcessor(PowerProfile powerProfile) {
+    BluetoothPowerStatsProcessor(PowerProfile powerProfile) {
         mRxPowerEstimator = new UsageBasedPowerEstimator(
                 powerProfile.getAveragePower(PowerProfile.POWER_BLUETOOTH_CONTROLLER_RX));
         mTxPowerEstimator = new UsageBasedPowerEstimator(
diff --git a/services/core/java/com/android/server/power/stats/CameraPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/CameraPowerStatsProcessor.java
similarity index 82%
rename from services/core/java/com/android/server/power/stats/CameraPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/CameraPowerStatsProcessor.java
index 15c3eb8..8309061 100644
--- a/services/core/java/com/android/server/power/stats/CameraPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/CameraPowerStatsProcessor.java
@@ -14,15 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.PowerStatsUidResolver;
 
-public class CameraPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
-    public CameraPowerStatsProcessor(PowerProfile powerProfile,
+class CameraPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
+    CameraPowerStatsProcessor(PowerProfile powerProfile,
             PowerStatsUidResolver uidResolver) {
         super(BatteryConsumer.POWER_COMPONENT_CAMERA, uidResolver,
                 powerProfile.getAveragePower(PowerProfile.POWER_CAMERA));
diff --git a/services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/CpuPowerStatsProcessor.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/CpuPowerStatsProcessor.java
index 6da0a8f..5f7a3da 100644
--- a/services/core/java/com/android/server/power/stats/CpuPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/CpuPowerStatsProcessor.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.util.ArraySet;
 import android.util.Log;
@@ -22,13 +22,14 @@
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.CpuPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
-public class CpuPowerStatsProcessor extends PowerStatsProcessor {
+class CpuPowerStatsProcessor extends PowerStatsProcessor {
     private static final String TAG = "CpuPowerStatsProcessor";
 
     private static final double HOUR_IN_MILLIS = TimeUnit.HOURS.toMillis(1);
@@ -72,7 +73,7 @@
     // Temp array for retrieval of UID power stats, to avoid repeated allocations
     private long[] mTmpUidStatsArray;
 
-    public CpuPowerStatsProcessor(PowerProfile powerProfile, CpuScalingPolicies scalingPolicies) {
+    CpuPowerStatsProcessor(PowerProfile powerProfile, CpuScalingPolicies scalingPolicies) {
         mCpuScalingPolicies = scalingPolicies;
         mCpuScalingStepCount = scalingPolicies.getScalingStepCount();
         mScalingStepToCluster = new int[mCpuScalingStepCount];
@@ -104,8 +105,7 @@
         }
 
         mLastUsedDescriptor = descriptor;
-        mStatsLayout = new CpuPowerStatsLayout();
-        mStatsLayout.fromExtras(descriptor.extras);
+        mStatsLayout = new CpuPowerStatsLayout(descriptor);
 
         mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
         mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
@@ -138,7 +138,7 @@
     }
 
     @Override
-    public void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
         if (stats.getPowerStatsDescriptor() == null) {
             return;
         }
diff --git a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsProcessor.java
similarity index 94%
rename from services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsProcessor.java
index a86242a..76adc47 100644
--- a/services/core/java/com/android/server/power/stats/CustomEnergyConsumerPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsProcessor.java
@@ -14,14 +14,15 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class CustomEnergyConsumerPowerStatsProcessor extends PowerStatsProcessor {
+class CustomEnergyConsumerPowerStatsProcessor extends PowerStatsProcessor {
     private static final EnergyConsumerPowerStatsLayout sLayout =
             new EnergyConsumerPowerStatsLayout();
     private long[] mTmpDeviceStatsArray;
diff --git a/services/core/java/com/android/server/power/stats/FlashlightPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/FlashlightPowerStatsProcessor.java
similarity index 82%
rename from services/core/java/com/android/server/power/stats/FlashlightPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/FlashlightPowerStatsProcessor.java
index f7216c9..b0bef69 100644
--- a/services/core/java/com/android/server/power/stats/FlashlightPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/FlashlightPowerStatsProcessor.java
@@ -14,15 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.PowerStatsUidResolver;
 
-public class FlashlightPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
-    public FlashlightPowerStatsProcessor(PowerProfile powerProfile,
+class FlashlightPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
+    FlashlightPowerStatsProcessor(PowerProfile powerProfile,
             PowerStatsUidResolver uidResolver) {
         super(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT, uidResolver,
                 powerProfile.getAveragePower(PowerProfile.POWER_FLASHLIGHT));
diff --git a/services/core/java/com/android/server/power/stats/GnssPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/GnssPowerStatsProcessor.java
similarity index 94%
rename from services/core/java/com/android/server/power/stats/GnssPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/GnssPowerStatsProcessor.java
index 0b28710..f1e3e90 100644
--- a/services/core/java/com/android/server/power/stats/GnssPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/GnssPowerStatsProcessor.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.location.GnssSignalQuality;
 import android.os.BatteryConsumer;
@@ -23,10 +23,13 @@
 
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.GnssPowerStatsLayout;
 
 import java.util.Arrays;
 
-public class GnssPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
+class GnssPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
     private static final GnssPowerStatsLayout sStatsLayout = new GnssPowerStatsLayout();
     private final UsageBasedPowerEstimator[] mSignalLevelEstimators =
             new UsageBasedPowerEstimator[GnssSignalQuality.NUM_GNSS_SIGNAL_QUALITY_LEVELS];
@@ -37,7 +40,7 @@
     private final long[] mGnssSignalDurations =
             new long[GnssSignalQuality.NUM_GNSS_SIGNAL_QUALITY_LEVELS];
 
-    public GnssPowerStatsProcessor(PowerProfile powerProfile, PowerStatsUidResolver uidResolver) {
+    GnssPowerStatsProcessor(PowerProfile powerProfile, PowerStatsUidResolver uidResolver) {
         super(BatteryConsumer.POWER_COMPONENT_GNSS, uidResolver,
                 powerProfile.getAveragePower(PowerProfile.POWER_GPS_ON),
                 sStatsLayout);
diff --git a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessor.java
similarity index 96%
rename from services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessor.java
index dcce562..b4c40de8 100644
--- a/services/core/java/com/android/server/power/stats/MobileRadioPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessor.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryStats;
 import android.telephony.CellSignalStrength;
@@ -26,13 +26,15 @@
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
 import com.android.internal.power.ModemPowerProfile;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-public class MobileRadioPowerStatsProcessor extends PowerStatsProcessor {
-    private static final String TAG = "MobileRadioPowerStatsProcessor";
+class MobileRadioPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "MobileRadioPowerStats";
     private static final boolean DEBUG = false;
 
     private static final int NUM_SIGNAL_STRENGTH_LEVELS =
@@ -61,7 +63,7 @@
     private long[] mTmpStateStatsArray;
     private long[] mTmpUidStatsArray;
 
-    public MobileRadioPowerStatsProcessor(PowerProfile powerProfile) {
+    MobileRadioPowerStatsProcessor(PowerProfile powerProfile) {
         final double sleepDrainRateMa = powerProfile.getAverageBatteryDrainOrDefaultMa(
                 PowerProfile.SUBSYSTEM_MODEM | ModemPowerProfile.MODEM_DRAIN_TYPE_SLEEP,
                 Double.NaN);
@@ -101,7 +103,7 @@
                     ? ServiceState.FREQUENCY_RANGE_COUNT : 1;
             for (int freqRange = 0; freqRange < freqCount; freqRange++) {
                 mRxTxPowerEstimators.put(
-                        MobileRadioPowerStatsCollector.makeStateKey(rat, freqRange),
+                        MobileRadioPowerStatsLayout.makeStateKey(rat, freqRange),
                         buildRxTxPowerEstimators(powerProfile, rat, freqRange));
             }
         }
@@ -114,8 +116,7 @@
                 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));
+            Log.w(TAG, "Unavailable Power Profile constant for key 0x" + Long.toHexString(rxKey));
             rxDrainRateMa = 0;
         }
         estimators.mRxPowerEstimator = new UsageBasedPowerEstimator(rxDrainRateMa);
diff --git a/services/core/java/com/android/server/power/stats/processor/MultiStatePowerAttributor.java b/services/core/java/com/android/server/power/stats/processor/MultiStatePowerAttributor.java
new file mode 100644
index 0000000..2ba4a52
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/processor/MultiStatePowerAttributor.java
@@ -0,0 +1,310 @@
+/*
+ * 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.processor;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.hardware.SensorManager;
+import android.os.BatteryConsumer;
+import android.os.BatteryUsageStats;
+import android.util.IndentingPrintWriter;
+import android.util.Slog;
+import android.util.SparseBooleanArray;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.BatteryStatsHistory;
+import com.android.internal.os.CpuScalingPolicies;
+import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.PowerAttributor;
+import com.android.server.power.stats.PowerStatsSpan;
+import com.android.server.power.stats.PowerStatsStore;
+import com.android.server.power.stats.PowerStatsUidResolver;
+
+import java.util.List;
+
+public class MultiStatePowerAttributor implements PowerAttributor {
+    private static final String TAG = "MultiStatePowerAttributor";
+
+    private final PowerStatsStore mPowerStatsStore;
+    private final PowerStatsExporter mPowerStatsExporter;
+    private final PowerStatsAggregator mPowerStatsAggregator;
+    private final SparseBooleanArray mPowerStatsExporterEnabled = new SparseBooleanArray();
+
+    // TODO(b/346371828): remove dependency on PowerStatsUidResolver. At the time of power
+    // attribution isolates UIDs are supposed to be long forgotten.
+    public MultiStatePowerAttributor(Context context, PowerStatsStore powerStatsStore,
+            @NonNull PowerProfile powerProfile, @NonNull CpuScalingPolicies cpuScalingPolicies,
+            @NonNull PowerStatsUidResolver powerStatsUidResolver) {
+        this(powerStatsStore, new PowerStatsAggregator(
+                createAggregatedPowerStatsConfig(context, powerProfile, cpuScalingPolicies,
+                        powerStatsUidResolver)));
+    }
+
+    @VisibleForTesting
+    MultiStatePowerAttributor(PowerStatsStore powerStatsStore,
+            PowerStatsAggregator powerStatsAggregator) {
+        mPowerStatsStore = powerStatsStore;
+        mPowerStatsAggregator = powerStatsAggregator;
+        mPowerStatsStore.addSectionReader(
+                new AggregatedPowerStatsSection.Reader(mPowerStatsAggregator.getConfig()));
+        mPowerStatsExporter = new PowerStatsExporter(mPowerStatsStore, mPowerStatsAggregator);
+    }
+
+    private static AggregatedPowerStatsConfig createAggregatedPowerStatsConfig(Context context,
+            PowerProfile powerProfile, CpuScalingPolicies cpuScalingPolicies,
+            PowerStatsUidResolver powerStatsUidResolver) {
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CPU)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new CpuPowerStatsProcessor(powerProfile, cpuScalingPolicies));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .setProcessorSupplier(
+                        () -> new ScreenPowerStatsProcessor(powerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .setProcessorSupplier(AmbientDisplayPowerStatsProcessor::new);
+
+        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)
+                .setProcessorSupplier(
+                        () -> new MobileRadioPowerStatsProcessor(powerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
+                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .setProcessorSupplier(PhoneCallPowerStatsProcessor::new);
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new WifiPowerStatsProcessor(powerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_BLUETOOTH)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new BluetoothPowerStatsProcessor(powerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AUDIO)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new AudioPowerStatsProcessor(powerProfile, powerStatsUidResolver));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_VIDEO)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new VideoPowerStatsProcessor(powerProfile, powerStatsUidResolver));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_FLASHLIGHT)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new FlashlightPowerStatsProcessor(powerProfile,
+                                powerStatsUidResolver));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CAMERA)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new CameraPowerStatsProcessor(powerProfile, powerStatsUidResolver));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_GNSS)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new GnssPowerStatsProcessor(powerProfile, powerStatsUidResolver));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SENSORS)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessorSupplier(() -> new SensorPowerStatsProcessor(
+                        () -> context.getSystemService(SensorManager.class)));
+
+        config.trackCustomPowerComponents(CustomEnergyConsumerPowerStatsProcessor::new)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
+        return config;
+    }
+
+    /**
+     * Marks the specified power component as supported by this PowerAttributor
+     */
+    public void setPowerComponentSupported(@BatteryConsumer.PowerComponentId int powerComponentId,
+            boolean enabled) {
+        mPowerStatsExporterEnabled.put(powerComponentId, enabled);
+        mPowerStatsExporter.setPowerComponentEnabled(powerComponentId, enabled);
+    }
+
+    @Override
+    public boolean isPowerComponentSupported(
+            @BatteryConsumer.PowerComponentId int powerComponentId) {
+        return mPowerStatsExporterEnabled.get(powerComponentId);
+    }
+
+    @Override
+    public void estimatePowerConsumption(BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            BatteryStatsHistory batteryHistory, long monotonicStartTime, long monotonicEndTime) {
+        mPowerStatsExporter.exportAggregatedPowerStats(batteryUsageStatsBuilder, batteryHistory,
+                monotonicStartTime, monotonicEndTime);
+    }
+
+    @Override
+    public void dumpEstimatedPowerConsumption(IndentingPrintWriter ipw,
+            BatteryStatsHistory batteryStatsHistory,
+            long startTime, long endTime) {
+        mPowerStatsAggregator.aggregatePowerStats(batteryStatsHistory, startTime, endTime,
+                stats -> {
+                    // Create a PowerStatsSpan for consistency of the textual output
+                    PowerStatsSpan span = createPowerStatsSpan(stats);
+                    if (span != null) {
+                        span.dump(ipw);
+                    }
+                });
+    }
+
+    @Override
+    public long storeEstimatedPowerConsumption(BatteryStatsHistory batteryStatsHistory,
+            long startTime, long endTimeMs) {
+        long[] lastSavedMonotonicTime = new long[1];
+        mPowerStatsAggregator.aggregatePowerStats(batteryStatsHistory, startTime, endTimeMs,
+                stats -> {
+                    storeAggregatedPowerStats(stats);
+                    lastSavedMonotonicTime[0] = stats.getStartTime() + stats.getDuration();
+                });
+        return lastSavedMonotonicTime[0];
+    }
+
+    @VisibleForTesting
+    void storeAggregatedPowerStats(AggregatedPowerStats stats) {
+        PowerStatsSpan span = createPowerStatsSpan(stats);
+        if (span == null) {
+            return;
+        }
+        mPowerStatsStore.storePowerStatsSpan(span);
+    }
+
+    private static PowerStatsSpan createPowerStatsSpan(AggregatedPowerStats stats) {
+        List<AggregatedPowerStats.ClockUpdate> clockUpdates = stats.getClockUpdates();
+        if (clockUpdates.isEmpty()) {
+            Slog.w(TAG, "No clock updates in aggregated power stats " + stats);
+            return null;
+        }
+
+        long monotonicTime = clockUpdates.get(0).monotonicTime;
+        long durationSum = 0;
+        PowerStatsSpan span = new PowerStatsSpan(monotonicTime);
+        for (int i = 0; i < clockUpdates.size(); i++) {
+            AggregatedPowerStats.ClockUpdate clockUpdate = clockUpdates.get(i);
+            long duration;
+            if (i == clockUpdates.size() - 1) {
+                duration = stats.getDuration() - durationSum;
+            } else {
+                duration = clockUpdate.monotonicTime - monotonicTime;
+            }
+            span.addTimeFrame(clockUpdate.monotonicTime, clockUpdate.currentTime, duration);
+            monotonicTime = clockUpdate.monotonicTime;
+            durationSum += duration;
+        }
+
+        span.addSection(new AggregatedPowerStatsSection(stats));
+        return span;
+    }
+
+    @Override
+    public long getLastSavedEstimatesPowerConsumptionTimestamp() {
+        long timestamp = -1;
+        for (PowerStatsSpan.Metadata metadata : mPowerStatsStore.getTableOfContents()) {
+            if (metadata.getSections().contains(AggregatedPowerStatsSection.TYPE)) {
+                for (PowerStatsSpan.TimeFrame timeFrame : metadata.getTimeFrames()) {
+                    long endMonotonicTime = timeFrame.startMonotonicTime + timeFrame.duration;
+                    if (endMonotonicTime > timestamp) {
+                        timestamp = endMonotonicTime;
+                    }
+                }
+            }
+        }
+        return timestamp;
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/MultiStateStats.java b/services/core/java/com/android/server/power/stats/processor/MultiStateStats.java
similarity index 94%
rename from services/core/java/com/android/server/power/stats/MultiStateStats.java
rename to services/core/java/com/android/server/power/stats/processor/MultiStateStats.java
index c3a0aeb..28474a5 100644
--- a/services/core/java/com/android/server/power/stats/MultiStateStats.java
+++ b/services/core/java/com/android/server/power/stats/processor/MultiStateStats.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.util.Slog;
 
@@ -37,7 +37,7 @@
  * CPU residency, Network packet counts etc.  All metrics must be represented as <code>long</code>
  * values;
  */
-public class MultiStateStats {
+class MultiStateStats {
     private static final String TAG = "MultiStateStats";
 
     private static final String XML_TAG_STATS = "stats";
@@ -47,12 +47,12 @@
      * A set of states, e.g. on-battery, screen-on, procstate.  The state values are integers
      * from 0 to States.mLabels.length
      */
-    public static class States {
+    static class States {
         final String mName;
         final boolean mTracked;
         final String[] mLabels;
 
-        public States(String name, boolean tracked, String... labels) {
+        States(String name, boolean tracked, String... labels) {
             mName = name;
             mTracked = tracked;
             mLabels = labels;
@@ -121,7 +121,7 @@
      * Factory for MultiStateStats containers. All generated containers retain their connection
      * to the Factory and the corresponding configuration.
      */
-    public static class Factory {
+    static class Factory {
         private static final int INVALID_SERIAL_STATE = -1;
         final int mDimensionCount;
         final States[] mStates;
@@ -147,7 +147,7 @@
         final int[] mCompositeToSerialState;
         final int mSerialStateCount;
 
-        public Factory(int dimensionCount, States... states) {
+        Factory(int dimensionCount, States... states) {
             mDimensionCount = dimensionCount;
             mStates = states;
 
@@ -250,7 +250,7 @@
         /**
          * Allocates a new stats container using this Factory's configuration.
          */
-        public MultiStateStats create() {
+        MultiStateStats create() {
             return new MultiStateStats(this, mDimensionCount);
         }
 
@@ -293,16 +293,16 @@
     private int mCompositeState;
     private boolean mTracking;
 
-    public MultiStateStats(Factory factory, int dimensionCount) {
+    MultiStateStats(Factory factory, int dimensionCount) {
         this.mFactory = factory;
         mCounter = new LongArrayMultiStateCounter(factory.mSerialStateCount, dimensionCount);
     }
 
-    public int getDimensionCount() {
+    int getDimensionCount() {
         return mFactory.mDimensionCount;
     }
 
-    public States[] getStates() {
+    States[] getStates() {
         return mFactory.mStates;
     }
 
@@ -310,7 +310,7 @@
      * Copies time-in-state and timestamps from the supplied prototype. Does not
      * copy accumulated counts.
      */
-    public void copyStatesFrom(MultiStateStats otherStats) {
+    void copyStatesFrom(MultiStateStats otherStats) {
         mCounter.copyStatesFrom(otherStats.mCounter);
     }
 
@@ -322,7 +322,7 @@
      * @param state       The new value of the state (e.g. 0 or 1 for "on-battery")
      * @param timestampMs The time when the state change occurred
      */
-    public void setState(int stateIndex, int state, long timestampMs) {
+    void setState(int stateIndex, int state, long timestampMs) {
         if (!mTracking) {
             mCounter.updateValues(new long[mCounter.getArrayLength()], timestampMs);
             mTracking = true;
@@ -335,7 +335,7 @@
      * Adds the delta to the metrics.  The number of values must correspond to the dimension count
      * supplied to the Factory constructor
      */
-    public void increment(long[] values, long timestampMs) {
+    void increment(long[] values, long timestampMs) {
         mCounter.incrementValues(values, timestampMs);
         mTracking = true;
     }
@@ -343,21 +343,21 @@
     /**
      * Returns accumulated stats for the specified composite state.
      */
-    public void getStats(long[] outValues, int[] states) {
+    void getStats(long[] outValues, int[] states) {
         mCounter.getCounts(outValues, mFactory.getSerialState(states));
     }
 
     /**
      * Updates the stats values for the provided combination of states.
      */
-    public void setStats(int[] states, long[] values) {
+    void setStats(int[] states, long[] values) {
         mCounter.setValues(mFactory.getSerialState(states), values);
     }
 
     /**
      * Resets the counters.
      */
-    public void reset() {
+    void reset() {
         mCounter.reset();
         mTracking = false;
     }
@@ -365,7 +365,7 @@
     /**
      * Stores contents in an XML doc.
      */
-    public void writeXml(TypedXmlSerializer serializer) throws IOException {
+    void writeXml(TypedXmlSerializer serializer) throws IOException {
         long[] tmpArray = new long[mCounter.getArrayLength()];
 
         try {
@@ -420,8 +420,7 @@
      * Populates the object with contents in an XML doc. The parser is expected to be
      * positioned on the opening tag of the corresponding element.
      */
-    public boolean readFromXml(TypedXmlPullParser parser) throws XmlPullParserException,
-            IOException {
+    boolean readFromXml(TypedXmlPullParser parser) throws XmlPullParserException, IOException {
         String outerTag = parser.getName();
         long[] tmpArray = new long[mCounter.getArrayLength()];
         int eventType = parser.getEventType();
diff --git a/services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessor.java
similarity index 88%
rename from services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessor.java
index ec23fa0..3957ae0 100644
--- a/services/core/java/com/android/server/power/stats/PhoneCallPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessor.java
@@ -13,24 +13,25 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.PersistableBundle;
 
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
+import com.android.server.power.stats.format.PhoneCallPowerStatsLayout;
 
-public class PhoneCallPowerStatsProcessor extends PowerStatsProcessor {
-    private final PowerStatsLayout mStatsLayout;
+class PhoneCallPowerStatsProcessor extends PowerStatsProcessor {
+    private final PhoneCallPowerStatsLayout 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();
+    PhoneCallPowerStatsProcessor() {
+        mStatsLayout = new PhoneCallPowerStatsLayout();
         PersistableBundle extras = new PersistableBundle();
         mStatsLayout.toExtras(extras);
         mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_PHONE,
diff --git a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/processor/PowerComponentAggregatedPowerStats.java
similarity index 98%
rename from services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
rename to services/core/java/com/android/server/power/stats/processor/PowerComponentAggregatedPowerStats.java
index a92a6fd3..d04c5ba 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerComponentAggregatedPowerStats.java
@@ -14,9 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
-
-import static com.android.server.power.stats.MultiStateStats.STATE_DOES_NOT_EXIST;
+package com.android.server.power.stats.processor;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -147,7 +145,8 @@
 
         int uidStateId = MultiStateStats.States
                 .findTrackedStateByName(mUidStateConfig, mDeviceStateConfig[stateId].getName());
-        if (uidStateId != STATE_DOES_NOT_EXIST && mUidStateConfig[uidStateId].isTracked()) {
+        if (uidStateId != MultiStateStats.STATE_DOES_NOT_EXIST
+                && mUidStateConfig[uidStateId].isTracked()) {
             for (int i = mUidStats.size() - 1; i >= 0; i--) {
                 PowerComponentAggregatedPowerStats.UidStats uidStats = mUidStats.valueAt(i);
                 if (uidStats.stats == null) {
@@ -271,7 +270,7 @@
                 if (mUidStateConfig[stateId].isTracked()) {
                     int deviceStateId = MultiStateStats.States.findTrackedStateByName(
                             mDeviceStateConfig, mUidStateConfig[stateId].getName());
-                    if (deviceStateId != STATE_DOES_NOT_EXIST
+                    if (deviceStateId != MultiStateStats.STATE_DOES_NOT_EXIST
                             && mDeviceStateConfig[deviceStateId].isTracked()) {
                         uidStats.states[stateId] = mDeviceStates[deviceStateId];
                     }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
similarity index 92%
rename from services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
rename to services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
index c734f68..32c1056 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsAggregator.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsAggregator.java
@@ -13,13 +13,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.annotation.NonNull;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 import android.util.SparseBooleanArray;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.BatteryStatsHistoryIterator;
 
@@ -33,20 +34,22 @@
 public class PowerStatsAggregator {
     private static final long UNINITIALIZED = -1;
     private final AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
-    private final BatteryStatsHistory mHistory;
     private final SparseBooleanArray mEnabledComponents =
             new SparseBooleanArray(BatteryConsumer.POWER_COMPONENT_COUNT + 10);
     private AggregatedPowerStats mStats;
     private int mCurrentBatteryState = AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
     private int mCurrentScreenState = AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
 
-    public PowerStatsAggregator(@NonNull AggregatedPowerStatsConfig aggregatedPowerStatsConfig,
-            @NonNull BatteryStatsHistory history) {
-        mAggregatedPowerStatsConfig = aggregatedPowerStatsConfig;
-        mHistory = history;
+    @VisibleForTesting
+    public PowerStatsAggregator() {
+        this(new AggregatedPowerStatsConfig());
     }
 
-    public AggregatedPowerStatsConfig getConfig() {
+    PowerStatsAggregator(@NonNull AggregatedPowerStatsConfig aggregatedPowerStatsConfig) {
+        mAggregatedPowerStatsConfig = aggregatedPowerStatsConfig;
+    }
+
+    AggregatedPowerStatsConfig getConfig() {
         return mAggregatedPowerStatsConfig;
     }
 
@@ -71,7 +74,7 @@
      * Note: the AggregatedPowerStats object is reused, so the consumer should fully consume
      * the stats in the <code>accept</code> method and never cache it.
      */
-    public void aggregatePowerStats(long startTimeMs, long endTimeMs,
+    public void aggregatePowerStats(BatteryStatsHistory history, long startTimeMs, long endTimeMs,
             Consumer<AggregatedPowerStats> consumer) {
         synchronized (this) {
             if (mStats == null) {
@@ -85,7 +88,7 @@
             long lastTime = 0;
             int lastStates = 0xFFFFFFFF;
             int lastStates2 = 0xFFFFFFFF;
-            try (BatteryStatsHistoryIterator iterator = mHistory.iterate(startTimeMs, endTimeMs)) {
+            try (BatteryStatsHistoryIterator iterator = history.iterate(startTimeMs, endTimeMs)) {
                 while (iterator.hasNext()) {
                     BatteryStats.HistoryItem item = iterator.next();
 
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
similarity index 94%
rename from services/core/java/com/android/server/power/stats/PowerStatsExporter.java
rename to services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
index c5bed24..fab87d6 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsExporter.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsExporter.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.annotation.Nullable;
 import android.os.AggregateBatteryConsumer;
@@ -24,7 +24,11 @@
 import android.util.Slog;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.PowerStatsSpan;
+import com.android.server.power.stats.PowerStatsStore;
+import com.android.server.power.stats.format.PowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -35,19 +39,18 @@
  * Given a time range, converts accumulated PowerStats to BatteryUsageStats.  Combines
  * stores spans of PowerStats with the yet-unprocessed tail of battery history.
  */
-public class PowerStatsExporter {
+class PowerStatsExporter {
     private static final String TAG = "PowerStatsExporter";
     private final PowerStatsStore mPowerStatsStore;
     private final PowerStatsAggregator mPowerStatsAggregator;
     private final long mBatterySessionTimeSpanSlackMillis;
     private static final long BATTERY_SESSION_TIME_SPAN_SLACK_MILLIS = TimeUnit.MINUTES.toMillis(2);
 
-    public PowerStatsExporter(PowerStatsStore powerStatsStore,
-            PowerStatsAggregator powerStatsAggregator) {
+    PowerStatsExporter(PowerStatsStore powerStatsStore, PowerStatsAggregator powerStatsAggregator) {
         this(powerStatsStore, powerStatsAggregator, BATTERY_SESSION_TIME_SPAN_SLACK_MILLIS);
     }
 
-    public PowerStatsExporter(PowerStatsStore powerStatsStore,
+    PowerStatsExporter(PowerStatsStore powerStatsStore,
             PowerStatsAggregator powerStatsAggregator,
             long batterySessionTimeSpanSlackMillis) {
         mPowerStatsStore = powerStatsStore;
@@ -59,8 +62,8 @@
      * Populates the provided BatteryUsageStats.Builder with power estimates from the accumulated
      * PowerStats, both stored in PowerStatsStore and not-yet processed.
      */
-    public void exportAggregatedPowerStats(BatteryUsageStats.Builder batteryUsageStatsBuilder,
-            long monotonicStartTime, long monotonicEndTime) {
+    void exportAggregatedPowerStats(BatteryUsageStats.Builder batteryUsageStatsBuilder,
+            BatteryStatsHistory history, long monotonicStartTime, long monotonicEndTime) {
         synchronized (mPowerStatsAggregator) {
             boolean hasStoredSpans = false;
             long maxEndTime = monotonicStartTime;
@@ -111,7 +114,7 @@
 
             if (!hasStoredSpans
                     || maxEndTime < monotonicEndTime - mBatterySessionTimeSpanSlackMillis) {
-                mPowerStatsAggregator.aggregatePowerStats(maxEndTime, monotonicEndTime,
+                mPowerStatsAggregator.aggregatePowerStats(history, maxEndTime, monotonicEndTime,
                         stats -> populateBatteryUsageStatsBuilder(batteryUsageStatsBuilder, stats));
             }
             mPowerStatsAggregator.reset();
@@ -140,9 +143,7 @@
             return;
         }
 
-        PowerStatsLayout layout = new PowerStatsLayout();
-        layout.fromExtras(descriptor.extras);
-
+        PowerStatsLayout layout = new PowerStatsLayout(descriptor);
         long[] deviceStats = new long[descriptor.statsArrayLength];
         for (int screenState = 0; screenState < BatteryConsumer.SCREEN_STATE_COUNT; screenState++) {
             if (batteryUsageStatsBuilder.isScreenStateDataNeeded()) {
@@ -328,8 +329,8 @@
         BatteryConsumer.Key key = getKeyForPartialTotal(batteryUsageStatsBuilder, allAppsScope,
                 powerComponentId, screenState, powerState);
         if (key != null) {
-                allAppsScope.addConsumedPower(key, powerAllApps,
-                        BatteryConsumer.POWER_MODEL_UNDEFINED);
+            allAppsScope.addConsumedPower(key, powerAllApps,
+                    BatteryConsumer.POWER_MODEL_UNDEFINED);
         }
         allAppsScope.addConsumedPower(powerComponentId, powerAllApps,
                 BatteryConsumer.POWER_MODEL_UNDEFINED);
@@ -393,7 +394,7 @@
         return true;
     }
 
-    void setPowerComponentEnabled(int powerComponentId, boolean enabled) {
+    public void setPowerComponentEnabled(int powerComponentId, boolean enabled) {
         mPowerStatsAggregator.setPowerComponentEnabled(powerComponentId, enabled);
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/PowerStatsProcessor.java
similarity index 97%
rename from services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/PowerStatsProcessor.java
index 6a8c6b12..838fc62 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/PowerStatsProcessor.java
@@ -13,10 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
-import static com.android.server.power.stats.MultiStateStats.STATE_DOES_NOT_EXIST;
-import static com.android.server.power.stats.MultiStateStats.States.findTrackedStateByName;
+import static com.android.server.power.stats.processor.MultiStateStats.STATE_DOES_NOT_EXIST;
+import static com.android.server.power.stats.processor.MultiStateStats.States.findTrackedStateByName;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -43,7 +43,7 @@
  * 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.
  */
-public abstract class PowerStatsProcessor {
+abstract class PowerStatsProcessor {
     private static final String TAG = "PowerStatsProcessor";
 
     private static final double MILLIAMPHOUR_PER_MICROCOULOMB = 1.0 / 1000.0 / 60.0 / 60.0;
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/ScreenPowerStatsProcessor.java
similarity index 93%
rename from services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/ScreenPowerStatsProcessor.java
index 8fb1fd6..b295e30 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/ScreenPowerStatsProcessor.java
@@ -14,28 +14,30 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static android.os.BatteryConsumer.PROCESS_STATE_UNSPECIFIED;
 
 import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_AMBIENT;
 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 com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
-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.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import android.os.BatteryStats;
 import android.util.Slog;
 
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.List;
 
-public class ScreenPowerStatsProcessor extends PowerStatsProcessor {
+class ScreenPowerStatsProcessor extends PowerStatsProcessor {
     private static final String TAG = "ScreenPowerStatsProcessor";
     private final int mDisplayCount;
     private final UsageBasedPowerEstimator[] mScreenOnPowerEstimators;
@@ -51,7 +53,7 @@
         public double power;
     }
 
-    public ScreenPowerStatsProcessor(PowerProfile powerProfile) {
+    ScreenPowerStatsProcessor(PowerProfile powerProfile) {
         mDisplayCount = powerProfile.getNumDisplays();
         mScreenOnPowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
         mScreenDozePowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
diff --git a/services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/SensorPowerStatsProcessor.java
similarity index 90%
rename from services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/SensorPowerStatsProcessor.java
index 79d8076..67013ea 100644
--- a/services/core/java/com/android/server/power/stats/SensorPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/SensorPowerStatsProcessor.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.hardware.Sensor;
 import android.hardware.SensorManager;
@@ -25,14 +25,16 @@
 import android.util.SparseArray;
 
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.PowerStatsLayout;
+import com.android.server.power.stats.format.SensorPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Comparator;
 import java.util.List;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
-public class SensorPowerStatsProcessor extends PowerStatsProcessor {
+class SensorPowerStatsProcessor extends PowerStatsProcessor {
     private static final String TAG = "SensorPowerStatsProcessor";
     private static final String ANDROID_SENSOR_TYPE_PREFIX = "android.sensor.";
 
@@ -64,7 +66,7 @@
     private long[] mTmpDeviceStatsArray;
     private long[] mTmpUidStatsArray;
 
-    public SensorPowerStatsProcessor(Supplier<SensorManager> sensorManagerSupplier) {
+    SensorPowerStatsProcessor(Supplier<SensorManager> sensorManagerSupplier) {
         mSensorManagerSupplier = sensorManagerSupplier;
     }
 
@@ -78,16 +80,9 @@
             return false;
         }
 
-        mStatsLayout = new SensorPowerStatsLayout();
-        List<Sensor> sensorList = new ArrayList<>(mSensorManager.getSensorList(Sensor.TYPE_ALL));
-        sensorList.sort(Comparator.comparingInt(Sensor::getId));
-        for (int i = 0; i < sensorList.size(); i++) {
-            Sensor sensor = sensorList.get(i);
-            String label = makeLabel(sensor, sensorList);
-            mStatsLayout.addUidSensorSection(sensor.getHandle(), label);
-        }
-        mStatsLayout.addUidSectionPowerEstimate();
-        mStatsLayout.addDeviceSectionPowerEstimate();
+        List<Sensor> sensorList = mSensorManager.getSensorList(Sensor.TYPE_ALL);
+        mStatsLayout = new SensorPowerStatsLayout(sensorList.stream().collect(
+                Collectors.toMap(Sensor::getHandle, sensor -> makeLabel(sensor, sensorList))));
 
         PersistableBundle extras = new PersistableBundle();
         mStatsLayout.toExtras(extras);
@@ -231,7 +226,8 @@
         sensorState.startTime = time;
     }
 
-    private void flushPowerStats(PowerComponentAggregatedPowerStats stats, long timestamp) {
+    private void flushPowerStats(
+            PowerComponentAggregatedPowerStats stats, long timestamp) {
         mPowerStats.durationMs = timestamp - mLastUpdateTimestamp;
         stats.addProcessedPowerStats(mPowerStats, timestamp);
 
@@ -240,7 +236,8 @@
         mLastUpdateTimestamp = timestamp;
     }
 
-    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats,
+    private void computeUidPowerEstimates(
+            PowerComponentAggregatedPowerStats stats,
             List<Integer> uids) {
         List<Sensor> sensorList = mSensorManager.getSensorList(Sensor.TYPE_ALL);
         int[] uidSensorDurationPositions = new int[sensorList.size()];
@@ -292,7 +289,8 @@
         }
     }
 
-    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats) {
+    private void computeDevicePowerEstimates(
+            PowerComponentAggregatedPowerStats stats) {
         for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
             CombinedDeviceStateEstimate estimation =
                     mPlan.combinedDeviceStateEstimations.get(i);
diff --git a/services/core/java/com/android/server/power/stats/VideoPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/VideoPowerStatsProcessor.java
similarity index 80%
rename from services/core/java/com/android/server/power/stats/VideoPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/VideoPowerStatsProcessor.java
index 48dac8a..a6c3807 100644
--- a/services/core/java/com/android/server/power/stats/VideoPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/VideoPowerStatsProcessor.java
@@ -14,16 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
 
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.PowerStatsUidResolver;
 
-public class VideoPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
-    public VideoPowerStatsProcessor(PowerProfile powerProfile,
-            PowerStatsUidResolver uidResolver) {
+class VideoPowerStatsProcessor extends BinaryStatePowerStatsProcessor {
+    VideoPowerStatsProcessor(PowerProfile powerProfile, PowerStatsUidResolver uidResolver) {
         super(BatteryConsumer.POWER_COMPONENT_VIDEO, uidResolver,
                 powerProfile.getAveragePower(PowerProfile.POWER_VIDEO));
     }
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/processor/WifiPowerStatsProcessor.java
similarity index 98%
rename from services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
rename to services/core/java/com/android/server/power/stats/processor/WifiPowerStatsProcessor.java
index 4e035c3..0df01cf 100644
--- a/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/processor/WifiPowerStatsProcessor.java
@@ -14,18 +14,20 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import android.util.Slog;
 
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.UsageBasedPowerEstimator;
+import com.android.server.power.stats.format.WifiPowerStatsLayout;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-public class WifiPowerStatsProcessor extends PowerStatsProcessor {
+class WifiPowerStatsProcessor extends PowerStatsProcessor {
     private static final String TAG = "WifiPowerStatsProcessor";
     private static final boolean DEBUG = false;
 
@@ -46,7 +48,7 @@
     private long[] mTmpUidStatsArray;
     private boolean mHasWifiPowerController;
 
-    public WifiPowerStatsProcessor(PowerProfile powerProfile) {
+    WifiPowerStatsProcessor(PowerProfile powerProfile) {
         mRxPowerEstimator = new UsageBasedPowerEstimator(
                 powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX));
         mTxPowerEstimator = new UsageBasedPowerEstimator(
diff --git a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
index 1c786e6..68026ea 100644
--- a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
+++ b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
@@ -18,6 +18,8 @@
 
 import static android.content.pm.Flags.provideInfoOfApkInApex;
 
+import static com.android.server.crashrecovery.CrashRecoveryUtils.logCrashRecoveryEvent;
+
 import android.annotation.AnyThread;
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -40,6 +42,7 @@
 import android.os.SystemProperties;
 import android.sysprop.CrashRecoveryProperties;
 import android.util.ArraySet;
+import android.util.Log;
 import android.util.Slog;
 import android.util.SparseArray;
 
@@ -532,11 +535,13 @@
     private void rollbackPackage(RollbackInfo rollback, VersionedPackage failedPackage,
             @FailureReasons int rollbackReason) {
         assertInWorkerThread();
+        String failedPackageName = (failedPackage == null ? null : failedPackage.getPackageName());
 
         Slog.i(TAG, "Rolling back package. RollbackId: " + rollback.getRollbackId()
-                + " failedPackage: "
-                + (failedPackage == null ? null : failedPackage.getPackageName())
+                + " failedPackage: " + failedPackageName
                 + " rollbackReason: " + rollbackReason);
+        logCrashRecoveryEvent(Log.DEBUG, String.format("Rolling back %s. Reason: %s",
+                failedPackageName, rollbackReason));
         final RollbackManager rollbackManager = mContext.getSystemService(RollbackManager.class);
         int reasonToLog = WatchdogRollbackLogger.mapFailureReasonToMetric(rollbackReason);
         final String failedPackageToLog;
@@ -724,6 +729,7 @@
         }
 
         Slog.i(TAG, "Rolling back all available low impact rollbacks");
+        logCrashRecoveryEvent(Log.DEBUG, "Rolling back all available. Reason: " + rollbackReason);
         // Add all rollback ids to mPendingStagedRollbackIds, so that we do not reboot before all
         // pending staged rollbacks are handled.
         for (RollbackInfo rollback : lowImpactRollbacks) {
diff --git a/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java b/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
index d6bf02f..6466519 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarShellCommand.java
@@ -190,6 +190,9 @@
                 case "notification-icons":
                     info.setNotificationIconsDisabled(true);
                     break;
+                case "quick-settings":
+                    info.setQuickSettingsDisabled(true);
+                    break;
                 default:
                     break;
             }
@@ -277,6 +280,7 @@
         pw.println("        system-icons        - disable system icons appearing in status bar");
         pw.println("        clock               - disable clock appearing in status bar");
         pw.println("        notification-icons  - disable notification icons from status bar");
+        pw.println("        quick-settings      - disable Quick Settings");
         pw.println("");
         pw.println("  tracing (start | stop)");
         pw.println("    Start or stop SystemUI tracing");
diff --git a/services/core/java/com/android/server/tv/TvInputManagerService.java b/services/core/java/com/android/server/tv/TvInputManagerService.java
index 6b3b5bd..67900f8 100644
--- a/services/core/java/com/android/server/tv/TvInputManagerService.java
+++ b/services/core/java/com/android/server/tv/TvInputManagerService.java
@@ -251,7 +251,7 @@
     }
 
     private void registerBroadcastReceivers() {
-        PackageMonitor monitor = new PackageMonitor() {
+        PackageMonitor monitor = new PackageMonitor(/* supportsPackageRestartQuery */ true) {
             private void buildTvInputList(String[] packages) {
                 int userId = getChangingUserId();
                 synchronized (mLock) {
diff --git a/services/core/java/com/android/server/tv/interactive/TvInteractiveAppManagerService.java b/services/core/java/com/android/server/tv/interactive/TvInteractiveAppManagerService.java
index edd2fa9..6a7fc6d 100644
--- a/services/core/java/com/android/server/tv/interactive/TvInteractiveAppManagerService.java
+++ b/services/core/java/com/android/server/tv/interactive/TvInteractiveAppManagerService.java
@@ -519,7 +519,7 @@
     }
 
     private void registerBroadcastReceivers() {
-        PackageMonitor monitor = new PackageMonitor() {
+        PackageMonitor monitor = new PackageMonitor(/* supportsPackageRestartQuery */ true) {
             private void buildTvInteractiveAppServiceList(String[] packages) {
                 int userId = getChangingUserId();
                 synchronized (mLock) {
diff --git a/services/core/java/com/android/server/tv/tunerresourcemanager/CasResource.java b/services/core/java/com/android/server/tv/tunerresourcemanager/CasResource.java
index 440d251..eb5361c 100644
--- a/services/core/java/com/android/server/tv/tunerresourcemanager/CasResource.java
+++ b/services/core/java/com/android/server/tv/tunerresourcemanager/CasResource.java
@@ -26,6 +26,11 @@
  * @hide
  */
 public class CasResource {
+    /**
+     * Handle of the current resource. Should not be changed and should be aligned with the driver
+     * level implementation.
+     */
+    final int mHandle;
 
     private final int mSystemId;
 
@@ -39,11 +44,16 @@
     private Map<Integer, Integer> mOwnerClientIdsToSessionNum = new HashMap<>();
 
     CasResource(Builder builder) {
+        this.mHandle = builder.mHandle;
         this.mSystemId = builder.mSystemId;
         this.mMaxSessionNum = builder.mMaxSessionNum;
         this.mAvailableSessionNum = builder.mMaxSessionNum;
     }
 
+    public int getHandle() {
+        return mHandle;
+    }
+
     public int getSystemId() {
         return mSystemId;
     }
@@ -136,10 +146,12 @@
      */
     public static class Builder {
 
+        private final int mHandle;
         private int mSystemId;
         protected int mMaxSessionNum;
 
-        Builder(int systemId) {
+        Builder(int handle, int systemId) {
+            this.mHandle = handle;
             this.mSystemId = systemId;
         }
 
diff --git a/services/core/java/com/android/server/tv/tunerresourcemanager/CiCamResource.java b/services/core/java/com/android/server/tv/tunerresourcemanager/CiCamResource.java
index 31149f3..5cef729 100644
--- a/services/core/java/com/android/server/tv/tunerresourcemanager/CiCamResource.java
+++ b/services/core/java/com/android/server/tv/tunerresourcemanager/CiCamResource.java
@@ -42,8 +42,8 @@
      * Builder class for {@link CiCamResource}.
      */
     public static class Builder extends CasResource.Builder {
-        Builder(int systemId) {
-            super(systemId);
+        Builder(int handle, int systemId) {
+            super(handle, systemId);
         }
 
         /**
diff --git a/services/core/java/com/android/server/tv/tunerresourcemanager/TunerResourceManagerService.java b/services/core/java/com/android/server/tv/tunerresourcemanager/TunerResourceManagerService.java
index 0afb049..9229f7f 100644
--- a/services/core/java/com/android/server/tv/tunerresourcemanager/TunerResourceManagerService.java
+++ b/services/core/java/com/android/server/tv/tunerresourcemanager/TunerResourceManagerService.java
@@ -203,13 +203,7 @@
         @Override
         public void unregisterClientProfile(int clientId) throws RemoteException {
             enforceTrmAccessPermission("unregisterClientProfile");
-            synchronized (mLock) {
-                if (!checkClientExists(clientId)) {
-                    Slog.e(TAG, "Unregistering non exists client:" + clientId);
-                    return;
-                }
-                unregisterClientProfileInternal(clientId);
-            }
+            unregisterClientProfileInternal(clientId);
         }
 
         @Override
@@ -291,20 +285,7 @@
                 Slog.e(TAG, "frontendHandle can't be null");
                 return false;
             }
-            synchronized (mLock) {
-                if (!checkClientExists(request.clientId)) {
-                    Slog.e(TAG, "Request frontend from unregistered client: "
-                            + request.clientId);
-                    return false;
-                }
-                // If the request client is holding or sharing a frontend, throw an exception.
-                if (!getClientProfile(request.clientId).getInUseFrontendHandles().isEmpty()) {
-                    Slog.e(TAG, "Release frontend before requesting another one. Client id: "
-                            + request.clientId);
-                    return false;
-                }
-                return requestFrontendInternal(request, frontendHandle);
-            }
+            return requestFrontendInternal(request, frontendHandle);
         }
 
         @Override
@@ -376,13 +357,7 @@
             if (demuxHandle == null) {
                 throw new RemoteException("demuxHandle can't be null");
             }
-            synchronized (mLock) {
-                if (!checkClientExists(request.clientId)) {
-                    throw new RemoteException("Request demux from unregistered client:"
-                            + request.clientId);
-                }
-                return requestDemuxInternal(request, demuxHandle);
-            }
+            return requestDemuxInternal(request, demuxHandle);
         }
 
         @Override
@@ -409,13 +384,7 @@
             if (casSessionHandle == null) {
                 throw new RemoteException("casSessionHandle can't be null");
             }
-            synchronized (mLock) {
-                if (!checkClientExists(request.clientId)) {
-                    throw new RemoteException("Request cas from unregistered client:"
-                            + request.clientId);
-                }
-                return requestCasSessionInternal(request, casSessionHandle);
-            }
+            return requestCasSessionInternal(request, casSessionHandle);
         }
 
         @Override
@@ -425,13 +394,7 @@
             if (ciCamHandle == null) {
                 throw new RemoteException("ciCamHandle can't be null");
             }
-            synchronized (mLock) {
-                if (!checkClientExists(request.clientId)) {
-                    throw new RemoteException("Request ciCam from unregistered client:"
-                            + request.clientId);
-                }
-                return requestCiCamInternal(request, ciCamHandle);
-            }
+            return requestCiCamInternal(request, ciCamHandle);
         }
 
         @Override
@@ -442,42 +405,14 @@
             if (lnbHandle == null) {
                 throw new RemoteException("lnbHandle can't be null");
             }
-            synchronized (mLock) {
-                if (!checkClientExists(request.clientId)) {
-                    throw new RemoteException("Request lnb from unregistered client:"
-                            + request.clientId);
-                }
-                return requestLnbInternal(request, lnbHandle);
-            }
+            return requestLnbInternal(request, lnbHandle);
         }
 
         @Override
         public void releaseFrontend(int frontendHandle, int clientId) throws RemoteException {
             enforceTunerAccessPermission("releaseFrontend");
             enforceTrmAccessPermission("releaseFrontend");
-            if (!validateResourceHandle(TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND,
-                    frontendHandle)) {
-                throw new RemoteException("frontendHandle can't be invalid");
-            }
-            synchronized (mLock) {
-                if (!checkClientExists(clientId)) {
-                    throw new RemoteException("Release frontend from unregistered client:"
-                            + clientId);
-                }
-                FrontendResource fe = getFrontendResource(frontendHandle);
-                if (fe == null) {
-                    throw new RemoteException("Releasing frontend does not exist.");
-                }
-                int ownerClientId = fe.getOwnerClientId();
-                ClientProfile ownerProfile = getClientProfile(ownerClientId);
-                if (ownerClientId != clientId
-                        && (ownerProfile != null
-                              && !ownerProfile.getShareFeClientIds().contains(clientId))) {
-                    throw new RemoteException(
-                            "Client is not the current owner of the releasing fe.");
-                }
-                releaseFrontendInternal(fe, clientId);
-            }
+            releaseFrontendInternal(frontendHandle, clientId);
         }
 
         @Override
@@ -746,17 +681,23 @@
 
     @VisibleForTesting
     protected void unregisterClientProfileInternal(int clientId) {
-        if (DEBUG) {
-            Slog.d(TAG, "unregisterClientProfile(clientId=" + clientId + ")");
-        }
-        removeClientProfile(clientId);
-        // Remove the Media Resource Manager callingPid to tvAppId mapping
-        if (mMediaResourceManager != null) {
-            try {
-                mMediaResourceManager.overridePid(Binder.getCallingPid(), -1);
-            } catch (RemoteException e) {
-                Slog.e(TAG, "Could not overridePid in resourceManagerSercice when unregister,"
-                        + " remote exception: " + e);
+        synchronized (mLock) {
+            if (!checkClientExists(clientId)) {
+                Slog.e(TAG, "Unregistering non exists client:" + clientId);
+                return;
+            }
+            if (DEBUG) {
+                Slog.d(TAG, "unregisterClientProfile(clientId=" + clientId + ")");
+            }
+            removeClientProfile(clientId);
+            // Remove the Media Resource Manager callingPid to tvAppId mapping
+            if (mMediaResourceManager != null) {
+                try {
+                    mMediaResourceManager.overridePid(Binder.getCallingPid(), -1);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Could not overridePid in resourceManagerSercice when unregister,"
+                            + " remote exception: " + e);
+                }
             }
         }
     }
@@ -992,10 +933,14 @@
             return;
         }
         // Add the new Cas Resource.
-        cas = new CasResource.Builder(casSystemId)
+        int casSessionHandle = generateResourceHandle(
+                TunerResourceManager.TUNER_RESOURCE_TYPE_CAS_SESSION, casSystemId);
+        cas = new CasResource.Builder(casSessionHandle, casSystemId)
                              .maxSessionNum(maxSessionNum)
                              .build();
-        ciCam = new CiCamResource.Builder(casSystemId)
+        int ciCamHandle = generateResourceHandle(
+                TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND_CICAM, casSystemId);
+        ciCam = new CiCamResource.Builder(ciCamHandle, casSystemId)
                              .maxSessionNum(maxSessionNum)
                              .build();
         addCasResource(cas);
@@ -1007,86 +952,120 @@
         if (DEBUG) {
             Slog.d(TAG, "requestFrontend(request=" + request + ")");
         }
-
-        frontendHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        ClientProfile requestClient = getClientProfile(request.clientId);
-        // TODO: check if this is really needed
-        if (requestClient == null) {
+        int[] reclaimOwnerId = new int[1];
+        if (!claimFrontend(request, frontendHandle, reclaimOwnerId)) {
             return false;
         }
-        clientPriorityUpdateOnRequest(requestClient);
-        int grantingFrontendHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        int inUseLowestPriorityFrHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        // Priority max value is 1000
-        int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
-        boolean isRequestFromSameProcess = false;
-        // If the desired frontend id was specified, we only need to check the frontend.
-        boolean hasDesiredFrontend = request.desiredId != TunerFrontendRequest.DEFAULT_DESIRED_ID;
-        for (FrontendResource fr : getFrontendResources().values()) {
-            int frontendId = getResourceIdFromHandle(fr.getHandle());
-            if (fr.getType() == request.frontendType
-                    && (!hasDesiredFrontend || frontendId == request.desiredId)) {
-                if (!fr.isInUse()) {
-                    // Unused resource cannot be acquired if the max is already reached, but
-                    // TRM still has to look for the reclaim candidate
-                    if (isFrontendMaxNumUseReached(request.frontendType)) {
-                        continue;
-                    }
-                    // Grant unused frontend with no exclusive group members first.
-                    if (fr.getExclusiveGroupMemberFeHandles().isEmpty()) {
-                        grantingFrontendHandle = fr.getHandle();
-                        break;
-                    } else if (grantingFrontendHandle
-                            == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
-                        // Grant the unused frontend with lower id first if all the unused
-                        // frontends have exclusive group members.
-                        grantingFrontendHandle = fr.getHandle();
-                    }
-                } else if (grantingFrontendHandle == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
-                    // Record the frontend id with the lowest client priority among all the
-                    // in use frontends when no available frontend has been found.
-                    int priority = getFrontendHighestClientPriority(fr.getOwnerClientId());
-                    if (currentLowestPriority > priority) {
-                        // we need to check the max used num if the target frontend type is not
-                        // currently in primary use (and simply blocked due to exclusive group)
-                        ClientProfile targetOwnerProfile = getClientProfile(fr.getOwnerClientId());
-                        int primaryFeId = targetOwnerProfile.getPrimaryFrontend();
-                        FrontendResource primaryFe = getFrontendResource(primaryFeId);
-                        if (fr.getType() != primaryFe.getType()
-                                && isFrontendMaxNumUseReached(fr.getType())) {
+        if (frontendHandle[0] == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
+            return false;
+        }
+        if (reclaimOwnerId[0] != INVALID_CLIENT_ID) {
+            if (!reclaimResource(reclaimOwnerId[0], TunerResourceManager
+                    .TUNER_RESOURCE_TYPE_FRONTEND)) {
+                return false;
+            }
+            synchronized (mLock) {
+                if (getFrontendResource(frontendHandle[0]).isInUse()) {
+                    Slog.e(TAG, "Reclaimed frontend still in use");
+                    return false;
+                }
+                updateFrontendClientMappingOnNewGrant(frontendHandle[0], request.clientId);
+            }
+        }
+        return true;
+    }
+
+    protected boolean claimFrontend(
+            TunerFrontendRequest request,
+            int[] frontendHandle,
+            int[] reclaimOwnerId
+    ) {
+        frontendHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
+        reclaimOwnerId[0] = INVALID_CLIENT_ID;
+        synchronized (mLock) {
+            if (!checkClientExists(request.clientId)) {
+                Slog.e(TAG, "Request frontend from unregistered client: "
+                        + request.clientId);
+                return false;
+            }
+            // If the request client is holding or sharing a frontend, throw an exception.
+            if (!getClientProfile(request.clientId).getInUseFrontendHandles().isEmpty()) {
+                Slog.e(TAG, "Release frontend before requesting another one. Client id: "
+                        + request.clientId);
+                return false;
+            }
+            ClientProfile requestClient = getClientProfile(request.clientId);
+            clientPriorityUpdateOnRequest(requestClient);
+            FrontendResource grantingFrontend = null;
+            FrontendResource inUseLowestPriorityFrontend = null;
+            // Priority max value is 1000
+            int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
+            boolean isRequestFromSameProcess = false;
+            // If the desired frontend id was specified, we only need to check the frontend.
+            boolean hasDesiredFrontend = request.desiredId != TunerFrontendRequest
+                    .DEFAULT_DESIRED_ID;
+            for (FrontendResource fr : getFrontendResources().values()) {
+                int frontendId = getResourceIdFromHandle(fr.getHandle());
+                if (fr.getType() == request.frontendType
+                        && (!hasDesiredFrontend || frontendId == request.desiredId)) {
+                    if (!fr.isInUse()) {
+                        // Unused resource cannot be acquired if the max is already reached, but
+                        // TRM still has to look for the reclaim candidate
+                        if (isFrontendMaxNumUseReached(request.frontendType)) {
                             continue;
                         }
-                        // update the target frontend
-                        inUseLowestPriorityFrHandle = fr.getHandle();
-                        currentLowestPriority = priority;
-                        isRequestFromSameProcess = (requestClient.getProcessId()
-                            == (getClientProfile(fr.getOwnerClientId())).getProcessId());
+                        // Grant unused frontend with no exclusive group members first.
+                        if (fr.getExclusiveGroupMemberFeHandles().isEmpty()) {
+                            grantingFrontend = fr;
+                            break;
+                        } else if (grantingFrontend == null) {
+                            // Grant the unused frontend with lower id first if all the unused
+                            // frontends have exclusive group members.
+                            grantingFrontend = fr;
+                        }
+                    } else if (grantingFrontend == null) {
+                        // Record the frontend id with the lowest client priority among all the
+                        // in use frontends when no available frontend has been found.
+                        int priority = getFrontendHighestClientPriority(fr.getOwnerClientId());
+                        if (currentLowestPriority > priority) {
+                            // we need to check the max used num if the target frontend type is not
+                            // currently in primary use (and simply blocked due to exclusive group)
+                            ClientProfile targetOwnerProfile =
+                                    getClientProfile(fr.getOwnerClientId());
+                            int primaryFeId = targetOwnerProfile.getPrimaryFrontend();
+                            FrontendResource primaryFe = getFrontendResource(primaryFeId);
+                            if (fr.getType() != primaryFe.getType()
+                                    && isFrontendMaxNumUseReached(fr.getType())) {
+                                continue;
+                            }
+                            // update the target frontend
+                            inUseLowestPriorityFrontend = fr;
+                            currentLowestPriority = priority;
+                            isRequestFromSameProcess = (requestClient.getProcessId()
+                                == (getClientProfile(fr.getOwnerClientId())).getProcessId());
+                        }
                     }
                 }
             }
-        }
 
-        // Grant frontend when there is unused resource.
-        if (grantingFrontendHandle != TunerResourceManager.INVALID_RESOURCE_HANDLE) {
-            frontendHandle[0] = grantingFrontendHandle;
-            updateFrontendClientMappingOnNewGrant(grantingFrontendHandle, request.clientId);
-            return true;
-        }
-
-        // When all the resources are occupied, grant the lowest priority resource if the
-        // request client has higher priority.
-        if (inUseLowestPriorityFrHandle != TunerResourceManager.INVALID_RESOURCE_HANDLE
-            && ((requestClient.getPriority() > currentLowestPriority) || (
-            (requestClient.getPriority() == currentLowestPriority) && isRequestFromSameProcess))) {
-            if (!reclaimResource(
-                    getFrontendResource(inUseLowestPriorityFrHandle).getOwnerClientId(),
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND)) {
-                return false;
+            // Grant frontend when there is unused resource.
+            if (grantingFrontend != null) {
+                updateFrontendClientMappingOnNewGrant(grantingFrontend.getHandle(),
+                        request.clientId);
+                frontendHandle[0] = grantingFrontend.getHandle();
+                return true;
             }
-            frontendHandle[0] = inUseLowestPriorityFrHandle;
-            updateFrontendClientMappingOnNewGrant(
-                    inUseLowestPriorityFrHandle, request.clientId);
-            return true;
+
+            // When all the resources are occupied, grant the lowest priority resource if the
+            // request client has higher priority.
+            if (inUseLowestPriorityFrontend != null
+                    && ((requestClient.getPriority() > currentLowestPriority)
+                    || ((requestClient.getPriority() == currentLowestPriority)
+                    && isRequestFromSameProcess))) {
+                frontendHandle[0] = inUseLowestPriorityFrontend.getHandle();
+                reclaimOwnerId[0] = inUseLowestPriorityFrontend.getOwnerClientId();
+                return true;
+            }
         }
 
         return false;
@@ -1192,165 +1171,257 @@
     }
 
     @VisibleForTesting
-    protected boolean requestLnbInternal(TunerLnbRequest request, int[] lnbHandle) {
+    protected boolean requestLnbInternal(TunerLnbRequest request, int[] lnbHandle)
+            throws RemoteException {
         if (DEBUG) {
             Slog.d(TAG, "requestLnb(request=" + request + ")");
         }
-
-        lnbHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        ClientProfile requestClient = getClientProfile(request.clientId);
-        clientPriorityUpdateOnRequest(requestClient);
-        int grantingLnbHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        int inUseLowestPriorityLnbHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        // Priority max value is 1000
-        int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
-        boolean isRequestFromSameProcess = false;
-        for (LnbResource lnb : getLnbResources().values()) {
-            if (!lnb.isInUse()) {
-                // Grant the unused lnb with lower handle first
-                grantingLnbHandle = lnb.getHandle();
-                break;
-            } else {
-                // Record the lnb id with the lowest client priority among all the
-                // in use lnb when no available lnb has been found.
-                int priority = updateAndGetOwnerClientPriority(lnb.getOwnerClientId());
-                if (currentLowestPriority > priority) {
-                    inUseLowestPriorityLnbHandle = lnb.getHandle();
-                    currentLowestPriority = priority;
-                    isRequestFromSameProcess = (requestClient.getProcessId()
-                        == (getClientProfile(lnb.getOwnerClientId())).getProcessId());
-                }
-            }
+        int[] reclaimOwnerId = new int[1];
+        if (!claimLnb(request, lnbHandle, reclaimOwnerId)) {
+            return false;
         }
-
-        // Grant Lnb when there is unused resource.
-        if (grantingLnbHandle > -1) {
-            lnbHandle[0] = grantingLnbHandle;
-            updateLnbClientMappingOnNewGrant(grantingLnbHandle, request.clientId);
-            return true;
+        if (lnbHandle[0] == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
+            return false;
         }
-
-        // When all the resources are occupied, grant the lowest priority resource if the
-        // request client has higher priority.
-        if (inUseLowestPriorityLnbHandle > TunerResourceManager.INVALID_RESOURCE_HANDLE
-            && ((requestClient.getPriority() > currentLowestPriority) || (
-            (requestClient.getPriority() == currentLowestPriority) && isRequestFromSameProcess))) {
-            if (!reclaimResource(getLnbResource(inUseLowestPriorityLnbHandle).getOwnerClientId(),
+        if (reclaimOwnerId[0] != INVALID_CLIENT_ID) {
+            if (!reclaimResource(reclaimOwnerId[0],
                     TunerResourceManager.TUNER_RESOURCE_TYPE_LNB)) {
                 return false;
             }
-            lnbHandle[0] = inUseLowestPriorityLnbHandle;
-            updateLnbClientMappingOnNewGrant(inUseLowestPriorityLnbHandle, request.clientId);
-            return true;
+            synchronized (mLock) {
+                if (getLnbResource(lnbHandle[0]).isInUse()) {
+                    Slog.e(TAG, "Reclaimed lnb still in use");
+                    return false;
+                }
+                updateLnbClientMappingOnNewGrant(lnbHandle[0], request.clientId);
+            }
+        }
+        return true;
+    }
+
+    protected boolean claimLnb(TunerLnbRequest request, int[] lnbHandle, int[] reclaimOwnerId)
+            throws RemoteException {
+        lnbHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
+        reclaimOwnerId[0] = INVALID_CLIENT_ID;
+        synchronized (mLock) {
+            if (!checkClientExists(request.clientId)) {
+                throw new RemoteException("Request lnb from unregistered client:"
+                        + request.clientId);
+            }
+            ClientProfile requestClient = getClientProfile(request.clientId);
+            clientPriorityUpdateOnRequest(requestClient);
+            LnbResource grantingLnb = null;
+            LnbResource inUseLowestPriorityLnb = null;
+            // Priority max value is 1000
+            int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
+            boolean isRequestFromSameProcess = false;
+            for (LnbResource lnb : getLnbResources().values()) {
+                if (!lnb.isInUse()) {
+                    // Grant the unused lnb with lower handle first
+                    grantingLnb = lnb;
+                    break;
+                } else {
+                    // Record the lnb id with the lowest client priority among all the
+                    // in use lnb when no available lnb has been found.
+                    int priority = updateAndGetOwnerClientPriority(lnb.getOwnerClientId());
+                    if (currentLowestPriority > priority) {
+                        inUseLowestPriorityLnb = lnb;
+                        currentLowestPriority = priority;
+                        isRequestFromSameProcess = (requestClient.getProcessId()
+                            == (getClientProfile(lnb.getOwnerClientId())).getProcessId());
+                    }
+                }
+            }
+
+            // Grant Lnb when there is unused resource.
+            if (grantingLnb != null) {
+                updateLnbClientMappingOnNewGrant(grantingLnb.getHandle(), request.clientId);
+                lnbHandle[0] = grantingLnb.getHandle();
+                return true;
+            }
+
+            // When all the resources are occupied, grant the lowest priority resource if the
+            // request client has higher priority.
+            if (inUseLowestPriorityLnb != null
+                    && ((requestClient.getPriority() > currentLowestPriority) || (
+                    (requestClient.getPriority() == currentLowestPriority)
+                        && isRequestFromSameProcess))) {
+                lnbHandle[0] = inUseLowestPriorityLnb.getHandle();
+                reclaimOwnerId[0] = inUseLowestPriorityLnb.getOwnerClientId();
+                return true;
+            }
         }
 
         return false;
     }
 
     @VisibleForTesting
-    protected boolean requestCasSessionInternal(CasSessionRequest request, int[] casSessionHandle) {
+    protected boolean requestCasSessionInternal(CasSessionRequest request, int[] casSessionHandle)
+            throws RemoteException {
         if (DEBUG) {
             Slog.d(TAG, "requestCasSession(request=" + request + ")");
         }
-        CasResource cas = getCasResource(request.casSystemId);
-        // Unregistered Cas System is treated as having unlimited sessions.
-        if (cas == null) {
-            cas = new CasResource.Builder(request.casSystemId)
-                                 .maxSessionNum(Integer.MAX_VALUE)
-                                 .build();
-            addCasResource(cas);
+        int[] reclaimOwnerId = new int[1];
+        if (!claimCasSession(request, casSessionHandle, reclaimOwnerId)) {
+            return false;
         }
-        casSessionHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        ClientProfile requestClient = getClientProfile(request.clientId);
-        clientPriorityUpdateOnRequest(requestClient);
-        int lowestPriorityOwnerId = -1;
-        // Priority max value is 1000
-        int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
-        boolean isRequestFromSameProcess = false;
-        if (!cas.isFullyUsed()) {
-            casSessionHandle[0] = generateResourceHandle(
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_CAS_SESSION, cas.getSystemId());
-            updateCasClientMappingOnNewGrant(request.casSystemId, request.clientId);
-            return true;
+        if (casSessionHandle[0] == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
+            return false;
         }
-        for (int ownerId : cas.getOwnerClientIds()) {
-            // Record the client id with lowest priority that is using the current Cas system.
-            int priority = updateAndGetOwnerClientPriority(ownerId);
-            if (currentLowestPriority > priority) {
-                lowestPriorityOwnerId = ownerId;
-                currentLowestPriority = priority;
-                isRequestFromSameProcess = (requestClient.getProcessId()
-                    == (getClientProfile(ownerId)).getProcessId());
-            }
-        }
-
-        // When all the Cas sessions are occupied, reclaim the lowest priority client if the
-        // request client has higher priority.
-        if (lowestPriorityOwnerId > -1 && ((requestClient.getPriority() > currentLowestPriority)
-        || ((requestClient.getPriority() == currentLowestPriority) && isRequestFromSameProcess))) {
-            if (!reclaimResource(lowestPriorityOwnerId,
+        if (reclaimOwnerId[0] != INVALID_CLIENT_ID) {
+            if (!reclaimResource(reclaimOwnerId[0],
                     TunerResourceManager.TUNER_RESOURCE_TYPE_CAS_SESSION)) {
                 return false;
             }
-            casSessionHandle[0] = generateResourceHandle(
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_CAS_SESSION, cas.getSystemId());
-            updateCasClientMappingOnNewGrant(request.casSystemId, request.clientId);
-            return true;
+            synchronized (mLock) {
+                if (getCasResource(request.casSystemId).isFullyUsed()) {
+                    Slog.e(TAG, "Reclaimed cas still fully used");
+                    return false;
+                }
+                updateCasClientMappingOnNewGrant(request.casSystemId, request.clientId);
+            }
         }
+        return true;
+    }
+
+    protected boolean claimCasSession(CasSessionRequest request, int[] casSessionHandle,
+            int[] reclaimOwnerId) throws RemoteException {
+        casSessionHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
+        reclaimOwnerId[0] = INVALID_CLIENT_ID;
+        synchronized (mLock) {
+            if (!checkClientExists(request.clientId)) {
+                throw new RemoteException("Request cas from unregistered client:"
+                        + request.clientId);
+            }
+            CasResource cas = getCasResource(request.casSystemId);
+            // Unregistered Cas System is treated as having unlimited sessions.
+            if (cas == null) {
+                int resourceHandle = generateResourceHandle(
+                        TunerResourceManager.TUNER_RESOURCE_TYPE_CAS_SESSION, request.clientId);
+                cas = new CasResource.Builder(resourceHandle, request.casSystemId)
+                                    .maxSessionNum(Integer.MAX_VALUE)
+                                    .build();
+                addCasResource(cas);
+            }
+            ClientProfile requestClient = getClientProfile(request.clientId);
+            clientPriorityUpdateOnRequest(requestClient);
+            int lowestPriorityOwnerId = INVALID_CLIENT_ID;
+            // Priority max value is 1000
+            int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
+            boolean isRequestFromSameProcess = false;
+            if (!cas.isFullyUsed()) {
+                updateCasClientMappingOnNewGrant(request.casSystemId, request.clientId);
+                casSessionHandle[0] = cas.getHandle();
+                return true;
+            }
+            for (int ownerId : cas.getOwnerClientIds()) {
+                // Record the client id with lowest priority that is using the current Cas system.
+                int priority = updateAndGetOwnerClientPriority(ownerId);
+                if (currentLowestPriority > priority) {
+                    lowestPriorityOwnerId = ownerId;
+                    currentLowestPriority = priority;
+                    isRequestFromSameProcess = (requestClient.getProcessId()
+                        == (getClientProfile(ownerId)).getProcessId());
+                }
+            }
+
+            // When all the Cas sessions are occupied, reclaim the lowest priority client if the
+            // request client has higher priority.
+            if (lowestPriorityOwnerId != INVALID_CLIENT_ID
+                    && ((requestClient.getPriority() > currentLowestPriority)
+                    || ((requestClient.getPriority() == currentLowestPriority)
+                    && isRequestFromSameProcess))) {
+                casSessionHandle[0] = cas.getHandle();
+                reclaimOwnerId[0] = lowestPriorityOwnerId;
+                return true;
+            }
+        }
+
         return false;
     }
 
     @VisibleForTesting
-    protected boolean requestCiCamInternal(TunerCiCamRequest request, int[] ciCamHandle) {
+    protected boolean requestCiCamInternal(TunerCiCamRequest request, int[] ciCamHandle)
+            throws RemoteException {
         if (DEBUG) {
             Slog.d(TAG, "requestCiCamInternal(TunerCiCamRequest=" + request + ")");
         }
-        CiCamResource ciCam = getCiCamResource(request.ciCamId);
-        // Unregistered Cas System is treated as having unlimited sessions.
-        if (ciCam == null) {
-            ciCam = new CiCamResource.Builder(request.ciCamId)
-                                     .maxSessionNum(Integer.MAX_VALUE)
-                                     .build();
-            addCiCamResource(ciCam);
+        int[] reclaimOwnerId = new int[1];
+        if (!claimCiCam(request, ciCamHandle, reclaimOwnerId)) {
+            return false;
         }
-        ciCamHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        ClientProfile requestClient = getClientProfile(request.clientId);
-        clientPriorityUpdateOnRequest(requestClient);
-        int lowestPriorityOwnerId = -1;
-        // Priority max value is 1000
-        int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
-        boolean isRequestFromSameProcess = false;
-        if (!ciCam.isFullyUsed()) {
-            ciCamHandle[0] = generateResourceHandle(
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND_CICAM, ciCam.getCiCamId());
-            updateCiCamClientMappingOnNewGrant(request.ciCamId, request.clientId);
-            return true;
+        if (ciCamHandle[0] == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
+            return false;
         }
-        for (int ownerId : ciCam.getOwnerClientIds()) {
-            // Record the client id with lowest priority that is using the current Cas system.
-            int priority = updateAndGetOwnerClientPriority(ownerId);
-            if (currentLowestPriority > priority) {
-                lowestPriorityOwnerId = ownerId;
-                currentLowestPriority = priority;
-                isRequestFromSameProcess = (requestClient.getProcessId()
-                    == (getClientProfile(ownerId)).getProcessId());
-            }
-        }
-
-        // When all the CiCam sessions are occupied, reclaim the lowest priority client if the
-        // request client has higher priority.
-        if (lowestPriorityOwnerId > -1 && ((requestClient.getPriority() > currentLowestPriority)
-            || ((requestClient.getPriority() == currentLowestPriority)
-                && isRequestFromSameProcess))) {
-            if (!reclaimResource(lowestPriorityOwnerId,
+        if (reclaimOwnerId[0] != INVALID_CLIENT_ID) {
+            if (!reclaimResource(reclaimOwnerId[0],
                     TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND_CICAM)) {
                 return false;
             }
-            ciCamHandle[0] = generateResourceHandle(
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND_CICAM, ciCam.getCiCamId());
-            updateCiCamClientMappingOnNewGrant(request.ciCamId, request.clientId);
-            return true;
+            synchronized (mLock) {
+                if (getCiCamResource(request.ciCamId).isFullyUsed()) {
+                    Slog.e(TAG, "Reclaimed ciCam still fully used");
+                    return false;
+                }
+                updateCiCamClientMappingOnNewGrant(request.ciCamId, request.clientId);
+            }
         }
+        return true;
+    }
+
+    protected boolean claimCiCam(TunerCiCamRequest request, int[] ciCamHandle,
+            int[] reclaimOwnerId) throws RemoteException {
+        ciCamHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
+        reclaimOwnerId[0] = INVALID_CLIENT_ID;
+        synchronized (mLock) {
+            if (!checkClientExists(request.clientId)) {
+                throw new RemoteException("Request ciCam from unregistered client:"
+                        + request.clientId);
+            }
+            CiCamResource ciCam = getCiCamResource(request.ciCamId);
+            // Unregistered CiCam is treated as having unlimited sessions.
+            if (ciCam == null) {
+                int resourceHandle = generateResourceHandle(
+                        TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND_CICAM, request.ciCamId);
+                ciCam = new CiCamResource.Builder(resourceHandle, request.ciCamId)
+                                    .maxSessionNum(Integer.MAX_VALUE)
+                                    .build();
+                addCiCamResource(ciCam);
+            }
+            ClientProfile requestClient = getClientProfile(request.clientId);
+            clientPriorityUpdateOnRequest(requestClient);
+            int lowestPriorityOwnerId = INVALID_CLIENT_ID;
+            // Priority max value is 1000
+            int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
+            boolean isRequestFromSameProcess = false;
+            if (!ciCam.isFullyUsed()) {
+                updateCiCamClientMappingOnNewGrant(request.ciCamId, request.clientId);
+                ciCamHandle[0] = ciCam.getHandle();
+                return true;
+            }
+            for (int ownerId : ciCam.getOwnerClientIds()) {
+                // Record the client id with lowest priority that is using the current CiCam.
+                int priority = updateAndGetOwnerClientPriority(ownerId);
+                if (currentLowestPriority > priority) {
+                    lowestPriorityOwnerId = ownerId;
+                    currentLowestPriority = priority;
+                    isRequestFromSameProcess = (requestClient.getProcessId()
+                        == (getClientProfile(ownerId)).getProcessId());
+                }
+            }
+
+            // When all the CiCam sessions are occupied, reclaim the lowest priority client if the
+            // request client has higher priority.
+            if (lowestPriorityOwnerId != INVALID_CLIENT_ID
+                    && ((requestClient.getPriority() > currentLowestPriority)
+                    || ((requestClient.getPriority() == currentLowestPriority)
+                    && isRequestFromSameProcess))) {
+                ciCamHandle[0] = ciCam.getHandle();
+                reclaimOwnerId[0] = lowestPriorityOwnerId;
+                return true;
+            }
+        }
+
         return false;
     }
 
@@ -1383,20 +1454,49 @@
     }
 
     @VisibleForTesting
-    protected void releaseFrontendInternal(FrontendResource fe, int clientId) {
+    protected void releaseFrontendInternal(int frontendHandle, int clientId)
+            throws RemoteException {
         if (DEBUG) {
-            Slog.d(TAG, "releaseFrontend(id=" + fe.getHandle() + ", clientId=" + clientId + " )");
+            Slog.d(TAG, "releaseFrontend(id=" + frontendHandle + ", clientId=" + clientId + " )");
         }
-        if (clientId == fe.getOwnerClientId()) {
-            ClientProfile ownerClient = getClientProfile(fe.getOwnerClientId());
-            if (ownerClient != null) {
-                for (int shareOwnerId : ownerClient.getShareFeClientIds()) {
-                    reclaimResource(shareOwnerId,
-                            TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND);
+        if (!validateResourceHandle(TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND,
+                frontendHandle)) {
+            throw new RemoteException("frontendHandle can't be invalid");
+        }
+        Set<Integer> reclaimedResourceOwnerIds = unclaimFrontend(frontendHandle, clientId);
+        if (reclaimedResourceOwnerIds != null) {
+            for (int shareOwnerId : reclaimedResourceOwnerIds) {
+                reclaimResource(shareOwnerId,
+                        TunerResourceManager.TUNER_RESOURCE_TYPE_FRONTEND);
+            }
+        }
+        synchronized (mLock) {
+            clearFrontendAndClientMapping(getClientProfile(clientId));
+        }
+    }
+
+    private Set<Integer> unclaimFrontend(int frontendHandle, int clientId) throws RemoteException {
+        Set<Integer> reclaimedResourceOwnerIds = null;
+        synchronized (mLock) {
+            if (!checkClientExists(clientId)) {
+                throw new RemoteException("Release frontend from unregistered client:"
+                        + clientId);
+            }
+            FrontendResource fe = getFrontendResource(frontendHandle);
+            if (fe == null) {
+                throw new RemoteException("Releasing frontend does not exist.");
+            }
+            int ownerClientId = fe.getOwnerClientId();
+            ClientProfile ownerProfile = getClientProfile(ownerClientId);
+            if (ownerClientId == clientId) {
+                reclaimedResourceOwnerIds = ownerProfile.getShareFeClientIds();
+            } else {
+                if (!ownerProfile.getShareFeClientIds().contains(clientId)) {
+                    throw new RemoteException("Client is not a sharee of the releasing fe.");
                 }
             }
         }
-        clearFrontendAndClientMapping(getClientProfile(clientId));
+        return reclaimedResourceOwnerIds;
     }
 
     @VisibleForTesting
@@ -1432,103 +1532,129 @@
     }
 
     @VisibleForTesting
-    protected boolean requestDemuxInternal(TunerDemuxRequest request, int[] demuxHandle) {
+    public boolean requestDemuxInternal(@NonNull TunerDemuxRequest request,
+                @NonNull int[] demuxHandle) throws RemoteException {
         if (DEBUG) {
             Slog.d(TAG, "requestDemux(request=" + request + ")");
         }
-
-        // For Tuner 2.0 and below or any HW constraint devices that are unable to support
-        // ITuner.openDemuxById(), demux resources are not really managed under TRM and
-        // mDemuxResources.size() will be zero
-        if (mDemuxResources.size() == 0) {
-            // There are enough Demux resources, so we don't manage Demux in R.
-            demuxHandle[0] =
-                    generateResourceHandle(TunerResourceManager.TUNER_RESOURCE_TYPE_DEMUX, 0);
-            return true;
-        }
-
-        demuxHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        ClientProfile requestClient = getClientProfile(request.clientId);
-
-        if (requestClient == null) {
+        int[] reclaimOwnerId = new int[1];
+        if (!claimDemux(request, demuxHandle, reclaimOwnerId)) {
             return false;
         }
+        if (demuxHandle[0] == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
+            return false;
+        }
+        if (reclaimOwnerId[0] != INVALID_CLIENT_ID) {
+            if (!reclaimResource(reclaimOwnerId[0],
+                    TunerResourceManager.TUNER_RESOURCE_TYPE_DEMUX)) {
+                return false;
+            }
+            synchronized (mLock) {
+                if (getDemuxResource(demuxHandle[0]).isInUse()) {
+                    Slog.e(TAG, "Reclaimed demux still in use");
+                    return false;
+                }
+                updateDemuxClientMappingOnNewGrant(demuxHandle[0], request.clientId);
+            }
+        }
+        return true;
+    }
 
-        clientPriorityUpdateOnRequest(requestClient);
-        int grantingDemuxHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        int inUseLowestPriorityDrHandle = TunerResourceManager.INVALID_RESOURCE_HANDLE;
-        // Priority max value is 1000
-        int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
-        boolean isRequestFromSameProcess = false;
-        // If the desired demux id was specified, we only need to check the demux.
-        boolean hasDesiredDemuxCap = request.desiredFilterTypes
-                != DemuxFilterMainType.UNDEFINED;
-        int smallestNumOfSupportedCaps = Integer.SIZE + 1;
-        int smallestNumOfSupportedCapsInUse = Integer.SIZE + 1;
-        for (DemuxResource dr : getDemuxResources().values()) {
-            if (!hasDesiredDemuxCap || dr.hasSufficientCaps(request.desiredFilterTypes)) {
-                if (!dr.isInUse()) {
-                    int numOfSupportedCaps = dr.getNumOfCaps();
+    protected boolean claimDemux(TunerDemuxRequest request, int[] demuxHandle, int[] reclaimOwnerId)
+            throws RemoteException {
+        demuxHandle[0] = TunerResourceManager.INVALID_RESOURCE_HANDLE;
+        reclaimOwnerId[0] = INVALID_CLIENT_ID;
+        synchronized (mLock) {
+            if (!checkClientExists(request.clientId)) {
+                throw new RemoteException("Request demux from unregistered client:"
+                        + request.clientId);
+            }
 
-                    // look for the demux with minimum caps supporting the desired caps
-                    if (smallestNumOfSupportedCaps > numOfSupportedCaps) {
-                        smallestNumOfSupportedCaps = numOfSupportedCaps;
-                        grantingDemuxHandle = dr.getHandle();
-                    }
-                } else if (grantingDemuxHandle == TunerResourceManager.INVALID_RESOURCE_HANDLE) {
-                    // Record the demux id with the lowest client priority among all the
-                    // in use demuxes when no availabledemux has been found.
-                    int priority = updateAndGetOwnerClientPriority(dr.getOwnerClientId());
-                    if (currentLowestPriority >= priority) {
+            // For Tuner 2.0 and below or any HW constraint devices that are unable to support
+            // ITuner.openDemuxById(), demux resources are not really managed under TRM and
+            // mDemuxResources.size() will be zero
+            if (mDemuxResources.size() == 0) {
+                // There are enough Demux resources, so we don't manage Demux in R.
+                demuxHandle[0] =
+                        generateResourceHandle(TunerResourceManager.TUNER_RESOURCE_TYPE_DEMUX, 0);
+                return true;
+            }
+
+            ClientProfile requestClient = getClientProfile(request.clientId);
+            if (requestClient == null) {
+                return false;
+            }
+            clientPriorityUpdateOnRequest(requestClient);
+            DemuxResource grantingDemux = null;
+            DemuxResource inUseLowestPriorityDemux = null;
+            // Priority max value is 1000
+            int currentLowestPriority = MAX_CLIENT_PRIORITY + 1;
+            boolean isRequestFromSameProcess = false;
+            // If the desired demux id was specified, we only need to check the demux.
+            boolean hasDesiredDemuxCap = request.desiredFilterTypes
+                    != DemuxFilterMainType.UNDEFINED;
+            int smallestNumOfSupportedCaps = Integer.SIZE + 1;
+            int smallestNumOfSupportedCapsInUse = Integer.SIZE + 1;
+            for (DemuxResource dr : getDemuxResources().values()) {
+                if (!hasDesiredDemuxCap || dr.hasSufficientCaps(request.desiredFilterTypes)) {
+                    if (!dr.isInUse()) {
                         int numOfSupportedCaps = dr.getNumOfCaps();
-                        boolean shouldUpdate = false;
-                        // update lowest priority
-                        if (currentLowestPriority > priority) {
-                            currentLowestPriority = priority;
-                            isRequestFromSameProcess = (requestClient.getProcessId()
-                                == (getClientProfile(dr.getOwnerClientId())).getProcessId());
 
-                            // reset the smallest caps when lower priority resource is found
-                            smallestNumOfSupportedCapsInUse = numOfSupportedCaps;
-
-                            shouldUpdate = true;
-                        } else {
-                            // This is the case when the priority is the same as previously found
-                            // one. Update smallest caps when priority.
-                            if (smallestNumOfSupportedCapsInUse > numOfSupportedCaps) {
-                                smallestNumOfSupportedCapsInUse = numOfSupportedCaps;
-                                shouldUpdate = true;
-                            }
+                        // look for the demux with minimum caps supporting the desired caps
+                        if (smallestNumOfSupportedCaps > numOfSupportedCaps) {
+                            smallestNumOfSupportedCaps = numOfSupportedCaps;
+                            grantingDemux = dr;
                         }
-                        if (shouldUpdate) {
-                            inUseLowestPriorityDrHandle = dr.getHandle();
+                    } else if (grantingDemux == null) {
+                        // Record the demux id with the lowest client priority among all the
+                        // in use demuxes when no availabledemux has been found.
+                        int priority = updateAndGetOwnerClientPriority(dr.getOwnerClientId());
+                        if (currentLowestPriority >= priority) {
+                            int numOfSupportedCaps = dr.getNumOfCaps();
+                            boolean shouldUpdate = false;
+                            // update lowest priority
+                            if (currentLowestPriority > priority) {
+                                currentLowestPriority = priority;
+                                isRequestFromSameProcess = (requestClient.getProcessId()
+                                    == (getClientProfile(dr.getOwnerClientId())).getProcessId());
+
+                                // reset the smallest caps when lower priority resource is found
+                                smallestNumOfSupportedCapsInUse = numOfSupportedCaps;
+
+                                shouldUpdate = true;
+                            } else {
+                                // This is the case when the priority is the same as previously
+                                // found one. Update smallest caps when priority.
+                                if (smallestNumOfSupportedCapsInUse > numOfSupportedCaps) {
+                                    smallestNumOfSupportedCapsInUse = numOfSupportedCaps;
+                                    shouldUpdate = true;
+                                }
+                            }
+                            if (shouldUpdate) {
+                                inUseLowestPriorityDemux = dr;
+                            }
                         }
                     }
                 }
             }
-        }
 
-        // Grant demux when there is unused resource.
-        if (grantingDemuxHandle != TunerResourceManager.INVALID_RESOURCE_HANDLE) {
-            demuxHandle[0] = grantingDemuxHandle;
-            updateDemuxClientMappingOnNewGrant(grantingDemuxHandle, request.clientId);
-            return true;
-        }
-
-        // When all the resources are occupied, grant the lowest priority resource if the
-        // request client has higher priority.
-        if (inUseLowestPriorityDrHandle != TunerResourceManager.INVALID_RESOURCE_HANDLE
-            && ((requestClient.getPriority() > currentLowestPriority) || (
-            (requestClient.getPriority() == currentLowestPriority) && isRequestFromSameProcess))) {
-            if (!reclaimResource(
-                    getDemuxResource(inUseLowestPriorityDrHandle).getOwnerClientId(),
-                    TunerResourceManager.TUNER_RESOURCE_TYPE_DEMUX)) {
-                return false;
+            // Grant demux when there is unused resource.
+            if (grantingDemux != null) {
+                updateDemuxClientMappingOnNewGrant(grantingDemux.getHandle(), request.clientId);
+                demuxHandle[0] = grantingDemux.getHandle();
+                return true;
             }
-            demuxHandle[0] = inUseLowestPriorityDrHandle;
-            updateDemuxClientMappingOnNewGrant(
-                    inUseLowestPriorityDrHandle, request.clientId);
-            return true;
+
+            // When all the resources are occupied, grant the lowest priority resource if the
+            // request client has higher priority.
+            if (inUseLowestPriorityDemux != null
+                    && ((requestClient.getPriority() > currentLowestPriority) || (
+                    (requestClient.getPriority() == currentLowestPriority)
+                        && isRequestFromSameProcess))) {
+                demuxHandle[0] = inUseLowestPriorityDemux.getHandle();
+                reclaimOwnerId[0] = inUseLowestPriorityDemux.getOwnerClientId();
+                return true;
+            }
         }
 
         return false;
@@ -1792,7 +1918,9 @@
             return;
         }
 
-        mListeners.put(clientId, record);
+        synchronized (mLock) {
+            mListeners.put(clientId, record);
+        }
     }
 
     @VisibleForTesting
@@ -1808,33 +1936,44 @@
 
         // Reclaim all the resources of the share owners of the frontend that is used by the current
         // resource reclaimed client.
-        ClientProfile profile = getClientProfile(reclaimingClientId);
-        // TODO: check if this check is really needed.
-        if (profile == null) {
-            return true;
-        }
-        Set<Integer> shareFeClientIds = profile.getShareFeClientIds();
-        for (int clientId : shareFeClientIds) {
-            try {
-                mListeners.get(clientId).getListener().onReclaimResources();
-            } catch (RemoteException e) {
-                Slog.e(TAG, "Failed to reclaim resources on client " + clientId, e);
-                return false;
+        Set<Integer> shareFeClientIds;
+        synchronized (mLock) {
+            ClientProfile profile = getClientProfile(reclaimingClientId);
+            if (profile == null) {
+                return true;
             }
-            clearAllResourcesAndClientMapping(getClientProfile(clientId));
+            shareFeClientIds = profile.getShareFeClientIds();
+        }
+        ResourcesReclaimListenerRecord listenerRecord = null;
+        for (int clientId : shareFeClientIds) {
+            synchronized (mLock) {
+                listenerRecord = mListeners.get(clientId);
+            }
+            if (listenerRecord != null) {
+                try {
+                    listenerRecord.getListener().onReclaimResources();
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Failed to reclaim resources on client " + clientId, e);
+                }
+            }
         }
 
         if (DEBUG) {
             Slog.d(TAG, "Reclaiming resources because higher priority client request resource type "
                     + resourceType + ", clientId:" + reclaimingClientId);
         }
-        try {
-            mListeners.get(reclaimingClientId).getListener().onReclaimResources();
-        } catch (RemoteException e) {
-            Slog.e(TAG, "Failed to reclaim resources on client " + reclaimingClientId, e);
-            return false;
+
+        synchronized (mLock) {
+            listenerRecord = mListeners.get(reclaimingClientId);
         }
-        clearAllResourcesAndClientMapping(profile);
+        if (listenerRecord != null) {
+            try {
+                listenerRecord.getListener().onReclaimResources();
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to reclaim resources on client " + reclaimingClientId, e);
+            }
+        }
+
         return true;
     }
 
@@ -2258,6 +2397,7 @@
         addResourcesReclaimListener(clientId, listener);
     }
 
+    @SuppressWarnings("GuardedBy") // Lock is held on mListeners
     private void removeClientProfile(int clientId) {
         for (int shareOwnerId : getClientProfile(clientId).getShareFeClientIds()) {
             clearFrontendAndClientMapping(getClientProfile(shareOwnerId));
@@ -2265,12 +2405,9 @@
         clearAllResourcesAndClientMapping(getClientProfile(clientId));
         mClientProfiles.remove(clientId);
 
-        // it may be called by unregisterClientProfileInternal under test
-        synchronized (mLock) {
-            ResourcesReclaimListenerRecord record = mListeners.remove(clientId);
-            if (record != null) {
-                record.getListener().asBinder().unlinkToDeath(record, 0);
-            }            
+        ResourcesReclaimListenerRecord record = mListeners.remove(clientId);
+        if (record != null) {
+            record.getListener().asBinder().unlinkToDeath(record, 0);
         }
     }
 
@@ -2304,7 +2441,8 @@
         profile.releaseFrontend();
     }
 
-    private void clearAllResourcesAndClientMapping(ClientProfile profile) {
+    @VisibleForTesting
+    protected void clearAllResourcesAndClientMapping(ClientProfile profile) {
         // TODO: check if this check is really needed. Maybe needed for reclaimResource path.
         if (profile == null) {
             return;
diff --git a/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java b/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java
new file mode 100644
index 0000000..b4d3862
--- /dev/null
+++ b/services/core/java/com/android/server/vibrator/ExternalVibrationSession.java
@@ -0,0 +1,146 @@
+/*
+ * 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.vibrator;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.os.ExternalVibration;
+import android.os.ExternalVibrationScale;
+import android.os.IBinder;
+import android.os.VibrationAttributes;
+import android.os.vibrator.Flags;
+
+import com.android.internal.util.FrameworkStatsLog;
+
+/**
+ * A vibration session holding a single {@link ExternalVibration} request.
+ */
+final class ExternalVibrationSession extends Vibration
+        implements VibrationSession, IBinder.DeathRecipient {
+
+    private final ExternalVibration mExternalVibration;
+    private final ExternalVibrationScale mScale = new ExternalVibrationScale();
+
+    @Nullable
+    private Runnable mBinderDeathCallback;
+
+    ExternalVibrationSession(ExternalVibration externalVibration) {
+        super(externalVibration.getToken(), new CallerInfo(
+                externalVibration.getVibrationAttributes(), externalVibration.getUid(),
+                // TODO(b/249785241): Find a way to link ExternalVibration to a VirtualDevice
+                // instead of using DEVICE_ID_INVALID here and relying on the UID checks.
+                Context.DEVICE_ID_INVALID, externalVibration.getPackage(), null));
+        mExternalVibration = externalVibration;
+    }
+
+    public ExternalVibrationScale getScale() {
+        return mScale;
+    }
+
+    @Override
+    public CallerInfo getCallerInfo() {
+        return callerInfo;
+    }
+
+    @Override
+    public VibrationSession.DebugInfo getDebugInfo() {
+        return new Vibration.DebugInfoImpl(getStatus(), stats, /* playedEffect= */ null,
+                /* originalEffect= */ null, mScale.scaleLevel, mScale.adaptiveHapticsScale,
+                callerInfo);
+    }
+
+    @Override
+    public VibrationStats.StatsInfo getStatsInfo(long completionUptimeMillis) {
+        return new VibrationStats.StatsInfo(
+                mExternalVibration.getUid(),
+                FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__EXTERNAL,
+                mExternalVibration.getVibrationAttributes().getUsage(), getStatus(), stats,
+                completionUptimeMillis);
+    }
+
+    @Override
+    public boolean isRepeating() {
+        // We don't currently know if the external vibration is repeating, so we just use a
+        // heuristic based on the usage. Ideally this would be propagated in the ExternalVibration.
+        int usage = mExternalVibration.getVibrationAttributes().getUsage();
+        return usage == VibrationAttributes.USAGE_RINGTONE
+                || usage == VibrationAttributes.USAGE_ALARM;
+    }
+
+    @Override
+    public void linkToDeath(Runnable callback) {
+        synchronized (this) {
+            mBinderDeathCallback = callback;
+        }
+        mExternalVibration.linkToDeath(this);
+    }
+
+    @Override
+    public void unlinkToDeath() {
+        mExternalVibration.unlinkToDeath(this);
+        synchronized (this) {
+            mBinderDeathCallback = null;
+        }
+    }
+
+    @Override
+    public void binderDied() {
+        synchronized (this) {
+            if (mBinderDeathCallback != null) {
+                mBinderDeathCallback.run();
+            }
+        }
+    }
+
+    @Override
+    void end(EndInfo endInfo) {
+        super.end(endInfo);
+        if (stats.hasStarted()) {
+            // External vibration doesn't have feedback from total time the vibrator was playing
+            // with non-zero amplitude, so we use the duration between start and end times of
+            // the vibration as the time the vibrator was ON, since the haptic channels are
+            // open for this duration and can receive vibration waveform data.
+            stats.reportVibratorOn(stats.getEndUptimeMillis() - stats.getStartUptimeMillis());
+        }
+    }
+
+    @Override
+    public void notifyEnded() {
+        // Notify external client that this vibration should stop sending data to the vibrator.
+        mExternalVibration.mute();
+    }
+
+    boolean isHoldingSameVibration(ExternalVibration vib) {
+        return mExternalVibration.equals(vib);
+    }
+
+    void muteScale() {
+        mScale.scaleLevel = ExternalVibrationScale.ScaleLevel.SCALE_MUTE;
+        if (Flags.hapticsScaleV2Enabled()) {
+            mScale.scaleFactor = 0;
+        }
+    }
+
+    void scale(VibrationScaler scaler, int usage) {
+        mScale.scaleLevel = scaler.getScaleLevel(usage);
+        if (Flags.hapticsScaleV2Enabled()) {
+            mScale.scaleFactor = scaler.getScaleFactor(usage);
+        }
+        mScale.adaptiveHapticsScale = scaler.getAdaptiveHapticsScale(usage);
+        stats.reportAdaptiveScale(mScale.adaptiveHapticsScale);
+    }
+}
diff --git a/services/core/java/com/android/server/vibrator/HalVibration.java b/services/core/java/com/android/server/vibrator/HalVibration.java
index fe0cf59..ea4bd01 100644
--- a/services/core/java/com/android/server/vibrator/HalVibration.java
+++ b/services/core/java/com/android/server/vibrator/HalVibration.java
@@ -51,37 +51,22 @@
     @NonNull
     private volatile CombinedVibration mEffectToPlay;
 
-    /** Vibration status. */
-    private Vibration.Status mStatus;
-
     /** Reported scale values applied to the vibration effects. */
     private int mScaleLevel;
     private float mAdaptiveScale;
 
     HalVibration(@NonNull IBinder token, @NonNull CombinedVibration effect,
-            @NonNull CallerInfo callerInfo) {
+            @NonNull VibrationSession.CallerInfo callerInfo) {
         super(token, callerInfo);
         mOriginalEffect = effect;
         mEffectToPlay = effect;
-        mStatus = Vibration.Status.RUNNING;
         mScaleLevel = VibrationScaler.SCALE_NONE;
         mAdaptiveScale = VibrationScaler.ADAPTIVE_SCALE_NONE;
     }
 
-    /**
-     * Set the {@link Status} of this vibration and reports the current system time as this
-     * vibration end time, for debugging purposes.
-     *
-     * <p>This method will only accept given value if the current status is {@link
-     * Status#RUNNING}.
-     */
-    public void end(EndInfo info) {
-        if (hasEnded()) {
-            // Vibration already ended, keep first ending status set and ignore this one.
-            return;
-        }
-        mStatus = info.status;
-        stats.reportEnded(info.endedBy);
+    @Override
+    public void end(EndInfo endInfo) {
+        super.end(endInfo);
         mCompletionLatch.countDown();
     }
 
@@ -144,11 +129,6 @@
         // No need to update fallback effects, they are already configured per device.
     }
 
-    /** Return true is current status is different from {@link Status#RUNNING}. */
-    public boolean hasEnded() {
-        return mStatus != Status.RUNNING;
-    }
-
     @Override
     public boolean isRepeating() {
         return mOriginalEffect.getDuration() == Long.MAX_VALUE;
@@ -159,16 +139,16 @@
         return mEffectToPlay;
     }
 
-    /** Return {@link Vibration.DebugInfo} with read-only debug information about this vibration. */
-    public Vibration.DebugInfo getDebugInfo() {
+    @Override
+    public VibrationSession.DebugInfo getDebugInfo() {
         // Clear the original effect if it's the same as the effect that was played, for simplicity
         CombinedVibration originalEffect =
                 Objects.equals(mOriginalEffect, mEffectToPlay) ? null : mOriginalEffect;
-        return new Vibration.DebugInfo(mStatus, stats, mEffectToPlay, originalEffect,
+        return new Vibration.DebugInfoImpl(getStatus(), stats, mEffectToPlay, originalEffect,
                 mScaleLevel, mAdaptiveScale, callerInfo);
     }
 
-    /** Return {@link VibrationStats.StatsInfo} with read-only metrics about this vibration. */
+    @Override
     public VibrationStats.StatsInfo getStatsInfo(long completionUptimeMillis) {
         int vibrationType = mEffectToPlay.hasVendorEffects()
                 ? FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__VENDOR
@@ -176,7 +156,7 @@
                         ? FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__REPEATED
                         : FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__SINGLE;
         return new VibrationStats.StatsInfo(
-                callerInfo.uid, vibrationType, callerInfo.attrs.getUsage(), mStatus,
+                callerInfo.uid, vibrationType, callerInfo.attrs.getUsage(), getStatus(),
                 stats, completionUptimeMillis);
     }
 
diff --git a/services/core/java/com/android/server/vibrator/InputDeviceDelegate.java b/services/core/java/com/android/server/vibrator/InputDeviceDelegate.java
index 4e58b9a..83e05f4 100644
--- a/services/core/java/com/android/server/vibrator/InputDeviceDelegate.java
+++ b/services/core/java/com/android/server/vibrator/InputDeviceDelegate.java
@@ -26,6 +26,7 @@
 import android.view.InputDevice;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.vibrator.VibrationSession.CallerInfo;
 
 /** Delegates vibrations to all connected {@link InputDevice} with one or more vibrators. */
 final class InputDeviceDelegate implements InputManager.InputDeviceListener {
@@ -93,7 +94,7 @@
      *
      * @return {@link #isAvailable()}
      */
-    public boolean vibrateIfAvailable(Vibration.CallerInfo callerInfo, CombinedVibration effect) {
+    public boolean vibrateIfAvailable(CallerInfo callerInfo, CombinedVibration effect) {
         synchronized (mLock) {
             for (int i = 0; i < mInputDeviceVibrators.size(); i++) {
                 mInputDeviceVibrators.valueAt(i).vibrate(callerInfo.uid, callerInfo.opPkg, effect,
diff --git a/services/core/java/com/android/server/vibrator/Vibration.java b/services/core/java/com/android/server/vibrator/Vibration.java
index aa4b9f3..9a04793 100644
--- a/services/core/java/com/android/server/vibrator/Vibration.java
+++ b/services/core/java/com/android/server/vibrator/Vibration.java
@@ -31,7 +31,6 @@
 import android.util.IndentingPrintWriter;
 import android.util.proto.ProtoOutputStream;
 
-import java.io.PrintWriter;
 import java.time.Instant;
 import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
@@ -52,131 +51,69 @@
     private static final AtomicInteger sNextVibrationId = new AtomicInteger(1); // 0 = no callback
 
     public final long id;
-    public final CallerInfo callerInfo;
+    public final VibrationSession.CallerInfo callerInfo;
     public final VibrationStats stats = new VibrationStats();
     public final IBinder callerToken;
 
-    /** Vibration status with reference to values from vibratormanagerservice.proto for logging. */
-    enum Status {
-        UNKNOWN(VibrationProto.UNKNOWN),
-        RUNNING(VibrationProto.RUNNING),
-        FINISHED(VibrationProto.FINISHED),
-        FINISHED_UNEXPECTED(VibrationProto.FINISHED_UNEXPECTED),
-        FORWARDED_TO_INPUT_DEVICES(VibrationProto.FORWARDED_TO_INPUT_DEVICES),
-        CANCELLED_BINDER_DIED(VibrationProto.CANCELLED_BINDER_DIED),
-        CANCELLED_BY_SCREEN_OFF(VibrationProto.CANCELLED_BY_SCREEN_OFF),
-        CANCELLED_BY_SETTINGS_UPDATE(VibrationProto.CANCELLED_BY_SETTINGS_UPDATE),
-        CANCELLED_BY_USER(VibrationProto.CANCELLED_BY_USER),
-        CANCELLED_BY_FOREGROUND_USER(VibrationProto.CANCELLED_BY_FOREGROUND_USER),
-        CANCELLED_BY_UNKNOWN_REASON(VibrationProto.CANCELLED_BY_UNKNOWN_REASON),
-        CANCELLED_SUPERSEDED(VibrationProto.CANCELLED_SUPERSEDED),
-        CANCELLED_BY_APP_OPS(VibrationProto.CANCELLED_BY_APP_OPS),
-        IGNORED_ERROR_APP_OPS(VibrationProto.IGNORED_ERROR_APP_OPS),
-        IGNORED_ERROR_CANCELLING(VibrationProto.IGNORED_ERROR_CANCELLING),
-        IGNORED_ERROR_SCHEDULING(VibrationProto.IGNORED_ERROR_SCHEDULING),
-        IGNORED_ERROR_TOKEN(VibrationProto.IGNORED_ERROR_TOKEN),
-        IGNORED_APP_OPS(VibrationProto.IGNORED_APP_OPS),
-        IGNORED_BACKGROUND(VibrationProto.IGNORED_BACKGROUND),
-        IGNORED_MISSING_PERMISSION(VibrationProto.IGNORED_MISSING_PERMISSION),
-        IGNORED_UNSUPPORTED(VibrationProto.IGNORED_UNSUPPORTED),
-        IGNORED_FOR_EXTERNAL(VibrationProto.IGNORED_FOR_EXTERNAL),
-        IGNORED_FOR_HIGHER_IMPORTANCE(VibrationProto.IGNORED_FOR_HIGHER_IMPORTANCE),
-        IGNORED_FOR_ONGOING(VibrationProto.IGNORED_FOR_ONGOING),
-        IGNORED_FOR_POWER(VibrationProto.IGNORED_FOR_POWER),
-        IGNORED_FOR_RINGER_MODE(VibrationProto.IGNORED_FOR_RINGER_MODE),
-        IGNORED_FOR_SETTINGS(VibrationProto.IGNORED_FOR_SETTINGS),
-        IGNORED_SUPERSEDED(VibrationProto.IGNORED_SUPERSEDED),
-        IGNORED_FROM_VIRTUAL_DEVICE(VibrationProto.IGNORED_FROM_VIRTUAL_DEVICE),
-        IGNORED_ON_WIRELESS_CHARGER(VibrationProto.IGNORED_ON_WIRELESS_CHARGER);
+    private VibrationSession.Status mStatus;
 
-        private final int mProtoEnumValue;
-
-        Status(int value) {
-            mProtoEnumValue = value;
-        }
-
-        public int getProtoEnumValue() {
-            return mProtoEnumValue;
-        }
-    }
-
-    Vibration(@NonNull IBinder token, @NonNull CallerInfo callerInfo) {
+    Vibration(@NonNull IBinder token, @NonNull VibrationSession.CallerInfo callerInfo) {
         Objects.requireNonNull(token);
         Objects.requireNonNull(callerInfo);
+        mStatus = VibrationSession.Status.RUNNING;
         this.id = sNextVibrationId.getAndIncrement();
         this.callerToken = token;
         this.callerInfo = callerInfo;
     }
 
+    VibrationSession.Status getStatus() {
+        return mStatus;
+    }
+
+    /** Return true is current status is different from {@link VibrationSession.Status#RUNNING}. */
+    boolean hasEnded() {
+        return mStatus != VibrationSession.Status.RUNNING;
+    }
+
+    /**
+     * Set the {@link VibrationSession} of this vibration and reports the current system time as
+     * this vibration end time, for debugging purposes.
+     *
+     * <p>This method will only accept given value if the current status is {@link
+     * VibrationSession.Status#RUNNING}.
+     */
+    void end(Vibration.EndInfo endInfo) {
+        if (hasEnded()) {
+            // Vibration already ended, keep first ending status set and ignore this one.
+            return;
+        }
+        mStatus = endInfo.status;
+        stats.reportEnded(endInfo.endedBy);
+    }
+
     /** Return true if vibration is a repeating vibration. */
     abstract boolean isRepeating();
 
-    /**
-     * Holds lightweight immutable info on the process that triggered the vibration. This data
-     * could potentially be kept in memory for a long time for bugreport dumpsys operations.
-     *
-     * Since CallerInfo can be kept in memory for a long time, it shouldn't hold any references to
-     * potentially expensive or resource-linked objects, such as {@link IBinder}.
-     */
-    static final class CallerInfo {
-        public final VibrationAttributes attrs;
-        public final int uid;
-        public final int deviceId;
-        public final String opPkg;
-        public final String reason;
+    /** Return {@link VibrationSession.DebugInfo} with read-only debug data about this vibration. */
+    abstract VibrationSession.DebugInfo getDebugInfo();
 
-        CallerInfo(@NonNull VibrationAttributes attrs, int uid, int deviceId, String opPkg,
-                String reason) {
-            Objects.requireNonNull(attrs);
-            this.attrs = attrs;
-            this.uid = uid;
-            this.deviceId = deviceId;
-            this.opPkg = opPkg;
-            this.reason = reason;
-        }
-
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) return true;
-            if (!(o instanceof CallerInfo)) return false;
-            CallerInfo that = (CallerInfo) o;
-            return Objects.equals(attrs, that.attrs)
-                    && uid == that.uid
-                    && deviceId == that.deviceId
-                    && Objects.equals(opPkg, that.opPkg)
-                    && Objects.equals(reason, that.reason);
-        }
-
-        @Override
-        public int hashCode() {
-            return Objects.hash(attrs, uid, deviceId, opPkg, reason);
-        }
-
-        @Override
-        public String toString() {
-            return "CallerInfo{"
-                    + " uid=" + uid
-                    + ", opPkg=" + opPkg
-                    + ", deviceId=" + deviceId
-                    + ", attrs=" + attrs
-                    + ", reason=" + reason
-                    + '}';
-        }
-    }
+    /** Return {@link VibrationStats.StatsInfo} with read-only metrics about this vibration. */
+    abstract VibrationStats.StatsInfo getStatsInfo(long completionUptimeMillis);
 
     /** Immutable info passed as a signal to end a vibration. */
     static final class EndInfo {
-        /** The {@link Status} to be set to the vibration when it ends with this info. */
+        /** The vibration status to be set when it ends with this info. */
         @NonNull
-        public final Status status;
+        public final VibrationSession.Status status;
         /** Info about the process that ended the vibration. */
-        public final CallerInfo endedBy;
+        public final VibrationSession.CallerInfo endedBy;
 
-        EndInfo(@NonNull Vibration.Status status) {
+        EndInfo(@NonNull VibrationSession.Status status) {
             this(status, null);
         }
 
-        EndInfo(@NonNull Vibration.Status status, @Nullable CallerInfo endedBy) {
+        EndInfo(@NonNull VibrationSession.Status status,
+                @Nullable VibrationSession.CallerInfo endedBy) {
             this.status = status;
             this.endedBy = endedBy;
         }
@@ -211,10 +148,10 @@
      * Since DebugInfo can be kept in memory for a long time, it shouldn't hold any references to
      * potentially expensive or resource-linked objects, such as {@link IBinder}.
      */
-    static final class DebugInfo {
-        final Status mStatus;
+    static final class DebugInfoImpl implements VibrationSession.DebugInfo {
+        final VibrationSession.Status mStatus;
         final long mCreateTime;
-        final CallerInfo mCallerInfo;
+        final VibrationSession.CallerInfo mCallerInfo;
         @Nullable
         final CombinedVibration mPlayedEffect;
 
@@ -226,9 +163,10 @@
         private final int mScaleLevel;
         private final float mAdaptiveScale;
 
-        DebugInfo(Status status, VibrationStats stats, @Nullable CombinedVibration playedEffect,
+        DebugInfoImpl(VibrationSession.Status status, VibrationStats stats,
+                @Nullable CombinedVibration playedEffect,
                 @Nullable CombinedVibration originalEffect, int scaleLevel,
-                float adaptiveScale, @NonNull CallerInfo callerInfo) {
+                float adaptiveScale, @NonNull VibrationSession.CallerInfo callerInfo) {
             Objects.requireNonNull(callerInfo);
             mCreateTime = stats.getCreateTimeDebug();
             mStartTime = stats.getStartTimeDebug();
@@ -243,6 +181,27 @@
         }
 
         @Override
+        public VibrationSession.Status getStatus() {
+            return mStatus;
+        }
+
+        @Override
+        public long getCreateUptimeMillis() {
+            return mCreateTime;
+        }
+
+        @Override
+        public VibrationSession.CallerInfo getCallerInfo() {
+            return mCallerInfo;
+        }
+
+        @Nullable
+        @Override
+        public Object getDumpAggregationKey() {
+            return mPlayedEffect;
+        }
+
+        @Override
         public String toString() {
             return "createTime: " + formatTime(mCreateTime, /*includeDate=*/ true)
                     + ", startTime: " + formatTime(mStartTime, /*includeDate=*/ true)
@@ -257,17 +216,13 @@
                     + ", callerInfo: " + mCallerInfo;
         }
 
-        void logMetrics(VibratorFrameworkStatsLogger statsLogger) {
+        @Override
+        public void logMetrics(VibratorFrameworkStatsLogger statsLogger) {
             statsLogger.logVibrationAdaptiveHapticScale(mCallerInfo.uid, mAdaptiveScale);
         }
 
-        /**
-         * Write this info in a compact way into given {@link PrintWriter}.
-         *
-         * <p>This is used by dumpsys to log multiple vibration records in single lines that are
-         * easy to skim through by the sorted created time.
-         */
-        void dumpCompact(IndentingPrintWriter pw) {
+        @Override
+        public void dumpCompact(IndentingPrintWriter pw) {
             boolean isExternalVibration = mPlayedEffect == null;
             String timingsStr = String.format(Locale.ROOT,
                     "%s | %8s | %20s | duration: %5dms | start: %12s | end: %12s",
@@ -299,8 +254,8 @@
             pw.println(timingsStr + paramStr + audioUsageStr + callerStr + effectStr);
         }
 
-        /** Write this info into given {@link PrintWriter}. */
-        void dump(IndentingPrintWriter pw) {
+        @Override
+        public void dump(IndentingPrintWriter pw) {
             pw.println("Vibration:");
             pw.increaseIndent();
             pw.println("status = " + mStatus.name().toLowerCase(Locale.ROOT));
@@ -317,8 +272,8 @@
             pw.decreaseIndent();
         }
 
-        /** Write this info into given {@code fieldId} on {@link ProtoOutputStream}. */
-        void dump(ProtoOutputStream proto, long fieldId) {
+        @Override
+        public void dump(ProtoOutputStream proto, long fieldId) {
             final long token = proto.start(fieldId);
             proto.write(VibrationProto.START_TIME, mStartTime);
             proto.write(VibrationProto.END_TIME, mEndTime);
diff --git a/services/core/java/com/android/server/vibrator/VibrationSession.java b/services/core/java/com/android/server/vibrator/VibrationSession.java
new file mode 100644
index 0000000..5640b49
--- /dev/null
+++ b/services/core/java/com/android/server/vibrator/VibrationSession.java
@@ -0,0 +1,208 @@
+/*
+ * 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.vibrator;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.CombinedVibration;
+import android.os.IBinder;
+import android.os.VibrationAttributes;
+import android.util.IndentingPrintWriter;
+import android.util.proto.ProtoOutputStream;
+
+import java.io.PrintWriter;
+import java.util.Objects;
+
+/**
+ * Represents a generic vibration session that plays one or more vibration requests.
+ *
+ * <p>This might represent:
+ *
+ * <ol>
+ *     <li>A single {@link CombinedVibration} playback.
+ *     <li>An {@link android.os.ExternalVibration} playback.
+ * </ol>
+ */
+interface VibrationSession {
+
+    /** Returns data about the client app that triggered this vibration session. */
+    CallerInfo getCallerInfo();
+
+    /** Returns debug data for logging and metric reports. */
+    DebugInfo getDebugInfo();
+
+    /**
+     * Links this session to the app process death with given callback to handle it.
+     *
+     * <p>This can be used by the service to end the vibration session when the app process dies.
+     */
+    void linkToDeath(Runnable callback);
+
+    /** Removes link to the app process death. */
+    void unlinkToDeath();
+
+    /** Notify the session end was requested, which might be acted upon asynchronously. */
+    void notifyEnded();
+
+    /**
+     * Session status with reference to values from vibratormanagerservice.proto for logging.
+     */
+    enum Status {
+        UNKNOWN(VibrationProto.UNKNOWN),
+        RUNNING(VibrationProto.RUNNING),
+        FINISHED(VibrationProto.FINISHED),
+        FINISHED_UNEXPECTED(VibrationProto.FINISHED_UNEXPECTED),
+        FORWARDED_TO_INPUT_DEVICES(VibrationProto.FORWARDED_TO_INPUT_DEVICES),
+        CANCELLED_BINDER_DIED(VibrationProto.CANCELLED_BINDER_DIED),
+        CANCELLED_BY_SCREEN_OFF(VibrationProto.CANCELLED_BY_SCREEN_OFF),
+        CANCELLED_BY_SETTINGS_UPDATE(VibrationProto.CANCELLED_BY_SETTINGS_UPDATE),
+        CANCELLED_BY_USER(VibrationProto.CANCELLED_BY_USER),
+        CANCELLED_BY_FOREGROUND_USER(VibrationProto.CANCELLED_BY_FOREGROUND_USER),
+        CANCELLED_BY_UNKNOWN_REASON(VibrationProto.CANCELLED_BY_UNKNOWN_REASON),
+        CANCELLED_SUPERSEDED(VibrationProto.CANCELLED_SUPERSEDED),
+        CANCELLED_BY_APP_OPS(VibrationProto.CANCELLED_BY_APP_OPS),
+        IGNORED_ERROR_APP_OPS(VibrationProto.IGNORED_ERROR_APP_OPS),
+        IGNORED_ERROR_CANCELLING(VibrationProto.IGNORED_ERROR_CANCELLING),
+        IGNORED_ERROR_SCHEDULING(VibrationProto.IGNORED_ERROR_SCHEDULING),
+        IGNORED_ERROR_TOKEN(VibrationProto.IGNORED_ERROR_TOKEN),
+        IGNORED_APP_OPS(VibrationProto.IGNORED_APP_OPS),
+        IGNORED_BACKGROUND(VibrationProto.IGNORED_BACKGROUND),
+        IGNORED_MISSING_PERMISSION(VibrationProto.IGNORED_MISSING_PERMISSION),
+        IGNORED_UNSUPPORTED(VibrationProto.IGNORED_UNSUPPORTED),
+        IGNORED_FOR_EXTERNAL(VibrationProto.IGNORED_FOR_EXTERNAL),
+        IGNORED_FOR_HIGHER_IMPORTANCE(VibrationProto.IGNORED_FOR_HIGHER_IMPORTANCE),
+        IGNORED_FOR_ONGOING(VibrationProto.IGNORED_FOR_ONGOING),
+        IGNORED_FOR_POWER(VibrationProto.IGNORED_FOR_POWER),
+        IGNORED_FOR_RINGER_MODE(VibrationProto.IGNORED_FOR_RINGER_MODE),
+        IGNORED_FOR_SETTINGS(VibrationProto.IGNORED_FOR_SETTINGS),
+        IGNORED_SUPERSEDED(VibrationProto.IGNORED_SUPERSEDED),
+        IGNORED_FROM_VIRTUAL_DEVICE(VibrationProto.IGNORED_FROM_VIRTUAL_DEVICE),
+        IGNORED_ON_WIRELESS_CHARGER(VibrationProto.IGNORED_ON_WIRELESS_CHARGER);
+
+        private final int mProtoEnumValue;
+
+        Status(int value) {
+            mProtoEnumValue = value;
+        }
+
+        public int getProtoEnumValue() {
+            return mProtoEnumValue;
+        }
+    }
+
+    /**
+     * Holds lightweight immutable info on the process that triggered the vibration session.
+     *
+     * <p>This data could potentially be kept in memory for a long time for bugreport dumpsys
+     * operations. It shouldn't hold any references to potentially expensive or resource-linked
+     * objects, such as {@link IBinder}.
+     */
+    final class CallerInfo {
+        public final VibrationAttributes attrs;
+        public final int uid;
+        public final int deviceId;
+        public final String opPkg;
+        public final String reason;
+
+        CallerInfo(@NonNull VibrationAttributes attrs, int uid, int deviceId, String opPkg,
+                String reason) {
+            Objects.requireNonNull(attrs);
+            this.attrs = attrs;
+            this.uid = uid;
+            this.deviceId = deviceId;
+            this.opPkg = opPkg;
+            this.reason = reason;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof CallerInfo)) return false;
+            CallerInfo that = (CallerInfo) o;
+            return Objects.equals(attrs, that.attrs)
+                    && uid == that.uid
+                    && deviceId == that.deviceId
+                    && Objects.equals(opPkg, that.opPkg)
+                    && Objects.equals(reason, that.reason);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(attrs, uid, deviceId, opPkg, reason);
+        }
+
+        @Override
+        public String toString() {
+            return "CallerInfo{"
+                    + " uid=" + uid
+                    + ", opPkg=" + opPkg
+                    + ", deviceId=" + deviceId
+                    + ", attrs=" + attrs
+                    + ", reason=" + reason
+                    + '}';
+        }
+    }
+
+    /**
+     * Interface for lightweight debug information about the vibration session for debugging.
+     *
+     * <p>This data could potentially be kept in memory for a long time for bugreport dumpsys
+     * operations. It shouldn't hold any references to potentially expensive or resource-linked
+     * objects, such as {@link IBinder}.
+     */
+    interface DebugInfo {
+
+        /** Return the vibration session status. */
+        Status getStatus();
+
+        /** Returns the session creation time from {@link android.os.SystemClock#uptimeMillis()}. */
+        long getCreateUptimeMillis();
+
+        /** Returns information about the process that created the session. */
+        CallerInfo getCallerInfo();
+
+        /**
+         * Returns the aggregation key for log records.
+         *
+         * <p>This is used to aggregate similar vibration sessions triggered in quick succession
+         * (e.g. multiple keyboard vibrations when the user is typing).
+         *
+         * <p>This does not need to include data from {@link CallerInfo} or {@link Status}.
+         *
+         * @see GroupedAggregatedLogRecords
+         */
+        @Nullable
+        Object getDumpAggregationKey();
+
+        /** Logs vibration session fields for metric reports. */
+        void logMetrics(VibratorFrameworkStatsLogger statsLogger);
+
+        /** Write this info into given {@code fieldId} on {@link ProtoOutputStream}. */
+        void dump(ProtoOutputStream proto, long fieldId);
+
+        /** Write this info into given {@link PrintWriter}. */
+        void dump(IndentingPrintWriter pw);
+
+        /**
+         * Write this info in a compact way into given {@link PrintWriter}.
+         *
+         * <p>This is used by dumpsys to log multiple records in single lines that are easy to skim
+         * through by the sorted created time.
+         */
+        void dumpCompact(IndentingPrintWriter pw);
+    }
+}
diff --git a/services/core/java/com/android/server/vibrator/VibrationSettings.java b/services/core/java/com/android/server/vibrator/VibrationSettings.java
index 0d6778c..69cdcf4 100644
--- a/services/core/java/com/android/server/vibrator/VibrationSettings.java
+++ b/services/core/java/com/android/server/vibrator/VibrationSettings.java
@@ -66,6 +66,8 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.LocalServices;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
+import com.android.server.vibrator.VibrationSession.CallerInfo;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -416,46 +418,46 @@
     /**
      * Check if given vibration should be ignored by the service.
      *
-     * @return One of Vibration.Status.IGNORED_* values if the vibration should be ignored,
+     * @return One of VibrationSession.Status.IGNORED_* values if the vibration should be ignored,
      * null otherwise.
      */
     @Nullable
-    public Vibration.Status shouldIgnoreVibration(@NonNull Vibration.CallerInfo callerInfo) {
+    public Status shouldIgnoreVibration(@NonNull CallerInfo callerInfo) {
         final int usage = callerInfo.attrs.getUsage();
         synchronized (mLock) {
             if (!mUidObserver.isUidForeground(callerInfo.uid)
                     && !BACKGROUND_PROCESS_USAGE_ALLOWLIST.contains(usage)) {
-                return Vibration.Status.IGNORED_BACKGROUND;
+                return Status.IGNORED_BACKGROUND;
             }
 
             if (callerInfo.deviceId != Context.DEVICE_ID_DEFAULT
                     && callerInfo.deviceId != Context.DEVICE_ID_INVALID) {
-                return Vibration.Status.IGNORED_FROM_VIRTUAL_DEVICE;
+                return Status.IGNORED_FROM_VIRTUAL_DEVICE;
             }
 
             if (callerInfo.deviceId == Context.DEVICE_ID_INVALID
                     && isAppRunningOnAnyVirtualDevice(callerInfo.uid)) {
-                return Vibration.Status.IGNORED_FROM_VIRTUAL_DEVICE;
+                return Status.IGNORED_FROM_VIRTUAL_DEVICE;
             }
 
             if (mBatterySaverMode && !BATTERY_SAVER_USAGE_ALLOWLIST.contains(usage)) {
-                return Vibration.Status.IGNORED_FOR_POWER;
+                return Status.IGNORED_FOR_POWER;
             }
 
             if (!callerInfo.attrs.isFlagSet(
                     VibrationAttributes.FLAG_BYPASS_USER_VIBRATION_INTENSITY_OFF)
                     && !shouldVibrateForUserSetting(callerInfo)) {
-                return Vibration.Status.IGNORED_FOR_SETTINGS;
+                return Status.IGNORED_FOR_SETTINGS;
             }
 
             if (!callerInfo.attrs.isFlagSet(VibrationAttributes.FLAG_BYPASS_INTERRUPTION_POLICY)) {
                 if (!shouldVibrateForRingerModeLocked(usage)) {
-                    return Vibration.Status.IGNORED_FOR_RINGER_MODE;
+                    return Status.IGNORED_FOR_RINGER_MODE;
                 }
             }
 
             if (mVibrationConfig.ignoreVibrationsOnWirelessCharger() && mOnWirelessCharger) {
-                return Vibration.Status.IGNORED_ON_WIRELESS_CHARGER;
+                return Status.IGNORED_ON_WIRELESS_CHARGER;
             }
         }
         return null;
@@ -471,7 +473,7 @@
      *
      * @return true if the vibration should be cancelled when the screen goes off, false otherwise.
      */
-    public boolean shouldCancelVibrationOnScreenOff(@NonNull Vibration.CallerInfo callerInfo,
+    public boolean shouldCancelVibrationOnScreenOff(@NonNull CallerInfo callerInfo,
             long vibrationStartUptimeMillis) {
         PowerManagerInternal pm;
         synchronized (mLock) {
@@ -483,8 +485,8 @@
             // ignored here and not cancel a vibration, and those are usually triggered by timeout
             // or inactivity, so it's unlikely that it will override a more active goToSleep reason.
             PowerManager.SleepData sleepData = pm.getLastGoToSleep();
-            if ((sleepData.goToSleepUptimeMillis < vibrationStartUptimeMillis)
-                    || POWER_MANAGER_SLEEP_REASON_ALLOWLIST.contains(sleepData.goToSleepReason)) {
+            if (sleepData != null && (sleepData.goToSleepUptimeMillis < vibrationStartUptimeMillis
+                    || POWER_MANAGER_SLEEP_REASON_ALLOWLIST.contains(sleepData.goToSleepReason))) {
                 // Ignore screen off events triggered before the vibration started, and all
                 // automatic "go to sleep" events from allowlist.
                 Slog.d(TAG, "Ignoring screen off event triggered at uptime "
@@ -522,7 +524,7 @@
      * {@code false} to ignore the vibration.
      */
     @GuardedBy("mLock")
-    private boolean shouldVibrateForUserSetting(Vibration.CallerInfo callerInfo) {
+    private boolean shouldVibrateForUserSetting(CallerInfo callerInfo) {
         final int usage = callerInfo.attrs.getUsage();
         if (!mVibrateOn && (VIBRATE_ON_DISABLED_USAGE_ALLOWED != usage)) {
             // Main setting disabled.
diff --git a/services/core/java/com/android/server/vibrator/VibrationStats.java b/services/core/java/com/android/server/vibrator/VibrationStats.java
index 8179d6a..fc0c6e7 100644
--- a/services/core/java/com/android/server/vibrator/VibrationStats.java
+++ b/services/core/java/com/android/server/vibrator/VibrationStats.java
@@ -166,7 +166,7 @@
      * @return true if the status was accepted. This method will only accept given values if
      * the end timestamp was never set.
      */
-    boolean reportEnded(@Nullable Vibration.CallerInfo endedBy) {
+    boolean reportEnded(@Nullable VibrationSession.CallerInfo endedBy) {
         if (hasEnded()) {
             // Vibration already ended, keep first ending stats set and ignore this one.
             return false;
@@ -187,7 +187,7 @@
      * <p>This method will only accept the first value as the one that was interrupted by this
      * vibration, and will ignore all successive calls.
      */
-    void reportInterruptedAnotherVibration(@NonNull Vibration.CallerInfo callerInfo) {
+    void reportInterruptedAnotherVibration(@NonNull VibrationSession.CallerInfo callerInfo) {
         if (mInterruptedUsage < 0) {
             mInterruptedUsage = callerInfo.attrs.getUsage();
         }
@@ -330,7 +330,7 @@
         public final int[] halUnsupportedEffectsUsed;
         private boolean mIsWritten;
 
-        StatsInfo(int uid, int vibrationType, int usage, Vibration.Status status,
+        StatsInfo(int uid, int vibrationType, int usage, VibrationSession.Status status,
                 VibrationStats stats, long completionUptimeMillis) {
             this.uid = uid;
             this.vibrationType = vibrationType;
diff --git a/services/core/java/com/android/server/vibrator/VibrationStepConductor.java b/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
index 7152844..5137d19 100644
--- a/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
+++ b/services/core/java/com/android/server/vibrator/VibrationStepConductor.java
@@ -32,6 +32,7 @@
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import java.util.ArrayList;
 import java.util.Iterator;
@@ -217,7 +218,7 @@
     }
 
     /**
-     * Calculate the {@link Vibration.Status} based on the current queue state and the expected
+     * Calculate the {@link Vibration.EndInfo} based on the current queue state and the expected
      * number of {@link StartSequentialEffectStep} to be played.
      */
     @Nullable
@@ -235,10 +236,10 @@
         }
         // No pending steps, and something happened.
         if (mSuccessfulVibratorOnSteps > 0) {
-            return new Vibration.EndInfo(Vibration.Status.FINISHED);
+            return new Vibration.EndInfo(Status.FINISHED);
         }
         // If no step was able to turn the vibrator ON successfully.
-        return new Vibration.EndInfo(Vibration.Status.IGNORED_UNSUPPORTED);
+        return new Vibration.EndInfo(Status.IGNORED_UNSUPPORTED);
     }
 
     /**
@@ -352,7 +353,7 @@
         if (DEBUG) {
             Slog.d(TAG, "Binder died, cancelling vibration...");
         }
-        notifyCancelled(new Vibration.EndInfo(Vibration.Status.CANCELLED_BINDER_DIED),
+        notifyCancelled(new Vibration.EndInfo(Status.CANCELLED_BINDER_DIED),
                 /* immediate= */ false);
     }
 
@@ -377,7 +378,7 @@
         if ((cancelInfo == null) || !cancelInfo.status.name().startsWith("CANCEL")) {
             Slog.w(TAG, "Vibration cancel requested with bad signal=" + cancelInfo
                     + ", using CANCELLED_UNKNOWN_REASON to ensure cancellation.");
-            cancelInfo = new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_UNKNOWN_REASON);
+            cancelInfo = new Vibration.EndInfo(Status.CANCELLED_BY_UNKNOWN_REASON);
         }
         synchronized (mLock) {
             if ((immediate && mSignalCancelImmediate) || (mSignalCancel != null)) {
diff --git a/services/core/java/com/android/server/vibrator/VibrationThread.java b/services/core/java/com/android/server/vibrator/VibrationThread.java
index cfb4c74..ab4a4d8 100644
--- a/services/core/java/com/android/server/vibrator/VibrationThread.java
+++ b/services/core/java/com/android/server/vibrator/VibrationThread.java
@@ -29,6 +29,7 @@
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import java.util.NoSuchElementException;
 import java.util.Objects;
@@ -240,7 +241,7 @@
                 runCurrentVibrationWithWakeLockAndDeathLink();
             } finally {
                 clientVibrationCompleteIfNotAlready(
-                        new Vibration.EndInfo(Vibration.Status.FINISHED_UNEXPECTED));
+                        new Vibration.EndInfo(Status.FINISHED_UNEXPECTED));
             }
         } finally {
             mWakeLock.release();
@@ -259,7 +260,7 @@
         } catch (RemoteException e) {
             Slog.e(TAG, "Error linking vibration to token death", e);
             clientVibrationCompleteIfNotAlready(
-                    new Vibration.EndInfo(Vibration.Status.IGNORED_ERROR_TOKEN));
+                    new Vibration.EndInfo(Status.IGNORED_ERROR_TOKEN));
             return;
         }
         // Ensure that the unlink always occurs now.
diff --git a/services/core/java/com/android/server/vibrator/VibratorControlService.java b/services/core/java/com/android/server/vibrator/VibratorControlService.java
index de5e662..3a814cd 100644
--- a/services/core/java/com/android/server/vibrator/VibratorControlService.java
+++ b/services/core/java/com/android/server/vibrator/VibratorControlService.java
@@ -559,8 +559,8 @@
     }
 
     /**
-     * Record for a single {@link Vibration.DebugInfo}, that can be grouped by usage and aggregated
-     * by UID, {@link VibrationAttributes} and {@link VibrationEffect}.
+     * Record for a single {@link VibrationSession.DebugInfo}, that can be grouped by usage and
+     * aggregated by UID, {@link VibrationAttributes} and {@link VibrationEffect}.
      */
     private static final class VibrationScaleParamRecord
             implements GroupedAggregatedLogRecords.SingleLogRecord {
diff --git a/services/core/java/com/android/server/vibrator/VibratorManagerService.java b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
index c143beb..799934a 100644
--- a/services/core/java/com/android/server/vibrator/VibratorManagerService.java
+++ b/services/core/java/com/android/server/vibrator/VibratorManagerService.java
@@ -73,9 +73,11 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.app.IBatteryStats;
 import com.android.internal.util.DumpUtils;
-import com.android.internal.util.FrameworkStatsLog;
 import com.android.server.SystemService;
 import com.android.server.pm.BackgroundUserSoundNotifier;
+import com.android.server.vibrator.VibrationSession.CallerInfo;
+import com.android.server.vibrator.VibrationSession.DebugInfo;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import libcore.util.NativeAllocationRegistry;
 
@@ -160,11 +162,12 @@
     @GuardedBy("mLock")
     private VibrationStepConductor mNextVibration;
     @GuardedBy("mLock")
-    private ExternalVibrationHolder mCurrentExternalVibration;
+    private ExternalVibrationSession mCurrentExternalVibration;
     @GuardedBy("mLock")
     private boolean mServiceReady;
 
-    private final VibrationSettings mVibrationSettings;
+    @VisibleForTesting
+    final VibrationSettings mVibrationSettings;
     private final VibrationScaler mVibrationScaler;
     private final VibratorControlService mVibratorControlService;
     private final InputDeviceDelegate mInputDeviceDelegate;
@@ -184,13 +187,12 @@
                     // When the system is entering a non-interactive state, we want to cancel
                     // vibrations in case a misbehaving app has abandoned them.
                     if (shouldCancelOnScreenOffLocked(mNextVibration)) {
-                        clearNextVibrationLocked(
-                                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF));
+                        clearNextVibrationLocked(new Vibration.EndInfo(
+                                Status.CANCELLED_BY_SCREEN_OFF));
                     }
                     if (shouldCancelOnScreenOffLocked(mCurrentVibration)) {
-                        mCurrentVibration.notifyCancelled(
-                                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
-                                /* immediate= */ false);
+                        mCurrentVibration.notifyCancelled(new Vibration.EndInfo(
+                                Status.CANCELLED_BY_SCREEN_OFF), /* immediate= */ false);
                     }
                 }
             } else if (android.multiuser.Flags.addUiForSoundsFromBackgroundUsers()
@@ -198,12 +200,11 @@
                 synchronized (mLock) {
                     if (shouldCancelOnFgUserRequest(mNextVibration)) {
                         clearNextVibrationLocked(new Vibration.EndInfo(
-                                Vibration.Status.CANCELLED_BY_FOREGROUND_USER));
+                                Status.CANCELLED_BY_FOREGROUND_USER));
                     }
                     if (shouldCancelOnFgUserRequest(mCurrentVibration)) {
                         mCurrentVibration.notifyCancelled(new Vibration.EndInfo(
-                                        Vibration.Status.CANCELLED_BY_FOREGROUND_USER),
-                                /* immediate= */ false);
+                                Status.CANCELLED_BY_FOREGROUND_USER), /* immediate= */ false);
                     }
                 }
             }
@@ -220,12 +221,12 @@
                     }
                     synchronized (mLock) {
                         if (shouldCancelAppOpModeChangedLocked(mNextVibration)) {
-                            clearNextVibrationLocked(
-                                    new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_APP_OPS));
+                            clearNextVibrationLocked(new Vibration.EndInfo(
+                                    Status.CANCELLED_BY_APP_OPS));
                         }
                         if (shouldCancelAppOpModeChangedLocked(mCurrentVibration)) {
                             mCurrentVibration.notifyCancelled(new Vibration.EndInfo(
-                                    Vibration.Status.CANCELLED_BY_APP_OPS), /* immediate= */ false);
+                                    Status.CANCELLED_BY_APP_OPS), /* immediate= */ false);
                         }
                     }
                 }
@@ -441,8 +442,8 @@
                     return false;
                 }
                 AlwaysOnVibration alwaysOnVibration = new AlwaysOnVibration(alwaysOnId,
-                        new Vibration.CallerInfo(attrs, uid, Context.DEVICE_ID_DEFAULT, opPkg,
-                                null), effects);
+                        new CallerInfo(attrs, uid, Context.DEVICE_ID_DEFAULT, opPkg, null),
+                        effects);
                 mAlwaysOnEffects.put(alwaysOnId, alwaysOnVibration);
                 updateAlwaysOnLocked(alwaysOnVibration);
             }
@@ -490,8 +491,7 @@
         // Make sure we report the constant id in the requested haptic feedback reason.
         reason = "performHapticFeedback(constant=" + constant + "): " + reason;
         HapticFeedbackVibrationProvider hapticVibrationProvider = getHapticVibrationProvider();
-        Vibration.Status ignoreStatus = shouldIgnoreHapticFeedback(constant, reason,
-                hapticVibrationProvider);
+        Status ignoreStatus = shouldIgnoreHapticFeedback(constant, reason, hapticVibrationProvider);
         if (ignoreStatus != null) {
             logAndRecordPerformHapticFeedbackAttempt(uid, deviceId, opPkg, reason, ignoreStatus);
             return null;
@@ -516,8 +516,7 @@
         reason = "performHapticFeedbackForInputDevice(constant=" + constant + ", inputDeviceId="
                 + inputDeviceId + ", inputSource=" + inputSource + "): " + reason;
         HapticFeedbackVibrationProvider hapticVibrationProvider = getHapticVibrationProvider();
-        Vibration.Status ignoreStatus = shouldIgnoreHapticFeedback(constant, reason,
-                hapticVibrationProvider);
+        Status ignoreStatus = shouldIgnoreHapticFeedback(constant, reason, hapticVibrationProvider);
         if (ignoreStatus != null) {
             logAndRecordPerformHapticFeedbackAttempt(uid, deviceId, opPkg, reason, ignoreStatus);
             return null;
@@ -533,7 +532,7 @@
             VibrationAttributes attrs) {
         if (effect == null) {
             logAndRecordPerformHapticFeedbackAttempt(uid, deviceId, opPkg, reason,
-                    Vibration.Status.IGNORED_UNSUPPORTED);
+                    Status.IGNORED_UNSUPPORTED);
             Slog.w(TAG,
                     "performHapticFeedbackWithEffect; vibration absent for constant " + constant);
             return null;
@@ -578,23 +577,21 @@
     private HalVibration vibrateInternal(int uid, int deviceId, String opPkg,
             @NonNull CombinedVibration effect, @NonNull VibrationAttributes attrs,
             String reason, IBinder token) {
-        Vibration.CallerInfo callerInfo =
-                new Vibration.CallerInfo(attrs, uid, deviceId, opPkg, reason);
+        CallerInfo callerInfo = new CallerInfo(attrs, uid, deviceId, opPkg, reason);
         if (token == null) {
             Slog.e(TAG, "token must not be null");
-            logAndRecordVibrationAttempt(effect, callerInfo, Vibration.Status.IGNORED_ERROR_TOKEN);
+            logAndRecordVibrationAttempt(effect, callerInfo, Status.IGNORED_ERROR_TOKEN);
             return null;
         }
         if (effect.hasVendorEffects()
                 && !hasPermission(android.Manifest.permission.VIBRATE_VENDOR_EFFECTS)) {
             Slog.e(TAG, "vibrate; no permission for vendor effects");
-            logAndRecordVibrationAttempt(effect, callerInfo,
-                    Vibration.Status.IGNORED_MISSING_PERMISSION);
+            logAndRecordVibrationAttempt(effect, callerInfo, Status.IGNORED_MISSING_PERMISSION);
             return null;
         }
         enforceUpdateAppOpsStatsPermission(uid);
         if (!isEffectValid(effect)) {
-            logAndRecordVibrationAttempt(effect, callerInfo, Vibration.Status.IGNORED_UNSUPPORTED);
+            logAndRecordVibrationAttempt(effect, callerInfo, Status.IGNORED_UNSUPPORTED);
             return null;
         }
         // Create Vibration.Stats as close to the received request as possible, for tracking.
@@ -625,11 +622,11 @@
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     if (mCurrentExternalVibration != null) {
-                        mCurrentExternalVibration.mute();
+                        mCurrentExternalVibration.notifyEnded();
                         vib.stats.reportInterruptedAnotherVibration(
                                 mCurrentExternalVibration.callerInfo);
                         endExternalVibrateLocked(
-                                new Vibration.EndInfo(Vibration.Status.CANCELLED_SUPERSEDED,
+                                new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
                                         vib.callerInfo),
                                 /* continueExternalControl= */ false);
                     } else if (mCurrentVibration != null) {
@@ -645,7 +642,7 @@
                             vib.stats.reportInterruptedAnotherVibration(
                                     mCurrentVibration.getVibration().callerInfo);
                             mCurrentVibration.notifyCancelled(
-                                    new Vibration.EndInfo(Vibration.Status.CANCELLED_SUPERSEDED,
+                                    new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
                                             vib.callerInfo),
                                     /* immediate= */ false);
                         }
@@ -677,7 +674,7 @@
                     Slog.d(TAG, "Canceling vibration");
                 }
                 Vibration.EndInfo cancelledByUserInfo =
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER);
+                        new Vibration.EndInfo(Status.CANCELLED_BY_USER);
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     if (mNextVibration != null
@@ -693,9 +690,9 @@
                     }
                     if (mCurrentExternalVibration != null
                             && shouldCancelVibration(
-                            mCurrentExternalVibration.externalVibration.getVibrationAttributes(),
+                            mCurrentExternalVibration.getCallerInfo().attrs,
                             usageFilter)) {
-                        mCurrentExternalVibration.mute();
+                        mCurrentExternalVibration.notifyEnded();
                         endExternalVibrateLocked(
                                 cancelledByUserInfo, /* continueExternalControl= */ false);
                     }
@@ -860,7 +857,7 @@
                             : vibrationEndInfo.status));
                 }
                 mCurrentVibration.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SETTINGS_UPDATE),
                         /* immediate= */ false);
             }
         }
@@ -911,7 +908,7 @@
             // Note that we don't consider pipelining here, because new pipelined ones should
             // replace pending non-executing pipelined ones anyway.
             clearNextVibrationLocked(
-                    new Vibration.EndInfo(Vibration.Status.IGNORED_SUPERSEDED, vib.callerInfo));
+                    new Vibration.EndInfo(Status.IGNORED_SUPERSEDED, vib.callerInfo));
             mNextVibration = conductor;
             return null;
         } finally {
@@ -934,15 +931,15 @@
                     if (!mVibrationThread.runVibrationOnVibrationThread(mCurrentVibration)) {
                         // Shouldn't happen. The method call already logs a wtf.
                         mCurrentVibration = null;  // Aborted.
-                        return new Vibration.EndInfo(Vibration.Status.IGNORED_ERROR_SCHEDULING);
+                        return new Vibration.EndInfo(Status.IGNORED_ERROR_SCHEDULING);
                     }
                     return null;
                 case AppOpsManager.MODE_ERRORED:
                     Slog.w(TAG, "Start AppOpsManager operation errored for uid "
                             + vib.callerInfo.uid);
-                    return new Vibration.EndInfo(Vibration.Status.IGNORED_ERROR_APP_OPS);
+                    return new Vibration.EndInfo(Status.IGNORED_ERROR_APP_OPS);
                 default:
-                    return new Vibration.EndInfo(Vibration.Status.IGNORED_APP_OPS);
+                    return new Vibration.EndInfo(Status.IGNORED_APP_OPS);
             }
         } finally {
             Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
@@ -950,7 +947,7 @@
     }
 
     @GuardedBy("mLock")
-    private void endVibrationLocked(HalVibration vib, Vibration.EndInfo vibrationEndInfo,
+    private void endVibrationLocked(Vibration vib, Vibration.EndInfo vibrationEndInfo,
             boolean shouldWriteStats) {
         vib.end(vibrationEndInfo);
         logAndRecordVibration(vib.getDebugInfo());
@@ -960,15 +957,6 @@
         }
     }
 
-    @GuardedBy("mLock")
-    private void endVibrationAndWriteStatsLocked(ExternalVibrationHolder vib,
-            Vibration.EndInfo vibrationEndInfo) {
-        vib.end(vibrationEndInfo);
-        logAndRecordVibration(vib.getDebugInfo());
-        mFrameworkStatsLogger.writeVibrationReportedAsync(
-                vib.getStatsInfo(/* completionUptimeMillis= */ SystemClock.uptimeMillis()));
-    }
-
     private VibrationStepConductor createVibrationStepConductor(HalVibration vib) {
         CompletableFuture<Void> requestVibrationParamsFuture = null;
 
@@ -990,33 +978,32 @@
         vib.scaleEffects(mVibrationScaler);
         mInputDeviceDelegate.vibrateIfAvailable(vib.callerInfo, vib.getEffectToPlay());
 
-        return new Vibration.EndInfo(Vibration.Status.FORWARDED_TO_INPUT_DEVICES);
+        return new Vibration.EndInfo(Status.FORWARDED_TO_INPUT_DEVICES);
     }
 
     private void logAndRecordPerformHapticFeedbackAttempt(int uid, int deviceId, String opPkg,
-            String reason, Vibration.Status status) {
-        Vibration.CallerInfo callerInfo = new Vibration.CallerInfo(
+            String reason, Status status) {
+        CallerInfo callerInfo = new CallerInfo(
                 VibrationAttributes.createForUsage(VibrationAttributes.USAGE_UNKNOWN),
                 uid, deviceId, opPkg, reason);
         logAndRecordVibrationAttempt(/* effect= */ null, callerInfo, status);
     }
 
     private void logAndRecordVibrationAttempt(@Nullable CombinedVibration effect,
-            Vibration.CallerInfo callerInfo, Vibration.Status status) {
+            CallerInfo callerInfo, Status status) {
         logAndRecordVibration(
-                new Vibration.DebugInfo(status, new VibrationStats(),
+                new Vibration.DebugInfoImpl(status, new VibrationStats(),
                         effect, /* originalEffect= */ null, VibrationScaler.SCALE_NONE,
                         VibrationScaler.ADAPTIVE_SCALE_NONE, callerInfo));
     }
 
-    private void logAndRecordVibration(Vibration.DebugInfo info) {
+    private void logAndRecordVibration(DebugInfo info) {
         info.logMetrics(mFrameworkStatsLogger);
-        logVibrationStatus(info.mCallerInfo.uid, info.mCallerInfo.attrs, info.mStatus);
+        logVibrationStatus(info.getCallerInfo().uid, info.getCallerInfo().attrs, info.getStatus());
         mVibratorManagerRecords.record(info);
     }
 
-    private void logVibrationStatus(int uid, VibrationAttributes attrs,
-            Vibration.Status status) {
+    private void logVibrationStatus(int uid, VibrationAttributes attrs, Status status) {
         switch (status) {
             case IGNORED_BACKGROUND:
                 Slog.e(TAG, "Ignoring incoming vibration as process with"
@@ -1161,15 +1148,14 @@
 
         if (ongoingVibrationImportance > newVibrationImportance) {
             // Existing vibration has higher importance and should not be cancelled.
-            return new Vibration.EndInfo(Vibration.Status.IGNORED_FOR_HIGHER_IMPORTANCE,
+            return new Vibration.EndInfo(Status.IGNORED_FOR_HIGHER_IMPORTANCE,
                     ongoingVibration.callerInfo);
         }
 
         // Same importance, use repeating as a tiebreaker.
         if (ongoingVibration.isRepeating() && !newVibration.isRepeating()) {
             // Ongoing vibration is repeating and new one is not, give priority to ongoing
-            return new Vibration.EndInfo(Vibration.Status.IGNORED_FOR_ONGOING,
-                    ongoingVibration.callerInfo);
+            return new Vibration.EndInfo(Status.IGNORED_FOR_ONGOING, ongoingVibration.callerInfo);
         }
         // New vibration is repeating or this is a complete tie between them,
         // give priority to new vibration.
@@ -1220,8 +1206,8 @@
      */
     @GuardedBy("mLock")
     @Nullable
-    private Vibration.EndInfo shouldIgnoreVibrationLocked(Vibration.CallerInfo callerInfo) {
-        Vibration.Status statusFromSettings = mVibrationSettings.shouldIgnoreVibration(callerInfo);
+    private Vibration.EndInfo shouldIgnoreVibrationLocked(CallerInfo callerInfo) {
+        Status statusFromSettings = mVibrationSettings.shouldIgnoreVibration(callerInfo);
         if (statusFromSettings != null) {
             return new Vibration.EndInfo(statusFromSettings);
         }
@@ -1231,9 +1217,9 @@
             if (mode == AppOpsManager.MODE_ERRORED) {
                 // We might be getting calls from within system_server, so we don't actually
                 // want to throw a SecurityException here.
-                return new Vibration.EndInfo(Vibration.Status.IGNORED_ERROR_APP_OPS);
+                return new Vibration.EndInfo(Status.IGNORED_ERROR_APP_OPS);
             } else {
-                return new Vibration.EndInfo(Vibration.Status.IGNORED_APP_OPS);
+                return new Vibration.EndInfo(Status.IGNORED_APP_OPS);
             }
         }
 
@@ -1241,16 +1227,16 @@
     }
 
     @Nullable
-    private Vibration.Status shouldIgnoreHapticFeedback(int constant, String reason,
+    private Status shouldIgnoreHapticFeedback(int constant, String reason,
             HapticFeedbackVibrationProvider hapticVibrationProvider) {
         if (hapticVibrationProvider == null) {
             Slog.e(TAG, reason + "; haptic vibration provider not ready.");
-            return Vibration.Status.IGNORED_ERROR_SCHEDULING;
+            return Status.IGNORED_ERROR_SCHEDULING;
         }
         if (hapticVibrationProvider.isRestrictedHapticFeedback(constant)
                 && !hasPermission(android.Manifest.permission.VIBRATE_SYSTEM_CONSTANTS)) {
             Slog.w(TAG, reason + "; no permission for system constant " + constant);
-            return Vibration.Status.IGNORED_MISSING_PERMISSION;
+            return Status.IGNORED_MISSING_PERMISSION;
         }
         return null;
     }
@@ -1291,7 +1277,7 @@
      * {@code attrs}. This will return one of the AppOpsManager.MODE_*.
      */
     @GuardedBy("mLock")
-    private int checkAppOpModeLocked(Vibration.CallerInfo callerInfo) {
+    private int checkAppOpModeLocked(CallerInfo callerInfo) {
         int mode = mAppOps.checkAudioOpNoThrow(AppOpsManager.OP_VIBRATE,
                 callerInfo.attrs.getAudioUsage(), callerInfo.uid, callerInfo.opPkg);
         int fixedMode = fixupAppOpModeLocked(mode, callerInfo.attrs);
@@ -1306,7 +1292,7 @@
 
     /** Start an operation in {@link AppOpsManager}, if allowed. */
     @GuardedBy("mLock")
-    private int startAppOpModeLocked(Vibration.CallerInfo callerInfo) {
+    private int startAppOpModeLocked(CallerInfo callerInfo) {
         return fixupAppOpModeLocked(
                 mAppOps.startOpNoThrow(AppOpsManager.OP_VIBRATE, callerInfo.uid, callerInfo.opPkg),
                 callerInfo.attrs);
@@ -1317,7 +1303,7 @@
      * operation with same uid was previously started.
      */
     @GuardedBy("mLock")
-    private void finishAppOpModeLocked(Vibration.CallerInfo callerInfo) {
+    private void finishAppOpModeLocked(CallerInfo callerInfo) {
         mAppOps.finishOp(AppOpsManager.OP_VIBRATE, callerInfo.uid, callerInfo.opPkg);
     }
 
@@ -1735,10 +1721,10 @@
      */
     private static final class AlwaysOnVibration {
         public final int alwaysOnId;
-        public final Vibration.CallerInfo callerInfo;
+        public final CallerInfo callerInfo;
         public final SparseArray<PrebakedSegment> effects;
 
-        AlwaysOnVibration(int alwaysOnId, Vibration.CallerInfo callerInfo,
+        AlwaysOnVibration(int alwaysOnId, CallerInfo callerInfo,
                 SparseArray<PrebakedSegment> effects) {
             this.alwaysOnId = alwaysOnId;
             this.callerInfo = callerInfo;
@@ -1746,113 +1732,6 @@
         }
     }
 
-    /** Holder for a {@link ExternalVibration}. */
-    private final class ExternalVibrationHolder extends Vibration implements
-            IBinder.DeathRecipient {
-
-        public final ExternalVibration externalVibration;
-        public final ExternalVibrationScale scale = new ExternalVibrationScale();
-
-        private Vibration.Status mStatus;
-
-        private ExternalVibrationHolder(ExternalVibration externalVibration) {
-            super(externalVibration.getToken(), new Vibration.CallerInfo(
-                    externalVibration.getVibrationAttributes(), externalVibration.getUid(),
-                    // TODO(b/249785241): Find a way to link ExternalVibration to a VirtualDevice
-                    // instead of using DEVICE_ID_INVALID here and relying on the UID checks.
-                    Context.DEVICE_ID_INVALID, externalVibration.getPackage(), null));
-            this.externalVibration = externalVibration;
-            mStatus = Vibration.Status.RUNNING;
-        }
-
-        public void muteScale() {
-            scale.scaleLevel = ExternalVibrationScale.ScaleLevel.SCALE_MUTE;
-            if (Flags.hapticsScaleV2Enabled()) {
-                scale.scaleFactor = 0;
-            }
-        }
-
-        public void scale(VibrationScaler scaler, int usage) {
-            scale.scaleLevel = scaler.getScaleLevel(usage);
-            if (Flags.hapticsScaleV2Enabled()) {
-                scale.scaleFactor = scaler.getScaleFactor(usage);
-            }
-            scale.adaptiveHapticsScale = scaler.getAdaptiveHapticsScale(usage);
-            stats.reportAdaptiveScale(scale.adaptiveHapticsScale);
-        }
-
-        public void mute() {
-            externalVibration.mute();
-        }
-
-        public void linkToDeath() {
-            externalVibration.linkToDeath(this);
-        }
-
-        public void unlinkToDeath() {
-            externalVibration.unlinkToDeath(this);
-        }
-
-        public boolean isHoldingSameVibration(ExternalVibration externalVibration) {
-            return this.externalVibration.equals(externalVibration);
-        }
-
-        public void end(Vibration.EndInfo info) {
-            if (mStatus != Vibration.Status.RUNNING) {
-                // Already ended, ignore this call
-                return;
-            }
-            mStatus = info.status;
-            stats.reportEnded(info.endedBy);
-
-            if (stats.hasStarted()) {
-                // External vibration doesn't have feedback from total time the vibrator was playing
-                // with non-zero amplitude, so we use the duration between start and end times of
-                // the vibration as the time the vibrator was ON, since the haptic channels are
-                // open for this duration and can receive vibration waveform data.
-                stats.reportVibratorOn(
-                        stats.getEndUptimeMillis() - stats.getStartUptimeMillis());
-            }
-        }
-
-        public void binderDied() {
-            synchronized (mLock) {
-                if (mCurrentExternalVibration != null) {
-                    if (DEBUG) {
-                        Slog.d(TAG, "External vibration finished because binder died");
-                    }
-                    endExternalVibrateLocked(
-                            new Vibration.EndInfo(Vibration.Status.CANCELLED_BINDER_DIED),
-                            /* continueExternalControl= */ false);
-                }
-            }
-        }
-
-        public Vibration.DebugInfo getDebugInfo() {
-            return new Vibration.DebugInfo(mStatus, stats, /* playedEffect= */ null,
-                    /* originalEffect= */ null, scale.scaleLevel, scale.adaptiveHapticsScale,
-                    callerInfo);
-        }
-
-        public VibrationStats.StatsInfo getStatsInfo(long completionUptimeMillis) {
-            return new VibrationStats.StatsInfo(
-                    externalVibration.getUid(),
-                    FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__EXTERNAL,
-                    externalVibration.getVibrationAttributes().getUsage(), mStatus, stats,
-                    completionUptimeMillis);
-        }
-
-        @Override
-        boolean isRepeating() {
-            // We don't currently know if the external vibration is repeating, so we just use a
-            // heuristic based on the usage. Ideally this would be propagated in the
-            // ExternalVibration.
-            int usage = externalVibration.getVibrationAttributes().getUsage();
-            return usage == VibrationAttributes.USAGE_RINGTONE
-                    || usage == VibrationAttributes.USAGE_ALARM;
-        }
-    }
-
     /** Wrapper around the static-native methods of {@link VibratorManagerService} for tests. */
     @VisibleForTesting
     public static class NativeWrapper {
@@ -1912,7 +1791,7 @@
                     new VibrationRecords(recentVibrationSizeLimit, /* aggregationTimeLimit= */ 0);
         }
 
-        synchronized void record(Vibration.DebugInfo info) {
+        synchronized void record(DebugInfo info) {
             GroupedAggregatedLogRecords.AggregatedLogRecord<VibrationRecord> droppedRecord =
                     mRecentVibrations.add(new VibrationRecord(info));
             if (droppedRecord != null) {
@@ -1969,25 +1848,25 @@
     }
 
     /**
-     * Record for a single {@link Vibration.DebugInfo}, that can be grouped by usage and aggregated
-     * by UID, {@link VibrationAttributes} and {@link VibrationEffect}.
+     * Record for a single {@link DebugInfo}, that can be grouped by usage and aggregated by UID,
+     * {@link VibrationAttributes} and {@link CombinedVibration}.
      */
     private static final class VibrationRecord
             implements GroupedAggregatedLogRecords.SingleLogRecord {
-        private final Vibration.DebugInfo mInfo;
+        private final DebugInfo mInfo;
 
-        VibrationRecord(Vibration.DebugInfo info) {
+        VibrationRecord(DebugInfo info) {
             mInfo = info;
         }
 
         @Override
         public int getGroupKey() {
-            return mInfo.mCallerInfo.attrs.getUsage();
+            return mInfo.getCallerInfo().attrs.getUsage();
         }
 
         @Override
         public long getCreateUptimeMs() {
-            return mInfo.mCreateTime;
+            return mInfo.getCreateUptimeMillis();
         }
 
         @Override
@@ -1995,10 +1874,10 @@
             if (!(record instanceof VibrationRecord)) {
                 return false;
             }
-            Vibration.DebugInfo info = ((VibrationRecord) record).mInfo;
-            return mInfo.mCallerInfo.uid == info.mCallerInfo.uid
-                    && Objects.equals(mInfo.mCallerInfo.attrs, info.mCallerInfo.attrs)
-                    && Objects.equals(mInfo.mPlayedEffect, info.mPlayedEffect);
+            DebugInfo info = ((VibrationRecord) record).mInfo;
+            return mInfo.getCallerInfo().uid == info.getCallerInfo().uid
+                    && Objects.equals(mInfo.getCallerInfo().attrs, info.getCallerInfo().attrs)
+                    && Objects.equals(mInfo.getDumpAggregationKey(), info.getDumpAggregationKey());
         }
 
         @Override
@@ -2048,7 +1927,8 @@
                 setExternalControl(false, mCurrentExternalVibration.stats);
             }
             // The external control was turned off, end it and report metrics right away.
-            endVibrationAndWriteStatsLocked(mCurrentExternalVibration, vibrationEndInfo);
+            endVibrationLocked(mCurrentExternalVibration, vibrationEndInfo,
+                    /* shouldWriteStats= */ true);
             mCurrentExternalVibration = null;
         } finally {
             Trace.traceEnd(Trace.TRACE_TAG_VIBRATOR);
@@ -2108,17 +1988,18 @@
         @Override
         public ExternalVibrationScale onExternalVibrationStart(ExternalVibration vib) {
             // Create Vibration.Stats as close to the received request as possible, for tracking.
-            ExternalVibrationHolder vibHolder = new ExternalVibrationHolder(vib);
+            ExternalVibrationSession externalVibration = new ExternalVibrationSession(vib);
             // Mute the request until we run all the checks and accept the vibration.
-            vibHolder.muteScale();
+            externalVibration.muteScale();
             boolean alreadyUnderExternalControl = false;
             boolean waitForCompletion = false;
 
             synchronized (mLock) {
                 if (!hasExternalControlCapability()) {
-                    endVibrationAndWriteStatsLocked(vibHolder,
-                            new Vibration.EndInfo(Vibration.Status.IGNORED_UNSUPPORTED));
-                    return vibHolder.scale;
+                    endVibrationLocked(externalVibration,
+                            new Vibration.EndInfo(Status.IGNORED_UNSUPPORTED),
+                            /* shouldWriteStats= */ true);
+                    return externalVibration.getScale();
                 }
 
                 if (ActivityManager.checkComponentPermission(android.Manifest.permission.VIBRATE,
@@ -2127,44 +2008,46 @@
                     Slog.w(TAG, "pkg=" + vib.getPackage() + ", uid=" + vib.getUid()
                             + " tried to play externally controlled vibration"
                             + " without VIBRATE permission, ignoring.");
-                    endVibrationAndWriteStatsLocked(vibHolder,
-                            new Vibration.EndInfo(Vibration.Status.IGNORED_MISSING_PERMISSION));
-                    return vibHolder.scale;
+                    endVibrationLocked(externalVibration,
+                            new Vibration.EndInfo(Status.IGNORED_MISSING_PERMISSION),
+                            /* shouldWriteStats= */ true);
+                    return externalVibration.getScale();
                 }
 
                 Vibration.EndInfo vibrationEndInfo = shouldIgnoreVibrationLocked(
-                        vibHolder.callerInfo);
+                        externalVibration.callerInfo);
 
                 if (vibrationEndInfo == null
                         && mCurrentExternalVibration != null
                         && mCurrentExternalVibration.isHoldingSameVibration(vib)) {
                     // We are already playing this external vibration, so we can return the same
                     // scale calculated in the previous call to this method.
-                    return mCurrentExternalVibration.scale;
+                    return mCurrentExternalVibration.getScale();
                 }
 
                 if (vibrationEndInfo == null) {
                     // Check if ongoing vibration is more important than this vibration.
-                    vibrationEndInfo = shouldIgnoreVibrationForOngoingLocked(vibHolder);
+                    vibrationEndInfo = shouldIgnoreVibrationForOngoingLocked(externalVibration);
                 }
 
                 if (vibrationEndInfo != null) {
-                    endVibrationAndWriteStatsLocked(vibHolder, vibrationEndInfo);
-                    return vibHolder.scale;
+                    endVibrationLocked(externalVibration, vibrationEndInfo,
+                            /* shouldWriteStats= */ true);
+                    return externalVibration.getScale();
                 }
 
                 if (mCurrentExternalVibration == null) {
                     // If we're not under external control right now, then cancel any normal
                     // vibration that may be playing and ready the vibrator for external control.
                     if (mCurrentVibration != null) {
-                        vibHolder.stats.reportInterruptedAnotherVibration(
+                        externalVibration.stats.reportInterruptedAnotherVibration(
                                 mCurrentVibration.getVibration().callerInfo);
                         clearNextVibrationLocked(
-                                new Vibration.EndInfo(Vibration.Status.IGNORED_FOR_EXTERNAL,
-                                        vibHolder.callerInfo));
+                                new Vibration.EndInfo(Status.IGNORED_FOR_EXTERNAL,
+                                        externalVibration.callerInfo));
                         mCurrentVibration.notifyCancelled(
-                                new Vibration.EndInfo(Vibration.Status.CANCELLED_SUPERSEDED,
-                                        vibHolder.callerInfo),
+                                new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
+                                        externalVibration.callerInfo),
                                 /* immediate= */ true);
                         waitForCompletion = true;
                     }
@@ -2178,12 +2061,12 @@
                     // Note that this doesn't support multiple concurrent external controls, as we
                     // would need to mute the old one still if it came from a different controller.
                     alreadyUnderExternalControl = true;
-                    mCurrentExternalVibration.mute();
-                    vibHolder.stats.reportInterruptedAnotherVibration(
+                    mCurrentExternalVibration.notifyEnded();
+                    externalVibration.stats.reportInterruptedAnotherVibration(
                             mCurrentExternalVibration.callerInfo);
                     endExternalVibrateLocked(
-                            new Vibration.EndInfo(Vibration.Status.CANCELLED_SUPERSEDED,
-                                    vibHolder.callerInfo),
+                            new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
+                                    externalVibration.callerInfo),
                             /* continueExternalControl= */ true);
                 }
 
@@ -2195,9 +2078,9 @@
                     mVibrationSettings.update();
                 }
 
-                mCurrentExternalVibration = vibHolder;
-                vibHolder.linkToDeath();
-                vibHolder.scale(mVibrationScaler, attrs.getUsage());
+                mCurrentExternalVibration = externalVibration;
+                externalVibration.linkToDeath(this::onExternalVibrationBinderDied);
+                externalVibration.scale(mVibrationScaler, attrs.getUsage());
             }
 
             if (waitForCompletion) {
@@ -2206,27 +2089,27 @@
                     synchronized (mLock) {
                         // Trigger endExternalVibrateLocked to unlink to death recipient.
                         endExternalVibrateLocked(
-                                new Vibration.EndInfo(Vibration.Status.IGNORED_ERROR_CANCELLING),
+                                new Vibration.EndInfo(Status.IGNORED_ERROR_CANCELLING),
                                 /* continueExternalControl= */ false);
                         // Mute the request, vibration will be ignored.
-                        vibHolder.muteScale();
+                        externalVibration.muteScale();
                     }
-                    return vibHolder.scale;
+                    return externalVibration.getScale();
                 }
             }
             if (!alreadyUnderExternalControl) {
                 if (DEBUG) {
                     Slog.d(TAG, "Vibrator going under external control.");
                 }
-                setExternalControl(true, vibHolder.stats);
+                setExternalControl(true, externalVibration.stats);
             }
             if (DEBUG) {
                 Slog.d(TAG, "Playing external vibration: " + vib);
             }
             // Vibrator will start receiving data from external channels after this point.
             // Report current time as the vibration start time, for debugging.
-            vibHolder.stats.reportStarted();
-            return vibHolder.scale;
+            externalVibration.stats.reportStarted();
+            return externalVibration.getScale();
         }
 
         @Override
@@ -2238,7 +2121,7 @@
                         Slog.d(TAG, "Stopping external vibration: " + vib);
                     }
                     endExternalVibrateLocked(
-                            new Vibration.EndInfo(Vibration.Status.FINISHED),
+                            new Vibration.EndInfo(Status.FINISHED),
                             /* continueExternalControl= */ false);
                 }
             }
@@ -2252,6 +2135,19 @@
             }
             return false;
         }
+
+        private void onExternalVibrationBinderDied() {
+            synchronized (mLock) {
+                if (mCurrentExternalVibration != null) {
+                    if (DEBUG) {
+                        Slog.d(TAG, "External vibration finished because binder died");
+                    }
+                    endExternalVibrateLocked(
+                            new Vibration.EndInfo(Status.CANCELLED_BINDER_DIED),
+                            /* continueExternalControl= */ false);
+                }
+            }
+        }
     }
 
     /** Provide limited functionality from {@link VibratorManagerService} as shell commands. */
diff --git a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
index ba2594a..f53dda6 100644
--- a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
+++ b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
@@ -42,6 +42,7 @@
 import static com.android.server.wallpaper.WallpaperUtils.WALLPAPER_LOCK_ORIG;
 import static com.android.server.wallpaper.WallpaperUtils.getWallpaperDir;
 import static com.android.server.wallpaper.WallpaperUtils.makeWallpaperIdLocked;
+import static com.android.window.flags.Flags.avoidRebindingIntentionallyDisconnectedWallpaper;
 import static com.android.window.flags.Flags.multiCrop;
 import static com.android.window.flags.Flags.offloadColorExtraction;
 
@@ -897,6 +898,12 @@
                     return;
                 }
 
+                if (avoidRebindingIntentionallyDisconnectedWallpaper()
+                        && mWallpaper.connection == null) {
+                    Slog.w(TAG, "Trying to reset an intentionally disconnected wallpaper!");
+                    return;
+                }
+
                 if (!mWallpaper.wallpaperUpdating && mWallpaper.userId == mCurrentUserId) {
                     Slog.w(TAG, "Wallpaper reconnect timed out for " + mWallpaper.wallpaperComponent
                             + ", reverting to built-in wallpaper!");
@@ -1066,6 +1073,13 @@
                 if (mWallpaper.wallpaperUpdating) {
                     return;
                 }
+
+                if (avoidRebindingIntentionallyDisconnectedWallpaper()
+                        && mWallpaper.connection == null) {
+                    Slog.w(TAG, "Trying to rebind an intentionally disconnected wallpaper!");
+                    return;
+                }
+
                 final ComponentName wpService = mWallpaper.wallpaperComponent;
                 // The broadcast of package update could be delayed after service disconnected. Try
                 // to re-bind the service for 10 seconds.
diff --git a/services/core/java/com/android/server/webkit/SystemImpl.java b/services/core/java/com/android/server/webkit/SystemImpl.java
index c4d601d..6740153 100644
--- a/services/core/java/com/android/server/webkit/SystemImpl.java
+++ b/services/core/java/com/android/server/webkit/SystemImpl.java
@@ -19,14 +19,12 @@
 import static android.webkit.Flags.updateServiceV2;
 
 import android.app.ActivityManager;
-import android.app.AppGlobals;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageInstaller;
 import android.content.pm.PackageManager;
 import android.content.pm.PackageManager.NameNotFoundException;
-import android.content.pm.UserInfo;
 import android.content.res.XmlResourceParser;
 import android.os.Build;
 import android.os.RemoteException;
@@ -79,7 +77,7 @@
         XmlResourceParser parser = null;
         List<WebViewProviderInfo> webViewProviders = new ArrayList<WebViewProviderInfo>();
         try {
-            parser = AppGlobals.getInitialApplication().getResources().getXml(
+            parser = mContext.getResources().getXml(
                     com.android.internal.R.xml.config_webview_packages);
             XmlUtils.beginDocument(parser, TAG_START);
             while(true) {
@@ -148,7 +146,7 @@
     }
 
     public long getFactoryPackageVersion(String packageName) throws NameNotFoundException {
-        PackageManager pm = AppGlobals.getInitialApplication().getPackageManager();
+        PackageManager pm = mContext.getPackageManager();
         return pm.getPackageInfo(packageName, PackageManager.MATCH_FACTORY_ONLY)
                 .getLongVersionCode();
     }
@@ -203,47 +201,48 @@
     @Override
     public void enablePackageForAllUsers(String packageName, boolean enable) {
         UserManager userManager = mContext.getSystemService(UserManager.class);
-        for(UserInfo userInfo : userManager.getUsers()) {
-            enablePackageForUser(packageName, enable, userInfo.id);
+        for (UserHandle user : userManager.getUserHandles(false)) {
+            enablePackageForUser(packageName, enable, user);
         }
     }
 
-    private void enablePackageForUser(String packageName, boolean enable, int userId) {
+    private void enablePackageForUser(String packageName, boolean enable, UserHandle user) {
+        Context contextAsUser = mContext.createContextAsUser(user, 0);
+        PackageManager pm = contextAsUser.getPackageManager();
         try {
-            AppGlobals.getPackageManager().setApplicationEnabledSetting(
+            pm.setApplicationEnabledSetting(
                     packageName,
                     enable ? PackageManager.COMPONENT_ENABLED_STATE_DEFAULT :
-                    PackageManager.COMPONENT_ENABLED_STATE_DISABLED_USER, 0,
-                    userId, null);
-        } catch (RemoteException | IllegalArgumentException e) {
+                    PackageManager.COMPONENT_ENABLED_STATE_DISABLED_USER, 0);
+        } catch (IllegalArgumentException e) {
             Log.w(TAG, "Tried to " + (enable ? "enable " : "disable ") + packageName
-                    + " for user " + userId + ": " + e);
+                    + " for user " + user + ": " + e);
         }
     }
 
     @Override
     public void installExistingPackageForAllUsers(String packageName) {
         UserManager userManager = mContext.getSystemService(UserManager.class);
-        for (UserInfo userInfo : userManager.getUsers()) {
-            installPackageForUser(packageName, userInfo.id);
+        for (UserHandle user : userManager.getUserHandles(false)) {
+            installPackageForUser(packageName, user);
         }
     }
 
-    private void installPackageForUser(String packageName, int userId) {
-        final Context contextAsUser = mContext.createContextAsUser(UserHandle.of(userId), 0);
-        final PackageInstaller installer = contextAsUser.getPackageManager().getPackageInstaller();
+    private void installPackageForUser(String packageName, UserHandle user) {
+        Context contextAsUser = mContext.createContextAsUser(user, 0);
+        PackageInstaller installer = contextAsUser.getPackageManager().getPackageInstaller();
         installer.installExistingPackage(packageName, PackageManager.INSTALL_REASON_UNKNOWN, null);
     }
 
     @Override
     public boolean systemIsDebuggable() {
-        return Build.IS_DEBUGGABLE;
+        return Build.isDebuggable();
     }
 
     @Override
     public PackageInfo getPackageInfoForProvider(WebViewProviderInfo configInfo)
             throws NameNotFoundException {
-        PackageManager pm = AppGlobals.getInitialApplication().getPackageManager();
+        PackageManager pm = mContext.getPackageManager();
         return pm.getPackageInfo(configInfo.packageName, PACKAGE_FLAGS);
     }
 
@@ -327,5 +326,5 @@
     // flags declaring we want extra info from the package manager for webview providers
     private final static int PACKAGE_FLAGS = PackageManager.GET_META_DATA
             | PackageManager.GET_SIGNATURES | PackageManager.GET_SHARED_LIBRARY_FILES
-            | PackageManager.MATCH_DEBUG_TRIAGED_MISSING | PackageManager.MATCH_ANY_USER;
+            | PackageManager.MATCH_ANY_USER;
 }
diff --git a/services/core/java/com/android/server/wm/ActionChain.java b/services/core/java/com/android/server/wm/ActionChain.java
new file mode 100644
index 0000000..d63044a
--- /dev/null
+++ b/services/core/java/com/android/server/wm/ActionChain.java
@@ -0,0 +1,240 @@
+/*
+ * 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.wm;
+
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.util.Slog;
+
+import com.android.window.flags.Flags;
+
+/**
+ * Represents a chain of WM actions where each action is "caused by" the prior action (except the
+ * first one of course). A whole chain is associated with one Transition (in fact, the purpose
+ * of this object is to communicate, to all callees, which transition they are part of).
+ *
+ * A single action is defined as "one logical thing requested of WM". This usually corresponds to
+ * each ingress-point into the process. For example, when starting an activity:
+ *   * the first action is to pause the current/top activity.
+ *       At this point, control leaves the process while the activity pauses.
+ *   * Then WM receives completePause (a new ingress). This is a new action that gets linked
+ *       to the prior action. This action involves resuming the next activity, at which point,
+ *       control leaves the process again.
+ *   * Eventually, when everything is done, we will have formed a chain of actions.
+ *
+ * We don't technically need to hold onto each prior action in the chain once a new action has
+ * been linked to the same transition; however, keeping the whole chain enables improved
+ * debugging and the ability to detect anomalies.
+ */
+public class ActionChain {
+    private static final String TAG = "TransitionChain";
+
+    /**
+     * Normal link type. This means the action was expected and is properly linked to the
+     * current chain.
+     */
+    static final int TYPE_NORMAL = 0;
+
+    /**
+     * This is the "default" link. It means we haven't done anything to properly track this case
+     * so it may or may not be correct. It represents the behavior as if there was no tracking.
+     *
+     * Any type that has "default" behavior uses the global "collecting transition" if it exists,
+     * otherwise it doesn't use any transition.
+     */
+    static final int TYPE_DEFAULT = 1;
+
+    /**
+     * This means the action was performed via a legacy code-path. These should be removed
+     * eventually. This will have the "default" behavior.
+     */
+    static final int TYPE_LEGACY = 2;
+
+    /** This is for a test. */
+    static final int TYPE_TEST = 3;
+
+    /** This is finishing a transition. Collection isn't supported during this. */
+    static final int TYPE_FINISH = 4;
+
+    /**
+     * Something unexpected happened so this action was started to recover from the unexpected
+     * state. This means that a "real" chain-link couldn't be determined. For now, the behavior of
+     * this is the same as "default".
+     */
+    static final int TYPE_FAILSAFE = 5;
+
+    /**
+     * Types of chain links (ie. how is this action associated with the chain it is linked to).
+     * @hide
+     */
+    @IntDef(prefix = { "TYPE_" }, value = {
+            TYPE_NORMAL,
+            TYPE_DEFAULT,
+            TYPE_LEGACY,
+            TYPE_TEST,
+            TYPE_FINISH,
+            TYPE_FAILSAFE
+    })
+    public @interface LinkType {}
+
+    /** Identifies the entry-point of this action. */
+    @NonNull
+    final String mSource;
+
+    /** Reference to ATMS. TEMPORARY! ONLY USE THIS WHEN tracker_plumbing flag is DISABLED! */
+    @Nullable
+    ActivityTaskManagerService mTmpAtm;
+
+    /** The transition that this chain's changes belong to. */
+    @Nullable
+    Transition mTransition;
+
+    /** The previous action in the chain. */
+    @Nullable
+    ActionChain mPrevious = null;
+
+    /** Classification of how this action is connected to the chain. */
+    @LinkType int mType = TYPE_NORMAL;
+
+    /** When this Action started. */
+    long mCreateTimeMs;
+
+    private ActionChain(String source, @LinkType int type, Transition transit) {
+        mSource = source;
+        mCreateTimeMs = System.currentTimeMillis();
+        mType = type;
+        mTransition = transit;
+        if (mTransition != null) {
+            mTransition.recordChain(this);
+        }
+    }
+
+    private Transition getTransition() {
+        if (!Flags.transitTrackerPlumbing()) {
+            return mTmpAtm.getTransitionController().getCollectingTransition();
+        }
+        return mTransition;
+    }
+
+    boolean isFinishing() {
+        return mType == TYPE_FINISH;
+    }
+
+    /**
+     * Some common checks to determine (and report) whether this chain has a collecting transition.
+     */
+    private boolean expectCollecting() {
+        final Transition transition = getTransition();
+        if (transition == null) {
+            Slog.e(TAG, "Can't collect into a chain with no transition");
+            return false;
+        }
+        if (isFinishing()) {
+            Slog.e(TAG, "Trying to collect into a finished transition");
+            return false;
+        }
+        if (transition.mController.getCollectingTransition() != mTransition) {
+            Slog.e(TAG, "Mismatch between current collecting ("
+                    + transition.mController.getCollectingTransition() + ") and chain ("
+                    + transition + ")");
+            return false;
+        }
+        return true;
+    }
+
+    /**
+     * Helper to collect a container into the associated transition. This will automatically do
+     * nothing if the chain isn't associated with a collecting transition.
+     */
+    void collect(@NonNull WindowContainer wc) {
+        if (!wc.mTransitionController.isShellTransitionsEnabled()) return;
+        if (!expectCollecting()) return;
+        getTransition().collect(wc);
+    }
+
+    /**
+     * An interface for creating and tracking action chains.
+     */
+    static class Tracker {
+        private final ActivityTaskManagerService mAtm;
+
+        Tracker(ActivityTaskManagerService atm) {
+            mAtm = atm;
+        }
+
+        private ActionChain makeChain(String source, @LinkType int type, Transition transit) {
+            final ActionChain out = new ActionChain(source, type, transit);
+            if (!Flags.transitTrackerPlumbing()) {
+                out.mTmpAtm = mAtm;
+            }
+            return out;
+        }
+
+        private ActionChain makeChain(String source, @LinkType int type) {
+            return makeChain(source, type,
+                    mAtm.getTransitionController().getCollectingTransition());
+        }
+
+        /**
+         * Starts tracking a normal action.
+         * @see #TYPE_NORMAL
+         */
+        @NonNull
+        ActionChain start(String source, Transition transit) {
+            return makeChain(source, TYPE_NORMAL, transit);
+        }
+
+        /** @see #TYPE_DEFAULT */
+        @NonNull
+        ActionChain startDefault(String source) {
+            return makeChain(source, TYPE_DEFAULT);
+        }
+
+        /**
+         * Starts tracking an action that finishes a transition.
+         * @see #TYPE_NORMAL
+         */
+        @NonNull
+        ActionChain startFinish(String source, Transition finishTransit) {
+            return makeChain(source, TYPE_FINISH, finishTransit);
+        }
+
+        /** @see #TYPE_LEGACY */
+        @NonNull
+        ActionChain startLegacy(String source) {
+            return makeChain(source, TYPE_LEGACY, null);
+        }
+
+        /** @see #TYPE_FAILSAFE */
+        @NonNull
+        ActionChain startFailsafe(String source) {
+            return makeChain(source, TYPE_FAILSAFE);
+        }
+    }
+
+    /** Helpers for usage in tests. */
+    @NonNull
+    static ActionChain test() {
+        return new ActionChain("test", TYPE_TEST, null /* transition */);
+    }
+
+    @NonNull
+    static ActionChain testFinish(Transition toFinish) {
+        return new ActionChain("test", TYPE_FINISH, toFinish);
+    }
+}
diff --git a/services/core/java/com/android/server/wm/ActivityClientController.java b/services/core/java/com/android/server/wm/ActivityClientController.java
index e27b2be..c1e859d 100644
--- a/services/core/java/com/android/server/wm/ActivityClientController.java
+++ b/services/core/java/com/android/server/wm/ActivityClientController.java
@@ -864,8 +864,9 @@
                 if (transition != null) {
                     if (changed) {
                         // Always set as scene transition because it expects to be a jump-cut.
-                        transition.setOverrideAnimation(TransitionInfo.AnimationOptions
-                                .makeSceneTransitionAnimOptions(), null, null);
+                        transition.setOverrideAnimation(
+                                TransitionInfo.AnimationOptions.makeSceneTransitionAnimOptions(), r,
+                                null, null);
                         r.mTransitionController.requestStartTransition(transition,
                                 null /*startTask */, null /* remoteTransition */,
                                 null /* displayChange */);
@@ -910,8 +911,9 @@
                                 && under.returningOptions.getAnimationType()
                                         == ANIM_SCENE_TRANSITION) {
                             // Pass along the scene-transition animation-type
-                            transition.setOverrideAnimation(TransitionInfo.AnimationOptions
-                                    .makeSceneTransitionAnimOptions(), null, null);
+                            transition.setOverrideAnimation(TransitionInfo
+                                            .AnimationOptions.makeSceneTransitionAnimOptions(), r,
+                                    null, null);
                         }
                     } else {
                         transition.abort();
@@ -1508,7 +1510,7 @@
                         r.mOverrideTaskTransition);
                 r.mTransitionController.setOverrideAnimation(
                         TransitionInfo.AnimationOptions.makeCustomAnimOptions(packageName,
-                                enterAnim, exitAnim, backgroundColor, r.mOverrideTaskTransition),
+                                enterAnim, exitAnim, backgroundColor, r.mOverrideTaskTransition), r,
                         null /* startCallback */, null /* finishCallback */);
             }
         }
diff --git a/services/core/java/com/android/server/wm/ActivityMetricsLogger.java b/services/core/java/com/android/server/wm/ActivityMetricsLogger.java
index 2ce1aa42..fb5c115 100644
--- a/services/core/java/com/android/server/wm/ActivityMetricsLogger.java
+++ b/services/core/java/com/android/server/wm/ActivityMetricsLogger.java
@@ -838,12 +838,13 @@
         }
 
         if (android.app.Flags.appStartInfoTimestamps()) {
+            final int pid = r.getPid();
             // Log here to match StatsD for time to first frame.
             mLoggerHandler.post(
                     () -> mSupervisor.mService.mWindowManager.mAmInternal.addStartInfoTimestamp(
                             ApplicationStartInfo.START_TIMESTAMP_FIRST_FRAME,
-                            timestampNs, r.getUid(), r.getPid(),
-                            info.mLastLaunchedActivity.mUserId));
+                            timestampNs, infoSnapshot.applicationInfo.uid, pid,
+                            infoSnapshot.userId));
         }
 
         return infoSnapshot;
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 4b0409b..e562ea8 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -353,6 +353,7 @@
 import android.window.SplashScreenView;
 import android.window.SplashScreenView.SplashScreenViewParcelable;
 import android.window.TaskSnapshot;
+import android.window.TransitionInfo;
 import android.window.TransitionInfo.AnimationOptions;
 import android.window.WindowContainerToken;
 import android.window.WindowOnBackInvokedDispatcher;
@@ -812,6 +813,8 @@
     /** The last set {@link DropInputMode} for this activity surface. */
     @DropInputMode
     private int mLastDropInputMode = DropInputMode.NONE;
+    /** Whether the input to this activity will be dropped during the current playing animation. */
+    private boolean mIsInputDroppedForAnimation;
 
     /**
      * Whether the application has desk mode resources. Calculated and cached when
@@ -1646,6 +1649,15 @@
         }
     }
 
+    /** Sets if all input will be dropped as a protection during the client-driven animation. */
+    void setDropInputForAnimation(boolean isInputDroppedForAnimation) {
+        if (mIsInputDroppedForAnimation == isInputDroppedForAnimation) {
+            return;
+        }
+        mIsInputDroppedForAnimation = isInputDroppedForAnimation;
+        updateUntrustedEmbeddingInputProtection();
+    }
+
     /**
      * Sets to drop input when obscured to activity if it is embedded in untrusted mode.
      *
@@ -1658,7 +1670,10 @@
         if (getSurfaceControl() == null) {
             return;
         }
-        if (isEmbeddedInUntrustedMode()) {
+        if (mIsInputDroppedForAnimation) {
+            // Disable all input during the animation.
+            setDropInputMode(DropInputMode.ALL);
+        } else if (isEmbeddedInUntrustedMode()) {
             // Set drop input to OBSCURED when untrusted embedded.
             setDropInputMode(DropInputMode.OBSCURED);
         } else {
@@ -5034,7 +5049,8 @@
                 // controller but don't clear the animation information from the options since they
                 // need to be sent to the animating activity.
                 mTransitionController.setOverrideAnimation(
-                        AnimationOptions.makeSceneTransitionAnimOptions(), null, null);
+                        TransitionInfo.AnimationOptions.makeSceneTransitionAnimOptions(), this,
+                        null, null);
                 return;
             }
             applyOptionsAnimation(mPendingOptions, intent);
@@ -5157,7 +5173,8 @@
         }
 
         if (options != null) {
-            mTransitionController.setOverrideAnimation(options, startCallback, finishCallback);
+            mTransitionController.setOverrideAnimation(options, this, startCallback,
+                    finishCallback);
         }
     }
 
@@ -8135,6 +8152,9 @@
      */
     @Override
     protected int getOverrideOrientation() {
+        if (mWmService.mConstants.mIgnoreActivityOrientationRequest) {
+            return ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
+        }
         return mAppCompatController.getOrientationPolicy()
                 .overrideOrientationIfNeeded(super.getOverrideOrientation());
     }
diff --git a/services/core/java/com/android/server/wm/ActivityStartController.java b/services/core/java/com/android/server/wm/ActivityStartController.java
index 1660ca9..35ec5ad 100644
--- a/services/core/java/com/android/server/wm/ActivityStartController.java
+++ b/services/core/java/com/android/server/wm/ActivityStartController.java
@@ -424,13 +424,19 @@
                 Intent intent = intents[i];
                 NeededUriGrants intentGrants = null;
 
-                intent.prepareToEnterSystemServer();
+                // Refuse possible leaked file descriptors.
+                if (intent.hasFileDescriptors()) {
+                    throw new IllegalArgumentException("File descriptors passed in Intent");
+                }
 
                 // Get the flag earlier because the intent may be modified in resolveActivity below.
                 final boolean componentSpecified = intent.getComponent() != null;
                 // Don't modify the client's object!
                 intent = new Intent(intent);
 
+                // Remove existing mismatch flag so it can be properly updated later
+                intent.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
+
                 // Collect information about the target of the Intent.
                 ActivityInfo aInfo = mSupervisor.resolveActivity(intent, resolvedTypes[i],
                         0 /* startFlags */, null /* profilerInfo */, userId, filterCallingUid,
diff --git a/services/core/java/com/android/server/wm/ActivityStarter.java b/services/core/java/com/android/server/wm/ActivityStarter.java
index 6479111..1822a80 100644
--- a/services/core/java/com/android/server/wm/ActivityStarter.java
+++ b/services/core/java/com/android/server/wm/ActivityStarter.java
@@ -721,7 +721,13 @@
             onExecutionStarted();
 
             if (mRequest.intent != null) {
-                mRequest.intent.prepareToEnterSystemServer();
+                // Refuse possible leaked file descriptors
+                if (mRequest.intent.hasFileDescriptors()) {
+                    throw new IllegalArgumentException("File descriptors passed in Intent");
+                }
+
+                // Remove existing mismatch flag so it can be properly updated later
+                mRequest.intent.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
             }
 
             final LaunchingState launchingState;
@@ -1945,7 +1951,7 @@
                 && sourceRecord != null && sourceRecord.getTask() == mStartActivity.getTask()
                 && balVerdict.allows()) {
             mRootWindowContainer.moveActivityToPinnedRootTask(mStartActivity,
-                    sourceRecord, "launch-into-pip");
+                    sourceRecord, "launch-into-pip", null /* bounds */);
         }
 
         mSupervisor.getBackgroundActivityLaunchController()
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 52447e8..3cfb9a07 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -795,6 +795,7 @@
     WindowOrganizerController mWindowOrganizerController;
     TaskOrganizerController mTaskOrganizerController;
     TaskFragmentOrganizerController mTaskFragmentOrganizerController;
+    ActionChain.Tracker mChainTracker;
 
     @Nullable
     private BackgroundActivityStartCallback mBackgroundActivityStartCallback;
@@ -869,6 +870,7 @@
         mInternal = new LocalService();
         GL_ES_VERSION = SystemProperties.getInt("ro.opengles.version", GL_ES_VERSION_UNDEFINED);
         mWindowOrganizerController = new WindowOrganizerController(this);
+        mChainTracker = new ActionChain.Tracker(this);
         mTaskOrganizerController = mWindowOrganizerController.mTaskOrganizerController;
         mTaskFragmentOrganizerController =
                 mWindowOrganizerController.mTaskFragmentOrganizerController;
@@ -1315,7 +1317,12 @@
             String resultWho, int requestCode, int flagsMask, int flagsValues, Bundle bOptions) {
         enforceNotIsolatedCaller("startActivityIntentSender");
         if (fillInIntent != null) {
-            fillInIntent.prepareToEnterSystemServer();
+            // Refuse possible leaked file descriptors
+            if (fillInIntent.hasFileDescriptors()) {
+                throw new IllegalArgumentException("File descriptors passed in Intent");
+            }
+            // Remove existing mismatch flag so it can be properly updated later
+            fillInIntent.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
         }
 
         if (!(target instanceof PendingIntentRecord)) {
@@ -1341,10 +1348,10 @@
     @Override
     public boolean startNextMatchingActivity(IBinder callingActivity, Intent intent,
             Bundle bOptions) {
-        if (intent != null) {
-            intent.prepareToEnterSystemServer();
+        // Refuse possible leaked file descriptors
+        if (intent != null && intent.hasFileDescriptors()) {
+            throw new IllegalArgumentException("File descriptors passed in Intent");
         }
-
         SafeActivityOptions options = SafeActivityOptions.fromBundle(bOptions);
 
         synchronized (mGlobalLock) {
@@ -1359,6 +1366,8 @@
                 return false;
             }
             intent = new Intent(intent);
+            // Remove existing mismatch flag so it can be properly updated later
+            intent.removeExtendedFlags(Intent.EXTENDED_FLAG_FILTER_MISMATCH);
             // The caller is not allowed to change the data.
             intent.setDataAndType(r.intent.getData(), r.intent.getType());
             // And we are resetting to find the next component...
@@ -3794,9 +3803,22 @@
                         r.shortComponentName, Boolean.toString(isAutoEnter));
                 r.setPictureInPictureParams(params);
                 r.mAutoEnteringPip = isAutoEnter;
-                mRootWindowContainer.moveActivityToPinnedRootTask(r,
-                        null /* launchIntoPipHostActivity */, "enterPictureInPictureMode",
-                        transition);
+
+                if (transition != null) {
+                    mRootWindowContainer.moveActivityToPinnedRootTaskAndRequestStart(r,
+                            "enterPictureInPictureMode");
+                } else if (getTransitionController().isCollecting()
+                        || !getTransitionController().isShellTransitionsEnabled()) {
+                    mRootWindowContainer.moveActivityToPinnedRootTask(r,
+                            null /* launchIntoPipHostActivity */, "enterPictureInPictureMode",
+                            null /* bounds */);
+                } else {
+                    // Need to make a transition just for this. This shouldn't really happen
+                    // though because if transition == null, we should be part of an existing one.
+                    getTransitionController().createTransition(TRANSIT_PIP);
+                    mRootWindowContainer.moveActivityToPinnedRootTaskAndRequestStart(r,
+                            "enterPictureInPictureMode");
+                }
                 // Continue the pausing process after entering pip.
                 if (r.isState(PAUSING) && r.mPauseSchedulePendingForPip) {
                     r.getTask().schedulePauseActivity(r, false /* userLeaving */,
diff --git a/services/core/java/com/android/server/wm/AppTransitionController.java b/services/core/java/com/android/server/wm/AppTransitionController.java
index 197bd5a..ab02d49 100644
--- a/services/core/java/com/android/server/wm/AppTransitionController.java
+++ b/services/core/java/com/android/server/wm/AppTransitionController.java
@@ -97,6 +97,7 @@
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 
 /**
@@ -253,10 +254,12 @@
                 getTopApp(mDisplayContent.mChangingContainers, false /* ignoreHidden */);
         final WindowManager.LayoutParams animLp = getAnimLp(animLpActivity);
 
-        // No AE remote animation with Shell transition.
-        // Unfreeze the windows that were previously frozen for TaskFragment animation.
-        unfreezeEmbeddedChangingWindows();
-        overrideWithRemoteAnimationIfSet(animLpActivity, transit, activityTypes);
+        // Check if there is any override
+        if (!overrideWithTaskFragmentRemoteAnimation(transit, activityTypes)) {
+            // Unfreeze the windows that were previously frozen for TaskFragment animation.
+            unfreezeEmbeddedChangingWindows();
+            overrideWithRemoteAnimationIfSet(animLpActivity, transit, activityTypes);
+        }
 
         final boolean voiceInteraction = containsVoiceInteraction(mDisplayContent.mClosingApps)
                 || containsVoiceInteraction(mDisplayContent.mOpeningApps);
@@ -687,6 +690,64 @@
     }
 
     /**
+     * Overrides the pending transition with the remote animation defined by the
+     * {@link ITaskFragmentOrganizer} if all windows in the transition are children of
+     * {@link TaskFragment} that are organized by the same organizer.
+     *
+     * @return {@code true} if the transition is overridden.
+     */
+    private boolean overrideWithTaskFragmentRemoteAnimation(@TransitionOldType int transit,
+            ArraySet<Integer> activityTypes) {
+        if (transitionMayContainNonAppWindows(transit)) {
+            return false;
+        }
+        if (!transitionContainsTaskFragmentWithBoundsOverride()) {
+            // No need to play TaskFragment remote animation if all embedded TaskFragment in the
+            // transition fill the Task.
+            return false;
+        }
+
+        final Task task = findParentTaskForAllEmbeddedWindows();
+        final ITaskFragmentOrganizer organizer = findTaskFragmentOrganizer(task);
+        final RemoteAnimationDefinition definition = organizer != null
+                ? mDisplayContent.mAtmService.mTaskFragmentOrganizerController
+                    .getRemoteAnimationDefinition(organizer)
+                : null;
+        final RemoteAnimationAdapter adapter = definition != null
+                ? definition.getAdapter(transit, activityTypes)
+                : null;
+        if (adapter == null) {
+            return false;
+        }
+        mDisplayContent.mAppTransition.overridePendingAppTransitionRemote(
+                adapter, false /* sync */, true /*isActivityEmbedding*/);
+        ProtoLog.v(WM_DEBUG_APP_TRANSITIONS,
+                "Override with TaskFragment remote animation for transit=%s",
+                AppTransition.appTransitionOldToString(transit));
+
+        final int organizerUid = mDisplayContent.mAtmService.mTaskFragmentOrganizerController
+                .getTaskFragmentOrganizerUid(organizer);
+        final boolean shouldDisableInputForRemoteAnimation = !task.isFullyTrustedEmbedding(
+                organizerUid);
+        final RemoteAnimationController remoteAnimationController =
+                mDisplayContent.mAppTransition.getRemoteAnimationController();
+        if (shouldDisableInputForRemoteAnimation && remoteAnimationController != null) {
+            // We are going to use client-driven animation, Disable all input on activity windows
+            // during the animation (unless it is fully trusted) to ensure it is safe to allow
+            // client to animate the surfaces.
+            // This is needed for all activity windows in the animation Task.
+            remoteAnimationController.setOnRemoteAnimationReady(() -> {
+                final Consumer<ActivityRecord> updateActivities =
+                        activity -> activity.setDropInputForAnimation(true);
+                task.forAllActivities(updateActivities);
+            });
+            ProtoLog.d(WM_DEBUG_APP_TRANSITIONS, "Task=%d contains embedded TaskFragment."
+                    + " Disabled all input during TaskFragment remote animation.", task.mTaskId);
+        }
+        return true;
+    }
+
+    /**
      * Overrides the pending transition with the remote animation defined for the transition in the
      * set of defined remote animations in the app window token.
      */
diff --git a/services/core/java/com/android/server/wm/BackNavigationController.java b/services/core/java/com/android/server/wm/BackNavigationController.java
index c44f838b..28dbc3a 100644
--- a/services/core/java/com/android/server/wm/BackNavigationController.java
+++ b/services/core/java/com/android/server/wm/BackNavigationController.java
@@ -574,10 +574,15 @@
 
     private void scheduleAnimation(@NonNull AnimationHandler.ScheduleAnimationBuilder builder) {
         mPendingAnimation = builder.build();
-        mWindowManagerService.mWindowPlacerLocked.requestTraversal();
-        if (mShowWallpaper) {
-            mWindowManagerService.getDefaultDisplayContentLocked().mWallpaperController
-                    .adjustWallpaperWindows();
+        if (mAnimationHandler.mOpenAnimAdaptor != null
+                && mAnimationHandler.mOpenAnimAdaptor.mPreparedOpenTransition != null) {
+            startAnimation();
+        } else {
+            mWindowManagerService.mWindowPlacerLocked.requestTraversal();
+            if (mShowWallpaper) {
+                mWindowManagerService.getDefaultDisplayContentLocked().mWallpaperController
+                        .adjustWallpaperWindows();
+            }
         }
     }
 
@@ -877,6 +882,7 @@
         } else {
             if (mAnimationHandler.mPrepareCloseTransition != null) {
                 Slog.e(TAG, "Gesture animation is applied on another transition?");
+                return;
             }
             mAnimationHandler.mPrepareCloseTransition = transition;
             if (!migratePredictToTransition) {
diff --git a/services/core/java/com/android/server/wm/BackgroundActivityStartController.java b/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
index 20c5f02..2259b5a 100644
--- a/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
+++ b/services/core/java/com/android/server/wm/BackgroundActivityStartController.java
@@ -21,10 +21,10 @@
 import static android.app.ActivityOptions.BackgroundActivityStartMode;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS;
+import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_COMPAT;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_DENIED;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_SYSTEM_DEFINED;
-import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE;
 import static android.content.Intent.FLAG_ACTIVITY_NEW_TASK;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.os.Build.VERSION_CODES.UPSIDE_DOWN_CAKE;
@@ -39,13 +39,15 @@
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.DEBUG_ACTIVITY_STARTS;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_ATM;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
+import static com.android.server.wm.ActivityTaskManagerService.ACTIVITY_BG_START_GRACE_PERIOD_MS;
 import static com.android.server.wm.ActivityTaskManagerService.APP_SWITCH_ALLOW;
 import static com.android.server.wm.ActivityTaskManagerService.APP_SWITCH_FG_ONLY;
 import static com.android.server.wm.ActivityTaskSupervisor.getApplicationLabel;
 import static com.android.server.wm.PendingRemoteAnimationRegistry.TIMEOUT_MS;
+import static com.android.window.flags.Flags.balAdditionalStartModes;
 import static com.android.window.flags.Flags.balDontBringExistingBackgroundTaskStackToFg;
-import static com.android.window.flags.Flags.balImprovedMetrics;
 import static com.android.window.flags.Flags.balImproveRealCallerVisibilityCheck;
+import static com.android.window.flags.Flags.balImprovedMetrics;
 import static com.android.window.flags.Flags.balRequireOptInByPendingIntentCreator;
 import static com.android.window.flags.Flags.balRequireOptInSameUid;
 import static com.android.window.flags.Flags.balRespectAppSwitchStateWhenCheckBoundByForegroundUid;
@@ -84,6 +86,7 @@
 import com.android.internal.util.Preconditions;
 import com.android.server.UiThread;
 import com.android.server.am.PendingIntentRecord;
+import com.android.server.wm.BackgroundLaunchProcessController.BalCheckConfiguration;
 
 import java.lang.annotation.Retention;
 import java.util.ArrayList;
@@ -107,6 +110,17 @@
     private static final int ASM_GRACEPERIOD_MAX_REPEATS = 5;
     private static final int NO_PROCESS_UID = -1;
 
+    private static final BalCheckConfiguration BAL_CHECK_FOREGROUND = new BalCheckConfiguration(
+            /* isCheckingForFgsStarts */ false,
+            /* checkVisibility */ true,
+            /* checkOtherExemptions */ false,
+            ACTIVITY_BG_START_GRACE_PERIOD_MS);
+    private static final BalCheckConfiguration BAL_CHECK_BACKGROUND = new BalCheckConfiguration(
+            /* isCheckingForFgsStarts */ false,
+            /* checkVisibility */ false,
+            /* checkOtherExemptions */ true,
+            ACTIVITY_BG_START_GRACE_PERIOD_MS);
+
     static final String AUTO_OPT_IN_NOT_PENDING_INTENT = "notPendingIntent";
     static final String AUTO_OPT_IN_CALL_FOR_RESULT = "callForResult";
     static final String AUTO_OPT_IN_SAME_UID = "sameUid";
@@ -412,6 +426,8 @@
                 int callingUid, String callingPackage, ActivityOptions checkedOptions) {
             switch (checkedOptions.getPendingIntentCreatorBackgroundActivityStartMode()) {
                 case MODE_BACKGROUND_ACTIVITY_START_ALLOWED:
+                case MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE:
+                case MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS:
                     return BackgroundStartPrivileges.ALLOW_BAL;
                 case MODE_BACKGROUND_ACTIVITY_START_DENIED:
                     return BackgroundStartPrivileges.NONE;
@@ -752,7 +768,7 @@
         // PendingIntents is null).
         BalVerdict resultForRealCaller = state.callerIsRealCaller() && resultForCaller.allows()
                 ? resultForCaller
-                : checkBackgroundActivityStartAllowedBySender(state)
+                : checkBackgroundActivityStartAllowedByRealCaller(state)
                         .setBasedOnRealCaller();
         state.setResultForRealCaller(resultForRealCaller);
 
@@ -827,6 +843,37 @@
      * or {@link #BAL_BLOCK} if the launch should be blocked
      */
     BalVerdict checkBackgroundActivityStartAllowedByCaller(BalState state) {
+        if (state.isPendingIntent()) {
+            // PendingIntents should mostly be allowed by the sender (real caller) or a permission
+            // the creator of the PendingIntent has. Visibility should be the exceptional case, so
+            // test it last (this does not change the result, just the bal code).
+            BalVerdict result = BalVerdict.BLOCK;
+            if (!(balAdditionalStartModes()
+                    && state.mCheckedOptions.getPendingIntentCreatorBackgroundActivityStartMode()
+                    == MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE)) {
+                result = checkBackgroundActivityStartAllowedByCallerInBackground(state);
+            }
+            if (result == BalVerdict.BLOCK) {
+                result = checkBackgroundActivityStartAllowedByCallerInForeground(state);
+
+            }
+            return result;
+        } else {
+            BalVerdict result = checkBackgroundActivityStartAllowedByCallerInForeground(state);
+            if (result == BalVerdict.BLOCK && !(balAdditionalStartModes()
+                    && state.mCheckedOptions.getPendingIntentCreatorBackgroundActivityStartMode()
+                    == MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE)) {
+                result = checkBackgroundActivityStartAllowedByCallerInBackground(state);
+            }
+            return result;
+        }
+    }
+
+    /**
+     * @return A code denoting which BAL rule allows an activity to be started,
+     * or {@link #BAL_BLOCK} if the launch should be blocked
+     */
+    BalVerdict checkBackgroundActivityStartAllowedByCallerInForeground(BalState state) {
         // This is used to block background activity launch even if the app is still
         // visible to user after user clicking home button.
 
@@ -842,7 +889,16 @@
             return new BalVerdict(BAL_ALLOW_NON_APP_VISIBLE_WINDOW,
                     /*background*/ false, "callingUid has non-app visible window");
         }
+        // Don't abort if the callerApp or other processes of that uid are considered to be in the
+        // foreground.
+        return checkProcessAllowsBal(state.mCallerApp, state, BAL_CHECK_FOREGROUND);
+    }
 
+    /**
+     * @return A code denoting which BAL rule allows an activity to be started,
+     * or {@link #BAL_BLOCK} if the launch should be blocked
+     */
+    BalVerdict checkBackgroundActivityStartAllowedByCallerInBackground(BalState state) {
         // don't abort for the most important UIDs
         final int callingAppId = UserHandle.getAppId(state.mCallingUid);
         if (state.mCallingUid == Process.ROOT_UID
@@ -922,25 +978,29 @@
                     "OP_SYSTEM_EXEMPT_FROM_ACTIVITY_BG_START_RESTRICTION appop is granted");
         }
 
-        // If we don't have callerApp at this point, no caller was provided to startActivity().
-        // That's the case for PendingIntent-based starts, since the creator's process might not be
-        // up and alive.
         // Don't abort if the callerApp or other processes of that uid are allowed in any way.
-        BalVerdict callerAppAllowsBal = checkProcessAllowsBal(state.mCallerApp, state);
-        if (callerAppAllowsBal.allows()) {
-            return callerAppAllowsBal;
-        }
-
-        // If we are here, it means all exemptions based on the creator failed
-        return BalVerdict.BLOCK;
+        return checkProcessAllowsBal(state.mCallerApp, state, BAL_CHECK_BACKGROUND);
     }
 
     /**
      * @return A code denoting which BAL rule allows an activity to be started,
      * or {@link #BAL_BLOCK} if the launch should be blocked
      */
-    BalVerdict checkBackgroundActivityStartAllowedBySender(BalState state) {
+    BalVerdict checkBackgroundActivityStartAllowedByRealCaller(BalState state) {
+        BalVerdict result = checkBackgroundActivityStartAllowedByRealCallerInForeground(state);
+        if (result == BalVerdict.BLOCK && !(balAdditionalStartModes()
+                && state.mCheckedOptions.getPendingIntentBackgroundActivityStartMode()
+                == MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE)) {
+            result = checkBackgroundActivityStartAllowedByRealCallerInBackground(state);
+        }
+        return result;
+    }
 
+    /**
+     * @return A code denoting which BAL rule allows an activity to be started,
+     * or {@link #BAL_BLOCK} if the launch should be blocked
+     */
+    BalVerdict checkBackgroundActivityStartAllowedByRealCallerInForeground(BalState state) {
         // Normal apps with visible app window will be allowed to start activity if app switching
         // is allowed, or apps like live wallpaper with non app visible window will be allowed.
         // The home app can start apps even if app switches are usually disallowed.
@@ -966,6 +1026,16 @@
             }
         }
 
+        // Don't abort if the realCallerApp or other processes of that uid are considered to be in
+        // the foreground.
+        return checkProcessAllowsBal(state.mRealCallerApp, state, BAL_CHECK_FOREGROUND);
+    }
+
+    /**
+     * @return A code denoting which BAL rule allows an activity to be started,
+     * or {@link #BAL_BLOCK} if the launch should be blocked
+     */
+    BalVerdict checkBackgroundActivityStartAllowedByRealCallerInBackground(BalState state) {
         if (state.mCheckedOptions.getPendingIntentBackgroundActivityStartMode()
                 == MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS
                 && hasBalPermission(state.mRealCallingUid, state.mRealCallingPid)) {
@@ -992,14 +1062,7 @@
         }
 
         // don't abort if the callerApp or other processes of that uid are allowed in any way
-        BalVerdict realCallerAppAllowsBal =
-                checkProcessAllowsBal(state.mRealCallerApp, state);
-        if (realCallerAppAllowsBal.allows()) {
-            return realCallerAppAllowsBal;
-        }
-
-        // If we are here, it means all exemptions based on PI sender failed
-        return BalVerdict.BLOCK;
+        return checkProcessAllowsBal(state.mRealCallerApp, state, BAL_CHECK_BACKGROUND);
     }
 
     @VisibleForTesting boolean hasBalPermission(int uid, int pid) {
@@ -1015,13 +1078,13 @@
      * exceptions.
      */
     @VisibleForTesting BalVerdict checkProcessAllowsBal(WindowProcessController app,
-            BalState state) {
+            BalState state, BalCheckConfiguration balCheckConfiguration) {
         if (app == null) {
             return BalVerdict.BLOCK;
         }
         // first check the original calling process
         final BalVerdict balAllowedForCaller = app
-                .areBackgroundActivityStartsAllowed(state.mAppSwitchState);
+                .areBackgroundActivityStartsAllowed(state.mAppSwitchState, balCheckConfiguration);
         if (balAllowedForCaller.allows()) {
             return balAllowedForCaller.withProcessInfo("callerApp process", app);
         } else {
@@ -1033,7 +1096,7 @@
                     final WindowProcessController proc = uidProcesses.valueAt(i);
                     if (proc != app) {
                         BalVerdict balAllowedForUid = proc.areBackgroundActivityStartsAllowed(
-                                state.mAppSwitchState);
+                                state.mAppSwitchState, balCheckConfiguration);
                         if (balAllowedForUid.allows()) {
                             return balAllowedForUid.withProcessInfo("process", proc);
                         }
@@ -1685,6 +1748,21 @@
                             (state.mOriginatingPendingIntent != null));
         }
 
+        if (finalVerdict.getRawCode() == BAL_ALLOW_GRACE_PERIOD) {
+            if (state.realCallerExplicitOptInOrAutoOptIn()
+                    && state.mResultForRealCaller.allows()
+                    && state.mResultForRealCaller.getRawCode() != BAL_ALLOW_GRACE_PERIOD) {
+                // real caller could allow with a different exemption
+            } else if (state.callerExplicitOptInOrAutoOptIn() && state.mResultForCaller.allows()
+                    && state.mResultForCaller.getRawCode() != BAL_ALLOW_GRACE_PERIOD) {
+                // caller could allow with a different exemption
+            } else {
+                // log to determine grace period length distribution
+                Slog.wtf(TAG, "Activity start ONLY allowed by BAL_ALLOW_GRACE_PERIOD "
+                        + finalVerdict.mMessage + ": " + state);
+            }
+        }
+
         if (balImprovedMetrics()) {
             if (shouldLogStats(finalVerdict, state)) {
                 String activityName;
diff --git a/services/core/java/com/android/server/wm/BackgroundLaunchProcessController.java b/services/core/java/com/android/server/wm/BackgroundLaunchProcessController.java
index 4a870a3..1073713 100644
--- a/services/core/java/com/android/server/wm/BackgroundLaunchProcessController.java
+++ b/services/core/java/com/android/server/wm/BackgroundLaunchProcessController.java
@@ -17,7 +17,6 @@
 package com.android.server.wm;
 
 import static com.android.internal.util.Preconditions.checkArgument;
-import static com.android.server.wm.ActivityTaskManagerDebugConfig.DEBUG_ACTIVITY_STARTS;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_ATM;
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.ActivityTaskManagerService.ACTIVITY_BG_START_GRACE_PERIOD_MS;
@@ -48,7 +47,6 @@
 import android.os.UserHandle;
 import android.util.ArrayMap;
 import android.util.IntArray;
-import android.util.Slog;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.server.wm.BackgroundActivityStartController.BalVerdict;
@@ -100,60 +98,75 @@
         mBackgroundActivityStartCallback = callback;
     }
 
+    record BalCheckConfiguration(
+            boolean isCheckingForFgsStart,
+            boolean checkVisibility,
+            boolean checkOtherExemptions,
+            long gracePeriod
+    ) {
+    }
+
+    /**
+     * Check configuration for foreground service starts.
+     *
+     * The check executes all parts of the BAL checks and uses the same grace period,
+     * so FGS is allowed whenever BAL is allowed.
+     */
+    static final BalCheckConfiguration CHECK_FOR_FGS_START = new BalCheckConfiguration(
+            /* isCheckingForFgsStarts */ true,
+            /* checkVisibility */ true,
+            /* checkOtherExemptions */ true,
+            ACTIVITY_BG_START_GRACE_PERIOD_MS);
+
     BalVerdict areBackgroundActivityStartsAllowed(
             int pid, int uid, String packageName,
-            int appSwitchState, boolean isCheckingForFgsStart,
+            int appSwitchState, BalCheckConfiguration checkConfiguration,
             boolean hasActivityInVisibleTask, boolean hasBackgroundActivityStartPrivileges,
             long lastStopAppSwitchesTime, long lastActivityLaunchTime,
             long lastActivityFinishTime) {
         // Allow if the proc is instrumenting with background activity starts privs.
-        if (hasBackgroundActivityStartPrivileges) {
+        if (checkConfiguration.checkOtherExemptions && hasBackgroundActivityStartPrivileges) {
             return new BalVerdict(BAL_ALLOW_PERMISSION, /*background*/ true,
                     "process instrumenting with background activity starts privileges");
         }
         // Allow if the flag was explicitly set.
-        if (isBackgroundStartAllowedByToken(uid, packageName, isCheckingForFgsStart)) {
+        if (checkConfiguration.checkOtherExemptions && isBackgroundStartAllowedByToken(uid,
+                packageName, checkConfiguration.isCheckingForFgsStart)) {
             return new BalVerdict(balImprovedMetrics() ? BAL_ALLOW_TOKEN : BAL_ALLOW_PERMISSION,
                     /*background*/ true, "process allowed by token");
         }
         // Allow if the caller is bound by a UID that's currently foreground.
         // But still respect the appSwitchState.
-        boolean allowBoundByForegroundUid =
+        if (checkConfiguration.checkVisibility && (
                 Flags.balRespectAppSwitchStateWhenCheckBoundByForegroundUid()
-                ? appSwitchState != APP_SWITCH_DISALLOW && isBoundByForegroundUid()
-                : isBoundByForegroundUid();
-        if (allowBoundByForegroundUid) {
+                        ? appSwitchState != APP_SWITCH_DISALLOW && isBoundByForegroundUid()
+                        : isBoundByForegroundUid())) {
             return new BalVerdict(balImprovedMetrics() ? BAL_ALLOW_BOUND_BY_FOREGROUND
                     : BAL_ALLOW_VISIBLE_WINDOW, /*background*/ false,
                     "process bound by foreground uid");
         }
         // Allow if the caller has an activity in any foreground task.
-        if (hasActivityInVisibleTask && appSwitchState != APP_SWITCH_DISALLOW) {
+        if (checkConfiguration.checkVisibility && hasActivityInVisibleTask
+                && appSwitchState != APP_SWITCH_DISALLOW) {
             return new BalVerdict(BAL_ALLOW_FOREGROUND, /*background*/ false,
                     "process has activity in foreground task");
         }
 
         // If app switching is not allowed, we ignore all the start activity grace period
         // exception so apps cannot start itself in onPause() after pressing home button.
-        if (appSwitchState == APP_SWITCH_ALLOW) {
+        if (checkConfiguration.checkOtherExemptions && appSwitchState == APP_SWITCH_ALLOW) {
             // Allow if any activity in the caller has either started or finished very recently, and
             // it must be started or finished after last stop app switches time.
-            final long now = SystemClock.uptimeMillis();
-            if (now - lastActivityLaunchTime < ACTIVITY_BG_START_GRACE_PERIOD_MS
-                    || now - lastActivityFinishTime < ACTIVITY_BG_START_GRACE_PERIOD_MS) {
-                // If activity is started and finished before stop app switch time, we should not
-                // let app to be able to start background activity even it's in grace period.
-                if (lastActivityLaunchTime > lastStopAppSwitchesTime
-                        || lastActivityFinishTime > lastStopAppSwitchesTime) {
+            if (lastActivityLaunchTime > lastStopAppSwitchesTime
+                    || lastActivityFinishTime > lastStopAppSwitchesTime) {
+                final long now = SystemClock.uptimeMillis();
+                long timeSinceLastStartOrFinish = now - Math.max(lastActivityLaunchTime,
+                        lastActivityFinishTime);
+                if (timeSinceLastStartOrFinish < checkConfiguration.gracePeriod) {
                     return new BalVerdict(BAL_ALLOW_GRACE_PERIOD, /*background*/ true,
-                            "within " + ACTIVITY_BG_START_GRACE_PERIOD_MS + "ms grace period");
+                            "within " + checkConfiguration.gracePeriod + "ms grace period ("
+                                    + timeSinceLastStartOrFinish + "ms)");
                 }
-                if (DEBUG_ACTIVITY_STARTS) {
-                    Slog.d(TAG, "[Process(" + pid + ")] Activity start within "
-                            + ACTIVITY_BG_START_GRACE_PERIOD_MS
-                            + "ms grace period but also within stop app switch window");
-                }
-
             }
         }
         return BalVerdict.BLOCK;
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 0597ed7..d495acb 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -235,7 +235,6 @@
 import android.view.Surface.Rotation;
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Transaction;
-import android.view.SurfaceSession;
 import android.view.WindowInsets;
 import android.view.WindowInsets.Type.InsetsType;
 import android.view.WindowManager;
@@ -572,8 +571,6 @@
 
     boolean mWallpaperMayChange = false;
 
-    private final SurfaceSession mSession = new SurfaceSession();
-
     /**
      * A perf hint session which will boost the refresh rate for the display and change sf duration
      * to handle larger workloads.
@@ -1284,7 +1281,7 @@
      * @param transaction as part of which to perform the configuration
      */
     private void configureSurfaces(Transaction transaction) {
-        final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder(mSession)
+        final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder()
                 .setOpaque(true)
                 .setContainerLayer()
                 .setCallsite("DisplayContent");
@@ -1835,7 +1832,7 @@
         if (mTransitionController.useShellTransitionsRotation()) {
             return ROTATION_UNDEFINED;
         }
-        final int activityOrientation = r.getOverrideOrientation();
+        int activityOrientation = r.getOverrideOrientation();
         if (!WindowManagerService.ENABLE_FIXED_ROTATION_TRANSFORM
                 || shouldIgnoreOrientationRequest(activityOrientation)) {
             return ROTATION_UNDEFINED;
@@ -1846,14 +1843,15 @@
                     r /* boundary */, false /* includeBoundary */, true /* traverseTopToBottom */);
             if (nextCandidate != null) {
                 r = nextCandidate;
+                activityOrientation = r.getOverrideOrientation();
             }
         }
-        if (r.inMultiWindowMode() || r.getRequestedConfigurationOrientation(true /* forDisplay */)
-                == getConfiguration().orientation) {
+        if (r.inMultiWindowMode() || r.getRequestedConfigurationOrientation(true /* forDisplay */,
+                activityOrientation) == getConfiguration().orientation) {
             return ROTATION_UNDEFINED;
         }
         final int currentRotation = getRotation();
-        final int rotation = mDisplayRotation.rotationForOrientation(r.getRequestedOrientation(),
+        final int rotation = mDisplayRotation.rotationForOrientation(activityOrientation,
                 currentRotation);
         if (rotation == currentRotation) {
             return ROTATION_UNDEFINED;
@@ -4596,7 +4594,7 @@
         // removed on the task.
         removeImeSurfaceImmediately();
         mImeScreenshot = new ImeScreenshot(
-                mWmService.mSurfaceControlFactory.apply(null), imeTarget);
+                mWmService.mSurfaceControlFactory.get(), imeTarget);
         // If the caller requests to hide IME, then allow to show IME snapshot for any target task.
         // So IME won't look like suddenly disappeared. It usually happens when turning off screen.
         mImeScreenshot.attachAndShow(t, hideImeWindow /* anyTargetTask */);
@@ -5429,14 +5427,8 @@
     }
 
     @Override
-    SurfaceSession getSession() {
-        return mSession;
-    }
-
-    @Override
     SurfaceControl.Builder makeChildSurface(WindowContainer child) {
-        SurfaceSession s = child != null ? child.getSession() : getSession();
-        final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder(s).setContainerLayer();
+        final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder().setContainerLayer();
         if (child == null) {
             return b;
         }
@@ -5452,12 +5444,12 @@
      * and other potpourii.
      */
     SurfaceControl.Builder makeOverlay() {
-        return mWmService.makeSurfaceBuilder(mSession).setParent(getOverlayLayer());
+        return mWmService.makeSurfaceBuilder().setParent(getOverlayLayer());
     }
 
     @Override
     public SurfaceControl.Builder makeAnimationLeash() {
-        return mWmService.makeSurfaceBuilder(mSession).setParent(mSurfaceControl)
+        return mWmService.makeSurfaceBuilder().setParent(mSurfaceControl)
                 .setContainerLayer();
     }
 
diff --git a/services/core/java/com/android/server/wm/DragState.java b/services/core/java/com/android/server/wm/DragState.java
index 59435b8..b09d63f 100644
--- a/services/core/java/com/android/server/wm/DragState.java
+++ b/services/core/java/com/android/server/wm/DragState.java
@@ -175,7 +175,7 @@
      */
     private CompletableFuture<Void> showInputSurface() {
         if (mInputSurface == null) {
-            mInputSurface = mService.makeSurfaceBuilder(mDisplayContent.getSession())
+            mInputSurface = mService.makeSurfaceBuilder()
                     .setContainerLayer()
                     .setName("Drag and Drop Input Consumer")
                     .setCallsite("DragState.showInputSurface")
diff --git a/services/core/java/com/android/server/wm/InputConsumerImpl.java b/services/core/java/com/android/server/wm/InputConsumerImpl.java
index 4204670..a288cc7 100644
--- a/services/core/java/com/android/server/wm/InputConsumerImpl.java
+++ b/services/core/java/com/android/server/wm/InputConsumerImpl.java
@@ -77,8 +77,7 @@
         mWindowHandle.scaleFactor = 1.0f;
         mWindowHandle.inputConfig = InputConfig.NOT_FOCUSABLE;
 
-        mInputSurface = mService.makeSurfaceBuilder(
-                        mService.mRoot.getDisplayContent(displayId).getSession())
+        mInputSurface = mService.makeSurfaceBuilder()
                 .setContainerLayer()
                 .setName("Input Consumer " + name)
                 .setCallsite("InputConsumerImpl")
diff --git a/services/core/java/com/android/server/wm/InputManagerCallback.java b/services/core/java/com/android/server/wm/InputManagerCallback.java
index 1e7de2b..232c3b6 100644
--- a/services/core/java/com/android/server/wm/InputManagerCallback.java
+++ b/services/core/java/com/android/server/wm/InputManagerCallback.java
@@ -281,7 +281,7 @@
                         + " - Input overlay layer is not initialized.");
                 return null;
             }
-            return mService.makeSurfaceBuilder(dc.getSession())
+            return mService.makeSurfaceBuilder()
                     .setContainerLayer()
                     .setName(name)
                     .setCallsite("createSurfaceForGestureMonitor")
diff --git a/services/core/java/com/android/server/wm/OWNERS b/services/core/java/com/android/server/wm/OWNERS
index 781023c..5d6d8bc 100644
--- a/services/core/java/com/android/server/wm/OWNERS
+++ b/services/core/java/com/android/server/wm/OWNERS
@@ -24,6 +24,7 @@
 per-file Background*Start* = set noparent
 per-file Background*Start* = file:/BAL_OWNERS
 per-file Background*Start* = ogunwale@google.com, louischang@google.com
+per-file BackgroundLaunchProcessController.java = file:/BAL_OWNERS
 
 # File related to activity callers
 per-file ActivityCallerState.java = file:/core/java/android/app/COMPONENT_CALLER_OWNERS
diff --git a/services/core/java/com/android/server/wm/RemoteAnimationController.java b/services/core/java/com/android/server/wm/RemoteAnimationController.java
index f8665c7..432089f 100644
--- a/services/core/java/com/android/server/wm/RemoteAnimationController.java
+++ b/services/core/java/com/android/server/wm/RemoteAnimationController.java
@@ -53,6 +53,7 @@
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.util.ArrayList;
+import java.util.function.Consumer;
 
 /**
  * Helper class to run app animations in a remote process.
@@ -348,6 +349,10 @@
             } finally {
                 mIsFinishing = false;
             }
+            // Reset input for all activities when the remote animation is finished.
+            final Consumer<ActivityRecord> updateActivities =
+                    activity -> activity.setDropInputForAnimation(false);
+            mDisplayContent.forAllActivities(updateActivities);
         }
         setRunningRemoteAnimation(false);
         ProtoLog.i(WM_DEBUG_REMOTE_ANIMATIONS, "Finishing remote animation");
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index a6d9659..866dcd5 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2032,33 +2032,39 @@
                 onTop);
     }
 
-    void moveActivityToPinnedRootTask(@NonNull ActivityRecord r,
-            @Nullable ActivityRecord launchIntoPipHostActivity, String reason) {
-        moveActivityToPinnedRootTask(r, launchIntoPipHostActivity, reason, null /* transition */);
+    /** Wrapper/Helper for tests */
+    void moveActivityToPinnedRootTask(@NonNull ActivityRecord r, String reason) {
+        Transition newTransit = (r.mTransitionController.isCollecting()
+                || !r.mTransitionController.isShellTransitionsEnabled())
+                ? null : r.mTransitionController.createTransition(TRANSIT_PIP);
+        moveActivityToPinnedRootTaskInner(r, null /* launchIntoPipHostActivity */, reason,
+                null /* bounds */, newTransit != null);
     }
 
     void moveActivityToPinnedRootTask(@NonNull ActivityRecord r,
             @Nullable ActivityRecord launchIntoPipHostActivity, String reason,
-            @Nullable Transition transition) {
-        moveActivityToPinnedRootTask(r, launchIntoPipHostActivity, reason, transition,
-                null /* bounds */);
+            @Nullable Rect bounds) {
+        moveActivityToPinnedRootTaskInner(r, launchIntoPipHostActivity, reason, bounds,
+                false /* requestStart */);
     }
 
-    void moveActivityToPinnedRootTask(@NonNull ActivityRecord r,
+    /**
+     * Moves activity to pinned in the provided transition and also requests start on that
+     * Transition at an appropriate time.
+     */
+    void moveActivityToPinnedRootTaskAndRequestStart(@NonNull ActivityRecord r, String reason) {
+        moveActivityToPinnedRootTaskInner(r, null /* launchIntoPipHostActivity */, reason,
+                null /* bounds */, true /* requestStart */);
+    }
+
+    private void moveActivityToPinnedRootTaskInner(@NonNull ActivityRecord r,
             @Nullable ActivityRecord launchIntoPipHostActivity, String reason,
-            @Nullable Transition transition, @Nullable Rect bounds) {
+            @Nullable Rect bounds, boolean requestStart) {
         final TaskDisplayArea taskDisplayArea = r.getDisplayArea();
         final Task task = r.getTask();
         final Task rootTask;
 
-        Transition newTransition = transition;
-        // Create a transition now (if not provided) to collect the current pinned Task dismiss.
-        // Only do the create here as the Task (trigger) to enter PIP is not ready yet.
         final TransitionController transitionController = task.mTransitionController;
-        if (newTransition == null && !transitionController.isCollecting()
-                && transitionController.getTransitionPlayer() != null) {
-            newTransition = transitionController.createTransition(TRANSIT_PIP);
-        }
 
         transitionController.deferTransitionReady();
         Transition.ReadyCondition pipChangesApplied = new Transition.ReadyCondition("movedToPip");
@@ -2273,14 +2279,16 @@
             }
         }
 
-        if (newTransition != null) {
+        // can be null (for now) if shell transitions are disabled or inactive at this time
+        final Transition transit = transitionController.getCollectingTransition();
+        if (requestStart && transit != null) {
             // Request at end since we want task-organizer events from ensureActivitiesVisible
             // to be recognized.
-            transitionController.requestStartTransition(newTransition, rootTask,
+            transitionController.requestStartTransition(transit, rootTask,
                     null /* remoteTransition */, null /* displayChange */);
             // A new transition was created just for this operations. Since the operation is
             // complete, mark it as ready.
-            newTransition.setReady(rootTask, true /* ready */);
+            transit.setReady(rootTask, true /* ready */);
         }
 
         resumeFocusedTasksTopActivities();
diff --git a/services/core/java/com/android/server/wm/Session.java b/services/core/java/com/android/server/wm/Session.java
index 7c875c1..2ea2aeb 100644
--- a/services/core/java/com/android/server/wm/Session.java
+++ b/services/core/java/com/android/server/wm/Session.java
@@ -38,7 +38,6 @@
 import static android.view.WindowManager.LayoutParams.isSystemAlertWindowType;
 
 import static com.android.internal.protolog.ProtoLogGroup.WM_DEBUG_IME;
-import static com.android.internal.protolog.ProtoLogGroup.WM_SHOW_TRANSACTIONS;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
@@ -72,7 +71,6 @@
 import android.view.InsetsSourceControl;
 import android.view.InsetsState;
 import android.view.SurfaceControl;
-import android.view.SurfaceSession;
 import android.view.View;
 import android.view.View.FocusDirection;
 import android.view.WindowInsets;
@@ -108,7 +106,6 @@
     @NonNull
     final WindowProcessController mProcess;
     private final String mStringName;
-    SurfaceSession mSurfaceSession;
     private final ArrayList<WindowState> mAddedWindows = new ArrayList<>();
     /** Set of visible alert/app-overlay windows connected to this session. */
     private final ArraySet<WindowState> mAlertWindows = new ArraySet<>();
@@ -719,12 +716,10 @@
             mPackageName = mProcess.mInfo.packageName;
             mRelayoutTag = "relayoutWindow: " + mPackageName;
         }
-        if (mSurfaceSession == null) {
+        if (mProcess.mWindowSession == null) {
             if (DEBUG) {
-                Slog.v(TAG_WM, "First window added to " + this + ", creating SurfaceSession");
+                Slog.v(TAG_WM, "First window added to " + mProcess);
             }
-            mSurfaceSession = new SurfaceSession();
-            ProtoLog.i(WM_SHOW_TRANSACTIONS, "  NEW SURFACE SESSION %s", mSurfaceSession);
             mService.mSessions.add(this);
             if (mLastReportedAnimatorScale != mService.getCurrentAnimatorScale()) {
                 mService.dispatchNewAnimatorScaleLocked(this);
@@ -821,18 +816,11 @@
         }
 
         mService.mSessions.remove(this);
-        if (mSurfaceSession == null) {
+        if (mProcess.mWindowSession == null) {
             return;
         }
 
-        ProtoLog.i(WM_SHOW_TRANSACTIONS, "  KILL SURFACE SESSION %s", mSurfaceSession);
-        try {
-            mSurfaceSession.kill();
-        } catch (Exception e) {
-            Slog.w(TAG_WM, "Exception thrown when killing surface session " + mSurfaceSession
-                    + " in session " + this + ": " + e.toString());
-        }
-        mSurfaceSession = null;
+        mProcess.mWindowSession = null;
         mAddedWindows.clear();
         mAlertWindows.clear();
         setHasOverlayUi(false);
@@ -857,7 +845,6 @@
                 pw.print(" mCanAddInternalSystemWindow="); pw.print(mCanAddInternalSystemWindow);
                 pw.print(" mAlertWindows="); pw.print(mAlertWindows);
                 pw.print(" mClientDead="); pw.print(mClientDead);
-                pw.print(" mSurfaceSession="); pw.println(mSurfaceSession);
         pw.print(prefix); pw.print("mPackageName="); pw.println(mPackageName);
         if (isSatellitePointingUiPackage()) {
             pw.print(prefix); pw.println("mIsSatellitePointingUiPackage=true");
diff --git a/services/core/java/com/android/server/wm/SnapshotController.java b/services/core/java/com/android/server/wm/SnapshotController.java
index 0f9c001..52994c7 100644
--- a/services/core/java/com/android/server/wm/SnapshotController.java
+++ b/services/core/java/com/android/server/wm/SnapshotController.java
@@ -31,6 +31,8 @@
 import android.view.WindowManager;
 import android.window.TaskSnapshot;
 
+import com.android.window.flags.Flags;
+
 import java.io.PrintWriter;
 import java.util.ArrayList;
 
@@ -150,6 +152,9 @@
                 if (mOpenActivities.isEmpty()) {
                     return false;
                 }
+                if (Flags.alwaysCaptureActivitySnapshot()) {
+                    return true;
+                }
                 for (int i = mOpenActivities.size() - 1; i >= 0; --i) {
                     if (!mOpenActivities.get(i).mOptInOnBackInvoked) {
                         return false;
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 21be0fc..8c93b4fe 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3507,7 +3507,7 @@
      * {@link android.window.TaskFragmentOrganizer}
      */
     TaskFragmentParentInfo getTaskFragmentParentInfo() {
-        return new TaskFragmentParentInfo(getConfiguration(), getDisplayId(),
+        return new TaskFragmentParentInfo(getConfiguration(), getDisplayId(), mTaskId,
                 shouldBeVisible(null /* starting */), hasNonFinishingDirectActivity(),
                 getDecorSurface());
     }
diff --git a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
index e4a3176..5aa34d2 100644
--- a/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
+++ b/services/core/java/com/android/server/wm/TaskFragmentOrganizerController.java
@@ -46,6 +46,7 @@
 import android.util.ArraySet;
 import android.util.Slog;
 import android.util.SparseArray;
+import android.view.RemoteAnimationDefinition;
 import android.view.WindowManager;
 import android.window.ITaskFragmentOrganizer;
 import android.window.ITaskFragmentOrganizerController;
@@ -156,6 +157,13 @@
         private final boolean mIsSystemOrganizer;
 
         /**
+         * {@link RemoteAnimationDefinition} for embedded activities transition animation that is
+         * organized by this organizer.
+         */
+        @Nullable
+        private RemoteAnimationDefinition mRemoteAnimationDefinition;
+
+        /**
          * Map from {@link TaskFragmentTransaction#getTransactionToken()} to the
          * {@link Transition#getSyncId()} that has been deferred. {@link TransitionController} will
          * wait until the organizer finished handling the {@link TaskFragmentTransaction}.
@@ -592,6 +600,50 @@
     }
 
     @Override
+    public void registerRemoteAnimations(@NonNull ITaskFragmentOrganizer organizer,
+            @NonNull RemoteAnimationDefinition definition) {
+        final int pid = Binder.getCallingPid();
+        final int uid = Binder.getCallingUid();
+        synchronized (mGlobalLock) {
+            ProtoLog.v(WM_DEBUG_WINDOW_ORGANIZER,
+                    "Register remote animations for organizer=%s uid=%d pid=%d",
+                    organizer.asBinder(), uid, pid);
+            final TaskFragmentOrganizerState organizerState =
+                    mTaskFragmentOrganizerState.get(organizer.asBinder());
+            if (organizerState == null) {
+                throw new IllegalStateException("The organizer hasn't been registered.");
+            }
+            if (organizerState.mRemoteAnimationDefinition != null) {
+                throw new IllegalStateException(
+                        "The organizer has already registered remote animations="
+                                + organizerState.mRemoteAnimationDefinition);
+            }
+
+            definition.setCallingPidUid(pid, uid);
+            organizerState.mRemoteAnimationDefinition = definition;
+        }
+    }
+
+    @Override
+    public void unregisterRemoteAnimations(@NonNull ITaskFragmentOrganizer organizer) {
+        final int pid = Binder.getCallingPid();
+        final long uid = Binder.getCallingUid();
+        synchronized (mGlobalLock) {
+            ProtoLog.v(WM_DEBUG_WINDOW_ORGANIZER,
+                    "Unregister remote animations for organizer=%s uid=%d pid=%d",
+                    organizer.asBinder(), uid, pid);
+            final TaskFragmentOrganizerState organizerState =
+                    mTaskFragmentOrganizerState.get(organizer.asBinder());
+            if (organizerState == null) {
+                Slog.e(TAG, "The organizer hasn't been registered.");
+                return;
+            }
+
+            organizerState.mRemoteAnimationDefinition = null;
+        }
+    }
+
+    @Override
     public void setSavedState(@NonNull ITaskFragmentOrganizer organizer, @Nullable Bundle state) {
         final int pid = Binder.getCallingPid();
         final int uid = Binder.getCallingUid();
@@ -649,6 +701,25 @@
         }
     }
 
+    /**
+     * Gets the {@link RemoteAnimationDefinition} set on the given organizer if exists. Returns
+     * {@code null} if it doesn't.
+     */
+    @Nullable
+    public RemoteAnimationDefinition getRemoteAnimationDefinition(
+            @NonNull ITaskFragmentOrganizer organizer) {
+        synchronized (mGlobalLock) {
+            final TaskFragmentOrganizerState organizerState =
+                    mTaskFragmentOrganizerState.get(organizer.asBinder());
+            if (organizerState == null) {
+                Slog.e(TAG, "TaskFragmentOrganizer has been unregistered or died when trying"
+                        + " to play animation on its organized windows.");
+                return null;
+            }
+            return organizerState.mRemoteAnimationDefinition;
+        }
+    }
+
     int getTaskFragmentOrganizerUid(@NonNull ITaskFragmentOrganizer organizer) {
         final TaskFragmentOrganizerState state = validateAndGetState(organizer);
         return state.mOrganizerUid;
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index e25db7e..7f6dc84 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -181,7 +181,7 @@
     final @TransitionType int mType;
     private int mSyncId = -1;
     private @TransitionFlags int mFlags;
-    private final TransitionController mController;
+    final TransitionController mController;
     private final BLASTSyncEngine mSyncEngine;
     private final Token mToken;
 
@@ -329,6 +329,9 @@
      */
     ArrayList<ActivityRecord> mConfigAtEndActivities = null;
 
+    /** The current head of the chain of actions related to this transition. */
+    ActionChain mChainHead = null;
+
     @VisibleForTesting
     Transition(@TransitionType int type, @TransitionFlags int flags,
             TransitionController controller, BLASTSyncEngine syncEngine) {
@@ -950,10 +953,13 @@
      * Set animation options for collecting transition by ActivityRecord.
      * @param options AnimationOptions captured from ActivityOptions
      */
-    void setOverrideAnimation(@Nullable AnimationOptions options,
+    void setOverrideAnimation(@Nullable AnimationOptions options, @NonNull ActivityRecord r,
             @Nullable IRemoteCallback startCallback, @Nullable IRemoteCallback finishCallback) {
         if (!isCollecting()) return;
         mOverrideOptions = options;
+        if (mOverrideOptions != null) {
+            mOverrideOptions.setUserId(r.mUserId);
+        }
         sendRemoteCallback(mClientAnimationStartCallback);
         mClientAnimationStartCallback = startCallback;
         mClientAnimationFinishCallback = finishCallback;
@@ -1051,9 +1057,8 @@
      * needs to be passed/applied in shell because until finish is called, shell owns the surfaces.
      * Additionally, this gives shell the ability to better deal with merged transitions.
      */
-    private void buildFinishTransaction(SurfaceControl.Transaction t, TransitionInfo info) {
-        // usually only size 1
-        final ArraySet<DisplayContent> displays = new ArraySet<>();
+    private void buildFinishTransaction(SurfaceControl.Transaction t, TransitionInfo info,
+            DisplayContent[] participantDisplays) {
         for (int i = mTargets.size() - 1; i >= 0; --i) {
             final WindowContainer<?> target = mTargets.get(i).mContainer;
             if (target.getParent() == null) continue;
@@ -1065,7 +1070,6 @@
             t.setCornerRadius(targetLeash, 0);
             t.setShadowRadius(targetLeash, 0);
             t.setAlpha(targetLeash, 1);
-            displays.add(target.getDisplayContent());
             // For config-at-end, the end-transform will be reset after the config is actually
             // applied in the client (since the transform depends on config). The other properties
             // remain here because shell might want to persistently override them.
@@ -1079,9 +1083,8 @@
         }
         // Need to update layers on involved displays since they were all paused while
         // the animation played. This puts the layers back into the correct order.
-        for (int i = displays.size() - 1; i >= 0; --i) {
-            if (displays.valueAt(i) == null) continue;
-            assignLayers(displays.valueAt(i), t);
+        for (int i = participantDisplays.length - 1; i >= 0; --i) {
+            assignLayers(participantDisplays[i], t);
         }
 
         for (int i = 0; i < info.getRootCount(); ++i) {
@@ -1207,10 +1210,14 @@
      * The transition has finished animating and is ready to finalize WM state. This should not
      * be called directly; use {@link TransitionController#finishTransition} instead.
      */
-    void finishTransition() {
+    void finishTransition(@NonNull ActionChain chain) {
         if (Trace.isTagEnabled(TRACE_TAG_WINDOW_MANAGER) && mIsPlayerEnabled) {
             asyncTraceEnd(System.identityHashCode(this));
         }
+        if (!chain.isFinishing()) {
+            throw new IllegalStateException("Can't finish on a non-finishing transition "
+                    + chain.mTransition);
+        }
         mLogger.mFinishTimeNs = SystemClock.elapsedRealtimeNanos();
         mController.mLoggerHandler.post(mLogger::logOnFinish);
         mController.mTransitionTracer.logFinishedTransition(this);
@@ -1790,6 +1797,8 @@
         mController.moveToPlaying(this);
 
         // Repopulate the displays based on the resolved targets.
+        final DisplayContent[] participantDisplays = mTargetDisplays.toArray(
+                new DisplayContent[mTargetDisplays.size()]);
         mTargetDisplays.clear();
         for (int i = 0; i < info.getRootCount(); ++i) {
             final DisplayContent dc = mController.mAtm.mRootWindowContainer.getDisplayContent(
@@ -1883,7 +1892,9 @@
                 controller.setupStartTransaction(transaction);
             }
         }
-        buildFinishTransaction(mFinishTransaction, info);
+        // Use participant displays here (rather than just targets) because it's possible for
+        // there to be order changes between non-top tasks in an otherwise no-op transition.
+        buildFinishTransaction(mFinishTransaction, info, participantDisplays);
         mCleanupTransaction = mController.mAtm.mWindowManager.mTransactionFactory.get();
         buildCleanupTransaction(mCleanupTransaction, info);
         if (mController.getTransitionPlayer() != null && mIsPlayerEnabled) {
@@ -2163,7 +2174,7 @@
         if (mFinishTransaction != null) {
             mFinishTransaction.apply();
         }
-        mController.finishTransition(this);
+        mController.finishTransition(mController.mAtm.mChainTracker.startFinish("clean-up", this));
     }
 
     private void cleanUpInternal() {
@@ -2811,7 +2822,7 @@
                 }
             }
             final SurfaceControl rootLeash = leashReference.makeAnimationLeash().setName(
-                    "Transition Root: " + leashReference.getName())
+                            "Transition Root: " + leashReference.getName())
                     .setCallsite("Transition.calculateTransitionRoots").build();
             rootLeash.setUnreleasedWarningCallSite("Transition.calculateTransitionRoots");
             // Update layers to start transaction because we prevent assignment during collect, so
@@ -2982,7 +2993,8 @@
                         // Create the options based on this change's custom animations and layout
                         // parameters
                         animOptions = getOptions(activityRecord /* customAnimActivity */,
-                                                 activityRecord /* animLpActivity */);
+                                activityRecord /* animLpActivity */);
+                        animOptions.setUserId(activityRecord.mUserId);
                         if (!change.hasFlags(FLAG_TRANSLUCENT)) {
                             // If this change is not translucent, its options are going to be
                             // inherited by the changes below
@@ -2992,6 +3004,7 @@
                 } else if (activityRecord != null && animOptionsForActivityTransition != null) {
                     // Use the same options from the top activity for all the activities
                     animOptions = animOptionsForActivityTransition;
+                    animOptions.setUserId(activityRecord.mUserId);
                 } else if (Flags.activityEmbeddingOverlayPresentationFlag()
                         && isEmbeddedTaskFragment) {
                     final TaskFragmentAnimationParams params = taskFragment.getAnimationParams();
@@ -3004,6 +3017,7 @@
                                 params.getOpenAnimationResId(), params.getChangeAnimationResId(),
                                 params.getCloseAnimationResId(), 0 /* backgroundColor */,
                                 false /* overrideTaskTransition */);
+                        animOptions.setUserId(taskFragment.getTask().mUserId);
                     }
                 }
                 if (animOptions != null) {
@@ -3067,6 +3081,9 @@
                     animOptions);
             animOptions = addCustomActivityTransition(customAnimActivity, false /* open */,
                     animOptions);
+            if (animOptions != null) {
+                animOptions.setUserId(customAnimActivity.mUserId);
+            }
         }
 
         // Layout parameters
@@ -3080,10 +3097,12 @@
             // are running an app starting animation, in which case we don't want the app to be
             // able to change its animation directly.
             if (animOptions != null) {
+                animOptions.setUserId(animLpActivity.mUserId);
                 animOptions.addOptionsFromLayoutParameters(animLp);
             } else {
                 animOptions = TransitionInfo.AnimationOptions
                         .makeAnimOptionsFromLayoutParameters(animLp);
+                animOptions.setUserId(animLpActivity.mUserId);
             }
         }
         return animOptions;
@@ -3109,6 +3128,7 @@
             if (animOptions == null) {
                 animOptions = TransitionInfo.AnimationOptions
                         .makeCommonAnimOptions(activity.packageName);
+                animOptions.setUserId(activity.mUserId);
             }
             animOptions.addCustomActivityTransition(open, customAnim.mEnterAnim,
                     customAnim.mExitAnim, customAnim.mBackgroundColor);
@@ -3237,7 +3257,7 @@
         // Remote animations always win, but fullscreen windows override non-fullscreen windows.
         ActivityRecord result = lookForTopWindowWithFilter(sortedTargets,
                 w -> w.getRemoteAnimationDefinition() != null
-                    && w.getRemoteAnimationDefinition().hasTransition(transit, activityTypes));
+                        && w.getRemoteAnimationDefinition().hasTransition(transit, activityTypes));
         if (result != null) {
             return result;
         }
@@ -3284,7 +3304,7 @@
     private void validateKeyguardOcclusion() {
         if ((mFlags & KEYGUARD_VISIBILITY_TRANSIT_FLAGS) != 0) {
             mController.mStateValidators.add(
-                mController.mAtm.mWindowManager.mPolicy::applyKeyguardOcclusionChange);
+                    mController.mAtm.mWindowManager.mPolicy::applyKeyguardOcclusionChange);
         }
     }
 
@@ -3379,6 +3399,11 @@
         return false;
     }
 
+    void recordChain(@NonNull ActionChain chain) {
+        chain.mPrevious = mChainHead;
+        mChainHead = chain;
+    }
+
     @VisibleForTesting
     static class ChangeInfo {
         private static final int FLAG_NONE = 0;
@@ -3888,9 +3913,9 @@
 
         /** @return true if all tracked subtrees are ready. */
         boolean allReady() {
-            ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS, " allReady query: used=%b "
-                    + "override=%b defer=%d states=[%s]", mUsed, mReadyOverride, mDeferReadyDepth,
-                    groupsToString());
+            ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS,
+                    " allReady query: used=%b " + "override=%b defer=%d states=[%s]", mUsed,
+                    mReadyOverride, mDeferReadyDepth, groupsToString());
             // If the readiness has never been touched, mUsed will be false. We never want to
             // consider a transition ready if nothing has been reported on it.
             if (!mUsed) return false;
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index 9bbf102..1d2a605 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -52,8 +52,8 @@
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
-import com.android.internal.protolog.ProtoLogGroup;
 import com.android.internal.protolog.ProtoLog;
+import com.android.internal.protolog.ProtoLogGroup;
 import com.android.server.FgThread;
 import com.android.window.flags.Flags;
 
@@ -880,10 +880,10 @@
     }
 
     /** @see Transition#setOverrideAnimation */
-    void setOverrideAnimation(TransitionInfo.AnimationOptions options,
+    void setOverrideAnimation(TransitionInfo.AnimationOptions options, ActivityRecord r,
             @Nullable IRemoteCallback startCallback, @Nullable IRemoteCallback finishCallback) {
         if (mCollectingTransition == null) return;
-        mCollectingTransition.setOverrideAnimation(options, startCallback, finishCallback);
+        mCollectingTransition.setOverrideAnimation(options, r, startCallback, finishCallback);
     }
 
     void setNoAnimation(WindowContainer wc) {
@@ -921,7 +921,12 @@
     }
 
     /** @see Transition#finishTransition */
-    void finishTransition(Transition record) {
+    void finishTransition(@NonNull ActionChain chain) {
+        if (!chain.isFinishing()) {
+            throw new IllegalStateException("Can't finish on a non-finishing transition "
+                    + chain.mTransition);
+        }
+        final Transition record = chain.mTransition;
         // It is usually a no-op but make sure that the metric consumer is removed.
         mTransitionMetricsReporter.reportAnimationStart(record.getToken(), 0 /* startTime */);
         // It is a no-op if the transition did not change the display.
@@ -937,7 +942,7 @@
             mTrackCount = 0;
         }
         updateRunningRemoteAnimation(record, false /* isPlaying */);
-        record.finishTransition();
+        record.finishTransition(chain);
         for (int i = mAnimatingExitWindows.size() - 1; i >= 0; i--) {
             final WindowState w = mAnimatingExitWindows.get(i);
             if (w.mAnimatingExit && w.mHasSurface && !w.inTransition()) {
diff --git a/services/core/java/com/android/server/wm/TrustedOverlayHost.java b/services/core/java/com/android/server/wm/TrustedOverlayHost.java
index 5f3c558..030e848 100644
--- a/services/core/java/com/android/server/wm/TrustedOverlayHost.java
+++ b/services/core/java/com/android/server/wm/TrustedOverlayHost.java
@@ -51,7 +51,7 @@
 
     void requireOverlaySurfaceControl() {
         if (mSurfaceControl == null) {
-            final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder(null)
+            final SurfaceControl.Builder b = mWmService.makeSurfaceBuilder()
                 .setContainerLayer()
                 .setHidden(true)
                 .setCallsite("TrustedOverlayHost.requireOverlaySurfaceControl")
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 6995027..1eeb3ec 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -100,7 +100,6 @@
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Builder;
 import android.view.SurfaceControlViewHost;
-import android.view.SurfaceSession;
 import android.view.WindowManager;
 import android.view.WindowManager.TransitionOldType;
 import android.view.animation.Animation;
@@ -706,7 +705,7 @@
         mLastSurfacePosition.set(0, 0);
         mLastDeltaRotation = Surface.ROTATION_0;
 
-        final Builder b = mWmService.makeSurfaceBuilder(null)
+        final Builder b = mWmService.makeSurfaceBuilder()
                 .setContainerLayer()
                 .setName(getName());
 
@@ -1731,13 +1730,13 @@
      *         last time {@link #getOrientation(int) was called.
      */
     @Nullable
-    WindowContainer getLastOrientationSource() {
-        final WindowContainer source = mLastOrientationSource;
-        if (source != null && source != this) {
-            final WindowContainer nextSource = source.getLastOrientationSource();
-            if (nextSource != null) {
-                return nextSource;
-            }
+    final WindowContainer<?> getLastOrientationSource() {
+        if (mLastOrientationSource == null) {
+            return null;
+        }
+        WindowContainer<?> source = this;
+        while (source != source.mLastOrientationSource && source.mLastOrientationSource != null) {
+            source = source.mLastOrientationSource;
         }
         return source;
     }
@@ -2662,13 +2661,6 @@
         return true;
     }
 
-    SurfaceSession getSession() {
-        if (getParent() != null) {
-            return getParent().getSession();
-        }
-        return null;
-    }
-
     void assignLayer(Transaction t, int layer) {
         // Don't assign layers while a transition animation is playing
         // TODO(b/173528115): establish robust best-practices around z-order fighting.
diff --git a/services/core/java/com/android/server/wm/WindowManagerConstants.java b/services/core/java/com/android/server/wm/WindowManagerConstants.java
index 1931be4..47c42f4 100644
--- a/services/core/java/com/android/server/wm/WindowManagerConstants.java
+++ b/services/core/java/com/android/server/wm/WindowManagerConstants.java
@@ -34,6 +34,10 @@
  */
 final class WindowManagerConstants {
 
+    /** The orientation of activity will be always "unspecified". */
+    private static final String KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST =
+            "ignore_activity_orientation_request";
+
     /**
      * The minimum duration between gesture exclusion logging for a given window in
      * milliseconds.
@@ -58,6 +62,9 @@
     /** @see AndroidDeviceConfig#KEY_SYSTEM_GESTURES_EXCLUDED_BY_PRE_Q_STICKY_IMMERSIVE */
     boolean mSystemGestureExcludedByPreQStickyImmersive;
 
+    /** @see #KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST */
+    boolean mIgnoreActivityOrientationRequest;
+
     private final WindowManagerGlobalLock mGlobalLock;
     private final Runnable mUpdateSystemGestureExclusionCallback;
     private final DeviceConfigInterface mDeviceConfig;
@@ -89,6 +96,7 @@
         updateSystemGestureExclusionLogDebounceMillis();
         updateSystemGestureExclusionLimitDp();
         updateSystemGestureExcludedByPreQStickyImmersive();
+        updateIgnoreActivityOrientationRequest();
     }
 
     private void onAndroidPropertiesChanged(DeviceConfig.Properties properties) {
@@ -127,6 +135,9 @@
                     case KEY_SYSTEM_GESTURE_EXCLUSION_LOG_DEBOUNCE_MILLIS:
                         updateSystemGestureExclusionLogDebounceMillis();
                         break;
+                    case KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST:
+                        updateIgnoreActivityOrientationRequest();
+                        break;
                     default:
                         break;
                 }
@@ -152,6 +163,12 @@
                 KEY_SYSTEM_GESTURES_EXCLUDED_BY_PRE_Q_STICKY_IMMERSIVE, false);
     }
 
+    private void updateIgnoreActivityOrientationRequest() {
+        mIgnoreActivityOrientationRequest = mDeviceConfig.getBoolean(
+                DeviceConfig.NAMESPACE_WINDOW_MANAGER,
+                KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST, false);
+    }
+
     void dump(PrintWriter pw) {
         pw.println("WINDOW MANAGER CONSTANTS (dumpsys window constants):");
 
@@ -161,6 +178,8 @@
         pw.print("="); pw.println(mSystemGestureExclusionLimitDp);
         pw.print("  "); pw.print(KEY_SYSTEM_GESTURES_EXCLUDED_BY_PRE_Q_STICKY_IMMERSIVE);
         pw.print("="); pw.println(mSystemGestureExcludedByPreQStickyImmersive);
+        pw.print("  "); pw.print(KEY_IGNORE_ACTIVITY_ORIENTATION_REQUEST);
+        pw.print("="); pw.println(mIgnoreActivityOrientationRequest);
         pw.println();
     }
 }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index dae5c8d..6b7ba66 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -288,7 +288,6 @@
 import android.view.Surface;
 import android.view.SurfaceControl;
 import android.view.SurfaceControlViewHost;
-import android.view.SurfaceSession;
 import android.view.View;
 import android.view.View.FocusDirection;
 import android.view.ViewDebug;
@@ -386,7 +385,6 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Function;
 import java.util.function.Supplier;
 
 /** {@hide} */
@@ -1111,7 +1109,7 @@
     static WindowManagerThreadPriorityBooster sThreadPriorityBooster =
             new WindowManagerThreadPriorityBooster();
 
-    Function<SurfaceSession, SurfaceControl.Builder> mSurfaceControlFactory;
+    Supplier<SurfaceControl.Builder> mSurfaceControlFactory;
     Supplier<SurfaceControl.Transaction> mTransactionFactory;
 
     private final SurfaceControl.Transaction mTransaction;
@@ -1202,7 +1200,7 @@
             final boolean showBootMsgs, WindowManagerPolicy policy, ActivityTaskManagerService atm,
             DisplayWindowSettingsProvider displayWindowSettingsProvider,
             Supplier<SurfaceControl.Transaction> transactionFactory,
-            Function<SurfaceSession, SurfaceControl.Builder> surfaceControlFactory) {
+            Supplier<SurfaceControl.Builder> surfaceControlFactory) {
         final WindowManagerService[] wms = new WindowManagerService[1];
         DisplayThread.getHandler().runWithScissors(() ->
                 wms[0] = new WindowManagerService(context, im, showBootMsgs, policy, atm,
@@ -1231,7 +1229,7 @@
             boolean showBootMsgs, WindowManagerPolicy policy, ActivityTaskManagerService atm,
             DisplayWindowSettingsProvider displayWindowSettingsProvider,
             Supplier<SurfaceControl.Transaction> transactionFactory,
-            Function<SurfaceSession, SurfaceControl.Builder> surfaceControlFactory) {
+            Supplier<SurfaceControl.Builder> surfaceControlFactory) {
         installLock(this, INDEX_WINDOW);
         mGlobalLock = atm.getGlobalLock();
         mAtmService = atm;
@@ -7920,7 +7918,7 @@
             }
             boolean allWindowsDrawn = false;
             synchronized (mGlobalLock) {
-                if (displayId == DEFAULT_DISPLAY
+                if (displayId == INVALID_DISPLAY
                         && mRoot.getDefaultDisplay().mDisplayUpdater.waitForTransition(message)) {
                     // Use the ready-to-play of transition as the signal.
                     return;
@@ -8521,7 +8519,7 @@
                     return null;
                 }
                 // TODO(b/210039666): Use a method like add/removeDisplayOverlay if available.
-                return makeSurfaceBuilder(dc.getSession())
+                return makeSurfaceBuilder()
                         .setContainerLayer()
                         .setName("IME Handwriting Surface")
                         .setCallsite("getHandwritingSurfaceForDisplay")
@@ -8829,8 +8827,8 @@
         }
     }
 
-    SurfaceControl.Builder makeSurfaceBuilder(SurfaceSession s) {
-        return mSurfaceControlFactory.apply(s);
+    SurfaceControl.Builder makeSurfaceBuilder() {
+        return mSurfaceControlFactory.get();
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index e1e64ee..476443a 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -223,7 +223,8 @@
         final long ident = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
-                applyTransaction(t, -1 /*syncId*/, null /*transition*/, caller);
+                final ActionChain chain = mService.mChainTracker.startLegacy("applyTransactLegacy");
+                applyTransaction(t, -1 /*syncId*/, chain, caller);
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
@@ -242,7 +243,8 @@
         try {
             synchronized (mGlobalLock) {
                 if (callback == null) {
-                    applyTransaction(t, -1 /* syncId*/, null /*transition*/, caller);
+                    final ActionChain chain = mService.mChainTracker.startLegacy("applySyncLegacy");
+                    applyTransaction(t, -1 /* syncId*/, chain, caller);
                     return -1;
                 }
 
@@ -262,13 +264,15 @@
                 final int syncId = syncGroup.mSyncId;
                 if (mTransitionController.isShellTransitionsEnabled()) {
                     mTransitionController.startLegacySyncOrQueue(syncGroup, (deferred) -> {
-                        applyTransaction(t, syncId, null /* transition */, caller, deferred);
+                        applyTransaction(t, syncId, mService.mChainTracker.startLegacy(
+                                "applySyncLegacy"), caller, deferred);
                         setSyncReady(syncId);
                     });
                 } else {
                     if (!mService.mWindowManager.mSyncEngine.hasActiveSync()) {
                         mService.mWindowManager.mSyncEngine.startSyncSet(syncGroup);
-                        applyTransaction(t, syncId, null /*transition*/, caller);
+                        applyTransaction(t, syncId, mService.mChainTracker.startLegacy(
+                                "applySyncLegacy"), caller);
                         setSyncReady(syncId);
                     } else {
                         // Because the BLAST engine only supports one sync at a time, queue the
@@ -276,7 +280,8 @@
                         mService.mWindowManager.mSyncEngine.queueSyncSet(
                                 () -> mService.mWindowManager.mSyncEngine.startSyncSet(syncGroup),
                                 () -> {
-                                    applyTransaction(t, syncId, null /*transition*/, caller);
+                                    applyTransaction(t, syncId, mService.mChainTracker.startLegacy(
+                                            "applySyncLegacy"), caller);
                                     setSyncReady(syncId);
                                 });
                     }
@@ -313,7 +318,8 @@
                         throw new IllegalArgumentException("Can't use legacy transitions in"
                                 + " compatibility mode with no WCT.");
                     }
-                    applyTransaction(t, -1 /* syncId */, null, caller);
+                    applyTransaction(t, -1 /* syncId */,
+                            mService.mChainTracker.startLegacy("wrongLegacyTransit"), caller);
                     return null;
                 }
                 final WindowContainerTransaction wct =
@@ -334,10 +340,11 @@
                     nextTransition.calcParallelCollectType(wct);
                     mTransitionController.startCollectOrQueue(nextTransition,
                             (deferred) -> {
+                                final ActionChain chain = mService.mChainTracker.start(
+                                        "startNewTransit", nextTransition);
                                 nextTransition.start();
                                 nextTransition.mLogger.mStartWCT = wct;
-                                applyTransaction(wct, -1 /* syncId */, nextTransition, caller,
-                                        deferred);
+                                applyTransaction(wct, -1 /* syncId */, chain, caller, deferred);
                                 wctApplied.meet();
                                 if (needsSetReady) {
                                     setAllReadyIfNeeded(nextTransition, wct);
@@ -351,7 +358,9 @@
                     Slog.e(TAG, "Trying to start a transition that isn't collecting. This probably"
                             + " means Shell took too long to respond to a request. WM State may be"
                             + " incorrect now, please file a bug");
-                    applyTransaction(wct, -1 /*syncId*/, null /*transition*/, caller);
+                    final ActionChain chain = mService.mChainTracker.startFailsafe("startTransit");
+                    chain.mTransition = null;
+                    applyTransaction(wct, -1 /*syncId*/, chain, caller);
                     return transition.getToken();
                 }
                 // Currently, application of wct can span multiple looper loops (ie.
@@ -367,16 +376,20 @@
                 if (transition.shouldApplyOnDisplayThread()) {
                     mService.mH.post(() -> {
                         synchronized (mService.mGlobalLock) {
+                            final ActionChain chain = mService.mChainTracker.start(
+                                    "startTransit", transition);
                             transition.start();
-                            applyTransaction(wct, -1 /* syncId */, transition, caller);
+                            applyTransaction(wct, -1 /* syncId */, chain, caller);
                             if (wctApplied != null) {
                                 wctApplied.meet();
                             }
                         }
                     });
                 } else {
+                    final ActionChain chain = mService.mChainTracker.start("startTransit",
+                            transition);
                     transition.start();
-                    applyTransaction(wct, -1 /* syncId */, transition, caller);
+                    applyTransaction(wct, -1 /* syncId */, chain, caller);
                     if (wctApplied != null) {
                         wctApplied.meet();
                     }
@@ -475,7 +488,8 @@
                 dc.mAppTransition.overridePendingAppTransitionRemote(adapter, true /* sync */,
                         false /* isActivityEmbedding */);
                 syncId = startSyncWithOrganizer(callback);
-                applyTransaction(t, syncId, null /* transition */, caller);
+                applyTransaction(t, syncId, mService.mChainTracker.startLegacy("legacyTransit"),
+                        caller);
                 setSyncReady(syncId);
             }
         } finally {
@@ -493,6 +507,8 @@
         try {
             synchronized (mGlobalLock) {
                 final Transition transition = Transition.fromBinder(transitionToken);
+                final ActionChain chain =
+                        mService.mChainTracker.startFinish("finishTransit", transition);
                 // apply the incoming transaction before finish in case it alters the visibility
                 // of the participants.
                 if (t != null) {
@@ -500,9 +516,9 @@
                     // changes of the transition participants will only set visible-requested
                     // and still let finishTransition handle the participants.
                     mTransitionController.mFinishingTransition = transition;
-                    applyTransaction(t, -1 /* syncId */, null /*transition*/, caller, transition);
+                    applyTransaction(t, -1 /* syncId */, chain, caller);
                 }
-                mTransitionController.finishTransition(transition);
+                mTransitionController.finishTransition(chain);
                 mTransitionController.mFinishingTransition = null;
             }
         } finally {
@@ -537,9 +553,10 @@
         final CallerInfo caller = new CallerInfo();
         final long ident = Binder.clearCallingIdentity();
         try {
-            if (mTransitionController.getTransitionPlayer() == null) {
+            if (!mTransitionController.isShellTransitionsEnabled()) {
                 // No need to worry about transition when Shell transition is not enabled.
-                applyTransaction(wct, -1 /* syncId */, null /* transition */, caller);
+                applyTransaction(wct, -1 /* syncId */,
+                        mService.mChainTracker.startLegacy("legacyTFTransact"), caller);
                 return;
             }
 
@@ -548,8 +565,8 @@
                 // Although there is an active sync, we want to apply the transaction now.
                 // TODO(b/232042367) Redesign the organizer update on activity callback so that we
                 // we will know about the transition explicitly.
-                final Transition transition = mTransitionController.getCollectingTransition();
-                if (transition == null) {
+                final ActionChain chain = mService.mChainTracker.startDefault("tfTransact");
+                if (chain.mTransition == null) {
                     // This should rarely happen, and we should try to avoid using
                     // {@link #applySyncTransaction} with Shell transition.
                     // We still want to apply and merge the transaction to the active sync
@@ -559,7 +576,7 @@
                                     + " because there is an ongoing sync for"
                                     + " applySyncTransaction().");
                 }
-                applyTransaction(wct, -1 /* syncId */, transition, caller);
+                applyTransaction(wct, -1 /* syncId */, chain, caller);
                 return;
             }
 
@@ -570,8 +587,9 @@
                     transition.abort();
                     return;
                 }
-                if (applyTransaction(wct, -1 /* syncId */, transition, caller, deferred)
-                        == TRANSACT_EFFECTS_NONE && transition.mParticipants.isEmpty()) {
+                final ActionChain chain = mService.mChainTracker.start("tfTransact", transition);
+                final int effects = applyTransaction(wct, -1 /* syncId */, chain, caller, deferred);
+                if (effects == TRANSACT_EFFECTS_NONE && transition.mParticipants.isEmpty()) {
                     transition.abort();
                     return;
                 }
@@ -586,15 +604,10 @@
     }
 
     private int applyTransaction(@NonNull WindowContainerTransaction t, int syncId,
-            @Nullable Transition transition, @NonNull CallerInfo caller) {
-        return applyTransaction(t, syncId, transition, caller, null /* finishTransition */);
-    }
-
-    private int applyTransaction(@NonNull WindowContainerTransaction t, int syncId,
-            @Nullable Transition transition, @NonNull CallerInfo caller, boolean deferred) {
+            @NonNull ActionChain chain, @NonNull CallerInfo caller, boolean deferred) {
         if (deferred) {
             try {
-                return applyTransaction(t, syncId, transition, caller);
+                return applyTransaction(t, syncId, chain, caller);
             } catch (RuntimeException e) {
                 // If the transaction is deferred, the caller could be from TransitionController
                 // #tryStartCollectFromQueue that executes on system's worker thread rather than
@@ -604,19 +617,17 @@
             }
             return TRANSACT_EFFECTS_NONE;
         }
-        return applyTransaction(t, syncId, transition, caller);
+        return applyTransaction(t, syncId, chain, caller);
     }
 
     /**
      * @param syncId If non-null, this will be a sync-transaction.
-     * @param transition A transition to collect changes into.
+     * @param chain A lifecycle-chain to acculumate changes into.
      * @param caller Info about the calling process.
-     * @param finishTransition The transition that is currently being finished.
      * @return The effects of the window container transaction.
      */
     private int applyTransaction(@NonNull WindowContainerTransaction t, int syncId,
-            @Nullable Transition transition, @NonNull CallerInfo caller,
-            @Nullable Transition finishTransition) {
+            @NonNull ActionChain chain, @NonNull CallerInfo caller) {
         int effects = TRANSACT_EFFECTS_NONE;
         ProtoLog.v(WM_DEBUG_WINDOW_ORGANIZER, "Apply window transaction, syncId=%d", syncId);
         mService.deferWindowLayout();
@@ -624,20 +635,21 @@
         boolean deferResume = true;
         mService.mTaskSupervisor.setDeferRootVisibilityUpdate(true /* deferUpdate */);
         boolean deferTransitionReady = false;
-        if (transition != null && !t.isEmpty()) {
-            if (transition.isCollecting()) {
+        if (chain.mTransition != null && !t.isEmpty() && !chain.isFinishing()) {
+            if (chain.mTransition.isCollecting()) {
                 deferTransitionReady = true;
-                transition.deferTransitionReady();
+                chain.mTransition.deferTransitionReady();
             } else {
                 Slog.w(TAG, "Transition is not collecting when applyTransaction."
-                        + " transition=" + transition + " state=" + transition.getState());
-                transition = null;
+                        + " transition=" + chain.mTransition + " state="
+                        + chain.mTransition.getState());
+                chain.mTransition = null;
             }
         }
         try {
             final ArraySet<WindowContainer<?>> haveConfigChanges = new ArraySet<>();
-            if (transition != null) {
-                transition.applyDisplayChangeIfNeeded(haveConfigChanges);
+            if (chain.mTransition != null) {
+                chain.mTransition.applyDisplayChangeIfNeeded(haveConfigChanges);
                 if (!haveConfigChanges.isEmpty()) {
                     effects |= TRANSACT_EFFECTS_CLIENT_CONFIG;
                 }
@@ -645,7 +657,7 @@
             final List<WindowContainerTransaction.HierarchyOp> hops = t.getHierarchyOps();
             final int hopSize = hops.size();
             Iterator<Map.Entry<IBinder, WindowContainerTransaction.Change>> entries;
-            if (transition != null) {
+            if (chain.mTransition != null) {
                 // Mark any config-at-end containers before applying config changes so that
                 // the config changes don't dispatch to client.
                 entries = t.getChanges().entrySet().iterator();
@@ -655,7 +667,7 @@
                     if (!entry.getValue().getConfigAtTransitionEnd()) continue;
                     final WindowContainer wc = WindowContainer.fromBinder(entry.getKey());
                     if (wc == null || !wc.isAttached()) continue;
-                    transition.setConfigAtEnd(wc);
+                    chain.mTransition.setConfigAtEnd(wc);
                 }
             }
             entries = t.getChanges().entrySet().iterator();
@@ -672,15 +684,13 @@
                 if (syncId >= 0) {
                     addToSyncSet(syncId, wc);
                 }
-                if (transition != null) transition.collect(wc);
+                chain.collect(wc);
 
                 if ((entry.getValue().getChangeMask()
                         & WindowContainerTransaction.Change.CHANGE_FORCE_NO_PIP) != 0) {
                     // Disable entering pip (eg. when recents pretends to finish itself)
-                    if (finishTransition != null) {
-                        finishTransition.setCanPipOnFinish(false /* canPipOnFinish */);
-                    } else if (transition != null) {
-                        transition.setCanPipOnFinish(false /* canPipOnFinish */);
+                    if (chain.mTransition != null) {
+                        chain.mTransition.setCanPipOnFinish(false /* canPipOnFinish */);
                     }
                 }
                 // A bit hacky, but we need to detect "remove PiP" so that we can "wrap" the
@@ -728,9 +738,9 @@
             if (hopSize > 0) {
                 final boolean isInLockTaskMode = mService.isInLockTaskMode();
                 for (int i = 0; i < hopSize; ++i) {
-                    effects |= applyHierarchyOp(hops.get(i), effects, syncId, transition,
+                    effects |= applyHierarchyOp(hops.get(i), effects, syncId, chain,
                             isInLockTaskMode, caller, t.getErrorCallbackToken(),
-                            t.getTaskFragmentOrganizer(), finishTransition);
+                            t.getTaskFragmentOrganizer());
                 }
             }
             // Queue-up bounds-change transactions for tasks which are now organized. Do
@@ -789,7 +799,7 @@
             }
         } finally {
             if (deferTransitionReady) {
-                transition.continueTransitionReady();
+                chain.mTransition.continueTransitionReady();
             }
             mService.mTaskSupervisor.setDeferRootVisibilityUpdate(false /* deferUpdate */);
             if (deferResume) {
@@ -1079,9 +1089,9 @@
     }
 
     private int applyHierarchyOp(WindowContainerTransaction.HierarchyOp hop, int effects,
-            int syncId, @Nullable Transition transition, boolean isInLockTaskMode,
+            int syncId, @NonNull ActionChain chain, boolean isInLockTaskMode,
             @NonNull CallerInfo caller, @Nullable IBinder errorCallbackToken,
-            @Nullable ITaskFragmentOrganizer organizer, @Nullable Transition finishTransition) {
+            @Nullable ITaskFragmentOrganizer organizer) {
         final int type = hop.getType();
         switch (type) {
             case HIERARCHY_OP_TYPE_REMOVE_TASK: {
@@ -1151,7 +1161,7 @@
                 break;
             }
             case HIERARCHY_OP_TYPE_CHILDREN_TASKS_REPARENT: {
-                effects |= reparentChildrenTasksHierarchyOp(hop, transition, syncId,
+                effects |= reparentChildrenTasksHierarchyOp(hop, chain.mTransition, syncId,
                         isInLockTaskMode);
                 break;
             }
@@ -1204,13 +1214,13 @@
                 if (syncId >= 0) {
                     addToSyncSet(syncId, wc);
                 }
-                if (transition != null) {
-                    transition.collect(wc);
+                if (chain.mTransition != null) {
+                    chain.mTransition.collect(wc);
                     if (hop.isReparent()) {
                         if (wc.getParent() != null) {
                             // Collect the current parent. It's visibility may change as
                             // a result of this reparenting.
-                            transition.collect(wc.getParent());
+                            chain.mTransition.collect(wc.getParent());
                         }
                         if (hop.getNewParent() != null) {
                             final WindowContainer parentWc =
@@ -1219,7 +1229,7 @@
                                 Slog.e(TAG, "Can't resolve parent window from token");
                                 break;
                             }
-                            transition.collect(parentWc);
+                            chain.mTransition.collect(parentWc);
                         }
                     }
                 }
@@ -1233,8 +1243,8 @@
                 break;
             }
             case HIERARCHY_OP_TYPE_ADD_TASK_FRAGMENT_OPERATION: {
-                effects |= applyTaskFragmentOperation(hop, transition, isInLockTaskMode, caller,
-                        errorCallbackToken, organizer);
+                effects |= applyTaskFragmentOperation(hop, chain, isInLockTaskMode,
+                        caller, errorCallbackToken, organizer);
                 break;
             }
             case HIERARCHY_OP_TYPE_PENDING_INTENT: {
@@ -1329,7 +1339,7 @@
                 Rect entryBounds = hop.getBounds();
                 mService.mRootWindowContainer.moveActivityToPinnedRootTask(
                         pipActivity, null /* launchIntoPipHostActivity */,
-                        "moveActivityToPinnedRootTask", null /* transition */, entryBounds);
+                        "moveActivityToPinnedRootTask", entryBounds);
 
                 if (pipActivity.isState(PAUSING) && pipActivity.mPauseSchedulePendingForPip) {
                     // Continue the pausing process. This must be done after moving PiP activity to
@@ -1348,13 +1358,13 @@
                 break;
             }
             case HIERARCHY_OP_TYPE_RESTORE_TRANSIENT_ORDER: {
-                if (finishTransition == null) break;
+                if (!chain.isFinishing()) break;
                 final WindowContainer container = WindowContainer.fromBinder(hop.getContainer());
                 if (container == null) break;
                 final Task thisTask = container.asActivityRecord() != null
                         ? container.asActivityRecord().getTask() : container.asTask();
                 if (thisTask == null) break;
-                final Task restoreAt = finishTransition.getTransientLaunchRestoreTarget(container);
+                final Task restoreAt = chain.mTransition.getTransientLaunchRestoreTarget(container);
                 if (restoreAt == null) break;
                 final TaskDisplayArea taskDisplayArea = thisTask.getTaskDisplayArea();
                 taskDisplayArea.moveRootTaskBehindRootTask(thisTask.getRootTask(), restoreAt);
@@ -1444,7 +1454,7 @@
      *         {@link #TRANSACT_EFFECTS_LIFECYCLE} or {@link #TRANSACT_EFFECTS_CLIENT_CONFIG}.
      */
     private int applyTaskFragmentOperation(@NonNull WindowContainerTransaction.HierarchyOp hop,
-            @Nullable Transition transition, boolean isInLockTaskMode, @NonNull CallerInfo caller,
+            @NonNull ActionChain chain, boolean isInLockTaskMode, @NonNull CallerInfo caller,
             @Nullable IBinder errorCallbackToken, @Nullable ITaskFragmentOrganizer organizer) {
         if (!validateTaskFragmentOperation(hop, errorCallbackToken, organizer)) {
             return TRANSACT_EFFECTS_NONE;
@@ -1467,7 +1477,7 @@
                     break;
                 }
                 createTaskFragment(taskFragmentCreationParams, errorCallbackToken, caller,
-                        transition);
+                        chain.mTransition);
                 break;
             }
             case OP_TYPE_DELETE_TASK_FRAGMENT: {
@@ -1484,7 +1494,7 @@
                         break;
                     }
                 }
-                effects |= deleteTaskFragment(taskFragment, transition);
+                effects |= deleteTaskFragment(taskFragment, chain.mTransition);
                 break;
             }
             case OP_TYPE_START_ACTIVITY_IN_TASK_FRAGMENT: {
@@ -1533,14 +1543,14 @@
                             opType, exception);
                     break;
                 }
-                if (transition != null) {
-                    transition.collect(activity);
+                if (chain.mTransition != null) {
+                    chain.collect(activity);
                     if (activity.getParent() != null) {
                         // Collect the current parent. Its visibility may change as a result of
                         // this reparenting.
-                        transition.collect(activity.getParent());
+                        chain.collect(activity.getParent());
                     }
-                    transition.collect(taskFragment);
+                    chain.collect(taskFragment);
                 }
                 activity.reparent(taskFragment, POSITION_TOP);
                 effects |= TRANSACT_EFFECTS_LIFECYCLE;
@@ -1696,8 +1706,8 @@
                 // If any TaskFragment in the Task is collected by the transition, we make the decor
                 // surface visible in sync with the TaskFragment transition. Otherwise, we make the
                 // decor surface visible immediately.
-                final TaskFragment syncTaskFragment = transition != null
-                        ? task.getTaskFragment(transition.mParticipants::contains)
+                final TaskFragment syncTaskFragment = chain.mTransition != null
+                        ? task.getTaskFragment(chain.mTransition.mParticipants::contains)
                         : null;
 
                 if (syncTaskFragment != null) {
@@ -1749,7 +1759,7 @@
                     // The decor surface boost/unboost must be applied after the transition is
                     // completed. Otherwise, the decor surface could be moved before Shell completes
                     // the transition, causing flicker.
-                    runAfterTransition(transition, task::commitDecorSurfaceBoostedState);
+                    runAfterTransition(chain.mTransition, task::commitDecorSurfaceBoostedState);
                 }
                 break;
             }
diff --git a/services/core/java/com/android/server/wm/WindowProcessController.java b/services/core/java/com/android/server/wm/WindowProcessController.java
index d96ebc6..b6b36c7 100644
--- a/services/core/java/com/android/server/wm/WindowProcessController.java
+++ b/services/core/java/com/android/server/wm/WindowProcessController.java
@@ -89,6 +89,7 @@
 import com.android.server.Watchdog;
 import com.android.server.grammaticalinflection.GrammaticalInflectionManagerInternal;
 import com.android.server.wm.ActivityTaskManagerService.HotPath;
+import com.android.server.wm.BackgroundLaunchProcessController.BalCheckConfiguration;
 
 import java.io.IOException;
 import java.io.PrintWriter;
@@ -695,20 +696,13 @@
     public boolean areBackgroundFgsStartsAllowed() {
         return areBackgroundActivityStartsAllowed(
                 mAtm.getBalAppSwitchesState(),
-                true /* isCheckingForFgsStart */).allows();
+                BackgroundLaunchProcessController.CHECK_FOR_FGS_START).allows();
     }
 
     BackgroundActivityStartController.BalVerdict areBackgroundActivityStartsAllowed(
-            int appSwitchState) {
-        return areBackgroundActivityStartsAllowed(
-                appSwitchState,
-                false /* isCheckingForFgsStart */);
-    }
-
-    private BackgroundActivityStartController.BalVerdict areBackgroundActivityStartsAllowed(
-            int appSwitchState, boolean isCheckingForFgsStart) {
+            int appSwitchState, BalCheckConfiguration checkConfiguration) {
         return mBgLaunchController.areBackgroundActivityStartsAllowed(mPid, mUid,
-                mInfo.packageName, appSwitchState, isCheckingForFgsStart,
+                mInfo.packageName, appSwitchState, checkConfiguration,
                 hasActivityInVisibleTask(), mInstrumentingWithBackgroundActivityStartPrivileges,
                 mAtm.getLastStopAppSwitchesTime(),
                 mLastActivityLaunchTime, mLastActivityFinishTime);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 7c05c29..256d0c6 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -233,7 +233,6 @@
 import android.view.Surface;
 import android.view.Surface.Rotation;
 import android.view.SurfaceControl;
-import android.view.SurfaceSession;
 import android.view.View;
 import android.view.ViewDebug;
 import android.view.ViewTreeObserver;
@@ -5162,15 +5161,6 @@
     }
 
     @Override
-    SurfaceSession getSession() {
-        if (mSession.mSurfaceSession != null) {
-            return mSession.mSurfaceSession;
-        } else {
-            return getParent().getSession();
-        }
-    }
-
-    @Override
     boolean needsZBoost() {
         final InsetsControlTarget target = getDisplayContent().getImeTarget(IME_TARGET_LAYERING);
         if (mIsImWindow && target != null) {
diff --git a/services/core/java/com/android/server/wm/WindowStateAnimator.java b/services/core/java/com/android/server/wm/WindowStateAnimator.java
index b40cf56..82fa9d4 100644
--- a/services/core/java/com/android/server/wm/WindowStateAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowStateAnimator.java
@@ -321,8 +321,7 @@
         }
 
         if (DEBUG_VISIBILITY) {
-            Slog.v(TAG, "Creating surface in session "
-                    + mSession.mSurfaceSession + " window " + this
+            Slog.v(TAG, "Creating surface " + this
                     + " format=" + attrs.format + " flags=" + flags);
         }
 
@@ -358,9 +357,8 @@
             w.mInputWindowHandle.forceChange();
 
             ProtoLog.i(WM_SHOW_SURFACE_ALLOC,
-                        "  CREATE SURFACE %s IN SESSION %s: pid=%d format=%d flags=0x%x / %s",
-                    mSurfaceControl, mSession.mSurfaceSession, mSession.mPid, attrs.format,
-                        flags, this);
+                    "  CREATE SURFACE %s: pid=%d format=%d flags=0x%x / %s",
+                    mSurfaceControl, mSession.mPid, attrs.format, flags, this);
         } catch (OutOfResourcesException e) {
             Slog.w(TAG, "OutOfResourcesException creating surface");
             mService.mRoot.reclaimSomeSurfaceMemory(this, "create", true);
diff --git a/services/core/java/com/android/server/wm/utils/DesktopModeFlagsUtil.java b/services/core/java/com/android/server/wm/utils/DesktopModeFlagsUtil.java
index 70c66de..d33313e 100644
--- a/services/core/java/com/android/server/wm/utils/DesktopModeFlagsUtil.java
+++ b/services/core/java/com/android/server/wm/utils/DesktopModeFlagsUtil.java
@@ -43,7 +43,7 @@
     // All desktop mode related flags to be overridden by developer option toggle will be added here
     DESKTOP_WINDOWING_MODE(
             Flags::enableDesktopWindowingMode, /* shouldOverrideByDevOption= */ true),
-    DYNAMIC_INITIAL_BOUNDS(Flags::enableWindowingDynamicInitialBounds, true);
+    DYNAMIC_INITIAL_BOUNDS(Flags::enableWindowingDynamicInitialBounds, false);
 
     private static final String TAG = "DesktopModeFlagsUtil";
     // Function called to obtain aconfig flag value.
diff --git a/services/core/xsd/display-device-config/display-device-config.xsd b/services/core/xsd/display-device-config/display-device-config.xsd
index 4231149..0eafb59 100644
--- a/services/core/xsd/display-device-config/display-device-config.xsd
+++ b/services/core/xsd/display-device-config/display-device-config.xsd
@@ -42,7 +42,7 @@
                     <xs:annotation name="nonnull"/>
                     <xs:annotation name="final"/>
                 </xs:element>
-                <xs:element type="thermalThrottling" name="thermalThrottling">
+                <xs:element type="thermalThrottling" name="thermalThrottling" minOccurs="0" maxOccurs="1">
                     <xs:annotation name="nonnull"/>
                     <xs:annotation name="final"/>
                 </xs:element>
@@ -464,7 +464,15 @@
             <xs:annotation name="nonnull"/>
             <xs:annotation name="final"/>
         </xs:element>
-        <xs:element name="pollingWindowMillis" type="xs:nonNegativeInteger">
+        <xs:element name="customAnimationRateSec" type="nonNegativeDecimal" minOccurs="0" maxOccurs="1">
+            <xs:annotation name="nonnull"/>
+            <xs:annotation name="final"/>
+        </xs:element>
+        <xs:element name="pollingWindowMaxMillis" type="xs:nonNegativeInteger">
+            <xs:annotation name="nonnull"/>
+            <xs:annotation name="final"/>
+        </xs:element>
+        <xs:element name="pollingWindowMinMillis" type="xs:nonNegativeInteger">
             <xs:annotation name="nonnull"/>
             <xs:annotation name="final"/>
         </xs:element>
diff --git a/services/core/xsd/display-device-config/schema/current.txt b/services/core/xsd/display-device-config/schema/current.txt
index cec2787..355b0ab 100644
--- a/services/core/xsd/display-device-config/schema/current.txt
+++ b/services/core/xsd/display-device-config/schema/current.txt
@@ -345,10 +345,14 @@
   public class PowerThrottlingConfig {
     ctor public PowerThrottlingConfig();
     method @NonNull public final java.math.BigDecimal getBrightnessLowestCapAllowed();
-    method @NonNull public final java.math.BigInteger getPollingWindowMillis();
+    method @NonNull public final java.math.BigDecimal getCustomAnimationRateSec();
+    method @NonNull public final java.math.BigInteger getPollingWindowMaxMillis();
+    method @NonNull public final java.math.BigInteger getPollingWindowMinMillis();
     method public final java.util.List<com.android.server.display.config.PowerThrottlingMap> getPowerThrottlingMap();
     method public final void setBrightnessLowestCapAllowed(@NonNull java.math.BigDecimal);
-    method public final void setPollingWindowMillis(@NonNull java.math.BigInteger);
+    method public final void setCustomAnimationRateSec(@NonNull java.math.BigDecimal);
+    method public final void setPollingWindowMaxMillis(@NonNull java.math.BigInteger);
+    method public final void setPollingWindowMinMillis(@NonNull java.math.BigInteger);
   }
 
   public class PowerThrottlingMap {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/ActiveAdmin.java b/services/devicepolicy/java/com/android/server/devicepolicy/ActiveAdmin.java
index 6a0dd5a..b982098 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/ActiveAdmin.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/ActiveAdmin.java
@@ -1325,31 +1325,6 @@
         pw.print("encryptionRequested=");
         pw.println(encryptionRequested);
 
-        if (!Flags.dumpsysPolicyEngineMigrationEnabled()) {
-            pw.print("disableCamera=");
-            pw.println(disableCamera);
-
-            pw.print("disableScreenCapture=");
-            pw.println(disableScreenCapture);
-
-            pw.print("requireAutoTime=");
-            pw.println(requireAutoTime);
-
-            if (permittedInputMethods != null) {
-                pw.print("permittedInputMethods=");
-                pw.println(permittedInputMethods);
-            }
-
-            pw.println("userRestrictions:");
-            UserRestrictionsUtils.dumpRestrictions(pw, "  ", userRestrictions);
-        }
-
-        if (!Flags.policyEngineMigrationV2Enabled()
-                || !Flags.dumpsysPolicyEngineMigrationEnabled()) {
-            pw.print("mUsbDataSignaling=");
-            pw.println(mUsbDataSignalingEnabled);
-        }
-
         pw.print("disableCallerId=");
         pw.println(disableCallerId);
 
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
index a08af72..4beb6a8 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyEngine.java
@@ -230,11 +230,9 @@
 
         synchronized (mLock) {
             PolicyState<V> localPolicyState = getLocalPolicyStateLocked(policyDefinition, userId);
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                if (!handleAdminPolicySizeLimit(localPolicyState, enforcingAdmin, value,
-                        policyDefinition, userId)) {
-                    return;
-                }
+            if (!handleAdminPolicySizeLimit(localPolicyState, enforcingAdmin, value,
+                    policyDefinition, userId)) {
+                return;
             }
 
             if (policyDefinition.isNonCoexistablePolicy()) {
@@ -354,9 +352,7 @@
             }
             PolicyState<V> localPolicyState = getLocalPolicyStateLocked(policyDefinition, userId);
 
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                decreasePolicySizeForAdmin(localPolicyState, enforcingAdmin);
-            }
+            decreasePolicySizeForAdmin(localPolicyState, enforcingAdmin);
 
             if (policyDefinition.isNonCoexistablePolicy()) {
                 setNonCoexistableLocalPolicyLocked(policyDefinition, localPolicyState,
@@ -500,11 +496,9 @@
 
         synchronized (mLock) {
             PolicyState<V> globalPolicyState = getGlobalPolicyStateLocked(policyDefinition);
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                if (!handleAdminPolicySizeLimit(globalPolicyState, enforcingAdmin, value,
-                        policyDefinition, UserHandle.USER_ALL)) {
-                    return;
-                }
+            if (!handleAdminPolicySizeLimit(globalPolicyState, enforcingAdmin, value,
+                    policyDefinition, UserHandle.USER_ALL)) {
+                return;
             }
             // TODO(b/270999567): Move error handling for DISALLOW_CELLULAR_2G into the code
             //  that honors the restriction once there's an API available
@@ -571,9 +565,7 @@
         synchronized (mLock) {
             PolicyState<V> policyState = getGlobalPolicyStateLocked(policyDefinition);
 
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                decreasePolicySizeForAdmin(policyState, enforcingAdmin);
-            }
+            decreasePolicySizeForAdmin(policyState, enforcingAdmin);
 
             boolean policyChanged = policyState.removePolicy(enforcingAdmin);
 
@@ -1739,25 +1731,23 @@
                 pw.println();
             }
             pw.decreaseIndent();
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                pw.println();
+            pw.println();
 
-                pw.println("Default admin policy size limit: " + DEFAULT_POLICY_SIZE_LIMIT);
-                pw.println("Current admin policy size limit: " + mPolicySizeLimit);
-                pw.println("Admin Policies size: ");
-                for (int i = 0; i < mAdminPolicySize.size(); i++) {
-                    int userId = mAdminPolicySize.keyAt(i);
-                    pw.printf("User %d:\n", userId);
-                    pw.increaseIndent();
-                    for (EnforcingAdmin admin : mAdminPolicySize.get(userId).keySet()) {
-                        pw.printf("Admin : " + admin + " : " + mAdminPolicySize.get(userId).get(
-                                admin));
-                        pw.println();
-                    }
-                    pw.decreaseIndent();
+            pw.println("Default admin policy size limit: " + DEFAULT_POLICY_SIZE_LIMIT);
+            pw.println("Current admin policy size limit: " + mPolicySizeLimit);
+            pw.println("Admin Policies size: ");
+            for (int i = 0; i < mAdminPolicySize.size(); i++) {
+                int userId = mAdminPolicySize.keyAt(i);
+                pw.printf("User %d:\n", userId);
+                pw.increaseIndent();
+                for (EnforcingAdmin admin : mAdminPolicySize.get(userId).keySet()) {
+                    pw.printf("Admin : " + admin + " : " + mAdminPolicySize.get(userId).get(
+                            admin));
+                    pw.println();
                 }
                 pw.decreaseIndent();
             }
+            pw.decreaseIndent();
         }
     }
 
@@ -2018,23 +2008,21 @@
 
         private void writeEnforcingAdminSizeInner(TypedXmlSerializer serializer)
                 throws IOException {
-            if (Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                if (mAdminPolicySize != null) {
-                    for (int i = 0; i < mAdminPolicySize.size(); i++) {
-                        int userId = mAdminPolicySize.keyAt(i);
-                        for (EnforcingAdmin admin : mAdminPolicySize.get(
-                                userId).keySet()) {
-                            serializer.startTag(/* namespace= */ null,
-                                    TAG_ENFORCING_ADMIN_AND_SIZE);
-                            serializer.startTag(/* namespace= */ null, TAG_ENFORCING_ADMIN);
-                            admin.saveToXml(serializer);
-                            serializer.endTag(/* namespace= */ null, TAG_ENFORCING_ADMIN);
-                            serializer.startTag(/* namespace= */ null, TAG_POLICY_SUM_SIZE);
-                            serializer.attributeInt(/* namespace= */ null, ATTR_POLICY_SUM_SIZE,
-                                    mAdminPolicySize.get(userId).get(admin));
-                            serializer.endTag(/* namespace= */ null, TAG_POLICY_SUM_SIZE);
-                            serializer.endTag(/* namespace= */ null, TAG_ENFORCING_ADMIN_AND_SIZE);
-                        }
+            if (mAdminPolicySize != null) {
+                for (int i = 0; i < mAdminPolicySize.size(); i++) {
+                    int userId = mAdminPolicySize.keyAt(i);
+                    for (EnforcingAdmin admin : mAdminPolicySize.get(
+                            userId).keySet()) {
+                        serializer.startTag(/* namespace= */ null,
+                                TAG_ENFORCING_ADMIN_AND_SIZE);
+                        serializer.startTag(/* namespace= */ null, TAG_ENFORCING_ADMIN);
+                        admin.saveToXml(serializer);
+                        serializer.endTag(/* namespace= */ null, TAG_ENFORCING_ADMIN);
+                        serializer.startTag(/* namespace= */ null, TAG_POLICY_SUM_SIZE);
+                        serializer.attributeInt(/* namespace= */ null, ATTR_POLICY_SUM_SIZE,
+                                mAdminPolicySize.get(userId).get(admin));
+                        serializer.endTag(/* namespace= */ null, TAG_POLICY_SUM_SIZE);
+                        serializer.endTag(/* namespace= */ null, TAG_ENFORCING_ADMIN_AND_SIZE);
                     }
                 }
             }
@@ -2042,9 +2030,6 @@
 
         private void writeMaxPolicySizeInner(TypedXmlSerializer serializer)
                 throws IOException {
-            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                return;
-            }
             serializer.startTag(/* namespace= */ null, TAG_MAX_POLICY_SIZE_LIMIT);
             serializer.attributeInt(
                     /* namespace= */ null, ATTR_POLICY_SUM_SIZE, mPolicySizeLimit);
@@ -2192,9 +2177,6 @@
 
         private void readMaxPolicySizeInner(TypedXmlPullParser parser)
                 throws XmlPullParserException, IOException {
-            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                return;
-            }
             mPolicySizeLimit = parser.getAttributeInt(/* namespace= */ null, ATTR_POLICY_SUM_SIZE);
         }
     }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index a80ee0f..fc61967 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -871,6 +871,16 @@
                 EXEMPT_FROM_POWER_RESTRICTIONS, OPSTR_SYSTEM_EXEMPT_FROM_POWER_RESTRICTIONS);
     }
 
+    private static final Set<String> METERED_DATA_RESTRICTION_EXEMPT_ROLES =
+            new ArraySet<>();
+    static {
+        // TODO(b/362545319): reference role name from role manager once it's exposed.
+        final String roleDeviceLockController =
+                "android.app.role.SYSTEM_FINANCED_DEVICE_CONTROLLER";
+        METERED_DATA_RESTRICTION_EXEMPT_ROLES.add(roleDeviceLockController);
+        METERED_DATA_RESTRICTION_EXEMPT_ROLES.add(RoleManager.ROLE_FINANCED_DEVICE_KIOSK);
+    }
+
     /**
      * Admin apps targeting Android S+ may not use
      * {@link android.app.admin.DevicePolicyManager#setPasswordQuality} to set password quality
@@ -1318,9 +1328,7 @@
                 Bundle prevRestrictions) {
             resetCrossProfileIntentFiltersIfNeeded(userId, newRestrictions, prevRestrictions);
             resetUserVpnIfNeeded(userId, newRestrictions, prevRestrictions);
-            if (Flags.deletePrivateSpaceUnderRestriction()) {
-                removePrivateSpaceIfRestrictionIsSet(userId, newRestrictions, prevRestrictions);
-            }
+            removePrivateSpaceIfRestrictionIsSet(userId, newRestrictions, prevRestrictions);
         }
 
         private void resetUserVpnIfNeeded(
@@ -1959,6 +1967,10 @@
             return UserManager.isHeadlessSystemUserMode();
         }
 
+        List<String> roleManagerGetRoleHoldersAsUser(String role, UserHandle userHandle) {
+            return getRoleManager().getRoleHoldersAsUser(role, userHandle);
+        }
+
         @SuppressWarnings("AndroidFrameworkPendingIntentMutability")
         PendingIntent pendingIntentGetActivityAsUser(Context context, int requestCode,
                 @NonNull Intent intent, int flags, Bundle options, UserHandle user) {
@@ -3681,9 +3693,6 @@
             }
 
             revertTransferOwnershipIfNecessaryLocked();
-            if (!Flags.policyEngineMigrationV2Enabled()) {
-                updateUsbDataSignal(mContext, isUsbDataSignalingEnabledInternalLocked());
-            }
         }
 
         // Check whether work apps were paused via suspension and unsuspend if necessary.
@@ -9363,8 +9372,7 @@
 
     void sendDeviceOwnerOrProfileOwnerCommand(String action, Bundle extras, int userId) {
         if (userId == UserHandle.USER_ALL) {
-            if (Flags.headlessDeviceOwnerDelegateSecurityLoggingBugFix()
-                    && getHeadlessDeviceOwnerModeForDeviceOwner()
+            if (getHeadlessDeviceOwnerModeForDeviceOwner()
                     == HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER) {
                 userId = mOwners.getDeviceOwnerUserId();
             } else {
@@ -11479,10 +11487,8 @@
                 pw.println();
                 mStatLogger.dump(pw);
                 pw.println();
-                if (Flags.dumpsysPolicyEngineMigrationEnabled()) {
-                    mDevicePolicyEngine.dump(pw);
-                    pw.println();
-                }
+                mDevicePolicyEngine.dump(pw);
+                pw.println();
                 pw.println("Encryption Status: " + getEncryptionStatusName(getEncryptionStatus()));
                 pw.println("Logout user: " + getLogoutUserIdUnchecked());
                 pw.println();
@@ -12440,12 +12446,6 @@
         }
 
         if (packageList != null) {
-            if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-                for (String pkg : packageList) {
-                    PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
-                }
-            }
-
             List<InputMethodInfo> enabledImes = mInjector.binderWithCleanCallingIdentity(() ->
                     InputMethodManagerInternal.get().getEnabledInputMethodListAsUser(userId));
             if (enabledImes != null) {
@@ -12682,14 +12682,12 @@
         Preconditions.checkCallAuthorization(isDefaultDeviceOwner(caller));
         checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_CREATE_AND_MANAGE_USER);
 
-        if (Flags.headlessDeviceOwnerSingleUserEnabled()) {
-            // Block this method if the device is in headless main user mode
-            Preconditions.checkCallAuthorization(
-                    !mInjector.userManagerIsHeadlessSystemUserMode()
-                            || getHeadlessDeviceOwnerModeForDeviceOwner()
-                            != HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER,
-                    "createAndManageUser was called while in headless single user mode");
-        }
+        // Block this method if the device is in headless main user mode
+        Preconditions.checkCallAuthorization(
+                !mInjector.userManagerIsHeadlessSystemUserMode()
+                        || getHeadlessDeviceOwnerModeForDeviceOwner()
+                        != HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER,
+                "createAndManageUser was called while in headless single user mode");
 
         // Only allow the system user to use this method
         Preconditions.checkCallAuthorization(caller.getUserHandle().isSystem(),
@@ -13976,11 +13974,9 @@
                     UserManager.DISALLOW_THREAD_NETWORK,
                     new String[]{MANAGE_DEVICE_POLICY_THREAD_NETWORK});
         }
-        if (Flags.assistContentUserRestrictionEnabled()) {
-            USER_RESTRICTION_PERMISSIONS.put(
-                    UserManager.DISALLOW_ASSIST_CONTENT,
-                    new String[]{MANAGE_DEVICE_POLICY_ASSIST_CONTENT});
-        }
+        USER_RESTRICTION_PERMISSIONS.put(
+                UserManager.DISALLOW_ASSIST_CONTENT,
+                new String[]{MANAGE_DEVICE_POLICY_ASSIST_CONTENT});
         USER_RESTRICTION_PERMISSIONS.put(
                 UserManager.DISALLOW_ULTRA_WIDEBAND_RADIO, new String[]{MANAGE_DEVICE_POLICY_NEARBY_COMMUNICATION});
         USER_RESTRICTION_PERMISSIONS.put(
@@ -14312,10 +14308,6 @@
             return;
         }
 
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            PolicySizeVerifier.enforceMaxStringLength(accountType, "account type");
-        }
-
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
         synchronized (getLockObject()) {
             int affectedUser = getAffectedUser(parent);
@@ -14926,11 +14918,6 @@
     public void setLockTaskPackages(ComponentName who, String callerPackageName, String[] packages)
             throws SecurityException {
         Objects.requireNonNull(packages, "packages is null");
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            for (String pkg : packages) {
-                PolicySizeVerifier.enforceMaxPackageNameLength(pkg);
-            }
-        }
 
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
         checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_SET_LOCK_TASK_PACKAGES);
@@ -15211,7 +15198,7 @@
         final CallerIdentity caller = getCallerIdentity(who);
         Preconditions.checkCallAuthorization(
                 isProfileOwner(caller) || isDefaultDeviceOwner(caller));
-        if (Flags.allowScreenBrightnessControlOnCope() && parent) {
+        if (parent) {
             Preconditions.checkCallAuthorization(isProfileOwnerOfOrganizationOwnedDevice(caller));
         }
         checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_SET_SYSTEM_SETTING);
@@ -15222,7 +15209,7 @@
                         "Permission denial: device owners cannot update %1$s", setting));
             }
             int affectedUser;
-            if (Flags.allowScreenBrightnessControlOnCope() && parent) {
+            if (parent) {
                 affectedUser = getProfileParentId(caller.getUserId());
             } else {
                 affectedUser = caller.getUserId();
@@ -16814,13 +16801,11 @@
                     mContext.sendBroadcastAsUser(intent, UserHandle.of(userId));
                 }
 
-                if (Flags.permissionMigrationForZeroTrustImplEnabled()) {
-                    final UserHandle user = UserHandle.of(userId);
-                    final String roleHolderPackage = getRoleHolderPackageNameOnUser(
-                            RoleManager.ROLE_DEVICE_POLICY_MANAGEMENT, userId);
-                    if (roleHolderPackage != null) {
-                        broadcastExplicitIntentToPackage(intent, roleHolderPackage, user);
-                    }
+                final UserHandle user = UserHandle.of(userId);
+                final String roleHolderPackage = getRoleHolderPackageNameOnUser(
+                        RoleManager.ROLE_DEVICE_POLICY_MANAGEMENT, userId);
+                if (roleHolderPackage != null) {
+                    broadcastExplicitIntentToPackage(intent, roleHolderPackage, user);
                 }
             }
         });
@@ -16828,18 +16813,10 @@
 
     @Override
     public SystemUpdateInfo getPendingSystemUpdate(ComponentName admin, String callerPackage) {
-        if (Flags.permissionMigrationForZeroTrustImplEnabled()) {
-            CallerIdentity caller = getCallerIdentity(admin, callerPackage);
-            enforcePermissions(new String[] {NOTIFY_PENDING_SYSTEM_UPDATE,
-                    MANAGE_DEVICE_POLICY_QUERY_SYSTEM_UPDATES}, caller.getPackageName(),
-                    caller.getUserId());
-        } else {
-            Objects.requireNonNull(admin, "ComponentName is null");
-
-            final CallerIdentity caller = getCallerIdentity(admin);
-            Preconditions.checkCallAuthorization(
-                    isDefaultDeviceOwner(caller) || isProfileOwner(caller));
-        }
+        CallerIdentity caller = getCallerIdentity(admin, callerPackage);
+        enforcePermissions(new String[] {NOTIFY_PENDING_SYSTEM_UPDATE,
+                MANAGE_DEVICE_POLICY_QUERY_SYSTEM_UPDATES}, caller.getPackageName(),
+                caller.getUserId());
         return mOwners.getSystemUpdateInfo();
     }
 
@@ -17315,7 +17292,7 @@
                     return STATUS_HEADLESS_SYSTEM_USER_MODE_NOT_SUPPORTED;
                 }
 
-                if (Flags.headlessDeviceOwnerSingleUserEnabled() && isHeadlessModeSingleUser) {
+                if (isHeadlessModeSingleUser) {
                     ensureSetUpUser = mUserManagerInternal.getMainUserId();
                     if (ensureSetUpUser == UserHandle.USER_NULL) {
                         return STATUS_HEADLESS_ONLY_SYSTEM_USER;
@@ -17383,17 +17360,10 @@
             @Nullable ComponentName componentName, @UserIdInt int callingUserId) {
         synchronized (getLockObject()) {
             int deviceOwnerUserId = -1;
-            if (Flags.headlessDeviceOwnerProvisioningFixEnabled()) {
-                deviceOwnerUserId = mInjector.userManagerIsHeadlessSystemUserMode()
-                        && getHeadlessDeviceOwnerModeForDeviceAdmin(componentName, callingUserId)
-                        == HEADLESS_DEVICE_OWNER_MODE_AFFILIATED
-                        ? UserHandle.USER_SYSTEM : callingUserId;
-            } else {
-                deviceOwnerUserId = mInjector.userManagerIsHeadlessSystemUserMode()
-                        && getHeadlessDeviceOwnerModeForDeviceOwner()
-                        == HEADLESS_DEVICE_OWNER_MODE_AFFILIATED
-                        ? UserHandle.USER_SYSTEM : callingUserId;
-            }
+            deviceOwnerUserId = mInjector.userManagerIsHeadlessSystemUserMode()
+                    && getHeadlessDeviceOwnerModeForDeviceAdmin(componentName, callingUserId)
+                    == HEADLESS_DEVICE_OWNER_MODE_AFFILIATED
+                    ? UserHandle.USER_SYSTEM : callingUserId;
             Slogf.i(LOG_TAG, "Calling user %d, device owner will be set on user %d",
                     callingUserId, deviceOwnerUserId);
             // hasIncompatibleAccountsOrNonAdb doesn't matter since the caller is not adb.
@@ -17906,15 +17876,28 @@
         });
     }
 
+    private Set<String> getMeteredDataRestrictionExemptPackages(int userId) {
+        final Set<String> exemptPkgs = new ArraySet<>();
+        for (String role: METERED_DATA_RESTRICTION_EXEMPT_ROLES) {
+            String pkg = getRoleHolderPackageNameOnUser(role, userId);
+            if (pkg != null) {
+                exemptPkgs.add(pkg);
+            }
+        }
+
+        return exemptPkgs;
+    }
+
     private List<String> removeInvalidPkgsForMeteredDataRestriction(
             int userId, List<String> pkgNames) {
+        final Set<String> exemptRolePkgs = getMeteredDataRestrictionExemptPackages(userId);
         synchronized (getLockObject()) {
             final Set<String> activeAdmins = getActiveAdminPackagesLocked(userId);
             final List<String> excludedPkgs = new ArrayList<>();
             for (int i = pkgNames.size() - 1; i >= 0; --i) {
                 final String pkgName = pkgNames.get(i);
-                // If the package is an active admin, don't restrict it.
-                if (activeAdmins.contains(pkgName)) {
+                // If the package is an active admin or exempt role, don't restrict it.
+                if (activeAdmins.contains(pkgName) || exemptRolePkgs.contains(pkgName)) {
                     excludedPkgs.add(pkgName);
                     continue;
                 }
@@ -19759,16 +19742,14 @@
     }
 
     private void transferSubscriptionOwnership(ComponentName admin, ComponentName target) {
-        if (Flags.esimManagementEnabled()) {
-            SubscriptionManager subscriptionManager = mContext.getSystemService(
-                    SubscriptionManager.class);
-            for (int subId : getSubscriptionIdsInternal(admin.getPackageName()).toArray()) {
-                try {
-                    subscriptionManager.setGroupOwner(subId, target.getPackageName());
-                } catch (Exception e) {
-                    // Shouldn't happen.
-                    Slogf.e(LOG_TAG, e, "Error setting group owner for subId: " + subId);
-                }
+        SubscriptionManager subscriptionManager = mContext.getSystemService(
+                SubscriptionManager.class);
+        for (int subId : getSubscriptionIdsInternal(admin.getPackageName()).toArray()) {
+            try {
+                subscriptionManager.setGroupOwner(subId, target.getPackageName());
+            } catch (Exception e) {
+                // Shouldn't happen.
+                Slogf.e(LOG_TAG, e, "Error setting group owner for subId: " + subId);
             }
         }
     }
@@ -20666,9 +20647,7 @@
                     // have OP_RUN_ANY_IN_BACKGROUND app op and won't execute in the background. The
                     // code below grants that app op, and once the exemption is in place, the user
                     // won't be able to disable background usage anymore.
-                    if (Flags.powerExemptionBgUsageFix()
-                            && exemption == EXEMPT_FROM_POWER_RESTRICTIONS
-                            && newMode == MODE_ALLOWED) {
+                    if (exemption == EXEMPT_FROM_POWER_RESTRICTIONS && newMode == MODE_ALLOWED) {
                         setBgUsageAppOp(appOpsMgr, appInfo);
                     }
                 }
@@ -21425,13 +21404,7 @@
 
         final CallerIdentity caller = getCallerIdentity(callerPackage);
 
-        if (Flags.permissionMigrationForZeroTrustImplEnabled()) {
-            enforcePermission(MANAGE_DEVICE_POLICY_CERTIFICATES, caller.getPackageName());
-        } else {
-            Preconditions.checkCallAuthorization(
-                    isDefaultDeviceOwner(caller) || isProfileOwner(caller)
-                            || isCallerDelegate(caller, DELEGATION_CERT_INSTALL));
-        }
+        enforcePermission(MANAGE_DEVICE_POLICY_CERTIFICATES, caller.getPackageName());
         synchronized (getLockObject()) {
             final ActiveAdmin requiredAdmin = getDeviceOrProfileOwnerAdminLocked(
                     caller.getUserId());
@@ -21759,8 +21732,6 @@
      */
     @Nullable
     private String getRoleHolderPackageNameOnUser(String role, int userId) {
-        RoleManager roleManager = mContext.getSystemService(RoleManager.class);
-
         // Clear calling identity as the RoleManager APIs require privileged permissions.
         return mInjector.binderWithCleanCallingIdentity(() -> {
             List<UserInfo> users;
@@ -21772,7 +21743,7 @@
             }
             for (UserInfo user : users) {
                 List<String> roleHolders =
-                        roleManager.getRoleHoldersAsUser(role, user.getUserHandle());
+                        mInjector.roleManagerGetRoleHoldersAsUser(role, user.getUserHandle());
                 if (!roleHolders.isEmpty()) {
                     return roleHolders.get(0);
                 }
@@ -22032,16 +22003,9 @@
         final long identity = Binder.clearCallingIdentity();
         try {
             boolean isSingleUserMode;
-            if (Flags.headlessDeviceOwnerProvisioningFixEnabled()) {
-                int headlessDeviceOwnerMode = getHeadlessDeviceOwnerModeForDeviceAdmin(
-                        deviceAdmin, caller.getUserId());
-                isSingleUserMode =
-                        headlessDeviceOwnerMode == HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER;
-            } else {
-                isSingleUserMode =
-                        getHeadlessDeviceOwnerModeForDeviceOwner()
-                                == HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER;
-            }
+            int headlessDeviceOwnerMode = getHeadlessDeviceOwnerModeForDeviceAdmin(
+                    deviceAdmin, caller.getUserId());
+            isSingleUserMode = headlessDeviceOwnerMode == HEADLESS_DEVICE_OWNER_MODE_SINGLE_USER;
 
             if (Flags.headlessSingleMinTargetSdk()
                     && mInjector.userManagerIsHeadlessSystemUserMode()
@@ -22065,8 +22029,8 @@
             setTimeAndTimezone(provisioningParams.getTimeZone(), provisioningParams.getLocalTime());
             setLocale(provisioningParams.getLocale());
 
-            int deviceOwnerUserId = Flags.headlessDeviceOwnerSingleUserEnabled()
-                    && isSingleUserMode && mInjector.userManagerIsHeadlessSystemUserMode()
+            int deviceOwnerUserId =
+                    isSingleUserMode && mInjector.userManagerIsHeadlessSystemUserMode()
                     ? mUserManagerInternal.getMainUserId() : UserHandle.USER_SYSTEM;
 
             if (!removeNonRequiredAppsForManagedDevice(
@@ -22440,35 +22404,17 @@
         Objects.requireNonNull(packageName, "Admin package name must be provided");
         final CallerIdentity caller = getCallerIdentity(packageName);
 
-        if (!Flags.policyEngineMigrationV2Enabled()) {
-            Preconditions.checkCallAuthorization(
-                    isDefaultDeviceOwner(caller) || isProfileOwnerOfOrganizationOwnedDevice(caller),
-                    "USB data signaling can only be controlled by a device owner or "
-                            + "a profile owner on an organization-owned device.");
+        synchronized (getLockObject()) {
+            EnforcingAdmin enforcingAdmin = enforcePermissionAndGetEnforcingAdmin(
+                    /* admin= */ null, MANAGE_DEVICE_POLICY_USB_DATA_SIGNALLING,
+                    caller.getPackageName(),
+                    caller.getUserId());
             Preconditions.checkState(canUsbDataSignalingBeDisabled(),
                     "USB data signaling cannot be disabled.");
-        }
-
-        synchronized (getLockObject()) {
-            if (Flags.policyEngineMigrationV2Enabled()) {
-                EnforcingAdmin enforcingAdmin = enforcePermissionAndGetEnforcingAdmin(
-                        /* admin= */ null, MANAGE_DEVICE_POLICY_USB_DATA_SIGNALLING,
-                        caller.getPackageName(),
-                        caller.getUserId());
-                Preconditions.checkState(canUsbDataSignalingBeDisabled(),
-                        "USB data signaling cannot be disabled.");
-                mDevicePolicyEngine.setGlobalPolicy(
-                        PolicyDefinition.USB_DATA_SIGNALING,
-                        enforcingAdmin,
-                        new BooleanPolicyValue(enabled));
-            } else {
-                ActiveAdmin admin = getProfileOwnerOrDeviceOwnerLocked(caller.getUserId());
-                if (admin.mUsbDataSignalingEnabled != enabled) {
-                    admin.mUsbDataSignalingEnabled = enabled;
-                    saveSettingsLocked(caller.getUserId());
-                    updateUsbDataSignal(mContext, isUsbDataSignalingEnabledInternalLocked());
-                }
-            }
+            mDevicePolicyEngine.setGlobalPolicy(
+                    PolicyDefinition.USB_DATA_SIGNALING,
+                    enforcingAdmin,
+                    new BooleanPolicyValue(enabled));
         }
         DevicePolicyEventLogger
                 .createEvent(DevicePolicyEnums.SET_USB_DATA_SIGNALING)
@@ -22490,24 +22436,10 @@
     @Override
     public boolean isUsbDataSignalingEnabled(String packageName) {
         final CallerIdentity caller = getCallerIdentity(packageName);
-        if (Flags.policyEngineMigrationV2Enabled()) {
-            Boolean enabled = mDevicePolicyEngine.getResolvedPolicy(
-                    PolicyDefinition.USB_DATA_SIGNALING,
-                    caller.getUserId());
-            return enabled == null || enabled;
-        } else {
-            synchronized (getLockObject()) {
-                // If the caller is an admin, return the policy set by itself. Otherwise
-                // return the device-wide policy.
-                if (isDefaultDeviceOwner(caller) || isProfileOwnerOfOrganizationOwnedDevice(
-                        caller)) {
-                    return getProfileOwnerOrDeviceOwnerLocked(
-                            caller.getUserId()).mUsbDataSignalingEnabled;
-                } else {
-                    return isUsbDataSignalingEnabledInternalLocked();
-                }
-            }
-        }
+        Boolean enabled = mDevicePolicyEngine.getResolvedPolicy(
+                PolicyDefinition.USB_DATA_SIGNALING,
+                caller.getUserId());
+        return enabled == null || enabled;
     }
 
     private boolean isUsbDataSignalingEnabledInternalLocked() {
@@ -24860,9 +24792,6 @@
 
     @Override
     public void setMaxPolicyStorageLimit(String callerPackageName, int storageLimit) {
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            return;
-        }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
                 caller.getUserId());
@@ -24876,9 +24805,6 @@
 
     @Override
     public int getMaxPolicyStorageLimit(String callerPackageName) {
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            return -1;
-        }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_PROFILE_AND_DEVICE_OWNERS, caller.getPackageName(),
                 caller.getUserId());
@@ -24888,9 +24814,6 @@
 
     @Override
     public void forceSetMaxPolicyStorageLimit(String callerPackageName, int storageLimit) {
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            return;
-        }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_DEVICE_POLICY_STORAGE_LIMIT, caller.getPackageName(),
                 caller.getUserId());
@@ -24901,9 +24824,6 @@
     @Override
     public int getPolicySizeForAdmin(
             String callerPackageName, android.app.admin.EnforcingAdmin admin) {
-        if (!Flags.devicePolicySizeTrackingInternalBugFixEnabled()) {
-            return -1;
-        }
         CallerIdentity caller = getCallerIdentity(callerPackageName);
         enforcePermission(MANAGE_DEVICE_POLICY_STORAGE_LIMIT, caller.getPackageName(),
                 caller.getUserId());
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyDefinition.java b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyDefinition.java
index 19a942c..24ee46f 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyDefinition.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyDefinition.java
@@ -536,7 +536,6 @@
             USER_RESTRICTION_FLAGS.put(
                     UserManager.DISALLOW_THREAD_NETWORK, POLICY_FLAG_GLOBAL_ONLY_POLICY);
         }
-        USER_RESTRICTION_FLAGS.put(UserManager.DISALLOW_ASSIST_CONTENT, /* flags= */ 0);
         for (String key : USER_RESTRICTION_FLAGS.keySet()) {
             createAndAddUserRestrictionPolicyDefinition(key, USER_RESTRICTION_FLAGS.get(key));
         }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
index e1cb37d..8068d46 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
@@ -238,9 +238,7 @@
             }
 
             for (int user : resolveUsers(userId)) {
-                if (Flags.disallowUserControlBgUsageFix()) {
-                    setBgUsageAppOp(packages, pmi, user, appOpsManager);
-                }
+                setBgUsageAppOp(packages, pmi, user, appOpsManager);
                 if (Flags.disallowUserControlStoppedStateFix()) {
                     for (String packageName : packages) {
                         pmi.setPackageStoppedState(packageName, false, user);
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index c5c371f..ab459df 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -16,7 +16,6 @@
 
 package com.android.server;
 
-import static android.app.appfunctions.flags.Flags.enableAppFunctionManager;
 import static android.net.NetworkStack.PERMISSION_MAINLINE_NETWORK_STACK;
 import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_CRITICAL;
 import static android.os.IServiceManager.DUMP_FLAG_PRIORITY_HIGH;
@@ -38,6 +37,7 @@
 import android.app.INotificationManager;
 import android.app.SystemServiceRegistry;
 import android.app.admin.DevicePolicySafetyChecker;
+import android.app.appfunctions.AppFunctionManagerConfiguration;
 import android.app.usage.UsageStatsManagerInternal;
 import android.content.ComponentName;
 import android.content.ContentResolver;
@@ -1743,12 +1743,11 @@
             mSystemServiceManager.startService(LogcatManagerService.class);
             t.traceEnd();
 
-            t.traceBegin("StartAppFunctionManager");
-            if (enableAppFunctionManager()) {
+            if (AppFunctionManagerConfiguration.isSupported(context)) {
+                t.traceBegin("StartAppFunctionManager");
                 mSystemServiceManager.startService(AppFunctionManagerService.class);
+                t.traceEnd();
             }
-            t.traceEnd();
-
         } catch (Throwable e) {
             Slog.e("System", "******************************************");
             Slog.e("System", "************ Failure starting core service");
@@ -2468,8 +2467,8 @@
                 reportWtf("starting RuntimeService", e);
             }
             t.traceEnd();
-
-            if (!isWatch && !disableNetworkTime) {
+            if (!disableNetworkTime && (!isWatch || (isWatch
+                    && android.server.Flags.allowNetworkTimeUpdateService()))) {
                 t.traceBegin("StartNetworkTimeUpdateService");
                 try {
                     networkTimeUpdater = new NetworkTimeUpdateService(context);
diff --git a/services/java/com/android/server/flags.aconfig b/services/java/com/android/server/flags.aconfig
index 29f3871..ec74ef19 100644
--- a/services/java/com/android/server/flags.aconfig
+++ b/services/java/com/android/server/flags.aconfig
@@ -28,4 +28,11 @@
      namespace: "wear_frameworks"
      description: "Allow removing VpnManagerService"
      bug: "340928692"
+}
+
+flag {
+    name: "allow_network_time_update_service"
+    namespace: "wear_systems"
+    description: "Allow NetworkTimeUpdateService on Wear"
+    bug: "327508176"
 }
\ No newline at end of file
diff --git a/services/supervision/OWNERS b/services/supervision/OWNERS
new file mode 100644
index 0000000..e5f4147
--- /dev/null
+++ b/services/supervision/OWNERS
@@ -0,0 +1 @@
+include /core/java/android/app/supervision/OWNERS
diff --git a/services/supervision/java/com/android/server/supervision/SupervisionService.java b/services/supervision/java/com/android/server/supervision/SupervisionService.java
index a4ef629..7ffd0ec 100644
--- a/services/supervision/java/com/android/server/supervision/SupervisionService.java
+++ b/services/supervision/java/com/android/server/supervision/SupervisionService.java
@@ -20,7 +20,9 @@
 import android.annotation.Nullable;
 import android.app.supervision.ISupervisionManager;
 import android.content.Context;
-
+import android.os.RemoteException;
+import android.os.ResultReceiver;
+import android.os.ShellCallback;
 
 import com.android.internal.util.DumpUtils;
 import com.android.server.SystemService;
@@ -28,7 +30,9 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 
-/** Service for handling system supervision. */
+/**
+ * Service for handling system supervision.
+ */
 public class SupervisionService extends ISupervisionManager.Stub {
     private static final String LOG_TAG = "SupervisionService";
 
@@ -44,8 +48,20 @@
     }
 
     @Override
-    protected void dump(@NonNull FileDescriptor fd,
-            @NonNull PrintWriter fout, @Nullable String[] args) {
+    public void onShellCommand(
+            @Nullable FileDescriptor in,
+            @Nullable FileDescriptor out,
+            @Nullable FileDescriptor err,
+            @NonNull String[] args,
+            @Nullable ShellCallback callback,
+            @NonNull ResultReceiver resultReceiver) throws RemoteException {
+        new SupervisionServiceShellCommand(this)
+                .exec(this, in, out, err, args, callback, resultReceiver);
+    }
+
+    @Override
+    protected void dump(
+            @NonNull FileDescriptor fd, @NonNull PrintWriter fout, @Nullable String[] args) {
         if (!DumpUtils.checkDumpPermission(mContext, LOG_TAG, fout)) return;
 
         fout.println("Supervision enabled: " + isSupervisionEnabled());
diff --git a/services/supervision/java/com/android/server/supervision/SupervisionServiceShellCommand.java b/services/supervision/java/com/android/server/supervision/SupervisionServiceShellCommand.java
new file mode 100644
index 0000000..3aba24a
--- /dev/null
+++ b/services/supervision/java/com/android/server/supervision/SupervisionServiceShellCommand.java
@@ -0,0 +1,61 @@
+/*
+ * 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.supervision;
+
+import android.os.ShellCommand;
+
+import java.io.PrintWriter;
+
+public class SupervisionServiceShellCommand extends ShellCommand {
+    private final SupervisionService mService;
+
+    public SupervisionServiceShellCommand(SupervisionService mService) {
+        this.mService = mService;
+    }
+
+    @Override
+    public int onCommand(String cmd) {
+        if (cmd == null) {
+            return handleDefaultCommands(null);
+        }
+        final PrintWriter pw = getOutPrintWriter();
+        switch (cmd) {
+            case "help": return help(pw);
+            case "is-enabled": return isEnabled(pw);
+            default: return handleDefaultCommands(cmd);
+        }
+    }
+
+    private int help(PrintWriter pw) {
+        pw.println("Supervision service commands:");
+        pw.println("  help");
+        pw.println("      Prints this help text");
+        pw.println("  is-enabled");
+        pw.println("      Is supervision enabled");
+        return 0;
+    }
+
+    private int isEnabled(PrintWriter pw) {
+        pw.println(mService.isSupervisionEnabled());
+        return 0;
+    }
+
+    @Override
+    public void onHelp() {
+        help(getOutPrintWriter());
+    }
+}
diff --git a/services/tests/appfunctions/OWNERS b/services/tests/appfunctions/OWNERS
new file mode 100644
index 0000000..7fa8917
--- /dev/null
+++ b/services/tests/appfunctions/OWNERS
@@ -0,0 +1,2 @@
+# Bug component: 1627156
+include platform/frameworks/base:/core/java/android/app/appfunctions/OWNERS
diff --git a/services/tests/displayservicetests/src/com/android/server/display/BrightnessRangeControllerTest.kt b/services/tests/displayservicetests/src/com/android/server/display/BrightnessRangeControllerTest.kt
index 1f3184d..f589a2c 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/BrightnessRangeControllerTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/BrightnessRangeControllerTest.kt
@@ -42,31 +42,31 @@
     private val mockToken = mock<IBinder>()
 
     @Test
-    fun `returns HBC max brightness if HBM supported and ON`() {
+    fun testMaxBrightness_HbmSupportedAndOn() {
         val controller = createController()
         assertThat(controller.currentBrightnessMax).isEqualTo(MAX_BRIGHTNESS)
     }
 
     @Test
-    fun `returns NBC max brightness if device does not support HBM`() {
+    fun testMaxBrightness_HbmNotSupported() {
         val controller = createController(hbmSupported = false)
         assertThat(controller.currentBrightnessMax).isEqualTo(NORMAL_BRIGHTNESS_LOW)
     }
 
     @Test
-    fun `returns NBC max brightness if HBM not allowed`() {
+    fun testMaxBrightness_HbmNotAllowed() {
         val controller = createController(hbmAllowed = false)
         assertThat(controller.currentBrightnessMax).isEqualTo(NORMAL_BRIGHTNESS_LOW)
     }
 
     @Test
-    fun `returns HBC max brightness if NBM is disabled`() {
+    fun testMaxBrightness_HbmDisabledAndNotAllowed() {
         val controller = createController(nbmEnabled = false, hbmAllowed = false)
         assertThat(controller.currentBrightnessMax).isEqualTo(MAX_BRIGHTNESS)
     }
 
     @Test
-    fun `returns HBC max brightness if lower than NBC max brightness`() {
+    fun testMaxBrightness_transitionPointLessThanCurrentNbmLimit() {
         val controller = createController(
             hbmAllowed = false,
             hbmMaxBrightness = TRANSITION_POINT,
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
index f690b1b..2d4a29b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
@@ -18,6 +18,8 @@
 
 import static org.junit.Assert.assertEquals;
 
+import android.hardware.display.BrightnessInfo;
+
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
@@ -112,7 +114,10 @@
                 .append("\n    mBrightnessAdjustmentFlag:")
                 .append(displayBrightnessState.getBrightnessAdjustmentFlag())
                 .append("\n    mIsUserInitiatedChange:")
-                .append(displayBrightnessState.isUserInitiatedChange());
+                .append(displayBrightnessState.isUserInitiatedChange())
+                .append("\n    mBrightnessMaxReason:")
+                .append(BrightnessInfo.briMaxReasonToString(
+                        displayBrightnessState.getBrightnessMaxReason()));
         return sb.toString();
     }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
index d450683..fd05b26 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayDeviceConfigTest.java
@@ -263,7 +263,9 @@
                 mDisplayDeviceConfig.getPowerThrottlingConfigData();
         assertNotNull(powerThrottlingConfigData);
         assertEquals(0.1f, powerThrottlingConfigData.brightnessLowestCapAllowed, SMALL_DELTA);
-        assertEquals(10, powerThrottlingConfigData.pollingWindowMillis);
+        assertEquals(15f, powerThrottlingConfigData.customAnimationRateSec, SMALL_DELTA);
+        assertEquals(20000, powerThrottlingConfigData.pollingWindowMaxMillis);
+        assertEquals(10000, powerThrottlingConfigData.pollingWindowMinMillis);
     }
 
     @Test
@@ -1295,7 +1297,9 @@
     private String getPowerThrottlingConfig() {
         return  "<powerThrottlingConfig >\n"
                 +       "<brightnessLowestCapAllowed>0.1</brightnessLowestCapAllowed>\n"
-                +       "<pollingWindowMillis>10</pollingWindowMillis>\n"
+                +       "<customAnimationRateSec>15</customAnimationRateSec>\n"
+                +       "<pollingWindowMaxMillis>20000</pollingWindowMaxMillis>\n"
+                +       "<pollingWindowMinMillis>10000</pollingWindowMinMillis>\n"
                 +       "<powerThrottlingMap>\n"
                 +           "<powerThrottlingPoint>\n"
                 +               "<thermalStatus>light</thermalStatus>\n"
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
index 30c384a..5e868a3 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayOffloadSessionImplTest.java
@@ -38,7 +38,7 @@
     private DisplayManagerInternal.DisplayOffloader mDisplayOffloader;
 
     @Mock
-    private DisplayPowerControllerInterface mDisplayPowerController;
+    private DisplayPowerController mDisplayPowerController;
 
     private DisplayOffloadSessionImpl mSession;
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
index 2166cb7..d0aec3b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerControllerTest.java
@@ -2587,7 +2587,7 @@
         BrightnessClamperController getBrightnessClamperController(Handler handler,
                 BrightnessClamperController.ClamperChangeListener clamperChangeListener,
                 BrightnessClamperController.DisplayDeviceData data, Context context,
-                DisplayManagerFlags flags, SensorManager sensorManager) {
+                DisplayManagerFlags flags, SensorManager sensorManager, float currentBrightness) {
             return mClamperController;
         }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerStateTest.kt b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerStateTest.kt
index 33d3020..a684fdb 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerStateTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayPowerStateTest.kt
@@ -60,7 +60,7 @@
     }
 
     @Test
-    fun `destroys ColorFade on stop`() {
+    fun destroysColorFadeOnStop() {
         displayPowerState.stop()
         val runnableCaptor = argumentCaptor<Runnable>()
 
@@ -71,13 +71,13 @@
     }
 
     @Test
-    fun `GIVEN not prepared WHEN draw runnable is called THEN colorFade not drawn`() {
+    fun testColorFadeDraw_notPrepared() {
         displayPowerState.mColorFadeDrawRunnable.run()
 
         verify(mockColorFade, never()).draw(anyFloat())
     }
     @Test
-    fun `GIVEN prepared WHEN draw runnable is called THEN colorFade is drawn`() {
+    fun testColorFadeDraw_prepared() {
         displayPowerState.prepareColorFade(mockContext, ColorFade.MODE_FADE)
         clearInvocations(mockColorFade)
 
@@ -87,7 +87,7 @@
     }
 
     @Test
-    fun `GIVEN prepared AND stopped WHEN draw runnable is called THEN colorFade is not drawn`() {
+    fun testColorFadeDraw_preparedAndStopped() {
         displayPowerState.prepareColorFade(mockContext, ColorFade.MODE_FADE)
         clearInvocations(mockColorFade)
         displayPowerState.stop()
diff --git a/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java b/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
index a7e0ebd..120cc84 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/LocalDisplayAdapterTest.java
@@ -118,7 +118,7 @@
     @Mock
     private DisplayManagerFlags mFlags;
     @Mock
-    private DisplayPowerControllerInterface mMockedDisplayPowerController;
+    private DisplayPowerController mMockedDisplayPowerController;
 
     private Handler mHandler;
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessClamperControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessClamperControllerTest.java
index 0ce9233..f9dc122 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessClamperControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessClamperControllerTest.java
@@ -31,7 +31,6 @@
 
 import android.content.Context;
 import android.hardware.SensorManager;
-import android.hardware.display.BrightnessInfo;
 import android.hardware.display.DisplayManagerInternal;
 import android.os.Handler;
 import android.os.PowerManager;
@@ -161,12 +160,6 @@
     }
 
     @Test
-    public void testMaxReasonIsNoneOnInit() {
-        assertEquals(BrightnessInfo.BRIGHTNESS_MAX_REASON_NONE,
-                mClamperController.getBrightnessMaxReason());
-    }
-
-    @Test
     public void testOnDisplayChanged_DelegatesToClamper() {
         mClamperController.onDisplayChanged(mMockDisplayDeviceData);
 
@@ -365,7 +358,7 @@
 
     private BrightnessClamperController createBrightnessClamperController() {
         return new BrightnessClamperController(mTestInjector, mTestHandler, mMockExternalListener,
-                mMockDisplayDeviceData, mMockContext, mFlags, mSensorManager);
+                mMockDisplayDeviceData, mMockContext, mFlags, mSensorManager, 0);
     }
 
     interface TestDisplayListenerModifier extends BrightnessStateModifier,
@@ -403,7 +396,7 @@
                 Handler handler,
                 BrightnessClamperController.ClamperChangeListener clamperChangeListener,
                 BrightnessClamperController.DisplayDeviceData data,
-                DisplayManagerFlags flags, Context context) {
+                DisplayManagerFlags flags, Context context, float currentBrightness) {
             mCapturedChangeListener = clamperChangeListener;
             return mClampers;
         }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessPowerClamperTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessPowerClamperTest.java
index b3f33ad..c4898da 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessPowerClamperTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/clamper/BrightnessPowerClamperTest.java
@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 
+import android.os.IThermalService;
 import android.os.PowerManager;
 import android.os.RemoteException;
 import android.os.Temperature;
@@ -58,12 +59,18 @@
     private final FakeDeviceConfigInterface mFakeDeviceConfigInterface =
             new FakeDeviceConfigInterface();
     private final TestHandler mTestHandler = new TestHandler(null);
+    private final TestInjector mTestInjector = new TestInjector();
     private BrightnessPowerClamper mClamper;
+    private final float mCurrentBrightness = 0.6f;
+    private PowerChangeListener mPowerChangeListener;
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mClamper = new BrightnessPowerClamper(new TestInjector(), mTestHandler,
-                mMockClamperChangeListener, new TestPowerData());
+        mClamper = new BrightnessPowerClamper(mTestInjector, mTestHandler,
+                mMockClamperChangeListener, new TestPowerData(), mCurrentBrightness);
+        mPowerChangeListener = mClamper.getPowerChangeListener();
+        mPmicMonitor = mTestInjector.getPmicMonitor(mPowerChangeListener, null, 5, 10);
+        mPmicMonitor.setPowerChangeListener(mPowerChangeListener);
         mTestHandler.flush();
     }
 
@@ -79,36 +86,27 @@
     }
 
     @Test
-    public void testPowerThrottlingNoOngoingAnimation() throws RemoteException {
-        mPmicMonitor.setThermalStatus(Temperature.THROTTLING_SEVERE);
+    public void testPowerThrottlingWithThermalLevelLight() throws RemoteException {
+        mPmicMonitor.setThermalStatus(Temperature.THROTTLING_LIGHT);
         mTestHandler.flush();
         assertFalse(mClamper.isActive());
         assertEquals(PowerManager.BRIGHTNESS_MAX, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
 
         // update a new device config for power-throttling.
         mClamper.onDisplayChanged(new TestPowerData(
-                List.of(new ThrottlingLevel(PowerManager.THERMAL_STATUS_SEVERE, 100f))));
+                List.of(new ThrottlingLevel(PowerManager.THERMAL_STATUS_LIGHT, 100f))));
 
         mPmicMonitor.setAvgPowerConsumed(200f);
         float expectedBrightness = 0.5f;
-        expectedBrightness = expectedBrightness * PowerManager.BRIGHTNESS_MAX;
+        expectedBrightness = expectedBrightness * mCurrentBrightness;
 
         mTestHandler.flush();
         // Assume current brightness as max, as there is no throttling.
         assertEquals(expectedBrightness, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
-        mPmicMonitor.setThermalStatus(Temperature.THROTTLING_CRITICAL);
-        // update a new device config for power-throttling.
-        mClamper.onDisplayChanged(new TestPowerData(
-                List.of(new ThrottlingLevel(PowerManager.THERMAL_STATUS_CRITICAL, 50f))));
-
-        mPmicMonitor.setAvgPowerConsumed(100f);
-        expectedBrightness = 0.5f * PowerManager.BRIGHTNESS_MAX;
-        mTestHandler.flush();
-        assertEquals(expectedBrightness, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
     }
 
     @Test
-    public void testPowerThrottlingWithOngoingAnimation() throws RemoteException {
+    public void testPowerThrottlingWithThermalLevelSevere() throws RemoteException {
         mPmicMonitor.setThermalStatus(Temperature.THROTTLING_SEVERE);
         mTestHandler.flush();
         assertEquals(PowerManager.BRIGHTNESS_MAX, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
@@ -119,20 +117,10 @@
 
         mPmicMonitor.setAvgPowerConsumed(200f);
         float expectedBrightness = 0.5f;
-        expectedBrightness = expectedBrightness * PowerManager.BRIGHTNESS_MAX;
-
+        expectedBrightness = expectedBrightness * mCurrentBrightness;
         mTestHandler.flush();
         // Assume current brightness as max, as there is no throttling.
         assertEquals(expectedBrightness, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
-        mPmicMonitor.setThermalStatus(Temperature.THROTTLING_CRITICAL);
-        // update a new device config for power-throttling.
-        mClamper.onDisplayChanged(new TestPowerData(
-                List.of(new ThrottlingLevel(PowerManager.THERMAL_STATUS_CRITICAL, 50f))));
-
-        mPmicMonitor.setAvgPowerConsumed(100f);
-        expectedBrightness = 0.5f * PowerManager.BRIGHTNESS_MAX;
-        mTestHandler.flush();
-        assertEquals(expectedBrightness, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
     }
 
     @Test
@@ -148,8 +136,7 @@
 
         mPmicMonitor.setAvgPowerConsumed(200f);
         float expectedBrightness = 0.5f;
-        expectedBrightness = expectedBrightness * PowerManager.BRIGHTNESS_MAX;
-
+        expectedBrightness = expectedBrightness * mCurrentBrightness;
         mTestHandler.flush();
 
         assertEquals(expectedBrightness, mClamper.getBrightnessCap(), FLOAT_TOLERANCE);
@@ -169,10 +156,11 @@
 
     private static class TestPmicMonitor extends PmicMonitor {
         private Temperature mCurrentTemperature;
-        private final PowerChangeListener mListener;
-        TestPmicMonitor(PowerChangeListener listener, int pollingTime) {
-            super(listener, pollingTime);
-            mListener = listener;
+        private PowerChangeListener mListener;
+        TestPmicMonitor(PowerChangeListener listener,
+                        IThermalService thermalService,
+                        int pollingTimeMax, int pollingTimeMin) {
+            super(listener, thermalService, pollingTimeMax, pollingTimeMin);
         }
         public void setAvgPowerConsumed(float power) {
             int status = mCurrentTemperature.getStatus();
@@ -181,13 +169,18 @@
         public void setThermalStatus(@Temperature.ThrottlingStatus int status) {
             mCurrentTemperature = new Temperature(100, Temperature.TYPE_SKIN, "test_temp", status);
         }
+        public void setPowerChangeListener(PowerChangeListener listener) {
+            mListener = listener;
+        }
     }
 
     private class TestInjector extends BrightnessPowerClamper.Injector {
         @Override
         TestPmicMonitor getPmicMonitor(PowerChangeListener listener,
-                int pollingTime) {
-            mPmicMonitor = new TestPmicMonitor(listener, pollingTime);
+                                       IThermalService thermalService,
+                                       int minPollingTimeMillis, int maxPollingTimeMillis) {
+            mPmicMonitor = new TestPmicMonitor(listener, thermalService, maxPollingTimeMillis,
+                    minPollingTimeMillis);
             return mPmicMonitor;
         }
 
@@ -216,7 +209,7 @@
             mUniqueDisplayId = uniqueDisplayId;
             mDataId = dataId;
             mData = PowerThrottlingData.create(data);
-            mConfigData = new PowerThrottlingConfigData(0.1f, 10);
+            mConfigData = new PowerThrottlingConfigData(0.1f, 10, 20, 10);
         }
 
         @NonNull
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/AppRequestObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/AppRequestObserverTest.kt
index 34c6ba9..1f3f19f 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/AppRequestObserverTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/AppRequestObserverTest.kt
@@ -50,7 +50,7 @@
     }
 
     @Test
-    fun `test app request votes`(@TestParameter testCase: AppRequestTestCase) {
+    fun testAppRequestVotes(@TestParameter testCase: AppRequestTestCase) {
         whenever(mockFlags.ignoreAppPreferredRefreshRateRequest())
                 .thenReturn(testCase.ignoreRefreshRateRequest)
         val displayModeDirector = DisplayModeDirector(
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/BaseModeRefreshRateVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/BaseModeRefreshRateVoteTest.kt
index bf2edfe..3841211 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/BaseModeRefreshRateVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/BaseModeRefreshRateVoteTest.kt
@@ -39,7 +39,7 @@
     }
 
     @Test
-    fun `updates summary with base mode refresh rate if not set`() {
+    fun updatesSummary_doesNotUpdateSummary_baseModeRefreshRateNotSet() {
         val summary = createVotesSummary()
 
         baseModeVote.updateSummary(summary)
@@ -48,7 +48,7 @@
     }
 
     @Test
-    fun `keeps summary base mode refresh rate if set`() {
+    fun doesNotUpdateSummary_baseModeRefreshRateSet() {
         val summary = createVotesSummary()
         summary.appRequestBaseModeRefreshRate = OTHER_BASE_REFRESH_RATE
 
@@ -58,7 +58,7 @@
     }
 
     @Test
-    fun `keeps summary with base mode refresh rate if vote refresh rate is negative`() {
+    fun doesNotUpdateSummary_baseModeRefreshRateNotSet_requestedRefreshRateInvalid() {
         val invalidBaseModeVote = BaseModeRefreshRateVote(-10f)
         val summary = createVotesSummary()
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/CombinedVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/CombinedVoteTest.kt
index 209e5a3..0a3c285 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/CombinedVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/CombinedVoteTest.kt
@@ -45,7 +45,7 @@
     }
 
     @Test
-    fun `delegates update to children`() {
+    fun delegatesUpdateToChildren() {
         val summary = createVotesSummary()
 
         combinedVote.updateSummary(summary)
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/DisableRefreshRateSwitchingVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/DisableRefreshRateSwitchingVoteTest.kt
index 38782c2..5b5ae65 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/DisableRefreshRateSwitchingVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/DisableRefreshRateSwitchingVoteTest.kt
@@ -28,7 +28,7 @@
 class DisableRefreshRateSwitchingVoteTest {
 
     @Test
-    fun `disabled refresh rate switching is not changed`(
+    fun testDisableRefreshRateSwitch_alreadyDisabled(
             @TestParameter voteDisableSwitching: Boolean
     ) {
         val summary = createVotesSummary()
@@ -41,7 +41,7 @@
     }
 
     @Test
-    fun `disables refresh rate switching if requested`() {
+    fun disablesRefreshRateSwitch_notDisabled_requested() {
         val summary = createVotesSummary()
         val vote = DisableRefreshRateSwitchingVote(true)
 
@@ -51,7 +51,7 @@
     }
 
     @Test
-    fun `does not disable refresh rate switching if not requested`() {
+    fun doesNotDisableRefreshRateSwitch_notDisabled_notRequested() {
         val summary = createVotesSummary()
         val vote = DisableRefreshRateSwitchingVote(false)
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/PhysicalVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/PhysicalVoteTest.kt
index 9edcc32..0968edb 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/PhysicalVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/PhysicalVoteTest.kt
@@ -37,7 +37,7 @@
     }
 
     @Test
-    fun `updates minPhysicalRefreshRate if summary has less`() {
+    fun updatesMinPhysicalRefreshRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.minPhysicalRefreshRate = 45f
 
@@ -47,7 +47,7 @@
     }
 
     @Test
-    fun `does not update minPhysicalRefreshRate if summary has more`() {
+    fun doesNotUpdateMinPhysicalRefreshRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.minPhysicalRefreshRate = 75f
 
@@ -57,7 +57,7 @@
     }
 
     @Test
-    fun `updates maxPhysicalRefreshRate if summary has more`() {
+    fun updatesMaxPhysicalRefreshRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.maxPhysicalRefreshRate = 120f
 
@@ -67,7 +67,7 @@
     }
 
     @Test
-    fun `does not update maxPhysicalRefreshRate if summary has less`() {
+    fun doesNotUpdateMaxPhysicalRefreshRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.maxPhysicalRefreshRate = 75f
 
@@ -77,7 +77,7 @@
     }
 
     @Test
-    fun `updates maxRenderFrameRate if summary has more`() {
+    fun updatesMaxRenderFrameRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.maxRenderFrameRate = 120f
 
@@ -87,7 +87,7 @@
     }
 
     @Test
-    fun `does not update maxRenderFrameRate if summary has less`() {
+    fun doesNotUpdateMaxRenderFrameRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.maxRenderFrameRate = 75f
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/RenderVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/RenderVoteTest.kt
index 2d65f1c..9fa1e1b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/RenderVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/RenderVoteTest.kt
@@ -38,7 +38,7 @@
     }
 
     @Test
-    fun `updates minRenderFrameRate if summary has less`() {
+    fun updatesMinRenderFrameRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.minRenderFrameRate = 45f
 
@@ -48,7 +48,7 @@
     }
 
     @Test
-    fun `does not update minRenderFrameRate if summary has more`() {
+    fun doesNotUpdateMinRenderFrameRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.minRenderFrameRate = 75f
 
@@ -58,7 +58,7 @@
     }
 
     @Test
-    fun `updates maxRenderFrameRate if summary has more`() {
+    fun updatesMaxPRenderFrameRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.maxRenderFrameRate = 120f
 
@@ -68,7 +68,7 @@
     }
 
     @Test
-    fun `does not update maxRenderFrameRate if summary has less`() {
+    fun doesNotUpdateMaxPRenderFrameRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.maxRenderFrameRate = 75f
 
@@ -78,7 +78,7 @@
     }
 
     @Test
-    fun `updates minPhysicalRefreshRate if summary has less`() {
+    fun updatesMinPhysicalRefreshRateWithBiggerValue() {
         val summary = createVotesSummary()
         summary.minPhysicalRefreshRate = 45f
 
@@ -88,7 +88,7 @@
     }
 
     @Test
-    fun `does not update minPhysicalRefreshRate if summary has more`() {
+    fun doesNotUpdateMinPhysicalRefreshRateWithSmallerValue() {
         val summary = createVotesSummary()
         summary.minPhysicalRefreshRate = 75f
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/RequestedRefreshRateVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/RequestedRefreshRateVoteTest.kt
index dbe9e4a..be9c563 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/RequestedRefreshRateVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/RequestedRefreshRateVoteTest.kt
@@ -28,7 +28,7 @@
 class RequestedRefreshRateVoteTest {
 
     @Test
-    fun `updates requestedRefreshRates`() {
+    fun testUpdatesRequestedRefreshRates() {
         val refreshRate = 90f
         val vote = RequestedRefreshRateVote(refreshRate)
         val summary = createVotesSummary()
@@ -40,7 +40,7 @@
     }
 
     @Test
-    fun `updates requestedRefreshRates with multiple refresh rates`() {
+    fun testUpdatesRequestedRefreshRates_multipleVotes() {
         val refreshRate1 = 90f
         val vote1 = RequestedRefreshRateVote(refreshRate1)
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
index 4fc574a..d7dcca7 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SettingsObserverTest.kt
@@ -103,7 +103,7 @@
     }
 
     @Test
-    fun `test low power mode`(@TestParameter testCase: LowPowerTestCase) {
+    fun testLowPowerMode(@TestParameter testCase: LowPowerTestCase) {
         whenever(mockFlags.isVsyncLowPowerVoteEnabled).thenReturn(testCase.vsyncLowPowerVoteEnabled)
         whenever(spyContext.contentResolver)
                 .thenReturn(settingsProviderRule.mockContentResolver(null))
@@ -151,7 +151,7 @@
     }
 
     @Test
-    fun `test settings refresh rates`(@TestParameter testCase: SettingsRefreshRateTestCase) {
+    fun testSettingsRefreshRates(@TestParameter testCase: SettingsRefreshRateTestCase) {
         whenever(mockFlags.isPeakRefreshRatePhysicalLimitEnabled)
                 .thenReturn(testCase.peakRefreshRatePhysicalLimitEnabled)
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SizeVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SizeVoteTest.kt
index 1be2fbf..319c21e 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SizeVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SizeVoteTest.kt
@@ -39,7 +39,7 @@
     }
 
     @Test
-    fun `updates size if width and height not set and display resolution voting disabled`() {
+    fun updatesSize_widthAndHeightNotSet_resolutionVotingDisabled() {
         val summary = createVotesSummary(isDisplayResolutionRangeVotingEnabled = false)
         summary.width = Vote.INVALID_SIZE
         summary.height = Vote.INVALID_SIZE
@@ -55,7 +55,7 @@
     }
 
     @Test
-    fun `does not update size if width set and display resolution voting disabled`() {
+    fun doesNotUpdateSiz_widthSet_resolutionVotingDisabled() {
         val summary = createVotesSummary(isDisplayResolutionRangeVotingEnabled = false)
         summary.width = 150
         summary.height = Vote.INVALID_SIZE
@@ -71,7 +71,7 @@
     }
 
     @Test
-    fun `does not update size if height set and display resolution voting disabled`() {
+    fun doesNotUpdateSize_heightSet_resolutionVotingDisabled() {
         val summary = createVotesSummary(isDisplayResolutionRangeVotingEnabled = false)
         summary.width = Vote.INVALID_SIZE
         summary.height = 250
@@ -87,7 +87,7 @@
     }
 
     @Test
-    fun `updates width if summary has more and display resolution voting enabled`() {
+    fun updatesWidthWithSmallerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 850
 
@@ -97,7 +97,7 @@
     }
 
     @Test
-    fun `does not update width if summary has less and display resolution voting enabled`() {
+    fun doesNotUpdateWidthWithBiggerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 750
 
@@ -107,7 +107,7 @@
     }
 
     @Test
-    fun `updates height if summary has more and display resolution voting enabled`() {
+    fun updatesHeightWithSmallerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.height = 1650
 
@@ -117,7 +117,7 @@
     }
 
     @Test
-    fun `does not update height if summary has less and display resolution voting enabled`() {
+    fun doesNotUpdateHeightWithBiggerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.height = 1550
 
@@ -127,7 +127,7 @@
     }
 
     @Test
-    fun `updates minWidth if summary has less and display resolution voting enabled`() {
+    fun updatesMinWidthWithSmallerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 150
         summary.minWidth = 350
@@ -138,7 +138,7 @@
     }
 
     @Test
-    fun `does not update minWidth if summary has more and display resolution voting enabled`() {
+    fun doesNotUpdateMinWidthWithBiggerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 150
         summary.minWidth = 450
@@ -149,7 +149,7 @@
     }
 
     @Test
-    fun `updates minHeight if summary has less and display resolution voting enabled`() {
+    fun updatesMinHeightWithSmallerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 150
         summary.minHeight = 1150
@@ -160,7 +160,7 @@
     }
 
     @Test
-    fun `does not update minHeight if summary has more and display resolution voting enabled`() {
+    fun doesNotUpdateMinHeightWithBiggerValue_resolutionVotingEnabled() {
         val summary = createVotesSummary()
         summary.width = 150
         summary.minHeight = 1250
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
index 6ce49b8..2a50a33 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedModesVoteTest.kt
@@ -39,7 +39,7 @@
     }
 
     @Test
-    fun `adds supported mode ids if supportedModeIds in summary is null`() {
+    fun addsSupportedModeIds_summaryHasNull() {
         val summary = createVotesSummary()
 
         supportedModesVote.updateSummary(summary)
@@ -48,7 +48,7 @@
     }
 
     @Test
-    fun `does not add supported mode ids if summary has empty list of modeIds`() {
+    fun doesNotAddSupportedModeIdes_summaryHasEmptyList() {
         val summary = createVotesSummary()
         summary.supportedModeIds = ArrayList()
 
@@ -58,7 +58,7 @@
     }
 
     @Test
-    fun `filters out modes that does not match vote`() {
+    fun filtersModeIdsThatDoesNotMatchVote() {
         val summary = createVotesSummary()
         summary.supportedModeIds = ArrayList(listOf(otherMode, supportedModes[0]))
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt
index d0c112b..0da6885 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SupportedRefreshRatesVoteTest.kt
@@ -42,7 +42,7 @@
     }
 
     @Test
-    fun `adds supported refresh rates if supportedModes in summary is null`() {
+    fun addsSupportedRefreshRates_summaryHasNull() {
         val summary = createVotesSummary()
 
         supportedRefreshRatesVote.updateSummary(summary)
@@ -51,7 +51,7 @@
     }
 
     @Test
-    fun `does not add supported refresh rates if summary has empty list of refresh rates`() {
+    fun doesNotAddSupportedRefreshRates_summaryHasEmptyList() {
         val summary = createVotesSummary()
         summary.supportedRefreshRates = ArrayList()
 
@@ -61,7 +61,7 @@
     }
 
     @Test
-    fun `filters out supported refresh rates that does not match vote`() {
+    fun filtersSupportedRefreshRatesThatDoesNotMatchVote() {
         val summary = createVotesSummary()
         summary.supportedRefreshRates = ArrayList(listOf(otherMode, refreshRates[0]))
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SyntheticModeManagerTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SyntheticModeManagerTest.kt
index 5cd3a33..b2d83d7 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SyntheticModeManagerTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SyntheticModeManagerTest.kt
@@ -41,7 +41,7 @@
     private val mockConfig = mock<DisplayDeviceConfig>()
 
     @Test
-    fun `test app supported modes`(@TestParameter testCase: AppSupportedModesTestCase) {
+    fun testAppSupportedModes(@TestParameter testCase: AppSupportedModesTestCase) {
         whenever(mockFlags.isSynthetic60HzModesEnabled).thenReturn(testCase.syntheticModesEnabled)
         whenever(mockConfig.isVrrSupportEnabled).thenReturn(testCase.vrrSupported)
         val syntheticModeManager = SyntheticModeManager(mockFlags)
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt
index c49205b..9ea7ea7 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/SystemRequestObserverTest.kt
@@ -51,7 +51,7 @@
     private val storage = VotesStorage({}, null)
 
     @Test
-    fun `requestDisplayModes adds vote to storage`() {
+    fun testRequestDisplayModes_voteAdded() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -69,7 +69,7 @@
     }
 
     @Test
-    fun `requestDisplayModes overrides votes in storage`() {
+    fun testRequestDisplayModes_voteReplaced() {
         val systemRequestObserver = SystemRequestObserver(storage)
 
         systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, intArrayOf(1, 2, 3))
@@ -89,7 +89,7 @@
     }
 
     @Test
-    fun `requestDisplayModes removes vote to storage`() {
+    fun testRequestDisplayModes_voteRemoved() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -101,7 +101,7 @@
     }
 
     @Test
-    fun `requestDisplayModes calls linkToDeath to token`() {
+    fun testTokenLinkToDeath() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -111,7 +111,7 @@
     }
 
     @Test
-    fun `does not add votes to storage if binder died when requestDisplayModes called`() {
+    fun testBinderDied_voteRemoved() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -123,7 +123,7 @@
     }
 
     @Test
-    fun `removes all votes from storage when binder dies`() {
+    fun testBinderDied_allVotesRemoved() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -138,7 +138,7 @@
     }
 
     @Test
-    fun `calls unlinkToDeath on token when no votes remaining`() {
+    fun testTokenUnlinkToDeath_noMoreVotes() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -149,7 +149,7 @@
     }
 
     @Test
-    fun `does not call unlinkToDeath on token when votes for other display in storage`() {
+    fun testTokenUnlinkToDeathNotCalled_votesForOtherDisplayInStorage() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
 
@@ -161,7 +161,7 @@
     }
 
     @Test
-    fun `requestDisplayModes subset modes from different tokens`() {
+    fun testRequestDisplayModes_differentToken_voteHasModesSubset() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
         systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
@@ -187,7 +187,7 @@
     }
 
     @Test
-    fun `recalculates vote if one binder dies`() {
+    fun testBinderDies_recalculatesVotes() {
         val systemRequestObserver = SystemRequestObserver(storage)
         val requestedModes = intArrayOf(1, 2, 3)
         systemRequestObserver.requestDisplayModes(mockToken, DISPLAY_ID, requestedModes)
diff --git a/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt b/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
index dd5e1be..239e59b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
+++ b/services/tests/displayservicetests/src/com/android/server/display/mode/VoteSummaryTest.kt
@@ -80,7 +80,7 @@
     }
 
     @Test
-    fun `filters modes for summary supportedRefreshRates`(
+    fun testFiltersModes_supportedRefreshRates(
             @TestParameter testCase: SupportedRefreshRatesTestCase
     ) {
         val summary = createSummary(testCase.supportedModesVoteEnabled)
@@ -142,9 +142,7 @@
     }
 
     @Test
-    fun `filters modes for summary supportedModes`(
-            @TestParameter testCase: SupportedModesTestCase
-    ) {
+    fun testFiltersModes_supportedModes(@TestParameter testCase: SupportedModesTestCase) {
         val summary = createSummary(testCase.supportedModesVoteEnabled)
         summary.supportedModeIds = testCase.summarySupportedModes
 
@@ -154,7 +152,7 @@
     }
 
     @Test
-    fun `summary invalid if has requestedRefreshRate less than minRenederRate`() {
+    fun testInvalidSummary_requestedRefreshRateLessThanMinRenderRate() {
         val summary = createSummary()
         summary.requestedRefreshRates = setOf(30f, 90f)
         summary.minRenderFrameRate = 60f
@@ -166,7 +164,7 @@
     }
 
     @Test
-    fun `summary invalid if has requestedRefreshRate more than maxRenderFrameRate`() {
+    fun testInvalidSummary_requestedRefreshRateMoreThanMaxRenderRate() {
         val summary = createSummary()
         summary.requestedRefreshRates = setOf(60f, 240f)
         summary.minRenderFrameRate = 60f
@@ -178,7 +176,7 @@
     }
 
     @Test
-    fun `summary valid if all requestedRefreshRates inside render rate limits`() {
+    fun testValidSummary_requestedRefreshRatesWithingRenderRateLimits() {
         val summary = createSummary()
         summary.requestedRefreshRates = setOf(60f, 90f)
         summary.minRenderFrameRate = 60f
diff --git a/services/tests/dreamservicetests/src/com/android/server/dreams/DreamOverlayServiceTest.java b/services/tests/dreamservicetests/src/com/android/server/dreams/DreamOverlayServiceTest.java
index 54f4607..1abc557 100644
--- a/services/tests/dreamservicetests/src/com/android/server/dreams/DreamOverlayServiceTest.java
+++ b/services/tests/dreamservicetests/src/com/android/server/dreams/DreamOverlayServiceTest.java
@@ -18,7 +18,9 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
@@ -27,7 +29,9 @@
 import android.content.Intent;
 import android.os.IBinder;
 import android.os.RemoteException;
+import android.platform.test.annotations.EnableFlags;
 import android.service.dreams.DreamOverlayService;
+import android.service.dreams.Flags;
 import android.service.dreams.IDreamOverlay;
 import android.service.dreams.IDreamOverlayCallback;
 import android.service.dreams.IDreamOverlayClient;
@@ -136,7 +140,7 @@
 
         // Start the dream.
         client.startDream(mLayoutParams, mOverlayCallback,
-                FIRST_DREAM_COMPONENT.flattenToString(), false);
+                FIRST_DREAM_COMPONENT.flattenToString(), false, false);
 
         // The callback should not have run yet.
         verify(monitor, never()).onStartDream();
@@ -194,22 +198,24 @@
         // Start a dream with the first client and ensure the dream is now active from the
         // overlay's perspective.
         firstClient.startDream(mLayoutParams, mOverlayCallback,
-                FIRST_DREAM_COMPONENT.flattenToString(), false);
+                FIRST_DREAM_COMPONENT.flattenToString(), true, false);
 
 
         verify(monitor).onStartDream();
         assertThat(service.getDreamComponent()).isEqualTo(FIRST_DREAM_COMPONENT);
+        assertThat(service.isDreamInPreviewMode()).isTrue();
 
         Mockito.clearInvocations(monitor);
 
         // Start a dream from the second client and verify that the overlay has both cycled to
         // the new dream (ended/started).
         secondClient.startDream(mLayoutParams, mOverlayCallback,
-                SECOND_DREAM_COMPONENT.flattenToString(), false);
+                SECOND_DREAM_COMPONENT.flattenToString(), false, false);
 
         verify(monitor).onEndDream();
         verify(monitor).onStartDream();
         assertThat(service.getDreamComponent()).isEqualTo(SECOND_DREAM_COMPONENT);
+        assertThat(service.isDreamInPreviewMode()).isFalse();
 
         Mockito.clearInvocations(monitor);
 
@@ -221,6 +227,47 @@
         verify(monitor, never()).onWakeUp();
     }
 
+    /**
+     * Verifies that only the currently started dream is able to affect the overlay.
+     */
+    @Test
+    @EnableFlags(Flags.FLAG_DREAM_WAKE_REDIRECT)
+    public void testRedirectToWakeAcrossClients() throws RemoteException {
+        doAnswer(invocation -> {
+            ((Runnable) invocation.getArgument(0)).run();
+            return null;
+        }).when(mExecutor).execute(any());
+
+        final TestDreamOverlayService.Monitor monitor = Mockito.mock(
+                TestDreamOverlayService.Monitor.class);
+        final TestDreamOverlayService service = new TestDreamOverlayService(monitor, mExecutor);
+        final IBinder binder = service.onBind(new Intent());
+        final IDreamOverlay overlay = IDreamOverlay.Stub.asInterface(binder);
+
+        service.redirectWake(true);
+
+        final IDreamOverlayClient client = getClient(overlay);
+
+        // Start the dream.
+        client.startDream(mLayoutParams, mOverlayCallback,
+                FIRST_DREAM_COMPONENT.flattenToString(), false, false);
+        // Make sure redirect state is set on dream.
+        verify(mOverlayCallback).onRedirectWake(eq(true));
+
+        // Make sure new changes are propagated.
+        clearInvocations(mOverlayCallback);
+        service.redirectWake(false);
+        verify(mOverlayCallback).onRedirectWake(eq(false));
+
+
+        // Start another dream, make sure new dream is informed of current state.
+        service.redirectWake(true);
+        clearInvocations(mOverlayCallback);
+        client.startDream(mLayoutParams, mOverlayCallback,
+                FIRST_DREAM_COMPONENT.flattenToString(), false, false);
+        verify(mOverlayCallback).onRedirectWake(eq(true));
+    }
+
     private static IDreamOverlayClient getClient(IDreamOverlay overlay) throws RemoteException {
         final OverlayClientCallback callback = new OverlayClientCallback();
         overlay.getClient(callback);
diff --git a/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java b/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
index 43aa7fe..7c239ef 100644
--- a/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
+++ b/services/tests/dreamservicetests/src/com/android/server/dreams/TestDreamEnvironment.java
@@ -385,7 +385,8 @@
             final ArgumentCaptor<IDreamOverlayCallback> overlayCallbackCaptor =
                     ArgumentCaptor.forClass(IDreamOverlayCallback.class);
             verify(mDreamOverlayClient, description("dream client not informed of dream start"))
-                    .startDream(any(), overlayCallbackCaptor.capture(), any(), anyBoolean());
+                    .startDream(any(), overlayCallbackCaptor.capture(), any(), anyBoolean(),
+                            anyBoolean());
 
             mDreamOverlayCallback = overlayCallbackCaptor.getValue();
         }
diff --git a/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java b/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
index 8656b99..51aa528 100644
--- a/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
+++ b/services/tests/mockingservicestests/src/com/android/server/am/MockingOomAdjusterTests.java
@@ -46,6 +46,7 @@
 
 import static com.android.server.am.ActivityManagerService.FOLLOW_UP_OOMADJUSTER_UPDATE_MSG;
 import static com.android.server.am.ProcessList.BACKUP_APP_ADJ;
+import static com.android.server.am.ProcessList.CACHED_APP_IMPORTANCE_LEVELS;
 import static com.android.server.am.ProcessList.CACHED_APP_MAX_ADJ;
 import static com.android.server.am.ProcessList.CACHED_APP_MIN_ADJ;
 import static com.android.server.am.ProcessList.FOREGROUND_APP_ADJ;
@@ -844,6 +845,49 @@
 
     @SuppressWarnings("GuardedBy")
     @Test
+    public void testUpdateOomAdj_DoAll_PreviousApp() {
+        final int numberOfApps = 15;
+        final ProcessRecord[] apps = new ProcessRecord[numberOfApps];
+        for (int i = 0; i < numberOfApps; i++) {
+            apps[i] = spy(makeDefaultProcessRecord(MOCKAPP_PID + i, MOCKAPP_UID + i,
+                    MOCKAPP_PROCESSNAME + i, MOCKAPP_PACKAGENAME + i, true));
+            final WindowProcessController wpc = apps[i].getWindowProcessController();
+            doReturn(true).when(wpc).isPreviousProcess();
+            doReturn(true).when(wpc).hasActivities();
+        }
+        mService.mWakefulness.set(PowerManagerInternal.WAKEFULNESS_AWAKE);
+        setProcessesToLru(apps);
+        mService.mOomAdjuster.updateOomAdjLocked(OOM_ADJ_REASON_NONE);
+
+        for (int i = 0; i < numberOfApps; i++) {
+            assertProcStates(apps[i], PROCESS_STATE_LAST_ACTIVITY, PREVIOUS_APP_ADJ,
+                    SCHED_GROUP_BACKGROUND, "previous");
+        }
+
+        if (!Flags.followUpOomadjUpdates()) return;
+
+        for (int i = 0; i < numberOfApps; i++) {
+            final ArgumentCaptor<Long> followUpTimeCaptor = ArgumentCaptor.forClass(Long.class);
+            verify(mService.mHandler).sendEmptyMessageAtTime(eq(FOLLOW_UP_OOMADJUSTER_UPDATE_MSG),
+                    followUpTimeCaptor.capture());
+            mInjector.jumpUptimeAheadTo(followUpTimeCaptor.getValue());
+        }
+
+        mService.mOomAdjuster.updateOomAdjFollowUpTargetsLocked();
+
+        for (int i = 0; i < numberOfApps; i++) {
+            final int mruIndex = numberOfApps - i - 1;
+            int expectedAdj = CACHED_APP_MIN_ADJ + (mruIndex * 2 * CACHED_APP_IMPORTANCE_LEVELS);
+            if (expectedAdj > CACHED_APP_MAX_ADJ) {
+                expectedAdj = CACHED_APP_MAX_ADJ;
+            }
+            assertProcStates(apps[i], PROCESS_STATE_LAST_ACTIVITY, expectedAdj,
+                    SCHED_GROUP_BACKGROUND, "previous-expired");
+        }
+    }
+
+    @SuppressWarnings("GuardedBy")
+    @Test
     public void testUpdateOomAdj_DoOne_Backup() {
         ProcessRecord app = spy(makeDefaultProcessRecord(MOCKAPP_PID, MOCKAPP_UID,
                 MOCKAPP_PROCESSNAME, MOCKAPP_PACKAGENAME, true));
diff --git a/services/tests/mockingservicestests/src/com/android/server/crashrecovery/CrashRecoveryUtilsTest.java b/services/tests/mockingservicestests/src/com/android/server/crashrecovery/CrashRecoveryUtilsTest.java
new file mode 100644
index 0000000..6f38fca
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/crashrecovery/CrashRecoveryUtilsTest.java
@@ -0,0 +1,145 @@
+/*
+ * 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.crashrecovery;
+
+
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.quality.Strictness.LENIENT;
+import static org.junit.Assert.fail;
+
+import android.content.Context;
+import android.os.Environment;
+import android.util.IndentingPrintWriter;
+import android.util.Log;
+
+import androidx.test.core.app.ApplicationProvider;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.dx.mockito.inline.extended.ExtendedMockito;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.MockitoSession;
+
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.StringWriter;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+
+
+/**
+ * Test CrashRecovery Utils.
+ */
+@RunWith(AndroidJUnit4.class)
+public class CrashRecoveryUtilsTest {
+
+    private MockitoSession mStaticMockSession;
+    private final String mLogMsg = "Logging from test";
+    private final String mCrashrecoveryEventTag = "CrashRecovery Events: ";
+    private File mCacheDir;
+
+    @Before
+    public void setup() throws IOException {
+        Context context = ApplicationProvider.getApplicationContext();
+        mCacheDir = context.getCacheDir();
+        mStaticMockSession = ExtendedMockito.mockitoSession()
+                .spyStatic(Environment.class)
+                .strictness(LENIENT)
+                .startMocking();
+        ExtendedMockito.doReturn(mCacheDir).when(() -> Environment.getDataDirectory());
+
+        createCrashRecoveryEventsTempDir();
+    }
+
+    @After
+    public void tearDown() throws IOException {
+        mStaticMockSession.finishMocking();
+        deleteCrashRecoveryEventsTempFile();
+    }
+
+    @Test
+    public void testCrashRecoveryUtils() {
+        testLogCrashRecoveryEvent();
+        testDumpCrashRecoveryEvents();
+    }
+
+    @Test
+    public void testDumpCrashRecoveryEventsWithoutAnyLogs() {
+        assertThat(getCrashRecoveryEventsTempFile().exists()).isFalse();
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter ipw = new IndentingPrintWriter(sw, "  ");
+        CrashRecoveryUtils.dumpCrashRecoveryEvents(ipw);
+        ipw.close();
+
+        String dump = sw.getBuffer().toString();
+        assertThat(dump).contains(mCrashrecoveryEventTag);
+        assertThat(dump).doesNotContain(mLogMsg);
+    }
+
+    private void testLogCrashRecoveryEvent() {
+        assertThat(getCrashRecoveryEventsTempFile().exists()).isFalse();
+        CrashRecoveryUtils.logCrashRecoveryEvent(Log.WARN, mLogMsg);
+
+        assertThat(getCrashRecoveryEventsTempFile().exists()).isTrue();
+        String fileContent = null;
+        try {
+            File file = getCrashRecoveryEventsTempFile();
+            FileInputStream fis = new FileInputStream(file);
+            byte[] data = new byte[(int) file.length()];
+            fis.read(data);
+            fis.close();
+            fileContent = new String(data, StandardCharsets.UTF_8);
+        } catch (Exception e) {
+            fail("Unable to read the events file");
+        }
+        assertThat(fileContent).contains(mLogMsg);
+    }
+
+    private void testDumpCrashRecoveryEvents() {
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter ipw = new IndentingPrintWriter(sw, "  ");
+        CrashRecoveryUtils.dumpCrashRecoveryEvents(ipw);
+        ipw.close();
+
+        String dump = sw.getBuffer().toString();
+        assertThat(dump).contains(mCrashrecoveryEventTag);
+        assertThat(dump).contains(mLogMsg);
+    }
+
+    private void createCrashRecoveryEventsTempDir() throws IOException {
+        Files.deleteIfExists(getCrashRecoveryEventsTempFile().toPath());
+        File mMockDirectory = new File(mCacheDir, "system");
+        if (!mMockDirectory.exists()) {
+            assertThat(mMockDirectory.mkdir()).isTrue();
+        }
+    }
+
+    private void deleteCrashRecoveryEventsTempFile() throws IOException {
+        Files.deleteIfExists(getCrashRecoveryEventsTempFile().toPath());
+    }
+
+    private File getCrashRecoveryEventsTempFile() {
+        File systemTempDir = new File(mCacheDir, "system");
+        return new File(systemTempDir, "crashrecovery-events.txt");
+    }
+}
diff --git a/services/tests/powerservicetests/src/com/android/server/power/WakeLockLogTest.java b/services/tests/powerservicetests/src/com/android/server/power/WakeLockLogTest.java
index 1c4db6a..c1d7c7b 100644
--- a/services/tests/powerservicetests/src/com/android/server/power/WakeLockLogTest.java
+++ b/services/tests/powerservicetests/src/com/android/server/power/WakeLockLogTest.java
@@ -25,6 +25,7 @@
 import android.content.Context;
 import android.content.pm.PackageManager;
 import android.os.PowerManager;
+import android.os.Process;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -54,6 +55,8 @@
 
         when(mPackageManager.getPackagesForUid(101)).thenReturn(new String[]{ "some.package1" });
         when(mPackageManager.getPackagesForUid(102)).thenReturn(new String[]{ "some.package2" });
+        when(mPackageManager.getPackagesForUid(Process.SYSTEM_UID))
+                .thenReturn(new String[]{ "some.package3" });
     }
 
     @Test
@@ -70,14 +73,20 @@
         log.onWakeLockAcquired("TagFull", 102,
                 PowerManager.FULL_WAKE_LOCK | PowerManager.ACQUIRE_CAUSES_WAKEUP, -1);
 
+        when(injectorSpy.currentTimeMillis()).thenReturn(1250L);
+        log.onWakeLockAcquired("TagSystem", 1000,
+                PowerManager.FULL_WAKE_LOCK | PowerManager.ACQUIRE_CAUSES_WAKEUP, -1);
+
         assertEquals("Wake Lock Log\n"
                         + "  01-01 00:00:01.000 - 101 (some.package1) - ACQ TagPartial "
                         + "(partial,on-after-release)\n"
                         + "  01-01 00:00:01.150 - 102 (some.package2) - ACQ TagFull "
                         + "(full,acq-causes-wake)\n"
+                        + "  01-01 00:00:01.250 - 1000 (" + WakeLockLog.SYSTEM_PACKAGE_NAME + ")"
+                        + " - ACQ TagSystem (full,acq-causes-wake)\n"
                         + "  -\n"
-                        + "  Events: 2, Time-Resets: 0\n"
-                        + "  Buffer, Bytes used: 6\n",
+                        + "  Events: 3, Time-Resets: 0\n"
+                        + "  Buffer, Bytes used: 9\n",
                 dumpLog(log, false));
     }
 
diff --git a/services/tests/powerstatstests/Android.bp b/services/tests/powerstatstests/Android.bp
index b2a5b02..91c62be 100644
--- a/services/tests/powerstatstests/Android.bp
+++ b/services/tests/powerstatstests/Android.bp
@@ -71,6 +71,8 @@
     ],
     srcs: [
         "src/com/android/server/power/stats/*.java",
+        "src/com/android/server/power/stats/format/*.java",
+        "src/com/android/server/power/stats/processor/*.java",
     ],
     java_resources: [
         "res/xml/power_profile*.xml",
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
index a1101cd..1d20538 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryStatsImplTest.java
@@ -117,7 +117,7 @@
     private PowerStatsStore mPowerStatsStore;
     private BatteryUsageStatsProvider mBatteryUsageStatsProvider;
     @Mock
-    private PowerStatsExporter mPowerStatsExporter;
+    private PowerAttributor mPowerAttributor;
 
     @Before
     public void setUp() throws IOException {
@@ -149,9 +149,8 @@
         } else {
             context = InstrumentationRegistry.getContext();
         }
-        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler,
-                new AggregatedPowerStatsConfig());
-        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mPowerStatsExporter,
+        mPowerStatsStore = new PowerStatsStore(systemDir, mHandler);
+        mBatteryUsageStatsProvider = new BatteryUsageStatsProvider(context, mPowerAttributor,
                 mPowerProfile, mBatteryStatsImpl.getCpuScalingPolicies(), mPowerStatsStore,
                 mMockClock);
     }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
index 17c7efa..fde84e9 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BatteryUsageStatsProviderTest.java
@@ -94,8 +94,9 @@
     public void test_getBatteryUsageStats() {
         BatteryStatsImpl batteryStats = prepareBatteryStats();
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats, BatteryUsageStatsQuery.DEFAULT);
@@ -130,8 +131,9 @@
     public void test_selectPowerComponents() {
         BatteryStatsImpl batteryStats = prepareBatteryStats();
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -235,8 +237,9 @@
             batteryStats.noteAlarmFinishLocked("foo", null, APP_UID, 3_001_000, 2_001_000);
         }
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -323,8 +326,9 @@
             }
         }
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
 
         final BatteryUsageStats batteryUsageStats =
                 provider.getBatteryUsageStats(batteryStats,
@@ -408,12 +412,12 @@
 
         PowerStatsStore powerStatsStore = new PowerStatsStore(
                 new File(mStatsRule.getHistoryDir(), "powerstatsstore"),
-                mStatsRule.getHandler(), null);
+                mStatsRule.getHandler());
         powerStatsStore.reset();
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), powerStatsStore,
-                mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), powerStatsStore, mMockClock);
 
         batteryStats.saveBatteryUsageStatsOnReset(provider, powerStatsStore);
         synchronized (batteryStats) {
@@ -522,8 +526,9 @@
             batteryStats.updateCustomEnergyConsumerStatsLocked(1, 200_000_000, uidEnergies);
         }
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), null, mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), mock(PowerStatsStore.class), mMockClock);
 
         PowerStatsStore powerStatsStore = mock(PowerStatsStore.class);
         doAnswer(invocation -> {
@@ -584,9 +589,9 @@
         when(powerStatsStore.loadPowerStatsSpan(1, BatteryUsageStatsSection.TYPE))
                 .thenReturn(span1);
 
-        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext, null,
-                mStatsRule.getPowerProfile(), mStatsRule.getCpuScalingPolicies(), powerStatsStore,
-                mMockClock);
+        BatteryUsageStatsProvider provider = new BatteryUsageStatsProvider(mContext,
+                mock(PowerAttributor.class), mStatsRule.getPowerProfile(),
+                mStatsRule.getCpuScalingPolicies(), powerStatsStore, mMockClock);
 
         BatteryUsageStatsQuery query = new BatteryUsageStatsQuery.Builder()
                 .aggregateSnapshots(0, 3000)
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
index 02c7b74..e392c5d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsCollectorTest.java
@@ -27,6 +27,7 @@
 import android.bluetooth.UidTraffic;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.Handler;
@@ -37,6 +38,7 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.BluetoothPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -48,7 +50,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.function.IntSupplier;
 
 public class BluetoothPowerStatsCollectorTest {
     private static final int APP_UID1 = 42;
@@ -132,11 +133,6 @@
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> 3500;
-                }
-
-                @Override
                 public BluetoothPowerStatsCollector.BluetoothStatsRetriever
                         getBluetoothStatsRetriever() {
                     return mBluetoothStatsRetriever;
@@ -232,6 +228,7 @@
         BluetoothPowerStatsCollector collector = new BluetoothPowerStatsCollector(mInjector);
         collector.setEnabled(true);
 
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH))
                 .thenReturn(new int[]{777});
 
@@ -242,8 +239,7 @@
 
         mUidScanTimes.put(APP_UID1, 100);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        mockConsumedEnergy(777, 10000);
 
         // Establish a baseline
         collector.collectStats();
@@ -258,13 +254,19 @@
         mUidScanTimes.put(APP_UID2, 300);
         mUidScanTimes.put(ISOLATED_UID, 400);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        mockConsumedEnergy(777, 64321);
 
         mStatsRule.setTime(20000, 20000);
         return collector.collectStats();
     }
 
+    private void mockConsumedEnergy(int consumerId, long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{consumerId})))
+                .thenReturn(new EnergyConsumerResult[]{ecr});
+    }
+
     private BluetoothActivityEnergyInfo mockBluetoothActivityEnergyInfo(long timestamp,
             long rxTimeMs, long txTimeMs, long idleTimeMs, UidTraffic... uidTraffic) {
         if (RavenwoodRule.isOnRavenwood()) {
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 d1105a4..1fea462 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
@@ -26,6 +26,7 @@
 import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.ConditionVariable;
@@ -42,6 +43,7 @@
 import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.CpuPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -53,7 +55,6 @@
 
 import java.io.IOException;
 import java.io.StringWriter;
-import java.util.function.IntSupplier;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
@@ -124,11 +125,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
             return 0;
         }
@@ -150,7 +146,9 @@
         mHandlerThread.start();
         mHandler = mHandlerThread.getThreadHandler();
         when(mMockKernelCpuStatsReader.isSupportedFeature()).thenReturn(true);
-        when(mConsumedEnergyRetriever.getEnergyConsumerIds(anyInt())).thenReturn(new int[0]);
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(anyInt()))
+                .thenReturn(new int[0]);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         mUidResolver.noteIsolatedUidAdded(ISOLATED_UID, UID_2);
     }
 
@@ -228,9 +226,7 @@
         assertThat(descriptor.name).isEqualTo("cpu");
         assertThat(descriptor.statsArrayLength).isEqualTo(13);
         assertThat(descriptor.uidStatsArrayLength).isEqualTo(5);
-        CpuPowerStatsLayout layout =
-                new CpuPowerStatsLayout();
-        layout.fromExtras(descriptor.extras);
+        CpuPowerStatsLayout layout = new CpuPowerStatsLayout(descriptor);
 
         long[] deviceStats = new long[descriptor.statsArrayLength];
         layout.setTimeByScalingStep(deviceStats, 2, 42);
@@ -267,8 +263,7 @@
         mockEnergyConsumers();
 
         CpuPowerStatsCollector collector = createCollector(8, 0);
-        CpuPowerStatsLayout layout = new CpuPowerStatsLayout();
-        layout.fromExtras(collector.getPowerStatsDescriptor().extras);
+        CpuPowerStatsLayout layout = new CpuPowerStatsLayout(collector.getPowerStatsDescriptor());
 
         mockKernelCpuStats(new long[]{1111, 2222, 3333},
                 new SparseArray<>() {{
@@ -338,8 +333,7 @@
         mockEnergyConsumers();
 
         CpuPowerStatsCollector collector = createCollector(8, 0);
-        CpuPowerStatsLayout layout = new CpuPowerStatsLayout();
-        layout.fromExtras(collector.getPowerStatsDescriptor().extras);
+        CpuPowerStatsLayout layout = new CpuPowerStatsLayout(collector.getPowerStatsDescriptor());
 
         mockKernelCpuStats(new long[]{1111, 2222, 3333},
                 new SparseArray<>() {{
@@ -462,17 +456,24 @@
 
     private void mockEnergyConsumers() {
         reset(mConsumedEnergyRetriever);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         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});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{1, 2})))
+                .thenReturn(new EnergyConsumerResult[]{
+                        mockEnergyConsumer(1000), mockEnergyConsumer(2000)})
+                .thenReturn(new EnergyConsumerResult[]{
+                        mockEnergyConsumer(1500), mockEnergyConsumer(2700)});
+    }
+
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
     }
 
     private static int[] getScalingStepToPowerBracketMap(CpuPowerStatsCollector collector) {
-        CpuPowerStatsLayout layout =
-                new CpuPowerStatsLayout();
-        layout.fromExtras(collector.getPowerStatsDescriptor().extras);
+        CpuPowerStatsLayout layout = new CpuPowerStatsLayout(collector.getPowerStatsDescriptor());
         return layout.getScalingStepToPowerBracketMap();
     }
 
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
index ef20946..00b911b 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsCollectorTest.java
@@ -33,6 +33,7 @@
 import android.annotation.NonNull;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.os.BatteryConsumer;
@@ -51,6 +52,7 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -63,7 +65,6 @@
 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;
 
@@ -139,11 +140,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
             return mNetworkStatsSupplier;
         }
@@ -178,6 +174,7 @@
                 return uid;
             }
         });
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         mBatteryStats = mStatsRule.getBatteryStats();
     }
 
@@ -242,8 +239,7 @@
         assertThat(powerStats.durationMs).isEqualTo(100);
 
         PowerStats.Descriptor descriptor = powerStats.descriptor;
-        MobileRadioPowerStatsLayout layout =
-                new MobileRadioPowerStatsLayout(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);
@@ -252,7 +248,7 @@
                 .isEqualTo((64321 - 10000) * 1000 / 3500);
 
         assertThat(powerStats.stateStats.size()).isEqualTo(2);
-        long[] state1 = powerStats.stateStats.get(MobileRadioPowerStatsCollector.makeStateKey(
+        long[] state1 = powerStats.stateStats.get(MobileRadioPowerStatsLayout.makeStateKey(
                 BatteryStats.RADIO_ACCESS_TECHNOLOGY_NR,
                 ServiceState.FREQUENCY_RANGE_MMWAVE
         ));
@@ -263,7 +259,7 @@
         assertThat(layout.getStateTxTime(state1, 3)).isEqualTo(4000);
         assertThat(layout.getStateTxTime(state1, 4)).isEqualTo(5000);
 
-        long[] state2 = powerStats.stateStats.get(MobileRadioPowerStatsCollector.makeStateKey(
+        long[] state2 = powerStats.stateStats.get(MobileRadioPowerStatsLayout.makeStateKey(
                 BatteryStats.RADIO_ACCESS_TECHNOLOGY_LTE,
                 ServiceState.FREQUENCY_RANGE_LOW
         ));
@@ -298,15 +294,14 @@
         assertThat(powerStats.durationMs).isEqualTo(100);
 
         PowerStats.Descriptor descriptor = powerStats.descriptor;
-        MobileRadioPowerStatsLayout layout =
-                new MobileRadioPowerStatsLayout(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(
+        long[] stateStats = powerStats.stateStats.get(MobileRadioPowerStatsLayout.makeStateKey(
                 AccessNetworkConstants.AccessNetworkType.UNKNOWN,
                 ServiceState.FREQUENCY_RANGE_UNKNOWN
         ));
@@ -416,8 +411,8 @@
                 4321, 321, 1234, 23,
                 4000, 40, 2000, 20);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10000)});
 
         when(mCallDurationSupplier.getAsLong()).thenReturn(10000L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(20000L);
@@ -439,8 +434,8 @@
                 5321, 421, 3234, 223,
                 8000, 80, 4000, 40);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(64321)});
         when(mCallDurationSupplier.getAsLong()).thenReturn(50000L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(80000L);
 
@@ -448,6 +443,12 @@
         return collector.collectStats();
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     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) {
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 89d6c1c..a04f721 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
@@ -114,13 +114,15 @@
         mockEnergyConsumers(powerStatsInternal);
 
         PowerStatsCollector.ConsumedEnergyRetrieverImpl retriever =
-                new PowerStatsCollector.ConsumedEnergyRetrieverImpl(powerStatsInternal);
+                new PowerStatsCollector.ConsumedEnergyRetrieverImpl(powerStatsInternal, ()-> 3500);
         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});
+        EnergyConsumerResult[] energy = retriever.getConsumedEnergy(energyConsumerIds);
+        assertThat(energy[0].energyUWs).isEqualTo(1000);
+        assertThat(energy[1].energyUWs).isEqualTo(2000);
+        energy = retriever.getConsumedEnergy(energyConsumerIds);
+        assertThat(energy[0].energyUWs).isEqualTo(1500);
+        assertThat(energy[1].energyUWs).isEqualTo(2700);
     }
 
     @SuppressWarnings("unchecked")
@@ -176,4 +178,11 @@
                 .thenReturn(future1)
                 .thenReturn(future2);
     }
+
+    private EnergyConsumerResult mockEnergyConsumerResult(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
index beec661..143d046 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsSchedulerTest.java
@@ -18,36 +18,18 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyLong;
-import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-
-import android.os.ConditionVariable;
-import android.os.Handler;
-import android.os.HandlerThread;
 import android.platform.test.ravenwood.RavenwoodRule;
 
 import androidx.test.runner.AndroidJUnit4;
 
-import com.android.internal.os.MonotonicClock;
-
-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;
 import java.time.Duration;
 import java.time.Instant;
-import java.util.ArrayList;
-import java.util.List;
 import java.util.TimeZone;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
 
 @RunWith(AndroidJUnit4.class)
 public class PowerStatsSchedulerTest {
@@ -56,134 +38,10 @@
             .setProvideMainThread(true)
             .build();
 
-    private PowerStatsStore mPowerStatsStore;
-    private Handler mHandler;
-    private MockClock mClock = new MockClock();
-    private MonotonicClock mMonotonicClock = new MonotonicClock(0, mClock);
-    private PowerStatsScheduler mPowerStatsScheduler;
-    private PowerStatsAggregator mPowerStatsAggregator;
-    private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
-    private List<Long> mScheduledAlarms = new ArrayList<>();
-    private boolean mPowerStatsCollectionOccurred;
-
-    private static final int START_REALTIME = 7654321;
-
-    @Before
-    @SuppressWarnings("GuardedBy")
-    public void setup() throws IOException {
-        TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
-
-        mClock.currentTime = Instant.parse("2023-01-02T03:04:05.00Z").toEpochMilli();
-        mClock.realtime = START_REALTIME;
-
-        HandlerThread bgThread = new HandlerThread("bg thread");
-        bgThread.start();
-        mHandler = new Handler(bgThread.getLooper());
-        mAggregatedPowerStatsConfig = new AggregatedPowerStatsConfig();
-        mPowerStatsStore = new PowerStatsStore(
-                Files.createTempDirectory("PowerStatsSchedulerTest").toFile(),
-                mHandler, mAggregatedPowerStatsConfig);
-        mPowerStatsAggregator = mock(PowerStatsAggregator.class);
-        mPowerStatsScheduler = new PowerStatsScheduler(
-                () -> mPowerStatsCollectionOccurred = true,
-                mPowerStatsAggregator, TimeUnit.MINUTES.toMillis(30), TimeUnit.HOURS.toMillis(1),
-                mPowerStatsStore,
-                ((triggerAtMillis, tag, onAlarmListener, handler) ->
-                        mScheduledAlarms.add(triggerAtMillis)),
-                mClock, mMonotonicClock, () -> 12345L, mHandler);
-    }
-
-    @Test
-    @SuppressWarnings("unchecked")
-    public void storeAggregatePowerStats() {
-        mPowerStatsStore.reset();
-
-        assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
-
-        mPowerStatsStore.storeAggregatedPowerStats(
-                createAggregatedPowerStats(mMonotonicClock.monotonicTime(), mClock.currentTime,
-                        123));
-
-        long delayBeforeAggregating = TimeUnit.MINUTES.toMillis(90);
-        mClock.realtime += delayBeforeAggregating;
-        mClock.currentTime += delayBeforeAggregating;
-
-        doAnswer(invocation -> {
-            // The first span is longer than 30 min, because the end time is being aligned with
-            // the wall clock.  Subsequent spans should be precisely 30 minutes.
-            long startTime = invocation.getArgument(0);
-            long endTime = invocation.getArgument(1);
-            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(2);
-
-            long startTimeWallClock =
-                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
-            long endTimeWallClock =
-                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
-
-            assertThat(startTime).isEqualTo(START_REALTIME + 123);
-            assertThat(endTime - startTime).isAtLeast(TimeUnit.MINUTES.toMillis(30));
-            assertThat(Instant.ofEpochMilli(endTimeWallClock))
-                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
-
-            consumer.accept(
-                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
-            return null;
-        }).doAnswer(invocation -> {
-            long startTime = invocation.getArgument(0);
-            long endTime = invocation.getArgument(1);
-            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(2);
-
-            long startTimeWallClock =
-                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
-            long endTimeWallClock =
-                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
-
-            assertThat(Instant.ofEpochMilli(startTimeWallClock))
-                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
-            assertThat(Instant.ofEpochMilli(endTimeWallClock))
-                    .isEqualTo(Instant.parse("2023-01-02T04:30:00Z"));
-
-            consumer.accept(
-                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
-            return null;
-        }).when(mPowerStatsAggregator).aggregatePowerStats(anyLong(), anyLong(),
-                any(Consumer.class));
-
-        mPowerStatsScheduler.start(/*enabled*/ true);
-        ConditionVariable done = new ConditionVariable();
-        mHandler.post(done::open);
-        done.block();
-
-        assertThat(mPowerStatsCollectionOccurred).isTrue();
-        assertThat(mScheduledAlarms).containsExactly(
-                START_REALTIME + TimeUnit.MINUTES.toMillis(90) + TimeUnit.HOURS.toMillis(1));
-
-        verify(mPowerStatsAggregator, times(2))
-                .aggregatePowerStats(anyLong(), anyLong(), any(Consumer.class));
-
-        List<PowerStatsSpan.Metadata> contents = mPowerStatsStore.getTableOfContents();
-        assertThat(contents).hasSize(3);
-        // Skip the first entry, which was placed in the store at the beginning of this test
-        PowerStatsSpan.TimeFrame timeFrame1 = contents.get(1).getTimeFrames().get(0);
-        PowerStatsSpan.TimeFrame timeFrame2 = contents.get(2).getTimeFrames().get(0);
-        assertThat(timeFrame1.startMonotonicTime).isEqualTo(START_REALTIME + 123);
-        assertThat(timeFrame2.startMonotonicTime)
-                .isEqualTo(timeFrame1.startMonotonicTime + timeFrame1.duration);
-        assertThat(Instant.ofEpochMilli(timeFrame2.startTime))
-                .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
-        assertThat(Duration.ofMillis(timeFrame2.duration)).isEqualTo(Duration.ofMinutes(30));
-    }
-
-    private AggregatedPowerStats createAggregatedPowerStats(long monotonicTime, long currentTime,
-            long duration) {
-        AggregatedPowerStats stats = new AggregatedPowerStats(mAggregatedPowerStatsConfig);
-        stats.addClockUpdate(monotonicTime, currentTime);
-        stats.setDuration(duration);
-        return stats;
-    }
-
     @Test
     public void alignToWallClock() {
+        TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
+
         // Expect the aligned value to be adjusted by 1 min 30 sec - rounded to the next 15 min
         assertThat(PowerStatsScheduler.alignToWallClock(123, TimeUnit.MINUTES.toMillis(15),
                 123 + TimeUnit.HOURS.toMillis(2),
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsStoreTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsStoreTest.java
index 36d7af5..dc8d920 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsStoreTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsStoreTest.java
@@ -59,14 +59,7 @@
         clearDirectory(mStoreDirectory);
 
         mPowerStatsStore = new PowerStatsStore(mStoreDirectory,
-                MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES,
-                new TestHandler(),
-                (sectionType, parser) -> {
-                    if (sectionType.equals(TestSection.TYPE)) {
-                        return TestSection.readXml(parser);
-                    }
-                    return null;
-                });
+                MAX_BATTERY_STATS_SNAPSHOT_STORAGE_BYTES, new TestHandler());
     }
 
     @Test
@@ -144,7 +137,7 @@
         }
 
         @Override
-        void write(TypedXmlSerializer serializer) throws IOException {
+        public void write(TypedXmlSerializer serializer) throws IOException {
             StringBuilder sb = new StringBuilder();
             for (int i = 0; i < mSize; i++) {
                 sb.append("X");
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
index 817fdcb..8c09d1d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsCollectorTest.java
@@ -23,6 +23,7 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -32,6 +33,7 @@
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
 import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -39,8 +41,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
-
 public class ScreenPowerStatsCollectorTest {
     private static final int APP_UID1 = 42;
     private static final int APP_UID2 = 24;
@@ -89,11 +89,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public int getDisplayCount() {
             return 2;
         }
@@ -115,6 +110,7 @@
                 return uid;
             }
         });
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
     }
 
     @Test
@@ -125,8 +121,8 @@
         // Establish a baseline
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                 .thenReturn(new int[]{77});
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                .thenReturn(new long[]{10_000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
 
         doAnswer(inv -> {
             ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
@@ -139,8 +135,8 @@
 
         collector.collectStats();
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                .thenReturn(new long[]{45_000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(45_000)});
         when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
                 .thenReturn(60_000L);
         when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
@@ -171,8 +167,7 @@
 
         PowerStats powerStats = collector.collectStats();
 
-        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout();
-        layout.fromExtras(powerStats.descriptor.extras);
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(powerStats.descriptor);
 
         // (45000 - 10000) / 3500
         assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
@@ -204,4 +199,10 @@
         assertThat(layout.getUidTopActivityDuration(powerStats.uidStats.get(APP_UID2)))
                 .isEqualTo(10000);
     }
+
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
 }
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
index b13fc53..8b5e6ee 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
@@ -32,6 +32,7 @@
 
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.net.wifi.WifiManager;
@@ -47,6 +48,7 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.format.WifiPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -57,7 +59,6 @@
 import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class WifiPowerStatsCollectorTest {
@@ -154,11 +155,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> 3500;
-        }
-
-        @Override
         public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
             return mNetworkStatsSupplier;
         }
@@ -368,6 +364,7 @@
         WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, null);
         collector.setEnabled(true);
 
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(3500);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
                 .thenReturn(new int[]{777});
 
@@ -385,8 +382,8 @@
         mockWifiScanTimes(APP_UID2, 3000, 4000);
         mockWifiScanTimes(ISOLATED_UID, 5000, 6000);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{10000});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
 
         collector.collectStats();
 
@@ -404,13 +401,19 @@
         mockWifiScanTimes(APP_UID2, 3100, 4200);
         mockWifiScanTimes(ISOLATED_UID, 5300, 6400);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
-                .thenReturn(new long[]{64321});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(eq(new int[]{777})))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(64_321)});
 
         mStatsRule.setTime(20000, 20000);
         return collector.collectStats();
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockWifiActivityInfo(long timestamp, long rxTimeMs, long txTimeMs, int scanTimeMs,
             int idleTimeMs) {
         int stackState = 0;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AggregatedPowerStatsTest.java
similarity index 93%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/AggregatedPowerStatsTest.java
index 04d53de..0e73329 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AggregatedPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AggregatedPowerStatsTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -49,7 +49,8 @@
     private static final int COMPONENT_STATE_1 = 1;
     private static final int COMPONENT_STATE_2 = 2;
 
-    private AggregatedPowerStatsConfig mAggregatedPowerStatsConfig;
+    private AggregatedPowerStatsConfig
+            mAggregatedPowerStatsConfig;
     private PowerStats.Descriptor mPowerComponentDescriptor;
 
     @Before
@@ -67,7 +68,8 @@
         mAggregatedPowerStatsConfig.trackCustomPowerComponents(
                         () -> new PowerStatsProcessor() {
                             @Override
-                            void finish(PowerComponentAggregatedPowerStats stats,
+                            void finish(
+                                    PowerComponentAggregatedPowerStats stats,
                                     long timestampMs) {
                             }
                         })
@@ -103,8 +105,8 @@
 
         TypedXmlPullParser parser = Xml.newFastPullParser();
         parser.setInput(new ByteArrayInputStream(baos.toByteArray()), "UTF-8");
-        AggregatedPowerStats actualStats = AggregatedPowerStats.createFromXml(parser,
-                mAggregatedPowerStatsConfig);
+        AggregatedPowerStats actualStats =
+                AggregatedPowerStats.createFromXml(parser, mAggregatedPowerStatsConfig);
 
         verifyAggregatedPowerStats(actualStats);
     }
@@ -163,7 +165,8 @@
         return stats;
     }
 
-    private void verifyAggregatedPowerStats(AggregatedPowerStats stats) {
+    private void verifyAggregatedPowerStats(
+            AggregatedPowerStats stats) {
         PowerStats.Descriptor descriptor = stats.getPowerComponentStats(TEST_POWER_COMPONENT)
                 .getPowerStatsDescriptor();
         assertThat(descriptor.powerComponentId).isEqualTo(TEST_POWER_COMPONENT);
@@ -277,7 +280,8 @@
                 .isEqualTo(new long[]{250, 300});
     }
 
-    private static long[] getDeviceStats(AggregatedPowerStats stats, int powerComponentId,
+    private static long[] getDeviceStats(
+            AggregatedPowerStats stats, int powerComponentId,
             int... states) {
         PowerComponentAggregatedPowerStats powerComponentStats =
                 stats.getPowerComponentStats(powerComponentId);
@@ -286,7 +290,8 @@
         return out;
     }
 
-    private static long[] getStateStats(AggregatedPowerStats stats, int key, int... states) {
+    private static long[] getStateStats(
+            AggregatedPowerStats stats, int key, int... states) {
         PowerComponentAggregatedPowerStats powerComponentStats =
                 stats.getPowerComponentStats(TEST_POWER_COMPONENT);
         long[] out = new long[powerComponentStats.getPowerStatsDescriptor().stateStatsArrayLength];
@@ -294,7 +299,8 @@
         return out;
     }
 
-    private static long[] getUidDeviceStats(AggregatedPowerStats stats, int powerComponentId,
+    private static long[] getUidDeviceStats(
+            AggregatedPowerStats stats, int powerComponentId,
             int uid, int... states) {
         PowerComponentAggregatedPowerStats powerComponentStats =
                 stats.getPowerComponentStats(powerComponentId);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
similarity index 68%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
index a2a7e00..21e615f 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/AmbientDisplayPowerStatsProcessorTest.java
@@ -14,14 +14,14 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
-import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
-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_SCREEN;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -35,7 +35,12 @@
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
-import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.ScreenPowerStatsCollector;
+import com.android.server.power.stats.ScreenPowerStatsCollector.ScreenUsageTimeRetriever;
+import com.android.server.power.stats.format.PowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -43,8 +48,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
-
 public class AmbientDisplayPowerStatsProcessorTest {
 
     @Rule(order = 0)
@@ -64,49 +67,45 @@
     @Mock
     private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
     @Mock
-    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+    private ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
 
-    private final Injector mInjector = new Injector() {
-        @Override
-        public Handler getHandler() {
-            return mStatsRule.getHandler();
-        }
+    private final ScreenPowerStatsCollector.Injector mInjector =
+            new ScreenPowerStatsCollector.Injector() {
+                @Override
+                public Handler getHandler() {
+                    return mStatsRule.getHandler();
+                }
 
-        @Override
-        public Clock getClock() {
-            return mStatsRule.getMockClock();
-        }
+                @Override
+                public Clock getClock() {
+                    return mStatsRule.getMockClock();
+                }
 
-        @Override
-        public PowerStatsUidResolver getUidResolver() {
-            return new PowerStatsUidResolver();
-        }
+                @Override
+                public PowerStatsUidResolver getUidResolver() {
+                    return new PowerStatsUidResolver();
+                }
 
-        @Override
-        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
-            return 0;
-        }
+                @Override
+                public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+                    return 0;
+                }
 
-        @Override
-        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
-            return mConsumedEnergyRetriever;
-        }
+                @Override
+                public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+                    return mConsumedEnergyRetriever;
+                }
 
-        @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> VOLTAGE_MV;
-        }
+                @Override
+                public int getDisplayCount() {
+                    return 2;
+                }
 
-        @Override
-        public int getDisplayCount() {
-            return 2;
-        }
-
-        @Override
-        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
-            return mScreenUsageTimeRetriever;
-        }
-    };
+                @Override
+                public ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+                    return mScreenUsageTimeRetriever;
+                }
+            };
 
     @Before
     public void setup() {
@@ -167,7 +166,8 @@
         return stats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY);
     }
 
-    private void assertPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+    private void assertPowerEstimate(
+            PowerComponentAggregatedPowerStats aggregatedStats,
             int powerState, int screenState, double expectedPowerEstimate) {
         PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
         PowerStatsLayout layout = new PowerStatsLayout(descriptor);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BinaryStatePowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessorTest.java
similarity index 92%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/BinaryStatePowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessorTest.java
index 4b40f68..b412ad6 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BinaryStatePowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BinaryStatePowerStatsProcessorTest.java
@@ -14,19 +14,19 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -40,6 +40,9 @@
 
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.MockClock;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.BinaryStatePowerStatsLayout;
 
 import org.junit.Rule;
 import org.junit.Test;
@@ -146,7 +149,8 @@
 
     @Test
     public void energyConsumerModel() {
-        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout();
+        BinaryStatePowerStatsLayout
+                statsLayout = new BinaryStatePowerStatsLayout();
         PersistableBundle extras = new PersistableBundle();
         statsLayout.toExtras(extras);
         PowerStats.Descriptor descriptor = new PowerStats.Descriptor(POWER_COMPONENT,
@@ -270,9 +274,8 @@
                 .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
                 .setProcessorSupplier(processorSupplier);
 
-        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
         PowerComponentAggregatedPowerStats powerComponentStats =
-                aggregatedPowerStats.getPowerComponentStats(POWER_COMPONENT);
+                new AggregatedPowerStats(config).getPowerComponentStats(POWER_COMPONENT);
         powerComponentStats.start(0);
 
         powerComponentStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
similarity index 91%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
index 4a8125f..6dfc220 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/BluetoothPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/BluetoothPowerStatsProcessorTest.java
@@ -14,19 +14,19 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -38,6 +38,7 @@
 import android.bluetooth.UidTraffic;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.Handler;
@@ -48,7 +49,12 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.BluetoothPowerStatsCollector;
 import com.android.server.power.stats.BluetoothPowerStatsCollector.BluetoothStatsRetriever;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.BluetoothPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -58,7 +64,6 @@
 
 import java.util.List;
 import java.util.concurrent.Executor;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class BluetoothPowerStatsProcessorTest {
@@ -143,11 +148,6 @@
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public BluetoothStatsRetriever getBluetoothStatsRetriever() {
                     return mBluetoothStatsRetriever;
                 }
@@ -363,7 +363,8 @@
 
     @Test
     public void consumedEnergyModel() {
-        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
+        // Power monitoring hardware exists
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.BLUETOOTH))
                 .thenReturn(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID});
 
@@ -378,8 +379,8 @@
 
         mUidScanTimes.put(APP_UID1, 100);
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{BLUETOOTH_ENERGY_CONSUMER_ID})).thenReturn(new long[]{0});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -404,8 +405,8 @@
 
         // 10 mAh represented as microWattSeconds
         long energyUws = 10 * 3600 * VOLTAGE_MV;
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{BLUETOOTH_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{BLUETOOTH_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
 
@@ -468,16 +469,15 @@
 
     private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
             Supplier<PowerStatsProcessor> processorSupplier) {
-        AggregatedPowerStatsConfig.PowerComponent config =
-                new AggregatedPowerStatsConfig.PowerComponent(
-                        BatteryConsumer.POWER_COMPONENT_BLUETOOTH)
-                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
-                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
-                        .setProcessorSupplier(processorSupplier);
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_BLUETOOTH)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                .setProcessorSupplier(processorSupplier);
 
         PowerComponentAggregatedPowerStats aggregatedStats =
-                new PowerComponentAggregatedPowerStats(
-                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+                new AggregatedPowerStats(config).getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_BLUETOOTH);
 
         aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
         aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
@@ -491,6 +491,12 @@
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private BluetoothActivityEnergyInfo mockBluetoothActivityEnergyInfo(long timestamp,
             long rxTimeMs, long txTimeMs, long idleTimeMs, UidTraffic... uidTraffic) {
         if (RavenwoodRule.isOnRavenwood()) {
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CameraPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
similarity index 88%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/CameraPowerStatsTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
index 88a4f5e..0afcbf1 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CameraPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CameraPowerStatsTest.java
@@ -14,23 +14,22 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
@@ -46,6 +45,12 @@
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.CameraPowerStatsCollector;
+import com.android.server.power.stats.EnergyConsumerPowerStatsCollector;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.BinaryStatePowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -53,7 +58,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class CameraPowerStatsTest {
@@ -103,11 +107,6 @@
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
     private MonotonicClock mMonotonicClock;
@@ -120,8 +119,9 @@
 
     @Test
     public void energyConsumerModel() {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.CAMERA), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.CAMERA)))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID});
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -161,8 +161,7 @@
         stats.finish(11_000);
 
         PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
-        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout();
-        statsLayout.fromExtras(descriptor.extras);
+        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout(descriptor);
 
         // Total estimated power = 3,600,000 uC = 1.0 mAh
         // of which 3,000,000 is distributed:
@@ -243,7 +242,8 @@
 
     private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
             Supplier<PowerStatsProcessor> processorSupplier) {
-        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        AggregatedPowerStatsConfig
+                config = new AggregatedPowerStatsConfig();
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_CAMERA)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -254,9 +254,8 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessorSupplier(processorSupplier);
 
-        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
-        PowerComponentAggregatedPowerStats powerComponentStats =
-                aggregatedPowerStats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_CAMERA);
+        PowerComponentAggregatedPowerStats powerComponentStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_CAMERA);
         powerComponentStats.start(0);
 
         powerComponentStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CpuPowerStatsProcessorTest.java
similarity index 86%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/CpuPowerStatsProcessorTest.java
index ab2e631..6938615 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CpuPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CpuPowerStatsProcessorTest.java
@@ -13,24 +13,23 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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 com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
-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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 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;
@@ -42,6 +41,8 @@
 
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.format.CpuPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -202,26 +203,18 @@
             PowerComponentAggregatedPowerStats {
         private final CpuPowerStatsLayout mStatsLayout;
         private final PowerStats.Descriptor mDescriptor;
-        private HashMap<String, long[]> mDeviceStats = new HashMap<>();
-        private HashMap<String, long[]> mUidStats = new HashMap<>();
-        private HashSet<Integer> mUids = new HashSet<>();
-        private HashMap<String, Double> mExpectedDevicePower = new HashMap<>();
-        private HashMap<String, Double> mExpectedUidPower = new HashMap<>();
+        private final HashMap<String, long[]> mDeviceStats = new HashMap<>();
+        private final HashMap<String, long[]> mUidStats = new HashMap<>();
+        private final HashSet<Integer> mUids = new HashSet<>();
+        private final HashMap<String, Double> mExpectedDevicePower = new HashMap<>();
+        private final HashMap<String, Double> mExpectedUidPower = new HashMap<>();
 
-        MockPowerComponentAggregatedPowerStats(AggregatedPowerStatsConfig.PowerComponent config,
+        MockPowerComponentAggregatedPowerStats(
+                AggregatedPowerStatsConfig.PowerComponent config,
                 boolean useEnergyConsumers) {
-            super(new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
-            mStatsLayout = new CpuPowerStatsLayout();
-            mStatsLayout.addDeviceSectionCpuTimeByScalingStep(3);
-            mStatsLayout.addDeviceSectionCpuTimeByCluster(2);
-            mStatsLayout.addDeviceSectionUsageDuration();
-            if (useEnergyConsumers) {
-                mStatsLayout.addDeviceSectionEnergyConsumers(2);
-            }
-            mStatsLayout.addDeviceSectionPowerEstimate();
-            mStatsLayout.addUidSectionCpuTimeByPowerBracket(new int[]{0, 1, 2});
-            mStatsLayout.addUidSectionPowerEstimate();
-
+            super(new AggregatedPowerStats(new AggregatedPowerStatsConfig()), config);
+            mStatsLayout = new CpuPowerStatsLayout(useEnergyConsumers ? 2 : 0, 2,
+                    new int[]{0, 1, 2});
             PersistableBundle extras = new PersistableBundle();
             mStatsLayout.toExtras(extras);
             mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/CustomEnergyConsumerPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
similarity index 90%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/CustomEnergyConsumerPowerStatsTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
index 8239fdb..42baba7 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/CustomEnergyConsumerPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/CustomEnergyConsumerPowerStatsTest.java
@@ -14,20 +14,20 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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_BATTERY;
-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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -42,6 +42,12 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.CustomEnergyConsumerPowerStatsCollector;
+import com.android.server.power.stats.EnergyConsumerPowerStatsCollector;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -54,7 +60,6 @@
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
-import java.util.function.IntSupplier;
 
 public class CustomEnergyConsumerPowerStatsTest {
     @Rule(order = 0)
@@ -105,11 +110,6 @@
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
 
@@ -237,6 +237,7 @@
 
     private List<PowerStats> collectPowerStats(long timestamp, int chargeUc1, int chargeUc2,
             EnergyConsumerAttribution... attributions2) throws Exception {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.OTHER))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID1, ENERGY_CONSUMER_ID2});
         when(mConsumedEnergyRetriever.getEnergyConsumerName(ENERGY_CONSUMER_ID1))
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/GnssPowerStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
similarity index 93%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/GnssPowerStatsTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
index f22279a..e6207d4 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/GnssPowerStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/GnssPowerStatsTest.java
@@ -14,23 +14,22 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.when;
 
@@ -47,6 +46,12 @@
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.EnergyConsumerPowerStatsCollector;
+import com.android.server.power.stats.GnssPowerStatsCollector;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.BinaryStatePowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -54,7 +59,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class GnssPowerStatsTest {
@@ -106,11 +110,6 @@
                 public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
                     return mConsumedEnergyRetriever;
                 }
-
-                @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
             };
 
     private MonotonicClock mMonotonicClock;
@@ -127,7 +126,7 @@
     public void powerProfileModel() {
         // ODPM unsupported
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[0]);
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -165,8 +164,7 @@
         stats.finish(START_TIME + 11_000);
 
         PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
-        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout();
-        statsLayout.fromExtras(descriptor.extras);
+        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout(descriptor);
 
         // scr-on, GNSS-good: 2500 * 100 = 250000 mA-ms = 0.06944 mAh
         // scr-off GNSS=good: 4500 * 100 = 0.12500 mAh
@@ -218,7 +216,7 @@
     public void initialStateGnssOn() {
         // ODPM unsupported
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[0]);
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -245,8 +243,7 @@
         stats.finish(START_TIME + 11_000);
 
         PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
-        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout();
-        statsLayout.fromExtras(descriptor.extras);
+        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout(descriptor);
 
         // scr-on, GNSS-good: 2500 * 100 = 250000 mA-ms = 0.06944 mAh
         // scr-off GNSS=good: 4500 * 100 = 0.12500 mAh
@@ -296,8 +293,9 @@
 
     @Test
     public void energyConsumerModel() {
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever
-                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS), any()))
+                .getEnergyConsumerIds(eq((int) EnergyConsumerType.GNSS)))
                 .thenReturn(new int[]{ENERGY_CONSUMER_ID});
 
         PowerComponentAggregatedPowerStats stats = createAggregatedPowerStats(
@@ -339,8 +337,7 @@
         stats.finish(START_TIME + 11_000);
 
         PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
-        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout();
-        statsLayout.fromExtras(descriptor.extras);
+        BinaryStatePowerStatsLayout statsLayout = new BinaryStatePowerStatsLayout(descriptor);
 
         // Total estimated power = 3,600,000 uC = 1.0 mAh
         // of which 3,000,000 is distributed:
@@ -442,7 +439,8 @@
 
     private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
             Supplier<PowerStatsProcessor> processorSupplier) {
-        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        AggregatedPowerStatsConfig
+                config = new AggregatedPowerStatsConfig();
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_GNSS)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -453,9 +451,8 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessorSupplier(processorSupplier);
 
-        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
-        PowerComponentAggregatedPowerStats powerComponentStats =
-                aggregatedPowerStats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_GNSS);
+        PowerComponentAggregatedPowerStats powerComponentStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_GNSS);
         powerComponentStats.start(START_TIME);
 
         powerComponentStats.setState(STATE_POWER, POWER_STATE_OTHER, START_TIME);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
similarity index 87%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
index 89d59a9..80358c5 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MobileRadioPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MobileRadioPowerStatsProcessorTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
 import static android.net.NetworkStats.METERED_NO;
@@ -23,12 +23,12 @@
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -41,6 +41,7 @@
 import android.annotation.Nullable;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.os.BatteryConsumer;
@@ -53,6 +54,11 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.MobileRadioPowerStatsCollector;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.MobileRadioPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -61,7 +67,6 @@
 import org.mockito.MockitoAnnotations;
 
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -129,11 +134,6 @@
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
@@ -172,17 +172,15 @@
 
         mStatsRule.setTestPowerProfile("power_profile_test_modem_calculator");
 
-        AggregatedPowerStatsConfig.PowerComponent config =
-                new AggregatedPowerStatsConfig.PowerComponent(
-                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
-                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
-                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
-                        .setProcessorSupplier(() -> new MobileRadioPowerStatsProcessor(
-                                mStatsRule.getPowerProfile()));
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile()));
 
-        PowerComponentAggregatedPowerStats aggregatedStats =
-                new PowerComponentAggregatedPowerStats(
-                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+        PowerComponentAggregatedPowerStats aggregatedStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
 
         aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
         aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
@@ -232,8 +230,7 @@
         aggregatedStats.finish(10_000);
 
         MobileRadioPowerStatsLayout statsLayout =
-                new MobileRadioPowerStatsLayout(
-                        aggregatedStats.getPowerStatsDescriptor());
+                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)
@@ -316,8 +313,7 @@
                 prepareAggregatedStats_energyConsumerModel();
 
         MobileRadioPowerStatsLayout statsLayout =
-                new MobileRadioPowerStatsLayout(
-                        aggregatedStats.getPowerStatsDescriptor());
+                new MobileRadioPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
 
         // 10_000_000 micro-Coulomb * 1/1000 milli/micro * 1/3600 hour/second = 2.77778 mAh
         double totalPower = 0;
@@ -406,23 +402,22 @@
 
     private PowerComponentAggregatedPowerStats prepareAggregatedStats_energyConsumerModel() {
         // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.MOBILE_RADIO))
                 .thenReturn(new int[] {MOBILE_RADIO_ENERGY_CONSUMER_ID});
 
         mStatsRule.setTestPowerProfile("power_profile_test_legacy_modem")
                 .initMeasuredEnergyStatsLocked();
 
-        AggregatedPowerStatsConfig.PowerComponent config =
-                new AggregatedPowerStatsConfig.PowerComponent(
-                        BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
-                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
-                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
-                        .setProcessorSupplier(() -> new MobileRadioPowerStatsProcessor(
-                                mStatsRule.getPowerProfile()));
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                .setProcessorSupplier(
+                        () -> new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile()));
 
-        PowerComponentAggregatedPowerStats aggregatedStats =
-                new PowerComponentAggregatedPowerStats(
-                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+        PowerComponentAggregatedPowerStats aggregatedStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
 
         aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
         aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
@@ -436,9 +431,9 @@
         // Initial empty ModemActivityInfo.
         mockModemActivityInfo(new ModemActivityInfo(0L, 0L, 0L, new int[5], 0L));
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+        when(mConsumedEnergyRetriever.getConsumedEnergy(
                 new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID}))
-                .thenReturn(new long[]{0});
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -467,8 +462,8 @@
         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(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{MOBILE_RADIO_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         when(mCallDurationSupplier.getAsLong()).thenReturn(200L);
         when(mScanDurationSupplier.getAsLong()).thenReturn(5555L);
@@ -485,6 +480,12 @@
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockModemActivityInfo(ModemActivityInfo emptyMai) {
         doAnswer(invocation -> {
             OutcomeReceiver<ModemActivityInfo, TelephonyManager.ModemActivityInfoException>
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStatePowerAttributorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStatePowerAttributorTest.java
new file mode 100644
index 0000000..704ee62
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStatePowerAttributorTest.java
@@ -0,0 +1,183 @@
+/*
+ * 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.processor;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.os.ConditionVariable;
+import android.os.Handler;
+import android.os.HandlerThread;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.BatteryStatsHistory;
+import com.android.internal.os.MonotonicClock;
+import com.android.server.power.stats.MockClock;
+import com.android.server.power.stats.PowerStatsScheduler;
+import com.android.server.power.stats.PowerStatsSpan;
+import com.android.server.power.stats.PowerStatsStore;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.TimeZone;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+public class MultiStatePowerAttributorTest {
+
+    @Rule
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    private PowerStatsStore mPowerStatsStore;
+    private Handler mHandler;
+    private final MockClock mClock = new MockClock();
+    private final MonotonicClock mMonotonicClock = new MonotonicClock(0, mClock);
+    private PowerStatsScheduler mPowerStatsScheduler;
+    private PowerStatsAggregator mPowerStatsAggregator;
+    private MultiStatePowerAttributor mPowerAttributor;
+    private final List<Long> mScheduledAlarms = new ArrayList<>();
+    private boolean mPowerStatsCollectionOccurred;
+
+    private static final int START_REALTIME = 7654321;
+
+    @Before
+    public void setup() throws IOException {
+        TimeZone.setDefault(TimeZone.getTimeZone("UTC"));
+
+        mClock.currentTime = Instant.parse("2023-01-02T03:04:05.00Z").toEpochMilli();
+        mClock.realtime = START_REALTIME;
+
+        HandlerThread bgThread = new HandlerThread("bg thread");
+        bgThread.start();
+        mHandler = new Handler(bgThread.getLooper());
+        mPowerStatsStore = new PowerStatsStore(
+                Files.createTempDirectory("MultiStatePowerAttributorTest").toFile(), mHandler);
+        mPowerStatsAggregator = mock(PowerStatsAggregator.class);
+        mPowerAttributor = new MultiStatePowerAttributor(mPowerStatsStore, mPowerStatsAggregator);
+        mPowerStatsScheduler = new PowerStatsScheduler(
+                () -> mPowerStatsCollectionOccurred = true,
+                mock(BatteryStatsHistory.class),
+                mPowerAttributor, TimeUnit.MINUTES.toMillis(30), TimeUnit.HOURS.toMillis(1),
+                mPowerStatsStore,
+                ((triggerAtMillis, tag, onAlarmListener, handler) ->
+                        mScheduledAlarms.add(triggerAtMillis)),
+                mClock, mMonotonicClock, () -> 12345L, mHandler);
+    }
+
+    @Test
+    public void storeAggregatedPowerStats() {
+        mPowerStatsStore.reset();
+
+        assertThat(mPowerStatsStore.getTableOfContents()).isEmpty();
+
+        mPowerAttributor.storeAggregatedPowerStats(
+                createAggregatedPowerStats(mMonotonicClock.monotonicTime(), mClock.currentTime,
+                        123));
+
+        long delayBeforeAggregating = TimeUnit.MINUTES.toMillis(90);
+        mClock.realtime += delayBeforeAggregating;
+        mClock.currentTime += delayBeforeAggregating;
+
+        doAnswer(invocation -> {
+            // The first span is longer than 30 min, because the end time is being aligned with
+            // the wall clock.  Subsequent spans should be precisely 30 minutes.
+            long startTime = invocation.getArgument(1);
+            long endTime = invocation.getArgument(2);
+            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(3);
+
+            long startTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
+            long endTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
+
+            assertThat(startTime).isEqualTo(START_REALTIME + 123);
+            assertThat(endTime - startTime).isAtLeast(TimeUnit.MINUTES.toMillis(30));
+            assertThat(Instant.ofEpochMilli(endTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+
+            consumer.accept(
+                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
+            return null;
+        }).doAnswer(invocation -> {
+            long startTime = invocation.getArgument(1);
+            long endTime = invocation.getArgument(2);
+            Consumer<AggregatedPowerStats> consumer = invocation.getArgument(3);
+
+            long startTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - startTime);
+            long endTimeWallClock =
+                    mClock.currentTime - (mMonotonicClock.monotonicTime() - endTime);
+
+            assertThat(Instant.ofEpochMilli(startTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+            assertThat(Instant.ofEpochMilli(endTimeWallClock))
+                    .isEqualTo(Instant.parse("2023-01-02T04:30:00Z"));
+
+            consumer.accept(
+                    createAggregatedPowerStats(startTime, startTimeWallClock, endTime - startTime));
+            return null;
+        }).when(mPowerStatsAggregator).aggregatePowerStats(any(BatteryStatsHistory.class),
+                anyLong(), anyLong(), any(Consumer.class));
+
+        mPowerStatsScheduler.start(/*enabled*/ true);
+        ConditionVariable done = new ConditionVariable();
+        mHandler.post(done::open);
+        done.block();
+
+        assertThat(mPowerStatsCollectionOccurred).isTrue();
+        assertThat(mScheduledAlarms).containsExactly(
+                START_REALTIME + TimeUnit.MINUTES.toMillis(90) + TimeUnit.HOURS.toMillis(1));
+
+        verify(mPowerStatsAggregator, times(2)).aggregatePowerStats(
+                any(BatteryStatsHistory.class), anyLong(), anyLong(), any(Consumer.class));
+
+        List<PowerStatsSpan.Metadata> contents = mPowerStatsStore.getTableOfContents();
+        assertThat(contents).hasSize(3);
+        // Skip the first entry, which was placed in the store at the beginning of this test
+        PowerStatsSpan.TimeFrame timeFrame1 = contents.get(1).getTimeFrames().get(0);
+        PowerStatsSpan.TimeFrame timeFrame2 = contents.get(2).getTimeFrames().get(0);
+        assertThat(timeFrame1.startMonotonicTime).isEqualTo(START_REALTIME + 123);
+        assertThat(timeFrame2.startMonotonicTime)
+                .isEqualTo(timeFrame1.startMonotonicTime + timeFrame1.duration);
+        assertThat(Instant.ofEpochMilli(timeFrame2.startTime))
+                .isEqualTo(Instant.parse("2023-01-02T04:00:00Z"));
+        assertThat(Duration.ofMillis(timeFrame2.duration)).isEqualTo(Duration.ofMinutes(30));
+    }
+
+    private AggregatedPowerStats createAggregatedPowerStats(long monotonicTime, long currentTime,
+            long duration) {
+        AggregatedPowerStats stats = new AggregatedPowerStats(new AggregatedPowerStatsConfig());
+        stats.addClockUpdate(monotonicTime, currentTime);
+        stats.setDuration(duration);
+        return stats;
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStateStatsTest.java
similarity index 99%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStateStatsTest.java
index ae258cd3..a232c0c 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/MultiStateStatsTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/MultiStateStatsTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -198,8 +198,7 @@
                 new MultiStateStats.States("scr", trackScreenState, "screen-off", "plugged-in"));
     }
 
-    private FactorySubject assertThatCpuPerformanceStatsFactory(
-            MultiStateStats.Factory factory) {
+    private FactorySubject assertThatCpuPerformanceStatsFactory(MultiStateStats.Factory factory) {
         FactorySubject subject = new FactorySubject();
         subject.mFactory = factory;
         return subject;
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
similarity index 84%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
index cb1bcfe..535f2da 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PhoneCallPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PhoneCallPowerStatsProcessorTest.java
@@ -13,14 +13,14 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
-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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -41,6 +41,11 @@
 import android.telephony.TelephonyManager;
 
 import com.android.internal.os.Clock;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.MobileRadioPowerStatsCollector;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.PowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -48,7 +53,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.LongSupplier;
 import java.util.function.Supplier;
 
@@ -113,11 +117,6 @@
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getMobileNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
@@ -156,19 +155,17 @@
 
     @Test
     public void copyEstimatesFromMobileRadioPowerStats() {
-
-        AggregatedPowerStatsConfig aggregatedPowerStatsConfig = new AggregatedPowerStatsConfig();
-        aggregatedPowerStatsConfig.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .trackDeviceStates(STATE_POWER, STATE_SCREEN)
                 .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
                 .setProcessorSupplier(
                         () -> new MobileRadioPowerStatsProcessor(mStatsRule.getPowerProfile()));
-        aggregatedPowerStatsConfig.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
                         BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .setProcessorSupplier(PhoneCallPowerStatsProcessor::new);
 
-        AggregatedPowerStats aggregatedPowerStats =
-                new AggregatedPowerStats(aggregatedPowerStatsConfig);
+        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
         PowerComponentAggregatedPowerStats mobileRadioStats =
                 aggregatedPowerStats.getPowerComponentStats(
                         BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO);
@@ -208,8 +205,7 @@
                 aggregatedPowerStats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_PHONE);
         stats.finish(10_000);
 
-        PowerStatsLayout statsLayout =
-                new PowerStatsLayout(stats.getPowerStatsDescriptor());
+        PowerStatsLayout statsLayout = new PowerStatsLayout(stats.getPowerStatsDescriptor());
 
         long[] deviceStats = new long[stats.getPowerStatsDescriptor().statsArrayLength];
         stats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsAggregatorTest.java
similarity index 95%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsAggregatorTest.java
index 3929137..f312bed 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsAggregatorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsAggregatorTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -32,6 +32,7 @@
 import com.android.internal.os.BatteryStatsHistory;
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.MockClock;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -53,7 +54,7 @@
     private final MockClock mClock = new MockClock();
     private final MonotonicClock mMonotonicClock = new MonotonicClock(START_TIME, mClock);
     private BatteryStatsHistory mHistory;
-    private PowerStatsAggregator mAggregator;
+    private com.android.server.power.stats.processor.PowerStatsAggregator mAggregator;
     private int mAggregatedStatsCount;
 
     @Before
@@ -71,7 +72,7 @@
                         AggregatedPowerStatsConfig.STATE_POWER,
                         AggregatedPowerStatsConfig.STATE_SCREEN,
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
-        mAggregator = new PowerStatsAggregator(config, mHistory);
+        mAggregator = new PowerStatsAggregator(config);
     }
 
     @Test
@@ -119,7 +120,7 @@
         powerStats.uidStats.put(TEST_UID, new long[]{4444});
         mHistory.recordPowerStats(mClock.realtime, mClock.uptime, powerStats);
 
-        mAggregator.aggregatePowerStats(0, MonotonicClock.UNDEFINED, stats -> {
+        mAggregator.aggregatePowerStats(mHistory, 0, MonotonicClock.UNDEFINED, stats -> {
             assertThat(mAggregatedStatsCount++).isEqualTo(0);
             assertThat(stats.getStartTime()).isEqualTo(START_TIME);
 
@@ -138,7 +139,8 @@
 
             long[] values = new long[1];
 
-            PowerComponentAggregatedPowerStats powerComponentStats = stats.getPowerComponentStats(
+            PowerComponentAggregatedPowerStats
+                    powerComponentStats = stats.getPowerComponentStats(
                     TEST_POWER_COMPONENT);
 
             assertThat(powerComponentStats.getDeviceStats(values, new int[]{
@@ -218,7 +220,7 @@
 
         mHistory.recordBatteryState(mClock.realtime, mClock.uptime, 50, /* plugged */ true);
 
-        mAggregator.aggregatePowerStats(0, MonotonicClock.UNDEFINED, stats -> {
+        mAggregator.aggregatePowerStats(mHistory, 0, MonotonicClock.UNDEFINED, stats -> {
             long[] values = new long[1];
 
             PowerComponentAggregatedPowerStats powerComponentStats =
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
similarity index 95%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
index 96203a5..024743d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsExporterTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsExporterTest.java
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
@@ -22,6 +22,7 @@
 import static org.mockito.Mockito.mock;
 
 import android.annotation.NonNull;
+import android.content.Context;
 import android.os.AggregateBatteryConsumer;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -37,9 +38,16 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.internal.os.BatteryStatsHistory;
+import com.android.internal.os.CpuScalingPolicies;
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.MockClock;
+import com.android.server.power.stats.PowerStatsStore;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.format.CpuPowerStatsLayout;
+import com.android.server.power.stats.format.EnergyConsumerPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -76,6 +84,7 @@
     private MonotonicClock mMonotonicClock = new MonotonicClock(0, mClock);
     private PowerStatsStore mPowerStatsStore;
     private PowerStatsAggregator mPowerStatsAggregator;
+    private MultiStatePowerAttributor mPowerAttributor;
     private BatteryStatsHistory mHistory;
     private CpuPowerStatsLayout mCpuStatsArrayLayout;
     private PowerStats.Descriptor mPowerStatsDescriptor;
@@ -108,25 +117,23 @@
                         AggregatedPowerStatsConfig.STATE_SCREEN,
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE);
 
-        mPowerStatsStore = new PowerStatsStore(storeDirectory, new TestHandler(), config);
+        mPowerStatsStore = new PowerStatsStore(storeDirectory, new TestHandler());
         mHistory = new BatteryStatsHistory(Parcel.obtain(), storeDirectory, 0, 10000,
                 mock(BatteryStatsHistory.HistoryStepDetailsCalculator.class), mClock,
                 mMonotonicClock, null, null);
-        mPowerStatsAggregator = new PowerStatsAggregator(config, mHistory);
+        mPowerStatsAggregator = new PowerStatsAggregator(config);
 
-        mCpuStatsArrayLayout = new CpuPowerStatsLayout();
-        mCpuStatsArrayLayout.addDeviceSectionCpuTimeByScalingStep(1);
-        mCpuStatsArrayLayout.addDeviceSectionCpuTimeByCluster(1);
-        mCpuStatsArrayLayout.addDeviceSectionUsageDuration();
-        mCpuStatsArrayLayout.addDeviceSectionPowerEstimate();
-        mCpuStatsArrayLayout.addUidSectionCpuTimeByPowerBracket(new int[]{0});
-        mCpuStatsArrayLayout.addUidSectionPowerEstimate();
+        mCpuStatsArrayLayout = new CpuPowerStatsLayout(0, 1, new int[]{0});
         PersistableBundle extras = new PersistableBundle();
         mCpuStatsArrayLayout.toExtras(extras);
 
         mPowerStatsDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_CPU,
                 mCpuStatsArrayLayout.getDeviceStatsArrayLength(),
                 null, 0, mCpuStatsArrayLayout.getUidStatsArrayLength(), extras);
+
+        mPowerAttributor = new MultiStatePowerAttributor(mock(Context.class), mPowerStatsStore,
+                mock(PowerProfile.class), mock(CpuScalingPolicies.class),
+                mock(PowerStatsUidResolver.class));
     }
 
     @Test
@@ -329,10 +336,12 @@
                 includeScreenStateData, includesPowerStateData);
     }
 
-    private @NonNull BatteryUsageStats exportToBatteryUsageStats(AggregatedPowerStats aps,
+    private @NonNull BatteryUsageStats exportToBatteryUsageStats(
+            AggregatedPowerStats aps,
             boolean includeProcessStateData, boolean includeScreenStateData,
             boolean includesPowerStateData) {
-        PowerStatsExporter exporter = new PowerStatsExporter(mPowerStatsStore,
+        PowerStatsExporter
+                exporter = new PowerStatsExporter(mPowerStatsStore,
                 mPowerStatsAggregator, /* batterySessionTimeSpanSlackMillis */ 0);
 
         BatteryUsageStats.Builder builder = new BatteryUsageStats.Builder(new String[0], false,
@@ -508,8 +517,8 @@
         mCpuStatsArrayLayout.setUidTimeByPowerBracket(uidStats2, 0, 2469);
         mHistory.recordPowerStats(3000, 3000, powerStats);
 
-        mPowerStatsAggregator.aggregatePowerStats(0, 3500,
-                stats -> mPowerStatsStore.storeAggregatedPowerStats(stats));
+        mPowerStatsAggregator.aggregatePowerStats(mHistory, 0, 3500,
+                stats -> mPowerAttributor.storeAggregatedPowerStats(stats));
 
         mHistory.recordProcessStateChange(4000, 4000, APP_UID1,
                 BatteryConsumer.PROCESS_STATE_BACKGROUND);
@@ -525,9 +534,8 @@
         mEnergyConsumerPowerStatsLayout.setUidConsumedEnergy(customUidStats, 0, 360_000);
         mHistory.recordPowerStats(6010, 6010, customPowerStats);
 
-        mPowerStatsAggregator.aggregatePowerStats(3500, 6500, stats -> {
-            mPowerStatsStore.storeAggregatedPowerStats(stats);
-        });
+        mPowerStatsAggregator.aggregatePowerStats(mHistory, 3500, 6500,
+                stats -> mPowerAttributor.storeAggregatedPowerStats(stats));
 
         mHistory.recordStateStartEvent(7000, 7000, BatteryStats.HistoryItem.STATE_SCREEN_ON_FLAG);
         mHistory.recordProcessStateChange(7000, 7000, APP_UID1,
@@ -548,7 +556,8 @@
         recordBatteryHistory();
         PowerStatsExporter exporter = new PowerStatsExporter(mPowerStatsStore,
                 mPowerStatsAggregator, /* batterySessionTimeSpanSlackMillis */ 0);
-        exporter.exportAggregatedPowerStats(builder, monotonicStartTime, monotonicEndTime);
+        exporter.exportAggregatedPowerStats(builder, mHistory, monotonicStartTime,
+                monotonicEndTime);
     }
 
     private void assertAggregatedPowerEstimate(String message, BatteryUsageStats bus, int scope,
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsProcessorTest.java
similarity index 94%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsProcessorTest.java
index 02e446a..13e0d9d 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/PowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/PowerStatsProcessorTest.java
@@ -13,11 +13,11 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
-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.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
similarity index 83%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
index 94f5662..1852165 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/ScreenPowerStatsProcessorTest.java
@@ -14,22 +14,22 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
-import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
-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_SCREEN;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.os.BatteryConsumer;
 import android.os.BatteryStats;
@@ -40,7 +40,12 @@
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerProfile;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.ScreenPowerStatsCollector;
 import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+import com.android.server.power.stats.format.ScreenPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -48,7 +53,6 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class ScreenPowerStatsProcessorTest {
@@ -106,11 +110,6 @@
         }
 
         @Override
-        public IntSupplier getVoltageSupplier() {
-            return () -> VOLTAGE_MV;
-        }
-
-        @Override
         public int getDisplayCount() {
             return 2;
         }
@@ -124,6 +123,7 @@
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
     }
 
     @Test
@@ -177,8 +177,8 @@
         if (energyConsumer) {
             when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                     .thenReturn(new int[]{77});
-            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                    .thenReturn(new long[]{10_000});
+            when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                    .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000)});
         } else {
             when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
                     .thenReturn(new int[0]);
@@ -200,8 +200,8 @@
         if (energyConsumer) {
             // 400 mAh represented as microWattSeconds
             long energyUws = 400L * 3600 * VOLTAGE_MV;
-            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
-                    .thenReturn(new long[]{10_000 + energyUws});
+            when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{77}))
+                    .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(10_000 + energyUws)});
         }
 
         when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
@@ -243,16 +243,14 @@
 
     private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
             Supplier<PowerStatsProcessor> processorSupplier) {
-        AggregatedPowerStatsConfig.PowerComponent config =
-                new AggregatedPowerStatsConfig.PowerComponent(
-                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
                         .trackDeviceStates(STATE_POWER, STATE_SCREEN)
                         .trackUidStates(STATE_POWER, STATE_SCREEN)
                         .setProcessorSupplier(processorSupplier);
 
-        PowerComponentAggregatedPowerStats aggregatedStats =
-                new PowerComponentAggregatedPowerStats(
-                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+        PowerComponentAggregatedPowerStats aggregatedStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_SCREEN);
 
         aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
         aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
@@ -260,7 +258,14 @@
         return aggregatedStats;
     }
 
-    private void assertDevicePowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
+    private void assertDevicePowerEstimate(
+            PowerComponentAggregatedPowerStats aggregatedStats,
             int powerState, int screenState, double expectedScreenPowerEstimate,
             double expectedDozePowerEstimate) {
         PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
@@ -273,7 +278,8 @@
                 .of(expectedDozePowerEstimate);
     }
 
-    private void assertUidPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats, int uid,
+    private void assertUidPowerEstimate(
+            PowerComponentAggregatedPowerStats aggregatedStats, int uid,
             int powerState, int screenState, double expectedScreenPowerEstimate) {
         PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
         ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(descriptor);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/SensorPowerStatsProcessorTest.java
similarity index 91%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/SensorPowerStatsProcessorTest.java
index 687d70b..d972604 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/SensorPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/SensorPowerStatsProcessorTest.java
@@ -14,19 +14,19 @@
  * limitations under the License.
  */
 
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -43,6 +43,8 @@
 
 import com.android.internal.os.MonotonicClock;
 import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.format.SensorPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -109,8 +111,7 @@
         stats.finish(10000);
 
         PowerStats.Descriptor descriptor = stats.getPowerStatsDescriptor();
-        SensorPowerStatsLayout statsLayout = new SensorPowerStatsLayout();
-        statsLayout.fromExtras(descriptor.extras);
+        SensorPowerStatsLayout statsLayout = new SensorPowerStatsLayout(descriptor);
 
         String dump = stats.toString();
         assertThat(dump).contains(" step_counter: ");
@@ -207,10 +208,8 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessorSupplier(processorSupplier);
 
-        AggregatedPowerStats aggregatedPowerStats = new AggregatedPowerStats(config);
-        PowerComponentAggregatedPowerStats powerComponentStats =
-                aggregatedPowerStats.getPowerComponentStats(
-                        BatteryConsumer.POWER_COMPONENT_SENSORS);
+        PowerComponentAggregatedPowerStats powerComponentStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_SENSORS);
         powerComponentStats.start(0);
 
         powerComponentStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
similarity index 91%
rename from services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
rename to services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
index 11c09bc..baf468e 100644
--- a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/processor/WifiPowerStatsProcessorTest.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.server.power.stats;
+package com.android.server.power.stats.processor;
 
 import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
 import static android.net.NetworkStats.METERED_NO;
@@ -23,12 +23,12 @@
 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.android.server.power.stats.processor.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.processor.AggregatedPowerStatsConfig.STATE_SCREEN;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -41,6 +41,7 @@
 import android.annotation.Nullable;
 import android.content.Context;
 import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerResult;
 import android.hardware.power.stats.EnergyConsumerType;
 import android.net.NetworkStats;
 import android.net.wifi.WifiManager;
@@ -53,6 +54,13 @@
 
 import com.android.internal.os.Clock;
 import com.android.internal.os.PowerProfile;
+import com.android.server.power.stats.BatteryUsageStatsRule;
+import com.android.server.power.stats.MockBatteryStatsImpl;
+import com.android.server.power.stats.PowerStatsCollector;
+import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.WifiPowerStatsCollector;
+import com.android.server.power.stats.WifiPowerStatsCollector.WifiStatsRetriever;
+import com.android.server.power.stats.format.WifiPowerStatsLayout;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -61,7 +69,6 @@
 import org.mockito.MockitoAnnotations;
 
 import java.util.List;
-import java.util.function.IntSupplier;
 import java.util.function.Supplier;
 
 public class WifiPowerStatsProcessorTest {
@@ -109,8 +116,7 @@
     private final SparseArray<ScanTimes> mScanTimes = new SparseArray<>();
     private long mWifiActiveDuration;
 
-    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
-            new WifiPowerStatsCollector.WifiStatsRetriever() {
+    private final WifiStatsRetriever mWifiStatsRetriever = new WifiStatsRetriever() {
         @Override
         public void retrieveWifiScanTimes(Callback callback) {
             for (int i = 0; i < mScanTimes.size(); i++) {
@@ -159,11 +165,6 @@
                 }
 
                 @Override
-                public IntSupplier getVoltageSupplier() {
-                    return () -> VOLTAGE_MV;
-                }
-
-                @Override
                 public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
                     return mNetworkStatsSupplier;
                 }
@@ -174,7 +175,7 @@
                 }
 
                 @Override
-                public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+                public WifiStatsRetriever getWifiStatsRetriever() {
                     return mWifiStatsRetriever;
                 }
             };
@@ -308,6 +309,7 @@
         when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
 
         // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getVoltageMv()).thenReturn(VOLTAGE_MV);
         when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
                 .thenReturn(new int[] {WIFI_ENERGY_CONSUMER_ID});
 
@@ -321,9 +323,9 @@
         mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
                 WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
 
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+        when(mConsumedEnergyRetriever.getConsumedEnergy(
                 new int[]{WIFI_ENERGY_CONSUMER_ID}))
-                .thenReturn(new long[]{0});
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(0)});
 
         aggregatedStats.start(0);
 
@@ -354,8 +356,8 @@
 
         // 10 mAh represented as microWattSeconds
         long energyUws = 10 * 3600 * VOLTAGE_MV;
-        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
-                new int[]{WIFI_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+        when(mConsumedEnergyRetriever.getConsumedEnergy(new int[]{WIFI_ENERGY_CONSUMER_ID}))
+                .thenReturn(new EnergyConsumerResult[]{mockEnergyConsumer(energyUws)});
 
         aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
 
@@ -525,15 +527,14 @@
 
     private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
             Supplier<PowerStatsProcessor> processorSupplier) {
-        AggregatedPowerStatsConfig.PowerComponent config =
-                new AggregatedPowerStatsConfig.PowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
-                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
-                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
-                        .setProcessorSupplier(processorSupplier);
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                .setProcessorSupplier(processorSupplier);
 
-        PowerComponentAggregatedPowerStats aggregatedStats =
-                new PowerComponentAggregatedPowerStats(
-                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+        PowerComponentAggregatedPowerStats aggregatedStats = new AggregatedPowerStats(config)
+                .getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_WIFI);
 
         aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
         aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
@@ -547,6 +548,12 @@
         return states;
     }
 
+    private EnergyConsumerResult mockEnergyConsumer(long energyUWs) {
+        EnergyConsumerResult ecr = new EnergyConsumerResult();
+        ecr.energyUWs = energyUWs;
+        return ecr;
+    }
+
     private void mockWifiActivityEnergyInfo(WifiActivityEnergyInfo waei) {
         doAnswer(invocation -> {
             WifiManager.OnWifiActivityEnergyInfoListener
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/AccessibilityManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/accessibility/AccessibilityManagerServiceTest.java
index c8cbbb5..2e6c93c 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/AccessibilityManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/AccessibilityManagerServiceTest.java
@@ -102,6 +102,7 @@
 import android.view.accessibility.AccessibilityManager;
 import android.view.accessibility.AccessibilityWindowAttributes;
 import android.view.accessibility.IAccessibilityManager;
+import android.view.accessibility.IUserInitializationCompleteCallback;
 
 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.filters.SmallTest;
@@ -137,6 +138,7 @@
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.mockito.Spy;
 import org.mockito.internal.util.reflection.FieldReader;
 import org.mockito.internal.util.reflection.FieldSetter;
 import org.mockito.stubbing.Answer;
@@ -209,6 +211,7 @@
     @Mock private FullScreenMagnificationController mMockFullScreenMagnificationController;
     @Mock private ProxyManager mProxyManager;
     @Mock private StatusBarManagerInternal mStatusBarManagerInternal;
+    @Spy private IUserInitializationCompleteCallback mUserInitializationCompleteCallback;
     @Captor private ArgumentCaptor<Intent> mIntentArgumentCaptor;
     private IAccessibilityManager mA11yManagerServiceOnDevice;
     private AccessibilityServiceConnection mAccessibilityServiceConnection;
@@ -2042,6 +2045,36 @@
                 .isEqualTo(ACCESSIBILITY_BUTTON_MODE_NAVIGATION_BAR);
     }
 
+    @Test
+    public void registerUserInitializationCompleteCallback_isRegistered() {
+        mA11yms.mUserInitializationCompleteCallbacks.clear();
+
+        mA11yms.registerUserInitializationCompleteCallback(mUserInitializationCompleteCallback);
+
+        assertThat(mA11yms.mUserInitializationCompleteCallbacks).containsExactly(
+                mUserInitializationCompleteCallback);
+    }
+
+    @Test
+    public void unregisterUserInitializationCompleteCallback_isUnregistered() {
+        mA11yms.mUserInitializationCompleteCallbacks.clear();
+        mA11yms.mUserInitializationCompleteCallbacks.add(mUserInitializationCompleteCallback);
+
+        mA11yms.unregisterUserInitializationCompleteCallback(mUserInitializationCompleteCallback);
+
+        assertThat(mA11yms.mUserInitializationCompleteCallbacks).isEmpty();
+    }
+
+    @Test
+    public void switchUser_callsUserInitializationCompleteCallback() throws RemoteException {
+        mA11yms.mUserInitializationCompleteCallbacks.add(mUserInitializationCompleteCallback);
+
+        mA11yms.switchUser(UserHandle.MIN_SECONDARY_USER_ID);
+
+        verify(mUserInitializationCompleteCallback).onUserInitializationComplete(
+                UserHandle.MIN_SECONDARY_USER_ID);
+    }
+
     private Set<String> readStringsFromSetting(String setting) {
         final Set<String> result = new ArraySet<>();
         mA11yms.readColonDelimitedSettingToSet(
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/MouseKeysInterceptorTest.kt b/services/tests/servicestests/src/com/android/server/accessibility/MouseKeysInterceptorTest.kt
index 8753b25..019ccf9 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/MouseKeysInterceptorTest.kt
+++ b/services/tests/servicestests/src/com/android/server/accessibility/MouseKeysInterceptorTest.kt
@@ -48,6 +48,7 @@
 import java.util.LinkedList
 import java.util.Queue
 import android.util.ArraySet
+import android.view.InputDevice
 
 /**
  * Tests for {@link MouseKeysInterceptor}
@@ -68,6 +69,8 @@
     }
 
     private lateinit var mouseKeysInterceptor: MouseKeysInterceptor
+    private lateinit var inputDevice: InputDevice
+
     private val clock = OffsettableClock()
     private val testLooper = TestLooper { clock.now() }
     private val nextInterceptor = TrackingInterceptor()
@@ -98,6 +101,10 @@
         testSession = InputManagerGlobal.createTestSession(iInputManager)
         mockInputManager = InputManager(context)
 
+        inputDevice = createInputDevice(DEVICE_ID)
+        Mockito.`when`(iInputManager.getInputDevice(DEVICE_ID))
+                .thenReturn(inputDevice)
+
         Mockito.`when`(mockVirtualDeviceManagerInternal.getDeviceIdsForUid(Mockito.anyInt()))
             .thenReturn(ArraySet(setOf(DEVICE_ID)))
         LocalServices.removeServiceForTest(VirtualDeviceManagerInternal::class.java)
@@ -115,7 +122,8 @@
         Mockito.`when`(iInputManager.inputDeviceIds).thenReturn(intArrayOf(DEVICE_ID))
         Mockito.`when`(mockAms.traceManager).thenReturn(mockTraceManager)
 
-        mouseKeysInterceptor = MouseKeysInterceptor(mockAms, testLooper.looper, DISPLAY_ID)
+        mouseKeysInterceptor = MouseKeysInterceptor(mockAms, mockInputManager,
+                testLooper.looper, DISPLAY_ID)
         mouseKeysInterceptor.next = nextInterceptor
     }
 
@@ -281,6 +289,17 @@
         }
     }
 
+    private fun createInputDevice(
+            deviceId: Int,
+            generation: Int = -1
+    ): InputDevice =
+            InputDevice.Builder()
+                    .setId(deviceId)
+                    .setName("Device $deviceId")
+                    .setDescriptor("descriptor $deviceId")
+                    .setGeneration(generation)
+                    .build()
+
     private class TrackingInterceptor : BaseEventStreamTransformation() {
         val events: Queue<KeyEvent> = LinkedList()
 
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
index 4ec2fb9..cdaeade 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/a11ychecker/AccessibilityCheckerUtilsTest.java
@@ -89,14 +89,15 @@
                         AccessibilityCheckResult.AccessibilityCheckResultType.NOT_RUN, null, 5,
                         null);
 
-        Set<AndroidAccessibilityCheckerResult> results =
-                AccessibilityCheckerUtils.processResults(
-                        mockNodeInfo,
-                        List.of(result1, result2, result3, result4),
-                        null,
+
+        AndroidAccessibilityCheckerResult.Builder resultBuilder =
+                AccessibilityCheckerUtils.getCommonResultBuilder(mockNodeInfo, null,
                         mMockPackageManager,
                         new ComponentName(TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME,
                                 TEST_A11Y_SERVICE_CLASS_NAME));
+        Set<AndroidAccessibilityCheckerResult> results =
+                AccessibilityCheckerUtils.processResults(mockNodeInfo,
+                        List.of(result1, result2, result3, result4), resultBuilder);
 
         assertThat(results).containsExactly(
                 createResult("TargetNode", "",
@@ -128,14 +129,14 @@
                         TouchTargetSizeCheck.class,
                         AccessibilityCheckResult.AccessibilityCheckResultType.ERROR, null, 2, null);
 
-        Set<AndroidAccessibilityCheckerResult> results =
-                AccessibilityCheckerUtils.processResults(
-                        mockNodeInfo,
-                        List.of(result1, result2),
-                        null,
+        AndroidAccessibilityCheckerResult.Builder resultBuilder =
+                AccessibilityCheckerUtils.getCommonResultBuilder(mockNodeInfo, null,
                         mMockPackageManager,
                         new ComponentName(TEST_A11Y_SERVICE_SOURCE_PACKAGE_NAME,
                                 TEST_A11Y_SERVICE_CLASS_NAME));
+        Set<AndroidAccessibilityCheckerResult> results =
+                AccessibilityCheckerUtils.processResults(mockNodeInfo,
+                        List.of(result1, result2), resultBuilder);
 
         assertThat(results).isEmpty();
     }
diff --git a/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java b/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
index 957ee06..598d3a3 100644
--- a/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
+++ b/services/tests/servicestests/src/com/android/server/accessibility/magnification/FullScreenMagnificationGestureHandlerTest.java
@@ -23,6 +23,8 @@
 import static android.view.MotionEvent.ACTION_POINTER_INDEX_SHIFT;
 import static android.view.MotionEvent.ACTION_POINTER_UP;
 import static android.view.MotionEvent.ACTION_UP;
+import static android.view.MotionEvent.CLASSIFICATION_TWO_FINGER_SWIPE;
+import static android.view.MotionEvent.TOOL_TYPE_FINGER;
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
@@ -1414,6 +1416,49 @@
     }
 
     @Test
+    public void testSynthesizedGestureEventsDoNotMoveMagnifierViewport() {
+        final EventCaptor eventCaptor = new EventCaptor();
+        mMgh.setNext(eventCaptor);
+
+        float centerX =
+                (INITIAL_MAGNIFICATION_BOUNDS.left + INITIAL_MAGNIFICATION_BOUNDS.width()) / 2.0f;
+        float centerY =
+                (INITIAL_MAGNIFICATION_BOUNDS.top + INITIAL_MAGNIFICATION_BOUNDS.height()) / 2.0f;
+        float scale = 5.6f; // value is unimportant but unique among tests to increase coverage.
+        mFullScreenMagnificationController.setScaleAndCenter(
+                DISPLAY_0, centerX, centerY, scale, /* animate= */ false, 1);
+        centerX = mFullScreenMagnificationController.getCenterX(DISPLAY_0);
+        centerY = mFullScreenMagnificationController.getCenterY(DISPLAY_0);
+
+        // Second finger down on trackpad starts a synthesized two-finger swipe with source
+        // mouse.
+        MotionEvent downEvent = motionEvent(centerX, centerY, ACTION_DOWN,
+                TOOL_TYPE_FINGER, CLASSIFICATION_TWO_FINGER_SWIPE);
+        send(downEvent, InputDevice.SOURCE_MOUSE);
+        fastForward(20);
+
+        // Two-finger swipe creates a synthesized move event, and shouldn't impact magnifier
+        // viewport.
+        MotionEvent moveEvent = motionEvent(centerX - 42, centerY - 42, ACTION_MOVE,
+                TOOL_TYPE_FINGER, CLASSIFICATION_TWO_FINGER_SWIPE);
+        send(moveEvent, InputDevice.SOURCE_MOUSE);
+        fastForward(20);
+
+        assertThat(mFullScreenMagnificationController.getCenterX(DISPLAY_0)).isEqualTo(centerX);
+        assertThat(mFullScreenMagnificationController.getCenterY(DISPLAY_0)).isEqualTo(centerY);
+
+        // The events were not consumed by magnifier.
+        assertThat(eventCaptor.mEvents.size()).isEqualTo(2);
+        assertThat(eventCaptor.mEvents.get(0).getSource()).isEqualTo(InputDevice.SOURCE_MOUSE);
+        assertThat(eventCaptor.mEvents.get(1).getSource()).isEqualTo(InputDevice.SOURCE_MOUSE);
+
+        final List<Integer> expectedActions = new ArrayList();
+        expectedActions.add(Integer.valueOf(ACTION_DOWN));
+        expectedActions.add(Integer.valueOf(ACTION_MOVE));
+        assertActionsInOrder(eventCaptor.mEvents, expectedActions);
+    }
+
+    @Test
     @RequiresFlagsDisabled(Flags.FLAG_ENABLE_MAGNIFICATION_FOLLOWS_MOUSE)
     public void testMouseHoverMoveEventsDoNotMoveMagnifierViewport() {
         runHoverMoveEventsDoNotMoveMagnifierViewport(InputDevice.SOURCE_MOUSE);
@@ -2130,6 +2175,30 @@
         return MotionEvent.obtain(mLastDownTime, mClock.now(), action, x, y, 0);
     }
 
+    private MotionEvent motionEvent(float x, float y, int action, int toolType,
+            int classification) {
+        // Create a generic motion event to populate the parameters.
+        MotionEvent event = motionEvent(x, y, action);
+        int pointerCount = event.getPointerCount();
+        MotionEvent.PointerCoords[] coords = new MotionEvent.PointerCoords[pointerCount];
+        MotionEvent.PointerProperties[] properties =
+                new MotionEvent.PointerProperties[pointerCount];
+        for (int i = 0; i < pointerCount; i++) {
+            properties[i] = new MotionEvent.PointerProperties();
+            event.getPointerProperties(i, properties[i]);
+            properties[i].toolType = toolType;
+            coords[i] = new MotionEvent.PointerCoords();
+            event.getPointerCoords(i, coords[i]);
+        }
+        // Apply the custom classification.
+        return MotionEvent.obtain(event.getDownTime(), event.getEventTime(), action,
+                /*pointerCount=*/1, properties, coords,
+                event.getMetaState(), event.getButtonState(),
+                event.getXPrecision(), event.getYPrecision(), event.getDeviceId(),
+                event.getEdgeFlags(), event.getSource(), event.getDisplayId(), event.getFlags(),
+                classification);
+    }
+
     private MotionEvent mouseEvent(float x, float y, int action) {
         return fromMouse(motionEvent(x, y, action));
     }
diff --git a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
index 1db46bf..2a55521 100644
--- a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
@@ -55,6 +55,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assume.assumeFalse;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyBoolean;
@@ -145,7 +146,6 @@
  */
 @SmallTest
 @Presubmit
-
 public class UserControllerTest {
     // Use big enough user id to avoid picking up already active user id.
     private static final int TEST_USER_ID = 100;
@@ -593,6 +593,7 @@
     @Test
     public void testScheduleStopOfBackgroundUser_switch() {
         mSetFlagsRule.enableFlags(android.multiuser.Flags.FLAG_SCHEDULE_STOP_OF_BACKGROUND_USER);
+        assumeFalse(UserManager.isVisibleBackgroundUsersEnabled());
 
         mUserController.setInitialConfig(/* userSwitchUiEnabled= */ true,
                 /* maxRunningUsers= */ 10, /* delayUserDataLocking= */ false,
@@ -642,6 +643,7 @@
     @Test
     public void testScheduleStopOfBackgroundUser_startInBackground() throws Exception {
         mSetFlagsRule.enableFlags(android.multiuser.Flags.FLAG_SCHEDULE_STOP_OF_BACKGROUND_USER);
+        assumeFalse(UserManager.isVisibleBackgroundUsersEnabled());
 
         mUserController.setInitialConfig(/* userSwitchUiEnabled= */ true,
                 /* maxRunningUsers= */ 10, /* delayUserDataLocking= */ false,
@@ -681,6 +683,7 @@
     @Test
     public void testScheduleStopOfBackgroundUser_rescheduleWhenGuest() throws Exception {
         mSetFlagsRule.enableFlags(android.multiuser.Flags.FLAG_SCHEDULE_STOP_OF_BACKGROUND_USER);
+        assumeFalse(UserManager.isVisibleBackgroundUsersEnabled());
 
         mUserController.setInitialConfig(/* userSwitchUiEnabled= */ true,
                 /* maxRunningUsers= */ 10, /* delayUserDataLocking= */ false,
@@ -736,6 +739,7 @@
     @Test
     public void testScheduleStopOfBackgroundUser_rescheduleIfAlarm() throws Exception {
         mSetFlagsRule.enableFlags(android.multiuser.Flags.FLAG_SCHEDULE_STOP_OF_BACKGROUND_USER);
+        assumeFalse(UserManager.isVisibleBackgroundUsersEnabled());
 
         mUserController.setInitialConfig(/* userSwitchUiEnabled= */ true,
                 /* maxRunningUsers= */ 10, /* delayUserDataLocking= */ false,
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/AuthSessionTest.java b/services/tests/servicestests/src/com/android/server/biometrics/AuthSessionTest.java
index bc2fd73..2f7b8d2 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/AuthSessionTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/AuthSessionTest.java
@@ -200,7 +200,8 @@
                     eq(TEST_REQUEST_ID),
                     eq(sensor.getCookie()),
                     anyBoolean() /* allowBackgroundAuthentication */,
-                    anyBoolean() /* isForLegacyFingerprintManager */);
+                    anyBoolean() /* isForLegacyFingerprintManager */,
+                    eq(false) /* isMandatoryBiometrics */);
         }
 
         final int cookie1 = session.mPreAuthInfo.eligibleSensors.get(0).getCookie();
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/BiometricServiceTest.java b/services/tests/servicestests/src/com/android/server/biometrics/BiometricServiceTest.java
index d2961bc..6b8e414 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/BiometricServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/BiometricServiceTest.java
@@ -636,7 +636,8 @@
                 eq(TEST_REQUEST_ID),
                 cookieCaptor.capture() /* cookie */,
                 anyBoolean() /* allowBackgroundAuthentication */,
-                anyBoolean() /* isForLegacyFingerprintManager */);
+                anyBoolean() /* isForLegacyFingerprintManager */,
+                eq(false) /* isMandatoryBiometrics */);
 
         // onReadyForAuthentication, mAuthSession state OK
         mBiometricService.mImpl.onReadyForAuthentication(TEST_REQUEST_ID, cookieCaptor.getValue());
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/AcquisitionClientTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/AcquisitionClientTest.java
index 4604b31..613cb20 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/AcquisitionClientTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/AcquisitionClientTest.java
@@ -98,7 +98,8 @@
                 @NonNull ClientMonitorCallbackConverter callback) {
             super(context, lazyDaemon, token, callback, 0 /* userId */, "Test", 0 /* cookie */,
                     TEST_SENSOR_ID /* sensorId */, true /* shouldVibrate */,
-                    mock(BiometricLogger.class), mock(BiometricContext.class));
+                    mock(BiometricLogger.class), mock(BiometricContext.class),
+                    false /* isMandatoryBiometrics */);
         }
 
         @Override
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerOperationTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerOperationTest.java
index ffc7811..4f07380 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerOperationTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerOperationTest.java
@@ -62,7 +62,8 @@
             extends HalClientMonitor<T> {
         public InterruptableMonitor() {
             super(null, null, null, null, 0, null, 0, 0,
-                    mock(BiometricLogger.class), mock(BiometricContext.class));
+                    mock(BiometricLogger.class), mock(BiometricContext.class),
+                    false /* isMandatoryBiometrics */);
         }
     }
 
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerTest.java
index 36a7b3d..90c07d4 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/BiometricSchedulerTest.java
@@ -1051,6 +1051,11 @@
         public String getAttributionTag() {
             return null;
         }
+
+        @Override
+        public boolean isMandatoryBiometrics() {
+            return false;
+        }
     }
 
     private static class TestAuthenticationClient
@@ -1176,7 +1181,7 @@
                 @NonNull Supplier<Object> lazyDaemon, int cookie, int protoEnum) {
             super(context, lazyDaemon, token /* token */, null /* listener */, 0 /* userId */, TAG,
                     cookie, TEST_SENSOR_ID, mock(BiometricLogger.class),
-                    mock(BiometricContext.class));
+                    mock(BiometricContext.class), false /* isMandatoryBiometrics */);
             mProtoEnum = protoEnum;
         }
 
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
index 5a78d9e..1a593dd 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
@@ -37,7 +37,6 @@
 
 import android.app.WindowConfiguration;
 import android.companion.virtual.IVirtualDeviceIntentInterceptor;
-import android.companion.virtual.VirtualDeviceManager;
 import android.content.AttributionSource;
 import android.content.ComponentName;
 import android.content.Context;
@@ -94,15 +93,9 @@
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     @Mock
-    private VirtualDeviceManager.ActivityListener mActivityListener;
-    @Mock
-    private GenericWindowPolicyController.IntentListenerCallback mIntentListenerCallback;
-    @Mock
-    private GenericWindowPolicyController.ActivityBlockedCallback mActivityBlockedCallback;
+    private GenericWindowPolicyController.ActivityListener mActivityListener;
     @Mock
     private GenericWindowPolicyController.RunningAppsChangedListener mRunningAppsChangedListener;
-    @Mock
-    private GenericWindowPolicyController.SecureWindowCallback mSecureWindowCallback;
 
     @Before
     public void setUp() throws Exception {
@@ -669,14 +662,14 @@
                 /* targetDisplayCategory */ null);
 
         // register interceptor and intercept intent
-        when(mIntentListenerCallback.shouldInterceptIntent(any(Intent.class))).thenReturn(true);
+        when(mActivityListener.shouldInterceptIntent(any(Intent.class))).thenReturn(true);
         assertThat(gwpc.canActivityBeLaunched(activityInfo, intent,
                 WindowConfiguration.WINDOWING_MODE_FULLSCREEN, DISPLAY_ID, /* isNewTask= */false,
                 /* isResultExpected = */ false, /* intentSender= */ null))
                 .isFalse();
 
         // unregister interceptor and launch activity
-        when(mIntentListenerCallback.shouldInterceptIntent(any(Intent.class))).thenReturn(false);
+        when(mActivityListener.shouldInterceptIntent(any(Intent.class))).thenReturn(false);
         assertThat(gwpc.canActivityBeLaunched(activityInfo, intent,
                 WindowConfiguration.WINDOWING_MODE_FULLSCREEN, DISPLAY_ID, /* isNewTask= */false,
                 /* isResultExpected = */ false, /* intentSender= */ null))
@@ -696,13 +689,12 @@
                 /* targetDisplayCategory */ null);
 
         // register interceptor with different filter
-        when(mIntentListenerCallback.shouldInterceptIntent(any(Intent.class))).thenReturn(false);
+        when(mActivityListener.shouldInterceptIntent(any(Intent.class))).thenReturn(false);
         assertThat(gwpc.canActivityBeLaunched(activityInfo, intent,
                 WindowConfiguration.WINDOWING_MODE_FULLSCREEN, DISPLAY_ID, /* isNewTask= */false,
                 /* isResultExpected = */ false, /* intentSender= */ null))
                 .isTrue();
-        verify(mIntentListenerCallback, timeout(TIMEOUT_MILLIS))
-                .shouldInterceptIntent(any(Intent.class));
+        verify(mActivityListener, timeout(TIMEOUT_MILLIS)).shouldInterceptIntent(any(Intent.class));
     }
 
     @Test
@@ -723,8 +715,8 @@
                 /* isResultExpected = */ true, /* intentSender= */ () -> intentSender))
                 .isFalse();
 
-        verify(mActivityBlockedCallback, timeout(TIMEOUT_MILLIS))
-                .onActivityBlocked(DISPLAY_ID, activityInfo, /* intentSender= */ null);
+        verify(mActivityListener, timeout(TIMEOUT_MILLIS))
+                .onActivityLaunchBlocked(DISPLAY_ID, activityInfo, /* intentSender= */ null);
     }
 
     @Test
@@ -761,10 +753,10 @@
 
         assertThat(gwpc.keepActivityOnWindowFlagsChanged(activityInfo, 0, 0)).isTrue();
 
-        verify(mSecureWindowCallback, after(TIMEOUT_MILLIS).never())
-                .onSecureWindowShown(DISPLAY_ID, activityInfo.applicationInfo.uid);
-        verify(mActivityBlockedCallback, never())
-                .onActivityBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
+                .onSecureWindowShown(eq(DISPLAY_ID), eq(activityInfo));
+        verify(mActivityListener, never())
+                .onActivityLaunchBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
     }
 
     @Test
@@ -780,10 +772,10 @@
 
         assertThat(gwpc.keepActivityOnWindowFlagsChanged(activityInfo, FLAG_SECURE, 0)).isTrue();
 
-        verify(mSecureWindowCallback, timeout(TIMEOUT_MILLIS)).onSecureWindowShown(DISPLAY_ID,
-                activityInfo.applicationInfo.uid);
-        verify(mActivityBlockedCallback, after(TIMEOUT_MILLIS).never())
-                .onActivityBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
+        verify(mActivityListener, timeout(TIMEOUT_MILLIS))
+                .onSecureWindowShown(eq(DISPLAY_ID), eq(activityInfo));
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
+                .onActivityLaunchBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
     }
 
     @Test
@@ -800,10 +792,10 @@
         assertThat(gwpc.keepActivityOnWindowFlagsChanged(activityInfo, 0,
                 SYSTEM_FLAG_HIDE_NON_SYSTEM_OVERLAY_WINDOWS)).isTrue();
 
-        verify(mSecureWindowCallback, after(TIMEOUT_MILLIS).never())
-                .onSecureWindowShown(DISPLAY_ID, activityInfo.applicationInfo.uid);
-        verify(mActivityBlockedCallback, never())
-                .onActivityBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
+                .onSecureWindowShown(eq(DISPLAY_ID), eq(activityInfo));
+        verify(mActivityListener, never())
+                .onActivityLaunchBlocked(eq(DISPLAY_ID), eq(activityInfo), any());
     }
 
     @Test
@@ -835,9 +827,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ mSecureWindowCallback,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -855,9 +844,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ mSecureWindowCallback,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -876,9 +862,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ homeComponent);
@@ -897,9 +880,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -918,9 +898,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -939,9 +916,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ Collections.singleton(displayCategory),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -960,9 +934,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ true,
                 /* crossTaskNavigationExemptions= */ Collections.singleton(blockedComponent),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -981,9 +952,6 @@
                 /* crossTaskNavigationAllowedByDefault= */ false,
                 /* crossTaskNavigationExemptions= */ Collections.singleton(allowedComponent),
                 /* activityListener= */ mActivityListener,
-                /* activityBlockedCallback= */ mActivityBlockedCallback,
-                /* secureWindowCallback= */ null,
-                /* intentListenerCallback= */ mIntentListenerCallback,
                 /* displayCategories= */ new ArraySet<>(),
                 /* showTasksInHostDeviceRecents= */ true,
                 /* customHomeComponent= */ null);
@@ -1029,9 +997,9 @@
         assertThat(gwpc.canActivityBeLaunched(activityInfo, null, windowingMode, fromDisplay,
                 isNewTask, /* isResultExpected= */ false, () -> intentSender)).isTrue();
 
-        verify(mActivityBlockedCallback, after(TIMEOUT_MILLIS).never())
-                .onActivityBlocked(fromDisplay, activityInfo, intentSender);
-        verify(mIntentListenerCallback, never()).shouldInterceptIntent(any(Intent.class));
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
+                .onActivityLaunchBlocked(fromDisplay, activityInfo, intentSender);
+        verify(mActivityListener, never()).shouldInterceptIntent(any(Intent.class));
     }
 
     private void assertActivityIsBlocked(GenericWindowPolicyController gwpc,
@@ -1046,9 +1014,9 @@
         assertThat(gwpc.canActivityBeLaunched(activityInfo, null, windowingMode, fromDisplay,
                 isNewTask, /* isResultExpected= */ false, () -> intentSender)).isFalse();
 
-        verify(mActivityBlockedCallback, timeout(TIMEOUT_MILLIS))
-                .onActivityBlocked(fromDisplay, activityInfo, intentSender);
-        verify(mIntentListenerCallback, after(TIMEOUT_MILLIS).never())
+        verify(mActivityListener, timeout(TIMEOUT_MILLIS))
+                .onActivityLaunchBlocked(fromDisplay, activityInfo, intentSender);
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
                 .shouldInterceptIntent(any(Intent.class));
     }
 
@@ -1060,8 +1028,8 @@
                 /* isResultExpected= */ false, () -> intentSender))
                 .isFalse();
 
-        verify(mActivityBlockedCallback, after(TIMEOUT_MILLIS).never())
-                .onActivityBlocked(eq(fromDisplay), eq(activityInfo), any());
-        verify(mIntentListenerCallback, never()).shouldInterceptIntent(any(Intent.class));
+        verify(mActivityListener, after(TIMEOUT_MILLIS).never())
+                .onActivityLaunchBlocked(eq(fromDisplay), eq(activityInfo), any());
+        verify(mActivityListener, never()).shouldInterceptIntent(any(Intent.class));
     }
 }
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/audio/VirtualAudioControllerTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/audio/VirtualAudioControllerTest.java
index 405929a..51c2ad1 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/audio/VirtualAudioControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/audio/VirtualAudioControllerTest.java
@@ -87,9 +87,6 @@
                         /* crossTaskNavigationAllowedByDefault= */ true,
                         /* crossTaskNavigationExemptions= */ new ArraySet<>(),
                         /* activityListener= */ null,
-                        /* activityBlockedCallback= */ null,
-                        /* secureWindowCallback= */ null,
-                        /* intentListenerCallback= */ null,
                         /* displayCategories= */ new ArraySet<>(),
                         /* showTasksInHostDeviceRecents= */ true,
                         /* customHomeComponent= */ null);
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
index b4cc343..698bda3 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
@@ -60,6 +60,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 /**
@@ -325,6 +326,11 @@
         }
 
         @Override
+        List<String> roleManagerGetRoleHoldersAsUser(String role, UserHandle userHandle) {
+            return services.roleManagerForMock.getRoleHoldersAsUser(role, userHandle);
+        }
+
+        @Override
         PendingIntent pendingIntentGetActivityAsUser(Context context, int requestCode,
                 Intent intent, int flags, Bundle options, UserHandle user) {
             return null;
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index b7483d6..cb4269a 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -109,6 +109,7 @@
 import android.app.admin.PreferentialNetworkServiceConfig;
 import android.app.admin.SystemUpdatePolicy;
 import android.app.admin.WifiSsidPolicy;
+import android.app.role.RoleManager;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Intent;
@@ -2889,6 +2890,52 @@
     }
 
     @Test
+    public void testSetMeteredDataDisabledPackagesExemptRoles() throws Exception {
+        // TODO(b/362545319): reference role name from role manager once it's exposed.
+        final String controllerRole = "android.app.role.SYSTEM_FINANCED_DEVICE_CONTROLLER";
+
+        setAsProfileOwner(admin1);
+
+        assertThat(dpm.getMeteredDataDisabledPackages(admin1)).isEmpty();
+
+        // Setup
+        final ArrayList<String> pkgsToRestrict = new ArrayList<>();
+        final ArrayList<String> pkgsExpectedAsNotRestricted = new ArrayList<>();
+        final String packageWithControllerRole = "com.example.controller";
+        final String packageWithKioskRole = "com.example.kiosk";
+        final String packageWithNotExemptRole = "com.example.notexempt";
+
+        pkgsToRestrict.add(packageWithControllerRole);
+        pkgsToRestrict.add(packageWithKioskRole);
+        pkgsToRestrict.add(packageWithNotExemptRole);
+
+        pkgsExpectedAsNotRestricted.add(packageWithControllerRole);
+        pkgsExpectedAsNotRestricted.add(packageWithKioskRole);
+
+        setupPackageInPackageManager(packageWithControllerRole, CALLER_USER_HANDLE, 123, 0);
+        setupPackageInPackageManager(packageWithKioskRole, CALLER_USER_HANDLE, 456, 0);
+        setupPackageInPackageManager(packageWithNotExemptRole, CALLER_USER_HANDLE, 789, 0);
+
+        when(getServices().roleManagerForMock.getRoleHoldersAsUser(controllerRole,
+                UserHandle.of(CALLER_USER_HANDLE)))
+                .thenReturn(new ArrayList<>(Arrays.asList(packageWithControllerRole)));
+        when(getServices().roleManagerForMock.getRoleHoldersAsUser(
+                RoleManager.ROLE_FINANCED_DEVICE_KIOSK,
+                UserHandle.of(CALLER_USER_HANDLE)))
+                .thenReturn(new ArrayList<>(Arrays.asList(packageWithKioskRole)));
+
+        List<String> excludedPkgs = dpm.setMeteredDataDisabledPackages(admin1, pkgsToRestrict);
+
+        // Verify
+        assertThat(excludedPkgs).containsExactlyElementsIn(pkgsExpectedAsNotRestricted);
+        assertThat(dpm.getMeteredDataDisabledPackages(admin1))
+                .isEqualTo(Arrays.asList(packageWithNotExemptRole));
+        verify(getServices().networkPolicyManagerInternal).setMeteredRestrictedPackages(
+                MockUtils.checkApps(packageWithNotExemptRole),
+                eq(CALLER_USER_HANDLE));
+    }
+
+    @Test
     public void testSetGetMeteredDataDisabledPackages_deviceAdmin() {
         mContext.callerPermissions.add(permission.MANAGE_DEVICE_ADMINS);
         dpm.setActiveAdmin(admin1, true);
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java b/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
index 76aa40c..2e200a9 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/MockSystemServices.java
@@ -142,6 +142,7 @@
     public final DevicePolicyManager devicePolicyManager;
     public final LocationManager locationManager;
     public final RoleManager roleManager;
+    public final RoleManagerForMock roleManagerForMock;
     public final SubscriptionManager subscriptionManager;
     /** Note this is a partial mock, not a real mock. */
     public final PackageManager packageManager;
@@ -200,6 +201,7 @@
         devicePolicyManager = mock(DevicePolicyManager.class);
         locationManager = mock(LocationManager.class);
         roleManager = realContext.getSystemService(RoleManager.class);
+        roleManagerForMock = mock(RoleManagerForMock.class);
         subscriptionManager = mock(SubscriptionManager.class);
 
         // Package manager is huge, so we use a partial mock instead.
@@ -495,6 +497,12 @@
         }
     }
 
+    public static class RoleManagerForMock {
+        public List<String> getRoleHoldersAsUser(String role, UserHandle userHandle) {
+            return new ArrayList<>();
+        }
+    }
+
     public static class SettingsForMock {
         public int settingsSecureGetIntForUser(String name, int def, int userHandle) {
             return 0;
diff --git a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
index 17b499e..d6f7e21 100644
--- a/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
+++ b/services/tests/servicestests/src/com/android/server/locksettings/RebootEscrowManagerTests.java
@@ -625,25 +625,10 @@
 
         // pretend reboot happens here
         when(mInjected.getBootCount()).thenReturn(1);
-        ArgumentCaptor<Boolean> metricsSuccessCaptor = ArgumentCaptor.forClass(Boolean.class);
-        ArgumentCaptor<Integer> metricsErrorCodeCaptor = ArgumentCaptor.forClass(Integer.class);
-        doNothing()
-                .when(mInjected)
-                .reportMetric(
-                        metricsSuccessCaptor.capture(),
-                        metricsErrorCodeCaptor.capture(),
-                        eq(2) /* Server based */,
-                        eq(1) /* attempt count */,
-                        anyInt(),
-                        eq(0) /* vbmeta status */,
-                        anyInt());
+
         mService.loadRebootEscrowDataIfAvailable(null);
         verify(mServiceConnection, never()).unwrap(any(), anyLong());
         verify(mCallbacks, never()).onRebootEscrowRestored(anyByte(), any(), anyInt());
-        assertFalse(metricsSuccessCaptor.getValue());
-        assertEquals(
-                Integer.valueOf(RebootEscrowManager.ERROR_NO_REBOOT_ESCROW_DATA),
-                metricsErrorCodeCaptor.getValue());
     }
 
     @Test
diff --git a/services/tests/servicestests/src/com/android/server/pm/ShortcutManagerTest7.java b/services/tests/servicestests/src/com/android/server/pm/ShortcutManagerTest7.java
index b2fd8aa..161b18c 100644
--- a/services/tests/servicestests/src/com/android/server/pm/ShortcutManagerTest7.java
+++ b/services/tests/servicestests/src/com/android/server/pm/ShortcutManagerTest7.java
@@ -165,7 +165,7 @@
 
         assertTrue(resultContains(
                 callShellCommand("reset-throttling", "--user", "10"),
-                "User 10 is not running or locked"));
+                "User (with userId=10) is not running or locked"));
 
         mRunningUsers.put(USER_10, true);
 
diff --git a/services/tests/servicestests/src/com/android/server/pm/UserRestrictionsUtilsTest.java b/services/tests/servicestests/src/com/android/server/pm/UserRestrictionsUtilsTest.java
index 55c48e0..f0a5f75 100644
--- a/services/tests/servicestests/src/com/android/server/pm/UserRestrictionsUtilsTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/UserRestrictionsUtilsTest.java
@@ -36,7 +36,6 @@
 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;
@@ -55,11 +54,6 @@
     @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
-    @Before
-    public void setUp() {
-        mSetFlagsRule.enableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
-    }
-
     @Test
     public void testNonNull() {
         Bundle out = UserRestrictionsUtils.nonNull(null);
@@ -144,7 +138,6 @@
 
     @Test
     public void testCanProfileOwnerChange_restrictionRequiresOrgOwnedDevice_orgOwned() {
-        mSetFlagsRule.enableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
         assertTrue(UserRestrictionsUtils.canProfileOwnerChange(
                 UserManager.DISALLOW_SIM_GLOBALLY,
                 false,
@@ -157,7 +150,6 @@
 
     @Test
     public void testCanProfileOwnerChange_restrictionRequiresOrgOwnedDevice_notOrgOwned() {
-        mSetFlagsRule.enableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
         assertFalse(UserRestrictionsUtils.canProfileOwnerChange(
                 UserManager.DISALLOW_SIM_GLOBALLY,
                 false,
@@ -169,22 +161,7 @@
     }
 
     @Test
-    public void
-            testCanProfileOwnerChange_disabled_restrictionRequiresOrgOwnedDevice_notOrgOwned() {
-        mSetFlagsRule.disableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
-        assertTrue(UserRestrictionsUtils.canProfileOwnerChange(
-                UserManager.DISALLOW_SIM_GLOBALLY,
-                false,
-                false));
-        assertTrue(UserRestrictionsUtils.canProfileOwnerChange(
-                UserManager.DISALLOW_SIM_GLOBALLY,
-                true,
-                false));
-    }
-
-    @Test
     public void testCanProfileOwnerChange_restrictionNotRequiresOrgOwnedDevice_orgOwned() {
-        mSetFlagsRule.enableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
         assertTrue(UserRestrictionsUtils.canProfileOwnerChange(
                 UserManager.DISALLOW_ADJUST_VOLUME,
                 false,
@@ -197,7 +174,6 @@
 
     @Test
     public void testCanProfileOwnerChange_restrictionNotRequiresOrgOwnedDevice_notOrgOwned() {
-        mSetFlagsRule.enableFlags(android.app.admin.flags.Flags.FLAG_ESIM_MANAGEMENT_ENABLED);
         assertTrue(UserRestrictionsUtils.canProfileOwnerChange(
                 UserManager.DISALLOW_ADJUST_VOLUME,
                 false,
diff --git a/services/tests/servicestests/src/com/android/server/tv/tunerresourcemanager/TunerResourceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/tv/tunerresourcemanager/TunerResourceManagerServiceTest.java
index 963b27e..bf58443 100644
--- a/services/tests/servicestests/src/com/android/server/tv/tunerresourcemanager/TunerResourceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/tv/tunerresourcemanager/TunerResourceManagerServiceTest.java
@@ -38,6 +38,7 @@
 import android.media.tv.tunerresourcemanager.TunerFrontendRequest;
 import android.media.tv.tunerresourcemanager.TunerLnbRequest;
 import android.media.tv.tunerresourcemanager.TunerResourceManager;
+import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.InstrumentationRegistry;
@@ -69,6 +70,61 @@
     private TunerResourceManagerService mTunerResourceManagerService;
     private boolean mIsForeground;
 
+    private final class TunerClient extends IResourcesReclaimListener.Stub {
+        int[] mClientId;
+        ClientProfile mProfile;
+        boolean mReclaimed;
+
+        TunerClient() {
+            mClientId = new int[1];
+            mClientId[0] = TunerResourceManagerService.INVALID_CLIENT_ID;
+        }
+
+        public void register(String sessionId, int useCase) {
+            ResourceClientProfile profile = new ResourceClientProfile();
+            profile.tvInputSessionId = sessionId;
+            profile.useCase = useCase;
+            mTunerResourceManagerService.registerClientProfileInternal(
+                    profile, this, mClientId);
+            assertThat(mClientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+            mProfile = mTunerResourceManagerService.getClientProfile(mClientId[0]);
+        }
+
+        public void register(String sessionId, int useCase, int priority, int niceValue) {
+            register(sessionId, useCase);
+            mTunerResourceManagerService.updateClientPriorityInternal(
+                    mClientId[0], priority, niceValue);
+        }
+
+        public void register(String sessionId, int useCase, int priority) {
+            register(sessionId, useCase, priority, 0);
+        }
+
+        public void unregister() {
+            mTunerResourceManagerService.unregisterClientProfileInternal(mClientId[0]);
+            mClientId[0] = TunerResourceManagerService.INVALID_CLIENT_ID;
+            mReclaimed = false;
+        }
+
+        public int getId() {
+            return mClientId[0];
+        }
+
+        public ClientProfile getProfile() {
+            return mProfile;
+        }
+
+        @Override
+        public void onReclaimResources() {
+            mTunerResourceManagerService.clearAllResourcesAndClientMapping(mProfile);
+            mReclaimed = true;
+        }
+
+        public boolean isReclaimed() {
+            return mReclaimed;
+        }
+    }
+
     private static final class TestResourcesReclaimListener extends IResourcesReclaimListener.Stub {
         boolean mReclaimed;
 
@@ -247,13 +303,11 @@
     }
 
     @Test
-    public void requestFrontendTest_NoFrontendWithGiveTypeAvailable() {
-        ResourceClientProfile profile = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile, null /*listener*/, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+    public void requestFrontendTest_NoFrontendWithGiveTypeAvailable() throws RemoteException {
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[1];
@@ -262,21 +316,20 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isFalse();
         assertThat(frontendHandle[0]).isEqualTo(TunerResourceManager.INVALID_RESOURCE_HANDLE);
+        client0.unregister();
     }
 
     @Test
-    public void requestFrontendTest_FrontendWithNoExclusiveGroupAvailable() {
-        ResourceClientProfile profile = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile, null /*listener*/, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+    public void requestFrontendTest_FrontendWithNoExclusiveGroupAvailable() throws RemoteException {
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[3];
@@ -295,27 +348,23 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
         assertThat(frontendHandle[0]).isEqualTo(0);
+        client0.unregister();
     }
 
     @Test
-    public void requestFrontendTest_FrontendWithExclusiveGroupAvailable() {
-        ResourceClientProfile profile0 = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        ResourceClientProfile profile1 = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile0, null /*listener*/, clientId0);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile1, null /*listener*/, clientId1);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+    public void requestFrontendTest_FrontendWithExclusiveGroupAvailable() throws RemoteException {
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[3];
@@ -335,13 +384,13 @@
 
         int[] frontendHandle = new int[1];
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId1[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client1.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
         assertThat(frontendHandle[0]).isEqualTo(infos[0].handle);
 
         request =
-                tunerFrontendRequest(clientId0[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
         assertThat(frontendHandle[0]).isEqualTo(infos[1].handle);
@@ -349,31 +398,20 @@
                 .isTrue();
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[2].handle).isInUse())
                 .isTrue();
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void requestFrontendTest_NoFrontendAvailable_RequestWithLowerPriority() {
+    public void requestFrontendTest_NoFrontendAvailable_RequestWithLowerPriority()
+            throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[2];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        profiles[1] = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientPriorities = {100, 50};
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[0], listener, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId0[0], clientPriorities[0], 0/*niceValue*/);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[1], new TestResourcesReclaimListener(), clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId1[0], clientPriorities[1], 0/*niceValue*/);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 100);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 50);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[2];
@@ -384,46 +422,36 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId0[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
 
         request =
-                tunerFrontendRequest(clientId1[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client1.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isFalse();
-        assertThat(listener.isReclaimed()).isFalse();
+        assertThat(client0.isReclaimed()).isFalse();
 
         request =
-                tunerFrontendRequest(clientId1[0] /*clientId*/, FrontendSettings.TYPE_DVBS);
+                tunerFrontendRequest(client1.getId() /*clientId*/, FrontendSettings.TYPE_DVBS);
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isFalse();
-        assertThat(listener.isReclaimed()).isFalse();
+        assertThat(client0.isReclaimed()).isFalse();
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void requestFrontendTest_NoFrontendAvailable_RequestWithHigherPriority() {
+    public void requestFrontendTest_NoFrontendAvailable_RequestWithHigherPriority()
+            throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[2];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        profiles[1] = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientPriorities = {100, 500};
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[0], listener, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId0[0], clientPriorities[0], 0/*niceValue*/);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[1], new TestResourcesReclaimListener(), clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId1[0], clientPriorities[1], 0/*niceValue*/);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 100);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 500);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[2];
@@ -434,17 +462,16 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId0[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
         assertThat(frontendHandle[0]).isEqualTo(infos[0].handle);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseFrontendHandles()).isEqualTo(new HashSet<Integer>(Arrays.asList(
-                        infos[0].handle, infos[1].handle)));
+        assertThat(client0.getProfile().getInUseFrontendHandles())
+                .isEqualTo(new HashSet<Integer>(Arrays.asList(infos[0].handle, infos[1].handle)));
 
         request =
-                tunerFrontendRequest(clientId1[0] /*clientId*/, FrontendSettings.TYPE_DVBS);
+                tunerFrontendRequest(client1.getId() /*clientId*/, FrontendSettings.TYPE_DVBS);
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
         assertThat(frontendHandle[0]).isEqualTo(infos[1].handle);
@@ -453,22 +480,20 @@
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
                 .isInUse()).isTrue();
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
-                .getOwnerClientId()).isEqualTo(clientId1[0]);
+                .getOwnerClientId()).isEqualTo(client1.getId());
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
-                .getOwnerClientId()).isEqualTo(clientId1[0]);
-        assertThat(listener.isReclaimed()).isTrue();
+                .getOwnerClientId()).isEqualTo(client1.getId());
+        assertThat(client0.isReclaimed()).isTrue();
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void releaseFrontendTest_UnderTheSameExclusiveGroup() {
+    public void releaseFrontendTest_UnderTheSameExclusiveGroup() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[1];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(profiles[0], listener, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[2];
@@ -479,7 +504,7 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
@@ -488,43 +513,29 @@
                 .getFrontendResource(infos[1].handle).isInUse()).isTrue();
 
         // Release frontend
-        mTunerResourceManagerService.releaseFrontendInternal(mTunerResourceManagerService
-                .getFrontendResource(frontendHandle[0]), clientId[0]);
+        mTunerResourceManagerService.releaseFrontendInternal(frontendHandle[0], client0.getId());
         assertThat(mTunerResourceManagerService
                 .getFrontendResource(frontendHandle[0]).isInUse()).isFalse();
         assertThat(mTunerResourceManagerService
                 .getFrontendResource(infos[1].handle).isInUse()).isFalse();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(clientId[0]).getInUseFrontendHandles().size()).isEqualTo(0);
+        assertThat(client0.getProfile().getInUseFrontendHandles().size()).isEqualTo(0);
+        client0.unregister();
     }
 
     @Test
-    public void requestCasTest_NoCasAvailable_RequestWithHigherPriority() {
+    public void requestCasTest_NoCasAvailable_RequestWithHigherPriority() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[2];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        profiles[1] = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientPriorities = {100, 500};
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[0], listener, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId0[0], clientPriorities[0], 0/*niceValue*/);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[1], new TestResourcesReclaimListener(), clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId1[0], clientPriorities[1], 0/*niceValue*/);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 100);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 500);
 
         // Init cas resources.
         mTunerResourceManagerService.updateCasInfoInternal(1 /*casSystemId*/, 2 /*maxSessionNum*/);
 
-        CasSessionRequest request = casSessionRequest(clientId0[0], 1 /*casSystemId*/);
+        CasSessionRequest request = casSessionRequest(client0.getId(), 1 /*casSystemId*/);
         int[] casSessionHandle = new int[1];
         // Request for 2 cas sessions.
         assertThat(mTunerResourceManagerService
@@ -533,54 +544,45 @@
                 .requestCasSessionInternal(request, casSessionHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(casSessionHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseCasSystemId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCasSystemId())
+                .isEqualTo(1);
         assertThat(mTunerResourceManagerService.getCasResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId0[0])));
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client0.getId())));
         assertThat(mTunerResourceManagerService.getCasResource(1).isFullyUsed()).isTrue();
 
-        request = casSessionRequest(clientId1[0], 1);
+        request = casSessionRequest(client1.getId(), 1);
         assertThat(mTunerResourceManagerService
                 .requestCasSessionInternal(request, casSessionHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(casSessionHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId1[0])
-                .getInUseCasSystemId()).isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseCasSystemId()).isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
+        assertThat(client1.getProfile().getInUseCasSystemId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCasSystemId())
+                .isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
         assertThat(mTunerResourceManagerService.getCasResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId1[0])));
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client1.getId())));
         assertThat(mTunerResourceManagerService.getCasResource(1).isFullyUsed()).isFalse();
-        assertThat(listener.isReclaimed()).isTrue();
+        assertThat(client0.isReclaimed()).isTrue();
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void requestCiCamTest_NoCiCamAvailable_RequestWithHigherPriority() {
+    public void requestCiCamTest_NoCiCamAvailable_RequestWithHigherPriority()
+            throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[2];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        profiles[1] = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientPriorities = {100, 500};
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[0], listener, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId0[0], clientPriorities[0], 0/*niceValue*/);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[1], new TestResourcesReclaimListener(), clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId1[0], clientPriorities[1], 0/*niceValue*/);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 100);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 500);
 
         // Init cicam/cas resources.
         mTunerResourceManagerService.updateCasInfoInternal(1 /*casSystemId*/, 2 /*maxSessionNum*/);
 
-        TunerCiCamRequest request = tunerCiCamRequest(clientId0[0], 1 /*ciCamId*/);
+        TunerCiCamRequest request = tunerCiCamRequest(client0.getId(), 1 /*ciCamId*/);
         int[] ciCamHandle = new int[1];
         // Request for 2 ciCam sessions.
         assertThat(mTunerResourceManagerService
@@ -589,139 +591,125 @@
                 .requestCiCamInternal(request, ciCamHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(ciCamHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseCiCamId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCiCamId()).isEqualTo(1);
         assertThat(mTunerResourceManagerService.getCiCamResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId0[0])));
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client0.getId())));
         assertThat(mTunerResourceManagerService.getCiCamResource(1).isFullyUsed()).isTrue();
 
-        request = tunerCiCamRequest(clientId1[0], 1);
+        request = tunerCiCamRequest(client1.getId(), 1);
         assertThat(mTunerResourceManagerService
                 .requestCiCamInternal(request, ciCamHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(ciCamHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId1[0])
-                .getInUseCiCamId()).isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseCiCamId()).isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
+        assertThat(client1.getProfile().getInUseCiCamId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCiCamId())
+                .isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
         assertThat(mTunerResourceManagerService.getCiCamResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId1[0])));
-        assertThat(mTunerResourceManagerService.getCiCamResource(1).isFullyUsed()).isFalse();
-        assertThat(listener.isReclaimed()).isTrue();
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client1.getId())));
+        assertThat(mTunerResourceManagerService
+                .getCiCamResource(1).isFullyUsed()).isFalse();
+        assertThat(client0.isReclaimed()).isTrue();
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void releaseCasTest() {
+    public void releaseCasTest() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[1];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(profiles[0], listener, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init cas resources.
         mTunerResourceManagerService.updateCasInfoInternal(1 /*casSystemId*/, 2 /*maxSessionNum*/);
 
-        CasSessionRequest request = casSessionRequest(clientId[0], 1 /*casSystemId*/);
+        CasSessionRequest request = casSessionRequest(client0.getId(), 1 /*casSystemId*/);
         int[] casSessionHandle = new int[1];
         // Request for 1 cas sessions.
         assertThat(mTunerResourceManagerService
                 .requestCasSessionInternal(request, casSessionHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(casSessionHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId[0])
-                .getInUseCasSystemId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCasSystemId()).isEqualTo(1);
         assertThat(mTunerResourceManagerService.getCasResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId[0])));
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client0.getId())));
         assertThat(mTunerResourceManagerService.getCasResource(1).isFullyUsed()).isFalse();
 
         // Release cas
         mTunerResourceManagerService.releaseCasSessionInternal(mTunerResourceManagerService
-                .getCasResource(1), clientId[0]);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId[0])
-                .getInUseCasSystemId()).isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
+                .getCasResource(1), client0.getId());
+        assertThat(client0.getProfile().getInUseCasSystemId())
+                .isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
         assertThat(mTunerResourceManagerService.getCasResource(1).isFullyUsed()).isFalse();
         assertThat(mTunerResourceManagerService.getCasResource(1)
                 .getOwnerClientIds()).isEmpty();
+        client0.unregister();
     }
 
     @Test
-    public void releaseCiCamTest() {
+    public void releaseCiCamTest() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[1];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(profiles[0], listener, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init cas resources.
         mTunerResourceManagerService.updateCasInfoInternal(1 /*casSystemId*/, 2 /*maxSessionNum*/);
 
-        TunerCiCamRequest request = tunerCiCamRequest(clientId[0], 1 /*ciCamId*/);
+        TunerCiCamRequest request = tunerCiCamRequest(client0.getId(), 1 /*ciCamId*/);
         int[] ciCamHandle = new int[1];
         // Request for 1 ciCam sessions.
         assertThat(mTunerResourceManagerService
                 .requestCiCamInternal(request, ciCamHandle)).isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(ciCamHandle[0]))
                 .isEqualTo(1);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId[0])
-                .getInUseCiCamId()).isEqualTo(1);
+        assertThat(client0.getProfile().getInUseCiCamId()).isEqualTo(1);
         assertThat(mTunerResourceManagerService.getCiCamResource(1)
-                .getOwnerClientIds()).isEqualTo(new HashSet<Integer>(Arrays.asList(clientId[0])));
-        assertThat(mTunerResourceManagerService.getCiCamResource(1).isFullyUsed()).isFalse();
+                .getOwnerClientIds()).isEqualTo(
+                        new HashSet<Integer>(Arrays.asList(client0.getId())));
+        assertThat(mTunerResourceManagerService
+                .getCiCamResource(1).isFullyUsed()).isFalse();
 
         // Release ciCam
         mTunerResourceManagerService.releaseCiCamInternal(mTunerResourceManagerService
-                .getCiCamResource(1), clientId[0]);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId[0])
-                .getInUseCiCamId()).isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
-        assertThat(mTunerResourceManagerService.getCiCamResource(1).isFullyUsed()).isFalse();
+                .getCiCamResource(1), client0.getId());
+        assertThat(client0.getProfile().getInUseCiCamId())
+                .isEqualTo(ClientProfile.INVALID_RESOURCE_ID);
+        assertThat(mTunerResourceManagerService
+                .getCiCamResource(1).isFullyUsed()).isFalse();
         assertThat(mTunerResourceManagerService.getCiCamResource(1)
                 .getOwnerClientIds()).isEmpty();
+        client0.unregister();
     }
 
     @Test
-    public void requestLnbTest_NoLnbAvailable_RequestWithHigherPriority() {
+    public void requestLnbTest_NoLnbAvailable_RequestWithHigherPriority() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[2];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        profiles[1] = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientPriorities = {100, 500};
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[0], listener, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId0[0], clientPriorities[0], 0/*niceValue*/);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profiles[1], new TestResourcesReclaimListener(), clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                clientId1[0], clientPriorities[1], 0/*niceValue*/);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 100);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK, 500);
 
         // Init lnb resources.
         int[] lnbHandles = {1};
         mTunerResourceManagerService.setLnbInfoListInternal(lnbHandles);
 
         TunerLnbRequest request = new TunerLnbRequest();
-        request.clientId = clientId0[0];
+        request.clientId = client0.getId();
         int[] lnbHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestLnbInternal(request, lnbHandle)).isTrue();
         assertThat(lnbHandle[0]).isEqualTo(lnbHandles[0]);
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0]).getInUseLnbHandles())
+        assertThat(client0.getProfile().getInUseLnbHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(lnbHandles[0])));
 
         request = new TunerLnbRequest();
-        request.clientId = clientId1[0];
+        request.clientId = client1.getId();
 
         assertThat(mTunerResourceManagerService
                 .requestLnbInternal(request, lnbHandle)).isTrue();
@@ -729,29 +717,26 @@
         assertThat(mTunerResourceManagerService.getLnbResource(lnbHandles[0])
                 .isInUse()).isTrue();
         assertThat(mTunerResourceManagerService.getLnbResource(lnbHandles[0])
-                .getOwnerClientId()).isEqualTo(clientId1[0]);
-        assertThat(listener.isReclaimed()).isTrue();
-        assertThat(mTunerResourceManagerService.getClientProfile(clientId0[0])
-                .getInUseLnbHandles().size()).isEqualTo(0);
+                .getOwnerClientId()).isEqualTo(client1.getId());
+        assertThat(client0.isReclaimed()).isTrue();
+        assertThat(client0.getProfile().getInUseLnbHandles().size()).isEqualTo(0);
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void releaseLnbTest() {
+    public void releaseLnbTest() throws RemoteException {
         // Register clients
-        ResourceClientProfile[] profiles = new ResourceClientProfile[1];
-        profiles[0] = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        TestResourcesReclaimListener listener = new TestResourcesReclaimListener();
-        mTunerResourceManagerService.registerClientProfileInternal(profiles[0], listener, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init lnb resources.
         int[] lnbHandles = {0};
         mTunerResourceManagerService.setLnbInfoListInternal(lnbHandles);
 
         TunerLnbRequest request = new TunerLnbRequest();
-        request.clientId = clientId[0];
+        request.clientId = client0.getId();
         int[] lnbHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestLnbInternal(request, lnbHandle)).isTrue();
@@ -762,19 +747,16 @@
                 .getLnbResource(lnbHandle[0]));
         assertThat(mTunerResourceManagerService
                 .getLnbResource(lnbHandle[0]).isInUse()).isFalse();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(clientId[0]).getInUseLnbHandles().size()).isEqualTo(0);
+        assertThat(client0.getProfile().getInUseLnbHandles().size()).isEqualTo(0);
+        client0.unregister();
     }
 
     @Test
-    public void unregisterClientTest_usingFrontend() {
-        // Register client
-        ResourceClientProfile profile = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile, null /*listener*/, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+    public void unregisterClientTest_usingFrontend() throws RemoteException {
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         // Init frontend resources.
         TunerFrontendInfo[] infos = new TunerFrontendInfo[2];
@@ -785,7 +767,7 @@
         mTunerResourceManagerService.setFrontendInfoListInternal(infos);
 
         TunerFrontendRequest request =
-                tunerFrontendRequest(clientId[0] /*clientId*/, FrontendSettings.TYPE_DVBT);
+                tunerFrontendRequest(client0.getId() /*clientId*/, FrontendSettings.TYPE_DVBT);
         int[] frontendHandle = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestFrontendInternal(request, frontendHandle)).isTrue();
@@ -796,26 +778,20 @@
                 .isInUse()).isTrue();
 
         // Unregister client when using frontend
-        mTunerResourceManagerService.unregisterClientProfileInternal(clientId[0]);
+        client0.unregister();
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
                 .isInUse()).isFalse();
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
                 .isInUse()).isFalse();
-        assertThat(mTunerResourceManagerService.checkClientExists(clientId[0])).isFalse();
-
+        assertThat(mTunerResourceManagerService.checkClientExists(client0.getId())).isFalse();
     }
 
     @Test
-    public void requestDemuxTest() {
-        // Register client
-        ResourceClientProfile profile0 = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        ResourceClientProfile profile1 = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId0 = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile0, null /*listener*/, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+    public void requestDemuxTest() throws RemoteException {
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         TunerDemuxInfo[] infos = new TunerDemuxInfo[3];
         infos[0] = tunerDemuxInfo(0 /* handle */, Filter.TYPE_TS | Filter.TYPE_IP);
@@ -825,7 +801,7 @@
 
         int[] demuxHandle0 = new int[1];
         // first with undefined type (should be the first one with least # of caps)
-        TunerDemuxRequest request = tunerDemuxRequest(clientId0[0], Filter.TYPE_UNDEFINED);
+        TunerDemuxRequest request = tunerDemuxRequest(client0.getId(), Filter.TYPE_UNDEFINED);
         assertThat(mTunerResourceManagerService.requestDemuxInternal(request, demuxHandle0))
                 .isTrue();
         assertThat(demuxHandle0[0]).isEqualTo(1);
@@ -846,16 +822,16 @@
         assertThat(demuxHandle0[0]).isEqualTo(2);
 
         // request for another TS
-        int[] clientId1 = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile1, null /*listener*/, clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client1 = new TunerClient();
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
+
         int[] demuxHandle1 = new int[1];
-        TunerDemuxRequest request1 = tunerDemuxRequest(clientId1[0], Filter.TYPE_TS);
+        TunerDemuxRequest request1 = tunerDemuxRequest(client1.getId(), Filter.TYPE_TS);
         assertThat(mTunerResourceManagerService.requestDemuxInternal(request1, demuxHandle1))
                 .isTrue();
         assertThat(demuxHandle1[0]).isEqualTo(0);
-        assertThat(mTunerResourceManagerService.getResourceIdFromHandle(demuxHandle1[0]))
+        assertThat(mTunerResourceManagerService.getResourceIdFromHandle(client1.getId()))
                 .isEqualTo(0);
 
         // release demuxes
@@ -863,33 +839,23 @@
         mTunerResourceManagerService.releaseDemuxInternal(dr);
         dr = mTunerResourceManagerService.getDemuxResource(demuxHandle1[0]);
         mTunerResourceManagerService.releaseDemuxInternal(dr);
+
+        client0.unregister();
+        client1.unregister();
     }
 
     @Test
-    public void requestDemuxTest_ResourceReclaim() {
+    public void requestDemuxTest_ResourceReclaim() throws RemoteException {
         // Register clients
-        ResourceClientProfile profile0 = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        ResourceClientProfile profile1 = resourceClientProfile("1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_SCAN);
-        ResourceClientProfile profile2 = resourceClientProfile("2" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_SCAN);
-        int[] clientId0 = new int[1];
-        int[] clientId1 = new int[1];
-        int[] clientId2 = new int[1];
-        TestResourcesReclaimListener listener0 = new TestResourcesReclaimListener();
-        TestResourcesReclaimListener listener1 = new TestResourcesReclaimListener();
-        TestResourcesReclaimListener listener2 = new TestResourcesReclaimListener();
-
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile0, listener0, clientId0);
-        assertThat(clientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile1, listener1, clientId1);
-        assertThat(clientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile2, listener2, clientId1);
-        assertThat(clientId2[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient client0 = new TunerClient();
+        TunerClient client1 = new TunerClient();
+        TunerClient client2 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
+        client1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_SCAN);
+        client2.register("2" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_SCAN);
 
         // Init demux resources.
         TunerDemuxInfo[] infos = new TunerDemuxInfo[2];
@@ -897,66 +863,67 @@
         infos[1] = tunerDemuxInfo(1 /*handle*/, Filter.TYPE_TS);
         mTunerResourceManagerService.setDemuxInfoListInternal(infos);
 
-        // let clientId0(prio:100) request for IP - should succeed
-        TunerDemuxRequest request0 = tunerDemuxRequest(clientId0[0], Filter.TYPE_IP);
+        // let client0(prio:100) request for IP - should succeed
+        TunerDemuxRequest request0 = tunerDemuxRequest(client0.getId(), Filter.TYPE_IP);
         int[] demuxHandle0 = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestDemuxInternal(request0, demuxHandle0)).isTrue();
         assertThat(demuxHandle0[0]).isEqualTo(0);
 
-        // let clientId1(prio:50) request for IP - should fail
-        TunerDemuxRequest request1 = tunerDemuxRequest(clientId1[0], Filter.TYPE_IP);
+        // let client1(prio:50) request for IP - should fail
+        TunerDemuxRequest request1 = tunerDemuxRequest(client1.getId(), Filter.TYPE_IP);
         int[] demuxHandle1 = new int[1];
         demuxHandle1[0] = -1;
         assertThat(mTunerResourceManagerService
                 .requestDemuxInternal(request1, demuxHandle1)).isFalse();
-        assertThat(listener0.isReclaimed()).isFalse();
+        assertThat(client0.isReclaimed()).isFalse();
         assertThat(demuxHandle1[0]).isEqualTo(-1);
 
-        // let clientId1(prio:50) request for TS - should succeed
+        // let client1(prio:50) request for TS - should succeed
         request1.desiredFilterTypes = Filter.TYPE_TS;
         assertThat(mTunerResourceManagerService
                 .requestDemuxInternal(request1, demuxHandle1)).isTrue();
         assertThat(demuxHandle1[0]).isEqualTo(1);
-        assertThat(listener0.isReclaimed()).isFalse();
+        assertThat(client0.isReclaimed()).isFalse();
 
-        // now release demux for the clientId0 (higher priority) and request demux
+        // now release demux for the client0 (higher priority) and request demux
         DemuxResource dr = mTunerResourceManagerService.getDemuxResource(demuxHandle0[0]);
         mTunerResourceManagerService.releaseDemuxInternal(dr);
 
-        // let clientId2(prio:50) request for TS - should succeed
-        TunerDemuxRequest request2 = tunerDemuxRequest(clientId2[0], Filter.TYPE_TS);
+        // let client2(prio:50) request for TS - should succeed
+        TunerDemuxRequest request2 = tunerDemuxRequest(client2.getId(), Filter.TYPE_TS);
         int[] demuxHandle2 = new int[1];
         assertThat(mTunerResourceManagerService
                 .requestDemuxInternal(request2, demuxHandle2)).isTrue();
         assertThat(demuxHandle2[0]).isEqualTo(0);
-        assertThat(listener1.isReclaimed()).isFalse();
+        assertThat(client1.isReclaimed()).isFalse();
 
-        // let clientId0(prio:100) request for TS - should reclaim from clientId2
+        // let client0(prio:100) request for TS - should reclaim from client1
         // , who has the smaller caps
         request0.desiredFilterTypes = Filter.TYPE_TS;
         assertThat(mTunerResourceManagerService
                 .requestDemuxInternal(request0, demuxHandle0)).isTrue();
-        assertThat(listener1.isReclaimed()).isFalse();
-        assertThat(listener2.isReclaimed()).isTrue();
+        assertThat(client1.isReclaimed()).isTrue();
+        assertThat(client2.isReclaimed()).isFalse();
+        client0.unregister();
+        client1.unregister();
+        client2.unregister();
     }
 
     @Test
     public void requestDescramblerTest() {
-        // Register client
-        ResourceClientProfile profile = resourceClientProfile("0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
-        int[] clientId = new int[1];
-        mTunerResourceManagerService.registerClientProfileInternal(
-                profile, null /*listener*/, clientId);
-        assertThat(clientId[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        // Register clients
+        TunerClient client0 = new TunerClient();
+        client0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_PLAYBACK);
 
         int[] desHandle = new int[1];
         TunerDescramblerRequest request = new TunerDescramblerRequest();
-        request.clientId = clientId[0];
+        request.clientId = client0.getId();
         assertThat(mTunerResourceManagerService.requestDescramblerInternal(request, desHandle))
                 .isTrue();
         assertThat(mTunerResourceManagerService.getResourceIdFromHandle(desHandle[0])).isEqualTo(0);
+        client0.unregister();
     }
 
     @Test
@@ -978,74 +945,26 @@
     }
 
     @Test
-    public void shareFrontendTest_FrontendWithExclusiveGroupReadyToShare() {
+    public void shareFrontendTest_FrontendWithExclusiveGroupReadyToShare() throws RemoteException {
         /**** Register Clients and Set Priority ****/
-
-        // Int array to save the returned client ids
-        int[] ownerClientId0 = new int[1];
-        int[] ownerClientId1 = new int[1];
-        int[] shareClientId0 = new int[1];
-        int[] shareClientId1 = new int[1];
-
-        // Predefined client profiles
-        ResourceClientProfile[] ownerProfiles = new ResourceClientProfile[2];
-        ResourceClientProfile[] shareProfiles = new ResourceClientProfile[2];
-        ownerProfiles[0] = resourceClientProfile(
-                "0" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_LIVE);
-        ownerProfiles[1] = resourceClientProfile(
-                "1" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_LIVE);
-        shareProfiles[0] = resourceClientProfile(
-                "2" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_RECORD);
-        shareProfiles[1] = resourceClientProfile(
-                "3" /*sessionId*/,
-                TvInputService.PRIORITY_HINT_USE_CASE_TYPE_RECORD);
-
-        // Predefined client reclaim listeners
-        TestResourcesReclaimListener ownerListener0 = new TestResourcesReclaimListener();
-        TestResourcesReclaimListener shareListener0 = new TestResourcesReclaimListener();
-        TestResourcesReclaimListener ownerListener1 = new TestResourcesReclaimListener();
-        TestResourcesReclaimListener shareListener1 = new TestResourcesReclaimListener();
-        // Register clients and validate the returned client ids
-        mTunerResourceManagerService
-                .registerClientProfileInternal(ownerProfiles[0], ownerListener0, ownerClientId0);
-        mTunerResourceManagerService
-                .registerClientProfileInternal(shareProfiles[0], shareListener0, shareClientId0);
-        mTunerResourceManagerService
-                .registerClientProfileInternal(ownerProfiles[1], ownerListener1, ownerClientId1);
-        mTunerResourceManagerService
-                .registerClientProfileInternal(shareProfiles[1], shareListener1, shareClientId1);
-        assertThat(ownerClientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        assertThat(shareClientId0[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        assertThat(ownerClientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
-        assertThat(shareClientId1[0]).isNotEqualTo(TunerResourceManagerService.INVALID_CLIENT_ID);
+        TunerClient ownerClient0 = new TunerClient();
+        TunerClient ownerClient1 = new TunerClient();
+        TunerClient shareClient0 = new TunerClient();
+        TunerClient shareClient1 = new TunerClient();
+        ownerClient0.register("0" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_LIVE, 100);
+        ownerClient1.register("1" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_LIVE, 300);
+        shareClient0.register("2" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_RECORD, 200);
+        shareClient1.register("3" /*sessionId*/,
+                        TvInputService.PRIORITY_HINT_USE_CASE_TYPE_RECORD, 400);
 
         mTunerResourceManagerService.updateClientPriorityInternal(
-                ownerClientId0[0],
-                100/*priority*/,
-                0/*niceValue*/);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                shareClientId0[0],
-                200/*priority*/,
-                0/*niceValue*/);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                ownerClientId1[0],
-                300/*priority*/,
-                0/*niceValue*/);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                shareClientId1[0],
-                400/*priority*/,
-                0/*niceValue*/);
-        mTunerResourceManagerService.updateClientPriorityInternal(
-                shareClientId1[0],
+                shareClient1.getId(),
                 -1/*invalid priority*/,
                 0/*niceValue*/);
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId1[0])
-                .getPriority())
-                .isEqualTo(400);
+        assertThat(shareClient1.getProfile().getPriority()).isEqualTo(400);
 
         /**** Init Frontend Resources ****/
 
@@ -1072,7 +991,7 @@
         // Predefined frontend request and array to save returned frontend handle
         int[] frontendHandle = new int[1];
         TunerFrontendRequest request = tunerFrontendRequest(
-                ownerClientId0[0] /*clientId*/,
+                ownerClient0.getId() /*clientId*/,
                 FrontendSettings.TYPE_DVBT);
 
         // Request call and validate granted resource and internal mapping
@@ -1080,9 +999,7 @@
                 .requestFrontendInternal(request, frontendHandle))
                 .isTrue();
         assertThat(frontendHandle[0]).isEqualTo(infos[0].handle);
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId0[0])
-                .getInUseFrontendHandles())
+        assertThat(ownerClient0.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
@@ -1091,11 +1008,11 @@
 
         // Share frontend call and validate the internal mapping
         mTunerResourceManagerService.shareFrontendInternal(
-                shareClientId0[0]/*selfClientId*/,
-                ownerClientId0[0]/*targetClientId*/);
+                shareClient0.getId()/*selfClientId*/,
+                ownerClient0.getId()/*targetClientId*/);
         mTunerResourceManagerService.shareFrontendInternal(
-                shareClientId1[0]/*selfClientId*/,
-                ownerClientId0[0]/*targetClientId*/);
+                shareClient1.getId()/*selfClientId*/,
+                ownerClient0.getId()/*targetClientId*/);
         // Verify fe in use status
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
                 .isInUse()).isTrue();
@@ -1103,31 +1020,24 @@
                 .isInUse()).isTrue();
         // Verify fe owner status
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
-                .getOwnerClientId()).isEqualTo(ownerClientId0[0]);
+                .getOwnerClientId()).isEqualTo(ownerClient0.getId());
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
-                .getOwnerClientId()).isEqualTo(ownerClientId0[0]);
+                .getOwnerClientId()).isEqualTo(ownerClient0.getId());
         // Verify share fe client status in the primary owner client
-        assertThat(mTunerResourceManagerService.getClientProfile(ownerClientId0[0])
-                .getShareFeClientIds())
+        assertThat(ownerClient0.getProfile().getShareFeClientIds())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
-                        shareClientId0[0],
-                        shareClientId1[0])));
+                        shareClient0.getId(),
+                        shareClient1.getId())));
         // Verify in use frontend list in all the primary owner and share owner clients
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId0[0])
-                .getInUseFrontendHandles())
+        assertThat(ownerClient0.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
-                .getInUseFrontendHandles())
+        assertThat(shareClient0.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId1[0])
-                .getInUseFrontendHandles())
+        assertThat(shareClient1.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
@@ -1135,21 +1045,17 @@
         /**** Remove Frontend Share Owner ****/
 
         // Unregister the second share fe client
-        mTunerResourceManagerService.unregisterClientProfileInternal(shareClientId1[0]);
+        shareClient1.unregister();
 
         // Validate the internal mapping
-        assertThat(mTunerResourceManagerService.getClientProfile(ownerClientId0[0])
-                .getShareFeClientIds())
+        assertThat(ownerClient0.getProfile().getShareFeClientIds())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
-                        shareClientId0[0])));
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId0[0])
-                .getInUseFrontendHandles())
+                        shareClient0.getId())));
+        assertThat(ownerClient0.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
+        assertThat(shareClient0.getProfile()
                 .getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
@@ -1159,7 +1065,7 @@
 
         // Predefined second frontend request
         request = tunerFrontendRequest(
-                ownerClientId1[0] /*clientId*/,
+                ownerClient1.getId() /*clientId*/,
                 FrontendSettings.TYPE_DVBT);
 
         // Second request call
@@ -1170,43 +1076,35 @@
         // Validate granted resource and internal mapping
         assertThat(frontendHandle[0]).isEqualTo(infos[0].handle);
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
-                .getOwnerClientId()).isEqualTo(ownerClientId1[0]);
+                .getOwnerClientId()).isEqualTo(ownerClient1.getId());
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
-                .getOwnerClientId()).isEqualTo(ownerClientId1[0]);
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId1[0])
-                .getInUseFrontendHandles())
+                .getOwnerClientId()).isEqualTo(ownerClient1.getId());
+        assertThat(ownerClient1.getProfile().getInUseFrontendHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         infos[0].handle,
                         infos[1].handle)));
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId0[0])
-                .getInUseFrontendHandles()
+        assertThat(ownerClient0.getProfile().getInUseFrontendHandles()
                 .isEmpty())
                 .isTrue();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
-                .getInUseFrontendHandles()
+        assertThat(shareClient0.getProfile().getInUseFrontendHandles()
                 .isEmpty())
                 .isTrue();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId0[0])
-                .getShareFeClientIds()
+        assertThat(ownerClient0.getProfile().getShareFeClientIds()
                 .isEmpty())
                 .isTrue();
-        assertThat(ownerListener0.isReclaimed()).isTrue();
-        assertThat(shareListener0.isReclaimed()).isTrue();
+        assertThat(ownerClient0.isReclaimed()).isTrue();
+        assertThat(shareClient0.isReclaimed()).isTrue();
 
         /**** Release Frontend Resource From Primary Owner ****/
 
         // Reshare the frontend
         mTunerResourceManagerService.shareFrontendInternal(
-                shareClientId0[0]/*selfClientId*/,
-                ownerClientId1[0]/*targetClientId*/);
+                shareClient0.getId()/*selfClientId*/,
+                ownerClient1.getId()/*targetClientId*/);
 
         // Release the frontend resource from the primary owner
-        mTunerResourceManagerService.releaseFrontendInternal(mTunerResourceManagerService
-                .getFrontendResource(infos[0].handle), ownerClientId1[0]);
+        mTunerResourceManagerService.releaseFrontendInternal(infos[0].handle,
+                ownerClient1.getId());
 
         // Validate the internal mapping
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
@@ -1214,19 +1112,13 @@
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
                 .isInUse()).isFalse();
         // Verify client status
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId1[0])
-                .getInUseFrontendHandles()
+        assertThat(ownerClient1.getProfile().getInUseFrontendHandles()
                 .isEmpty())
                 .isTrue();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
-                .getInUseFrontendHandles()
+        assertThat(shareClient0.getProfile().getInUseFrontendHandles()
                 .isEmpty())
                 .isTrue();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(ownerClientId1[0])
-                .getShareFeClientIds()
+        assertThat(ownerClient1.getProfile().getShareFeClientIds()
                 .isEmpty())
                 .isTrue();
 
@@ -1234,7 +1126,7 @@
 
         // Predefined Lnb request and handle array
         TunerLnbRequest requestLnb = new TunerLnbRequest();
-        requestLnb.clientId = shareClientId0[0];
+        requestLnb.clientId = shareClient0.getId();
         int[] lnbHandle = new int[1];
 
         // Request for an Lnb
@@ -1247,11 +1139,11 @@
                 .requestFrontendInternal(request, frontendHandle))
                 .isTrue();
         mTunerResourceManagerService.shareFrontendInternal(
-                shareClientId0[0]/*selfClientId*/,
-                ownerClientId1[0]/*targetClientId*/);
+                shareClient0.getId()/*selfClientId*/,
+                ownerClient1.getId()/*targetClientId*/);
 
         // Unregister the primary owner of the shared frontend
-        mTunerResourceManagerService.unregisterClientProfileInternal(ownerClientId1[0]);
+        ownerClient1.unregister();
 
         // Validate the internal mapping
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[0].handle)
@@ -1259,16 +1151,15 @@
         assertThat(mTunerResourceManagerService.getFrontendResource(infos[1].handle)
                 .isInUse()).isFalse();
         // Verify client status
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
-                .getInUseFrontendHandles()
+        assertThat(shareClient0.getProfile().getInUseFrontendHandles()
                 .isEmpty())
                 .isTrue();
-        assertThat(mTunerResourceManagerService
-                .getClientProfile(shareClientId0[0])
-                .getInUseLnbHandles())
+        assertThat(shareClient0.getProfile().getInUseLnbHandles())
                 .isEqualTo(new HashSet<Integer>(Arrays.asList(
                         lnbHandles[0])));
+
+        ownerClient0.unregister();
+        shareClient0.unregister();
     }
 
     private TunerFrontendInfo tunerFrontendInfo(
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index 0a52238..6a1140c 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -14995,7 +14995,6 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void enqueueNotification_acceptsCorrectToken() throws RemoteException {
         Notification sent = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setContentIntent(createPendingIntent("content"))
@@ -15014,7 +15013,6 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void enqueueNotification_acceptsNullToken_andPopulatesIt() throws RemoteException {
         Notification receivedWithoutParceling = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setContentIntent(createPendingIntent("content"))
@@ -15031,7 +15029,6 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void enqueueNotification_directlyThroughRunnable_populatesAllowlistToken() {
         Notification receivedWithoutParceling = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setContentIntent(createPendingIntent("content"))
@@ -15054,7 +15051,6 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void enqueueNotification_rejectsOtherToken() throws RemoteException {
         Notification sent = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setContentIntent(createPendingIntent("content"))
@@ -15072,7 +15068,6 @@
     }
 
     @Test
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void enqueueNotification_customParcelingWithFakeInnerToken_hasCorrectTokenInIntents()
             throws RemoteException {
         Notification sentFromApp = new Notification.Builder(mContext, TEST_CHANNEL_ID)
@@ -15278,7 +15273,6 @@
 
     @Test
     @SuppressWarnings("unchecked")
-    @EnableFlags(android.app.Flags.FLAG_SECURE_ALLOWLIST_TOKEN)
     public void getActiveNotifications_doesNotLeakAllowlistToken() throws RemoteException {
         Notification sentFromApp = new Notification.Builder(mContext, TEST_CHANNEL_ID)
                 .setContentIntent(createPendingIntent("content"))
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
index e70ed5f..f8ff1f4 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeConfigTest.java
@@ -788,6 +788,19 @@
     }
 
     @Test
+    public void testRuleXml_invalidInterruptionFilter_readsDefault() throws Exception {
+        ZenModeConfig.ZenRule rule = new ZenModeConfig.ZenRule();
+        rule.zenMode = 1979;
+
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        writeRuleXml(rule, baos);
+        ByteArrayInputStream bais = new ByteArrayInputStream(baos.toByteArray());
+        ZenModeConfig.ZenRule fromXml = readRuleXml(bais);
+
+        assertThat(fromXml.zenMode).isEqualTo(ZEN_MODE_IMPORTANT_INTERRUPTIONS);
+    }
+
+    @Test
     public void testZenPolicyXml_allUnset() throws Exception {
         ZenPolicy policy = new ZenPolicy.Builder().build();
 
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
index c1e3f47..baa633f 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ZenModeHelperTest.java
@@ -18,6 +18,8 @@
 
 import static android.app.AutomaticZenRule.TYPE_BEDTIME;
 import static android.app.AutomaticZenRule.TYPE_IMMERSIVE;
+import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
+import static android.app.AutomaticZenRule.TYPE_UNKNOWN;
 import static android.app.Flags.FLAG_MODES_API;
 import static android.app.Flags.FLAG_MODES_UI;
 import static android.app.NotificationManager.AUTOMATIC_RULE_STATUS_ACTIVATED;
@@ -6875,6 +6877,52 @@
                 "Didn't find rule with id %s", ruleId);
     }
 
+    @Test
+    @DisableFlags({FLAG_MODES_API, FLAG_MODES_UI})
+    public void testDefaultConfig_preModesApi_rulesAreBare() {
+        // Create a new user, which should get a copy of the default policy.
+        mZenModeHelper.onUserSwitched(101);
+
+        ZenRule eventsRule = mZenModeHelper.mConfig.automaticRules.get(
+                ZenModeConfig.EVENTS_DEFAULT_RULE_ID);
+
+        assertThat(eventsRule).isNotNull();
+        assertThat(eventsRule.zenPolicy).isNull();
+        assertThat(eventsRule.type).isEqualTo(TYPE_UNKNOWN);
+        assertThat(eventsRule.triggerDescription).isNull();
+    }
+
+    @Test
+    @EnableFlags(FLAG_MODES_API)
+    @DisableFlags(FLAG_MODES_UI)
+    public void testDefaultConfig_modesApi_rulesHaveFullPolicy() {
+        // Create a new user, which should get a copy of the default policy.
+        mZenModeHelper.onUserSwitched(201);
+
+        ZenRule eventsRule = mZenModeHelper.mConfig.automaticRules.get(
+                ZenModeConfig.EVENTS_DEFAULT_RULE_ID);
+
+        assertThat(eventsRule).isNotNull();
+        assertThat(eventsRule.zenPolicy).isEqualTo(mZenModeHelper.getDefaultZenPolicy());
+        assertThat(eventsRule.type).isEqualTo(TYPE_UNKNOWN);
+        assertThat(eventsRule.triggerDescription).isNull();
+    }
+
+    @Test
+    @EnableFlags({FLAG_MODES_API, FLAG_MODES_UI})
+    public void testDefaultConfig_modesUi_rulesHaveFullPolicy() {
+        // Create a new user, which should get a copy of the default policy.
+        mZenModeHelper.onUserSwitched(301);
+
+        ZenRule eventsRule = mZenModeHelper.mConfig.automaticRules.get(
+                ZenModeConfig.EVENTS_DEFAULT_RULE_ID);
+
+        assertThat(eventsRule).isNotNull();
+        assertThat(eventsRule.zenPolicy).isEqualTo(mZenModeHelper.getDefaultZenPolicy());
+        assertThat(eventsRule.type).isEqualTo(TYPE_SCHEDULE_CALENDAR);
+        assertThat(eventsRule.triggerDescription).isNotEmpty();
+    }
+
     private static void addZenRule(ZenModeConfig config, String id, String ownerPkg, int zenMode,
             @Nullable ZenPolicy zenPolicy) {
         ZenRule rule = new ZenRule();
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/InputDeviceDelegateTest.java b/services/tests/vibrator/src/com/android/server/vibrator/InputDeviceDelegateTest.java
index f3ecfcc..66788b6 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/InputDeviceDelegateTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/InputDeviceDelegateTest.java
@@ -45,6 +45,8 @@
 
 import androidx.test.InstrumentationRegistry;
 
+import com.android.server.vibrator.VibrationSession.CallerInfo;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -262,7 +264,7 @@
     public void vibrateIfAvailable_withNoInputDevice_returnsFalse() {
         assertFalse(mInputDeviceDelegate.isAvailable());
         assertFalse(mInputDeviceDelegate.vibrateIfAvailable(
-                new Vibration.CallerInfo(VIBRATION_ATTRIBUTES, UID, -1, PACKAGE_NAME, REASON),
+                new CallerInfo(VIBRATION_ATTRIBUTES, UID, -1, PACKAGE_NAME, REASON),
                 SYNCED_EFFECT));
     }
 
@@ -277,7 +279,7 @@
         mInputDeviceDelegate.updateInputDeviceVibrators(/* vibrateInputDevices= */ true);
 
         assertTrue(mInputDeviceDelegate.vibrateIfAvailable(
-                new Vibration.CallerInfo(VIBRATION_ATTRIBUTES, UID, -1, PACKAGE_NAME, REASON),
+                new CallerInfo(VIBRATION_ATTRIBUTES, UID, -1, PACKAGE_NAME, REASON),
                 SYNCED_EFFECT));
         verify(mIInputManagerMock).vibrateCombined(eq(1), same(SYNCED_EFFECT), any());
         verify(mIInputManagerMock).vibrateCombined(eq(2), same(SYNCED_EFFECT), any());
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibrationTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibrationSessionTest.java
similarity index 80%
rename from services/tests/vibrator/src/com/android/server/vibrator/VibrationTest.java
rename to services/tests/vibrator/src/com/android/server/vibrator/VibrationSessionTest.java
index 84f8412..f69d1c4 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibrationTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibrationSessionTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2022 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.
@@ -24,13 +24,13 @@
 
 import java.util.Arrays;
 
-public class VibrationTest {
+public class VibrationSessionTest {
 
     @Test
     public void status_hasUniqueProtoEnumValues() {
         assertThat(
-                Arrays.stream(Vibration.Status.values())
-                        .map(Vibration.Status::getProtoEnumValue)
+                Arrays.stream(VibrationSession.Status.values())
+                        .map(VibrationSession.Status::getProtoEnumValue)
                         .collect(toList()))
                 .containsNoDuplicates();
     }
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibrationSettingsTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibrationSettingsTest.java
index 38cd49d..c7a136a 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibrationSettingsTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibrationSettingsTest.java
@@ -81,6 +81,8 @@
 import com.android.internal.util.test.FakeSettingsProviderRule;
 import com.android.server.LocalServices;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
+import com.android.server.vibrator.VibrationSession.CallerInfo;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import org.junit.After;
 import org.junit.Before;
@@ -292,7 +294,7 @@
             if (expectedAllowedVibrations.contains(usage)) {
                 assertVibrationNotIgnoredForUsage(usage);
             } else {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_BACKGROUND);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_BACKGROUND);
             }
         }
     }
@@ -350,7 +352,7 @@
         createSystemReadyVibrationSettings();
 
         for (int usage : ALL_USAGES) {
-            assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_ON_WIRELESS_CHARGER);
+            assertVibrationIgnoredForUsage(usage, Status.IGNORED_ON_WIRELESS_CHARGER);
         }
     }
 
@@ -365,7 +367,7 @@
         mRegisteredBatteryBroadcastReceiver.onReceive(mContextSpy, wirelessChargingIntent);
 
         for (int usage : ALL_USAGES) {
-            assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_ON_WIRELESS_CHARGER);
+            assertVibrationIgnoredForUsage(usage, Status.IGNORED_ON_WIRELESS_CHARGER);
         }
     }
 
@@ -377,7 +379,7 @@
         createSystemReadyVibrationSettings();
         // Check that initially, all usages are ignored due to the wireless charging.
         for (int usage : ALL_USAGES) {
-            assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_ON_WIRELESS_CHARGER);
+            assertVibrationIgnoredForUsage(usage, Status.IGNORED_ON_WIRELESS_CHARGER);
         }
 
         Intent nonWirelessChargingIntent = getBatteryChangedIntent(BATTERY_PLUGGED_USB);
@@ -404,7 +406,7 @@
             if (expectedAllowedVibrations.contains(usage)) {
                 assertVibrationNotIgnoredForUsage(usage);
             } else {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_POWER);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_POWER);
             }
         }
     }
@@ -426,7 +428,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_RINGTONE || usage == USAGE_NOTIFICATION) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_RINGER_MODE);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_RINGER_MODE);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -470,7 +472,7 @@
             if (usage == USAGE_ACCESSIBILITY) {
                 assertVibrationNotIgnoredForUsage(usage);
             } else {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             }
             assertVibrationNotIgnoredForUsageAndFlags(usage,
                     VibrationAttributes.FLAG_BYPASS_USER_VIBRATION_INTENSITY_OFF);
@@ -512,7 +514,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_TOUCH) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -527,7 +529,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_TOUCH || usage == USAGE_IME_FEEDBACK) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -542,7 +544,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_HARDWARE_FEEDBACK || usage == USAGE_PHYSICAL_EMULATION) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -557,7 +559,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_NOTIFICATION) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -574,7 +576,7 @@
 
         for (int usage : ALL_USAGES) {
             if (usage == USAGE_RINGTONE) {
-                assertVibrationIgnoredForUsage(usage, Vibration.Status.IGNORED_FOR_SETTINGS);
+                assertVibrationIgnoredForUsage(usage, Status.IGNORED_FOR_SETTINGS);
             } else {
                 assertVibrationNotIgnoredForUsage(usage);
             }
@@ -597,7 +599,7 @@
         mVibrationSettings.mSettingChangeReceiver.onReceive(mContextSpy,
                 new Intent(AudioManager.INTERNAL_RINGER_MODE_CHANGED_ACTION));
 
-        assertVibrationIgnoredForUsage(USAGE_RINGTONE, Vibration.Status.IGNORED_FOR_RINGER_MODE);
+        assertVibrationIgnoredForUsage(USAGE_RINGTONE, Status.IGNORED_FOR_RINGER_MODE);
     }
 
     @Test
@@ -611,7 +613,7 @@
                 new VibrationAttributes.Builder()
                         .setUsage(USAGE_IME_FEEDBACK)
                         .build(),
-                Vibration.Status.IGNORED_FOR_SETTINGS);
+                Status.IGNORED_FOR_SETTINGS);
 
         // General touch and keyboard touch with bypass flag not ignored.
         assertVibrationNotIgnoredForUsage(USAGE_TOUCH);
@@ -629,7 +631,7 @@
         setUserSetting(Settings.System.KEYBOARD_VIBRATION_ENABLED, 1 /* ON */);
 
         // General touch ignored.
-        assertVibrationIgnoredForUsage(USAGE_TOUCH, Vibration.Status.IGNORED_FOR_SETTINGS);
+        assertVibrationIgnoredForUsage(USAGE_TOUCH, Status.IGNORED_FOR_SETTINGS);
 
         // Keyboard touch not ignored.
         assertVibrationNotIgnoredForAttributes(
@@ -645,14 +647,14 @@
         setUserSetting(Settings.System.KEYBOARD_VIBRATION_ENABLED, 1 /* ON */);
 
         // General touch ignored.
-        assertVibrationIgnoredForUsage(USAGE_TOUCH, Vibration.Status.IGNORED_FOR_SETTINGS);
+        assertVibrationIgnoredForUsage(USAGE_TOUCH, Status.IGNORED_FOR_SETTINGS);
 
         // Keyboard touch ignored.
         assertVibrationIgnoredForAttributes(
                 new VibrationAttributes.Builder()
                         .setUsage(USAGE_IME_FEEDBACK)
                         .build(),
-                Vibration.Status.IGNORED_FOR_SETTINGS);
+                Status.IGNORED_FOR_SETTINGS);
     }
 
     @Test
@@ -668,7 +670,7 @@
         // Ignore the vibration when the coming device id represents a virtual device.
         for (int usage : ALL_USAGES) {
             assertVibrationIgnoredForUsageAndDevice(usage, VIRTUAL_DEVICE_ID,
-                    Vibration.Status.IGNORED_FROM_VIRTUAL_DEVICE);
+                    Status.IGNORED_FROM_VIRTUAL_DEVICE);
         }
     }
 
@@ -911,22 +913,21 @@
     }
 
     private void assertVibrationIgnoredForUsage(@VibrationAttributes.Usage int usage,
-            Vibration.Status expectedStatus) {
+            Status expectedStatus) {
         assertVibrationIgnoredForUsageAndDevice(usage, Context.DEVICE_ID_DEFAULT, expectedStatus);
     }
 
     private void assertVibrationIgnoredForUsageAndDevice(@VibrationAttributes.Usage int usage,
-            int deviceId, Vibration.Status expectedStatus) {
-        Vibration.CallerInfo callerInfo = new Vibration.CallerInfo(
+            int deviceId, Status expectedStatus) {
+        CallerInfo callerInfo = new CallerInfo(
                 VibrationAttributes.createForUsage(usage), UID, deviceId, null, null);
         assertEquals(errorMessageForUsage(usage), expectedStatus,
                 mVibrationSettings.shouldIgnoreVibration(callerInfo));
     }
 
     private void assertVibrationIgnoredForAttributes(VibrationAttributes attrs,
-            Vibration.Status expectedStatus) {
-        Vibration.CallerInfo callerInfo = new Vibration.CallerInfo(attrs, UID,
-                Context.DEVICE_ID_DEFAULT, null, null);
+            Status expectedStatus) {
+        CallerInfo callerInfo = new CallerInfo(attrs, UID, Context.DEVICE_ID_DEFAULT, null, null);
         assertEquals(errorMessageForAttributes(attrs), expectedStatus,
                 mVibrationSettings.shouldIgnoreVibration(callerInfo));
     }
@@ -948,7 +949,7 @@
     private void assertVibrationNotIgnoredForUsageAndFlagsAndDevice(
             @VibrationAttributes.Usage int usage, int deviceId,
             @VibrationAttributes.Flag int flags) {
-        Vibration.CallerInfo callerInfo = new Vibration.CallerInfo(
+        CallerInfo callerInfo = new CallerInfo(
                 new VibrationAttributes.Builder().setUsage(usage).setFlags(flags).build(), UID,
                 deviceId, null, null);
         assertNull(errorMessageForUsage(usage),
@@ -956,7 +957,7 @@
     }
 
     private void assertVibrationNotIgnoredForAttributes(VibrationAttributes attrs) {
-        Vibration.CallerInfo callerInfo = new Vibration.CallerInfo(attrs, UID,
+        CallerInfo callerInfo = new CallerInfo(attrs, UID,
                 Context.DEVICE_ID_DEFAULT, null, null);
         assertNull(errorMessageForAttributes(attrs),
                 mVibrationSettings.shouldIgnoreVibration(callerInfo));
@@ -1017,10 +1018,10 @@
                 new PowerManager.SleepData(sleepTime, reason));
     }
 
-    private Vibration.CallerInfo createCallerInfo(int uid, String opPkg,
+    private CallerInfo createCallerInfo(int uid, String opPkg,
             @VibrationAttributes.Usage int usage) {
         VibrationAttributes attrs = VibrationAttributes.createForUsage(usage);
-        return new Vibration.CallerInfo(attrs, uid, VIRTUAL_DEVICE_ID, opPkg, null);
+        return new CallerInfo(attrs, uid, VIRTUAL_DEVICE_ID, opPkg, null);
     }
 
     private void setBatteryReceiverRegistrationResult(Intent result) {
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
index bfdaa78..31cc50f1 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibrationThreadTest.java
@@ -78,6 +78,8 @@
 import com.android.internal.util.test.FakeSettingsProvider;
 import com.android.internal.util.test.FakeSettingsProviderRule;
 import com.android.server.LocalServices;
+import com.android.server.vibrator.VibrationSession.CallerInfo;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import org.junit.After;
 import org.junit.Before;
@@ -189,7 +191,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.IGNORED_UNSUPPORTED);
+        verifyCallbacksTriggered(vibrationId, Status.IGNORED_UNSUPPORTED);
     }
 
     @Test
@@ -202,7 +204,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks, never()).onComplete(anyInt(), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.IGNORED_UNSUPPORTED);
+        verifyCallbacksTriggered(vibrationId, Status.IGNORED_UNSUPPORTED);
     }
 
     @Test
@@ -216,7 +218,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedOneShot(10)),
@@ -233,7 +235,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedOneShot(10)),
@@ -253,7 +255,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(15L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedOneShot(15)),
@@ -326,13 +328,10 @@
         assertTrue(mThread.isRunningVibrationId(vibrationId));
         assertTrue(mControllers.get(VIBRATOR_ID).isVibrating());
 
-        Vibration.EndInfo cancelVibrationInfo = new Vibration.EndInfo(
-                Vibration.Status.CANCELLED_SUPERSEDED, new Vibration.CallerInfo(
-                VibrationAttributes.createForUsage(VibrationAttributes.USAGE_ALARM), /* uid= */
-                1, /* deviceId= */ -1, /* opPkg= */ null, /* reason= */ null));
-        mVibrationConductor.notifyCancelled(
-                cancelVibrationInfo,
-                /* immediate= */ false);
+        Vibration.EndInfo cancelVibrationInfo = new Vibration.EndInfo(Status.CANCELLED_SUPERSEDED,
+                new CallerInfo(VibrationAttributes.createForUsage(VibrationAttributes.USAGE_ALARM),
+                        /* uid= */ 1, /* deviceId= */ -1, /* opPkg= */ null, /* reason= */ null));
+        mVibrationConductor.notifyCancelled(cancelVibrationInfo, /* immediate= */ false);
         waitForCompletion();
         assertFalse(mThread.isRunningVibrationId(vibrationId));
 
@@ -363,11 +362,10 @@
 
         assertTrue(waitUntil(() -> !fakeVibrator.getAmplitudes().isEmpty(), TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(expectedOneShot(5000)),
                 fakeVibrator.getEffectSegments(vibrationId));
@@ -385,7 +383,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(300L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId))
@@ -406,7 +404,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(350L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId))
@@ -433,11 +431,10 @@
         assertTrue(waitUntil(() -> fakeVibrator.getEffectSegments(vibrationId).size() >= 5,
                 5000L + TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId).subList(0, 5))
@@ -466,12 +463,11 @@
         assertTrue(waitUntil(() -> !fakeVibrator.getEffectSegments(vibrationId).isEmpty(),
                 TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
         // PWLE size max was used to generate a single vibrate call with 10 segments.
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(10, fakeVibrator.getEffectSegments(vibrationId).size());
     }
@@ -496,12 +492,11 @@
         assertTrue(waitUntil(() -> !fakeVibrator.getEffectSegments(vibrationId).isEmpty(),
                 TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF), /* immediate= */ false);
         waitForCompletion();
 
         // Composition size max was used to generate a single vibrate call with 10 primitives.
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SCREEN_OFF);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(10, fakeVibrator.getEffectSegments(vibrationId).size());
     }
@@ -519,11 +514,10 @@
 
         assertTrue(waitUntil(() -> !fakeVibrator.getAmplitudes().isEmpty(), TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(expectedOneShot(5550)),
                 fakeVibrator.getEffectSegments(vibrationId));
@@ -545,11 +539,10 @@
         assertTrue(waitUntil(() -> fakeVibrator.getEffectSegments(vibrationId).size() > 1,
                 expectedOnDuration + TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         List<VibrationEffectSegment> effectSegments = fakeVibrator.getEffectSegments(vibrationId);
         // First time, turn vibrator ON for the expected fixed duration.
@@ -584,14 +577,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread =
                 new Thread(() -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SETTINGS_UPDATE),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SETTINGS_UPDATE);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
     }
 
@@ -614,14 +607,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread =
                 new Thread(() -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SETTINGS_UPDATE),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SETTINGS_UPDATE);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
     }
 
@@ -641,14 +634,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread =
                 new Thread(() -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SCREEN_OFF);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
     }
 
@@ -663,7 +656,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_THUD)),
@@ -684,7 +677,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedOneShot(10)),
@@ -701,7 +694,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(0L));
         verify(mManagerHooks, never()).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks, never()).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.IGNORED_UNSUPPORTED);
+        verifyCallbacksTriggered(vibrationId, Status.IGNORED_UNSUPPORTED);
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId).isEmpty());
     }
 
@@ -718,7 +711,7 @@
                 eq(PerformVendorEffectVibratorStep.VENDOR_EFFECT_MAX_DURATION_MS));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertThat(mControllers.get(VIBRATOR_ID).isVibrating()).isFalse();
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getVendorEffects(vibrationId))
@@ -743,7 +736,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(40L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
                 expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0),
@@ -762,7 +755,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(0L));
         verify(mManagerHooks, never()).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks, never()).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.IGNORED_UNSUPPORTED);
+        verifyCallbacksTriggered(vibrationId, Status.IGNORED_UNSUPPORTED);
         assertTrue(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId).isEmpty());
     }
 
@@ -784,7 +777,7 @@
         long vibrationId = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         // Vibrator compose called twice.
         verify(mControllerCallbacks, times(2)).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
         assertEquals(3, fakeVibrator.getEffectSegments(vibrationId).size());
@@ -824,7 +817,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(10L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks, times(5)).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
                 expectedOneShot(10),
@@ -865,7 +858,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), anyLong());
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks, times(4)).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         List<VibrationEffectSegment> segments =
@@ -904,7 +897,7 @@
         verify(mManagerHooks).noteVibratorOn(eq(UID), eq(100L));
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
         assertEquals(Arrays.asList(
                 expectedRamp(/* amplitude= */ 1, /* frequencyHz= */ 150, /* duration= */ 10),
@@ -942,7 +935,7 @@
         long vibrationId = startThreadAndDispatcher(effect);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         // Vibrator compose called 3 times with 2 segments instead of 2 times with 3 segments.
         // Using best split points instead of max-packing PWLEs.
@@ -967,7 +960,7 @@
         waitForCompletion();
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BINDER_DIED);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BINDER_DIED);
     }
 
     @Test
@@ -977,7 +970,7 @@
         long vibrationId = startThreadAndDispatcher(VibrationEffect.createOneShot(10, 100));
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         verify(mManagerHooks, never()).prepareSyncedVibration(anyLong(), any());
         verify(mManagerHooks, never()).triggerSyncedVibration(anyLong());
         verify(mManagerHooks, never()).cancelSyncedVibration();
@@ -998,7 +991,7 @@
         verify(mManagerHooks).noteVibratorOff(eq(UID));
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
         verify(mControllerCallbacks, never()).onComplete(eq(2), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_TICK)),
@@ -1022,7 +1015,7 @@
         verify(mControllerCallbacks).onComplete(eq(1), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(2), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(3), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
         assertFalse(mControllers.get(3).isVibrating());
@@ -1065,7 +1058,7 @@
         verify(mControllerCallbacks).onComplete(eq(2), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(3), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(4), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
         assertFalse(mControllers.get(3).isVibrating());
@@ -1117,7 +1110,7 @@
         batteryVerifier.verify(mManagerHooks).noteVibratorOn(eq(UID), eq(20L));
         batteryVerifier.verify(mManagerHooks).noteVibratorOff(eq(UID));
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
         assertFalse(mControllers.get(3).isVibrating());
@@ -1166,7 +1159,7 @@
         verify(mManagerHooks).prepareSyncedVibration(eq(expectedCap), eq(vibratorIds));
         verify(mManagerHooks).triggerSyncedVibration(eq(vibrationId));
         verify(mManagerHooks, never()).cancelSyncedVibration();
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         VibrationEffectSegment expected = expectedPrimitive(
                 VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 100);
@@ -1213,7 +1206,7 @@
         verify(mManagerHooks).prepareSyncedVibration(eq(expectedCap), eq(vibratorIds));
         verify(mManagerHooks).triggerSyncedVibration(eq(vibrationId));
         verify(mManagerHooks, never()).cancelSyncedVibration();
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
     }
 
     @Test
@@ -1305,7 +1298,7 @@
         verify(mControllerCallbacks).onComplete(eq(1), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(2), eq(vibrationId));
         verify(mControllerCallbacks).onComplete(eq(3), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
         assertFalse(mControllers.get(3).isVibrating());
@@ -1478,8 +1471,7 @@
         // fail at waitForCompletion(cancellingThread).
         Thread cancellingThread = new Thread(
                 () -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                        /* immediate= */ false));
+                        new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false));
         cancellingThread.start();
 
         // Cancelling the vibration should be fast and return right away, even if the thread is
@@ -1488,7 +1480,7 @@
 
         // After the vibrator call ends the vibration is cancelled and the vibrator is turned off.
         waitForCompletion(/* timeout= */ latency + TEST_TIMEOUT_MILLIS);
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
     }
 
@@ -1517,14 +1509,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread = new Thread(
                 () -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SCREEN_OFF);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
     }
@@ -1551,14 +1543,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread = new Thread(
                 () -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SCREEN_OFF);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
     }
@@ -1585,15 +1577,14 @@
         // fail at waitForCompletion(vibrationThread) if the vibration not cancelled immediately.
         Thread cancellingThread = new Thread(
                 () -> mVibrationConductor.notifyCancelled(
-                        new Vibration.EndInfo(
-                                Vibration.Status.CANCELLED_BY_SCREEN_OFF),
+                        new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF),
                         /* immediate= */ false));
         cancellingThread.start();
 
         waitForCompletion(/* timeout= */ 50);
         cancellingThread.join();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_SCREEN_OFF);
         assertFalse(mControllers.get(1).isVibrating());
         assertFalse(mControllers.get(2).isVibrating());
     }
@@ -1612,7 +1603,7 @@
 
         verify(mVibrationToken).linkToDeath(same(mVibrationConductor), eq(0));
         verify(mVibrationToken).unlinkToDeath(same(mVibrationConductor), eq(0));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BINDER_DIED);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BINDER_DIED);
         assertFalse(mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId).isEmpty());
         assertFalse(mControllers.get(VIBRATOR_ID).isVibrating());
     }
@@ -1628,7 +1619,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         // Duration extended for 5 + 5 + 5 + 15.
         assertEquals(Arrays.asList(expectedOneShot(30)),
@@ -1651,7 +1642,7 @@
 
         // Vibration completed but vibrator not yet released.
         verify(mManagerHooks, timeout(TEST_TIMEOUT_MILLIS)).onVibrationCompleted(eq(vibrationId),
-                eq(new Vibration.EndInfo(Vibration.Status.FINISHED)));
+                eq(new Vibration.EndInfo(Status.FINISHED)));
         verify(mManagerHooks, never()).onVibrationThreadReleased(anyLong());
 
         // Thread still running ramp down.
@@ -1663,13 +1654,12 @@
 
         // Will stop the ramp down right away.
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE),
-                /* immediate= */ true);
+                new Vibration.EndInfo(Status.CANCELLED_BY_SETTINGS_UPDATE), /* immediate= */ true);
         waitForCompletion();
 
         // Does not cancel already finished vibration, but releases vibrator.
         verify(mManagerHooks, never()).onVibrationCompleted(eq(vibrationId),
-                eq(new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SETTINGS_UPDATE)));
+                eq(new Vibration.EndInfo(Status.CANCELLED_BY_SETTINGS_UPDATE)));
         verify(mManagerHooks).onVibrationThreadReleased(vibrationId);
     }
 
@@ -1684,11 +1674,10 @@
         assertTrue(waitUntil(() -> mControllers.get(VIBRATOR_ID).isVibrating(),
                 TEST_TIMEOUT_MILLIS));
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId, Status.CANCELLED_BY_USER);
 
         // Duration extended for 10000 + 15.
         assertEquals(Arrays.asList(expectedOneShot(10_015)),
@@ -1711,7 +1700,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
                 mVibratorProviders.get(VIBRATOR_ID).getEffectSegments(vibrationId));
@@ -1729,7 +1718,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         assertThat(mVibratorProviders.get(VIBRATOR_ID).getVendorEffects(vibrationId))
                 .containsExactly(effect)
@@ -1752,7 +1741,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         assertEquals(
                 Arrays.asList(expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
@@ -1779,7 +1768,7 @@
         waitForCompletion();
 
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId));
-        verifyCallbacksTriggered(vibrationId, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedRamp(0, 1, 150, 150, 1)),
                 fakeVibrator.getEffectSegments(vibrationId));
@@ -1810,14 +1799,13 @@
         long vibrationId1 = startThreadAndDispatcher(effect1);
         waitForCompletion();
         verify(mControllerCallbacks).onComplete(VIBRATOR_ID, vibrationId1);
-        verifyCallbacksTriggered(vibrationId1, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId1, Status.FINISHED);
 
         long vibrationId2 = startThreadAndDispatcher(effect2);
         // Effect2 won't complete on its own. Cancel it after a couple of repeats.
         Thread.sleep(150);  // More than two TICKs.
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_USER),
-                /* immediate= */ false);
+                new Vibration.EndInfo(Status.CANCELLED_BY_USER), /* immediate= */ false);
         waitForCompletion();
 
         long vibrationId3 = startThreadAndDispatcher(effect3);
@@ -1827,8 +1815,7 @@
         long start4 = System.currentTimeMillis();
         long vibrationId4 = startThreadAndDispatcher(effect4);
         mVibrationConductor.notifyCancelled(
-                new Vibration.EndInfo(Vibration.Status.CANCELLED_BY_SCREEN_OFF),
-                /* immediate= */ true);
+                new Vibration.EndInfo(Status.CANCELLED_BY_SCREEN_OFF), /* immediate= */ true);
         waitForCompletion();
         long duration4 = System.currentTimeMillis() - start4;
 
@@ -1841,14 +1828,14 @@
 
         // Effect1
         verify(mControllerCallbacks).onComplete(VIBRATOR_ID, vibrationId1);
-        verifyCallbacksTriggered(vibrationId1, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId1, Status.FINISHED);
 
         assertEquals(Arrays.asList(expectedPrebaked(VibrationEffect.EFFECT_CLICK)),
                 fakeVibrator.getEffectSegments(vibrationId1));
 
         // Effect2: repeating, cancelled.
         verify(mControllerCallbacks, atLeast(2)).onComplete(VIBRATOR_ID, vibrationId2);
-        verifyCallbacksTriggered(vibrationId2, Vibration.Status.CANCELLED_BY_USER);
+        verifyCallbacksTriggered(vibrationId2, Status.CANCELLED_BY_USER);
 
         // The exact count of segments might vary, so just check that there's more than 2 and
         // all elements are the same segment.
@@ -1860,13 +1847,13 @@
 
         // Effect3
         verify(mControllerCallbacks).onComplete(eq(VIBRATOR_ID), eq(vibrationId3));
-        verifyCallbacksTriggered(vibrationId3, Vibration.Status.FINISHED);
+        verifyCallbacksTriggered(vibrationId3, Status.FINISHED);
         assertEquals(Arrays.asList(
                         expectedPrimitive(VibrationEffect.Composition.PRIMITIVE_CLICK, 1, 0)),
                 fakeVibrator.getEffectSegments(vibrationId3));
 
         // Effect4: cancelled quickly.
-        verifyCallbacksTriggered(vibrationId4, Vibration.Status.CANCELLED_BY_SCREEN_OFF);
+        verifyCallbacksTriggered(vibrationId4, Status.CANCELLED_BY_SCREEN_OFF);
         assertTrue("Tested duration=" + duration4, duration4 < 2000);
 
         // Effect5: played normally after effect4, which may or may not have played.
@@ -1907,7 +1894,7 @@
                 .build();
         HalVibration vib = new HalVibration(mVibrationToken,
                 CombinedVibration.createParallel(effect),
-                new Vibration.CallerInfo(attrs, UID, DEVICE_ID, PACKAGE_NAME, "reason"));
+                new CallerInfo(attrs, UID, DEVICE_ID, PACKAGE_NAME, "reason"));
         return startThreadAndDispatcher(vib, requestVibrationParamsFuture);
     }
 
@@ -1944,7 +1931,7 @@
 
     private HalVibration createVibration(CombinedVibration effect) {
         return new HalVibration(mVibrationToken, effect,
-                new Vibration.CallerInfo(ATTRS, UID, DEVICE_ID, PACKAGE_NAME, "reason"));
+                new CallerInfo(ATTRS, UID, DEVICE_ID, PACKAGE_NAME, "reason"));
     }
 
     private SparseArray<VibratorController> createVibratorControllers() {
@@ -2007,7 +1994,7 @@
                 .collect(Collectors.toList());
     }
 
-    private void verifyCallbacksTriggered(long vibrationId, Vibration.Status expectedStatus) {
+    private void verifyCallbacksTriggered(long vibrationId, Status expectedStatus) {
         verifyCallbacksTriggered(vibrationId, new Vibration.EndInfo(expectedStatus));
     }
 
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibratorFrameworkStatsLoggerTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibratorFrameworkStatsLoggerTest.java
index 3466bbb..cd057b6 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibratorFrameworkStatsLoggerTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibratorFrameworkStatsLoggerTest.java
@@ -22,6 +22,8 @@
 import android.os.Handler;
 import android.os.test.TestLooper;
 
+import com.android.server.vibrator.VibrationSession.Status;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -113,7 +115,6 @@
     }
 
     private static VibrationStats.StatsInfo newEmptyStatsInfo() {
-        return new VibrationStats.StatsInfo(
-                0, 0, 0, Vibration.Status.FINISHED, new VibrationStats(), 0L);
+        return new VibrationStats.StatsInfo(0, 0, 0, Status.FINISHED, new VibrationStats(), 0L);
     }
 }
diff --git a/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java b/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
index 4afb562..4012575 100644
--- a/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
+++ b/services/tests/vibrator/src/com/android/server/vibrator/VibratorManagerServiceTest.java
@@ -110,6 +110,7 @@
 import com.android.server.LocalServices;
 import com.android.server.companion.virtual.VirtualDeviceManagerInternal;
 import com.android.server.pm.BackgroundUserSoundNotifier;
+import com.android.server.vibrator.VibrationSession.Status;
 
 import org.junit.After;
 import org.junit.Before;
@@ -803,8 +804,8 @@
         mockVibrators(1);
         VibratorManagerService service = createSystemReadyService();
 
-        var vib = vibrate(service,
-                VibrationEffect.createWaveform(new long[]{100, 100, 100, 100}, 0), RINGTONE_ATTRS);
+        HalVibration vib = vibrate(service, VibrationEffect.createWaveform(
+                new long[]{0, TEST_TIMEOUT_MILLIS, TEST_TIMEOUT_MILLIS}, 0), RINGTONE_ATTRS);
 
         assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
 
@@ -815,14 +816,80 @@
         service.mAppOpsChangeListener.onOpChanged(AppOpsManager.OP_VIBRATE, null);
 
         assertTrue(waitUntil(s -> vib.hasEnded(), service, TEST_TIMEOUT_MILLIS));
+        assertThat(vib.getStatus()).isEqualTo(Status.CANCELLED_BY_APP_OPS);
+    }
 
-        var statsInfoCaptor = ArgumentCaptor.forClass(VibrationStats.StatsInfo.class);
-        verify(mVibratorFrameworkStatsLoggerMock, timeout(TEST_TIMEOUT_MILLIS))
-                .writeVibrationReportedAsync(statsInfoCaptor.capture());
+    @Test
+    public void vibrate_thenPowerModeChanges_getsCancelled() throws Exception {
+        mockVibrators(1, 2);
+        VibratorManagerService service = createSystemReadyService();
 
-        VibrationStats.StatsInfo touchMetrics = statsInfoCaptor.getAllValues().get(0);
-        assertEquals(Vibration.Status.CANCELLED_BY_APP_OPS.getProtoEnumValue(),
-                touchMetrics.status);
+        HalVibration vib = vibrate(service,
+                CombinedVibration.startParallel()
+                        .addVibrator(1, VibrationEffect.createOneShot(2 * TEST_TIMEOUT_MILLIS, 100))
+                        .combine(),
+                HAPTIC_FEEDBACK_ATTRS);
+
+        // VibrationThread will start this vibration async, so wait until vibration is triggered.
+        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
+
+        mRegisteredPowerModeListener.onLowPowerModeChanged(LOW_POWER_STATE);
+
+        assertTrue(waitUntil(s -> vib.hasEnded(), service, TEST_TIMEOUT_MILLIS));
+        assertThat(vib.getStatus()).isEqualTo(Status.CANCELLED_BY_SETTINGS_UPDATE);
+    }
+
+    @Test
+    public void vibrate_thenSettingsRefreshedWithoutChange_doNotCancelVibration() throws Exception {
+        mockVibrators(1);
+        VibratorManagerService service = createSystemReadyService();
+
+        vibrate(service, VibrationEffect.createOneShot(2 * TEST_TIMEOUT_MILLIS, 100),
+                HAPTIC_FEEDBACK_ATTRS);
+
+        // VibrationThread will start this vibration async, so wait until vibration is triggered.
+        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
+
+        service.updateServiceState();
+
+        // Vibration is not stopped nearly after updating service.
+        assertFalse(waitUntil(s -> !s.isVibrating(1), service, 50));
+    }
+
+    @Test
+    public void vibrate_thenSettingsChange_getsCancelled() throws Exception {
+        mockVibrators(1);
+        VibratorManagerService service = createSystemReadyService();
+
+        HalVibration vib = vibrate(service,
+                VibrationEffect.createOneShot(2 * TEST_TIMEOUT_MILLIS, 100),
+                HAPTIC_FEEDBACK_ATTRS);
+
+        // VibrationThread will start this vibration async, so wait until vibration is triggered.
+        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
+
+        setUserSetting(Settings.System.HAPTIC_FEEDBACK_INTENSITY, Vibrator.VIBRATION_INTENSITY_OFF);
+        service.mVibrationSettings.mSettingObserver.onChange(false);
+        service.updateServiceState();
+
+        assertTrue(waitUntil(s -> vib.hasEnded(), service, TEST_TIMEOUT_MILLIS));
+        assertThat(vib.getStatus()).isEqualTo(Status.CANCELLED_BY_SETTINGS_UPDATE);
+    }
+
+    @Test
+    public void vibrate_thenScreenTurnsOff_getsCancelled() throws Throwable {
+        mockVibrators(1);
+        VibratorManagerService service = createSystemReadyService();
+
+        HalVibration vib = vibrate(service, VibrationEffect.createWaveform(
+                new long[]{0, TEST_TIMEOUT_MILLIS, TEST_TIMEOUT_MILLIS}, 0), ALARM_ATTRS);
+
+        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
+
+        service.mIntentReceiver.onReceive(mContextSpy, new Intent(Intent.ACTION_SCREEN_OFF));
+
+        assertTrue(waitUntil(s -> vib.hasEnded(), service, TEST_TIMEOUT_MILLIS));
+        assertThat(vib.getStatus()).isEqualTo(Status.CANCELLED_BY_SCREEN_OFF);
     }
 
     @Test
@@ -831,8 +898,8 @@
         mockVibrators(1);
         VibratorManagerService service = createSystemReadyService();
 
-        var vib = vibrate(service,
-                VibrationEffect.createWaveform(new long[]{100, 100, 100, 100}, 0), ALARM_ATTRS);
+        HalVibration vib = vibrate(service, VibrationEffect.createWaveform(
+                new long[]{0, TEST_TIMEOUT_MILLIS, TEST_TIMEOUT_MILLIS}, 0), ALARM_ATTRS);
 
         assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
 
@@ -841,14 +908,7 @@
                 BackgroundUserSoundNotifier.ACTION_MUTE_SOUND));
 
         assertTrue(waitUntil(s -> vib.hasEnded(), service, TEST_TIMEOUT_MILLIS));
-
-        var statsInfoCaptor = ArgumentCaptor.forClass(VibrationStats.StatsInfo.class);
-        verify(mVibratorFrameworkStatsLoggerMock, timeout(TEST_TIMEOUT_MILLIS))
-                .writeVibrationReportedAsync(statsInfoCaptor.capture());
-
-        VibrationStats.StatsInfo touchMetrics = statsInfoCaptor.getAllValues().get(0);
-        assertEquals(Vibration.Status.CANCELLED_BY_FOREGROUND_USER.getProtoEnumValue(),
-                touchMetrics.status);
+        assertThat(vib.getStatus()).isEqualTo(Status.CANCELLED_BY_FOREGROUND_USER);
     }
 
     @Test
@@ -1314,7 +1374,7 @@
     }
 
     @Test
-    public void vibrate_withriggerCallback_finishesVibration() throws Exception {
+    public void vibrate_withTriggerCallback_finishesVibration() throws Exception {
         mockCapabilities(IVibratorManager.CAP_SYNC, IVibratorManager.CAP_PREPARE_COMPOSE);
         mockVibrators(1, 2);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_COMPOSE_EFFECTS);
@@ -1947,41 +2007,6 @@
     }
 
     @Test
-    public void vibrate_withPowerModeChange_cancelVibrationIfNotAllowed() throws Exception {
-        mockVibrators(1, 2);
-        VibratorManagerService service = createSystemReadyService();
-        vibrate(service,
-                CombinedVibration.startParallel()
-                        .addVibrator(1, VibrationEffect.createOneShot(1000, 100))
-                        .combine(),
-                HAPTIC_FEEDBACK_ATTRS);
-
-        // VibrationThread will start this vibration async, so wait until vibration is triggered.
-        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
-
-        mRegisteredPowerModeListener.onLowPowerModeChanged(LOW_POWER_STATE);
-
-        // Haptic feedback cancelled on low power mode.
-        assertTrue(waitUntil(s -> !s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
-    }
-
-    @Test
-    public void vibrate_withSettingsChange_doNotCancelVibration() throws Exception {
-        mockVibrators(1);
-        VibratorManagerService service = createSystemReadyService();
-
-        vibrate(service, VibrationEffect.createOneShot(1000, 100), HAPTIC_FEEDBACK_ATTRS);
-
-        // VibrationThread will start this vibration async, so wait until vibration is triggered.
-        assertTrue(waitUntil(s -> s.isVibrating(1), service, TEST_TIMEOUT_MILLIS));
-
-        service.updateServiceState();
-
-        // Vibration is not stopped nearly after updating service.
-        assertFalse(waitUntil(s -> !s.isVibrating(1), service, 50));
-    }
-
-    @Test
     public void vibrate_ignoreVibrationFromVirtualDevice() throws Exception {
         mockVibrators(1);
         VibratorManagerService service = createSystemReadyService();
@@ -2475,6 +2500,113 @@
     }
 
     @Test
+    public void onExternalVibration_thenDeniedAppOps_doNotCancelVibration() throws Throwable {
+        mockVibrators(1);
+        mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
+        VibratorManagerService service = createSystemReadyService();
+
+        IExternalVibrationController externalVibrationControllerMock =
+                mock(IExternalVibrationController.class);
+        ExternalVibration externalVibration = new ExternalVibration(UID, PACKAGE_NAME,
+                AUDIO_ALARM_ATTRS, externalVibrationControllerMock, mock(IBinder.class));
+        ExternalVibrationScale scale = mExternalVibratorService.onExternalVibrationStart(
+                externalVibration);
+
+        assertThat(scale.scaleLevel).isNotEqualTo(ExternalVibrationScale.ScaleLevel.SCALE_MUTE);
+
+        when(mAppOpsManagerMock.checkAudioOpNoThrow(eq(AppOpsManager.OP_VIBRATE),
+                eq(AudioAttributes.USAGE_ALARM), anyInt(), anyString()))
+                .thenReturn(AppOpsManager.MODE_IGNORED);
+        service.mAppOpsChangeListener.onOpChanged(AppOpsManager.OP_VIBRATE, null);
+
+        verify(externalVibrationControllerMock, never()).mute();
+    }
+
+    @Test
+    public void onExternalVibration_thenPowerModeChanges_doNotCancelVibration() throws Exception {
+        mockVibrators(1);
+        mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
+        createSystemReadyService();
+
+        IExternalVibrationController externalVibrationControllerMock =
+                mock(IExternalVibrationController.class);
+        ExternalVibration externalVibration = new ExternalVibration(UID, PACKAGE_NAME,
+                AUDIO_ALARM_ATTRS, externalVibrationControllerMock, mock(IBinder.class));
+        ExternalVibrationScale scale = mExternalVibratorService.onExternalVibrationStart(
+                externalVibration);
+
+        assertThat(scale.scaleLevel).isNotEqualTo(ExternalVibrationScale.ScaleLevel.SCALE_MUTE);
+
+        mRegisteredPowerModeListener.onLowPowerModeChanged(LOW_POWER_STATE);
+
+        verify(externalVibrationControllerMock, never()).mute();
+    }
+
+    @Test
+    public void onExternalVibration_thenSettingsChange_doNotCancelVibration() throws Exception {
+        mockVibrators(1);
+        mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
+        VibratorManagerService service = createSystemReadyService();
+
+        IExternalVibrationController externalVibrationControllerMock =
+                mock(IExternalVibrationController.class);
+        ExternalVibration externalVibration = new ExternalVibration(UID, PACKAGE_NAME,
+                AUDIO_ALARM_ATTRS, externalVibrationControllerMock, mock(IBinder.class));
+        ExternalVibrationScale scale = mExternalVibratorService.onExternalVibrationStart(
+                externalVibration);
+
+        assertThat(scale.scaleLevel).isNotEqualTo(ExternalVibrationScale.ScaleLevel.SCALE_MUTE);
+
+        setUserSetting(Settings.System.ALARM_VIBRATION_INTENSITY, Vibrator.VIBRATION_INTENSITY_OFF);
+        service.mVibrationSettings.mSettingObserver.onChange(false);
+        service.updateServiceState();
+
+        verify(externalVibrationControllerMock, never()).mute();
+    }
+
+    @Test
+    public void onExternalVibration_thenScreenTurnsOff_doNotCancelVibration() throws Throwable {
+        mockVibrators(1);
+        mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
+        VibratorManagerService service = createSystemReadyService();
+
+        IExternalVibrationController externalVibrationControllerMock =
+                mock(IExternalVibrationController.class);
+        ExternalVibration externalVibration = new ExternalVibration(UID, PACKAGE_NAME,
+                AUDIO_ALARM_ATTRS, externalVibrationControllerMock, mock(IBinder.class));
+        ExternalVibrationScale scale = mExternalVibratorService.onExternalVibrationStart(
+                externalVibration);
+
+        assertThat(scale.scaleLevel).isNotEqualTo(ExternalVibrationScale.ScaleLevel.SCALE_MUTE);
+
+        service.mIntentReceiver.onReceive(mContextSpy, new Intent(Intent.ACTION_SCREEN_OFF));
+
+        verify(externalVibrationControllerMock, never()).mute();
+    }
+
+    @Test
+    @RequiresFlagsEnabled(android.multiuser.Flags.FLAG_ADD_UI_FOR_SOUNDS_FROM_BACKGROUND_USERS)
+    public void onExternalVibration_thenFgUserRequestsMute_doNotCancelVibration() throws Throwable {
+        mockVibrators(1);
+        mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
+        VibratorManagerService service = createSystemReadyService();
+
+        IExternalVibrationController externalVibrationControllerMock =
+                mock(IExternalVibrationController.class);
+        ExternalVibration externalVibration = new ExternalVibration(UID, PACKAGE_NAME,
+                AUDIO_ALARM_ATTRS, externalVibrationControllerMock, mock(IBinder.class));
+        ExternalVibrationScale scale = mExternalVibratorService.onExternalVibrationStart(
+                externalVibration);
+
+        assertThat(scale.scaleLevel).isNotEqualTo(ExternalVibrationScale.ScaleLevel.SCALE_MUTE);
+
+        service.mIntentReceiver.onReceive(mContextSpy, new Intent(
+                BackgroundUserSoundNotifier.ACTION_MUTE_SOUND));
+
+        verify(externalVibrationControllerMock, never()).mute();
+    }
+
+    @Test
     public void frameworkStats_externalVibration_reportsAllMetrics() throws Exception {
         mockVibrators(1);
         mVibratorProviders.get(1).setCapabilities(IVibrator.CAP_EXTERNAL_CONTROL);
@@ -2501,7 +2633,7 @@
         assertEquals(FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__EXTERNAL,
                 statsInfo.vibrationType);
         assertEquals(VibrationAttributes.USAGE_ALARM, statsInfo.usage);
-        assertEquals(Vibration.Status.FINISHED.getProtoEnumValue(), statsInfo.status);
+        assertEquals(Status.FINISHED.getProtoEnumValue(), statsInfo.status);
         assertTrue(statsInfo.totalDurationMillis > 0);
         assertTrue(
                 "Expected vibrator ON for at least 10ms, got " + statsInfo.vibratorOnMillis + "ms",
@@ -2533,7 +2665,7 @@
         assertEquals(FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__SINGLE,
                 metrics.vibrationType);
         assertEquals(VibrationAttributes.USAGE_RINGTONE, metrics.usage);
-        assertEquals(Vibration.Status.FINISHED.getProtoEnumValue(), metrics.status);
+        assertEquals(Status.FINISHED.getProtoEnumValue(), metrics.status);
         assertTrue("Total duration was too low, " + metrics.totalDurationMillis + "ms",
                 metrics.totalDurationMillis >= 20);
         assertTrue("Vibrator ON duration was too low, " + metrics.vibratorOnMillis + "ms",
@@ -2586,7 +2718,7 @@
         assertEquals(FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__REPEATED,
                 metrics.vibrationType);
         assertEquals(VibrationAttributes.USAGE_RINGTONE, metrics.usage);
-        assertEquals(Vibration.Status.CANCELLED_BY_USER.getProtoEnumValue(), metrics.status);
+        assertEquals(Status.CANCELLED_BY_USER.getProtoEnumValue(), metrics.status);
         assertTrue("Total duration was too low, " + metrics.totalDurationMillis + "ms",
                 metrics.totalDurationMillis >= 100);
         assertTrue("Vibrator ON duration was too low, " + metrics.vibratorOnMillis + "ms",
@@ -2647,7 +2779,7 @@
         assertEquals(FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__SINGLE,
                 metrics.vibrationType);
         assertEquals(VibrationAttributes.USAGE_ALARM, metrics.usage);
-        assertEquals(Vibration.Status.FINISHED.getProtoEnumValue(), metrics.status);
+        assertEquals(Status.FINISHED.getProtoEnumValue(), metrics.status);
 
         // At least 4 effect/primitive played, 20ms each, plus configured fallback.
         assertTrue("Total duration was too low, " + metrics.totalDurationMillis + "ms",
@@ -2703,7 +2835,7 @@
         VibrationStats.StatsInfo touchMetrics = argumentCaptor.getAllValues().get(0);
         assertEquals(UID, touchMetrics.uid);
         assertEquals(VibrationAttributes.USAGE_TOUCH, touchMetrics.usage);
-        assertEquals(Vibration.Status.CANCELLED_SUPERSEDED.getProtoEnumValue(),
+        assertEquals(Status.CANCELLED_SUPERSEDED.getProtoEnumValue(),
                 touchMetrics.status);
         assertTrue(touchMetrics.endedBySameUid);
         assertEquals(VibrationAttributes.USAGE_ALARM, touchMetrics.endedByUsage);
@@ -2712,7 +2844,7 @@
         VibrationStats.StatsInfo alarmMetrics = argumentCaptor.getAllValues().get(1);
         assertEquals(UID, alarmMetrics.uid);
         assertEquals(VibrationAttributes.USAGE_ALARM, alarmMetrics.usage);
-        assertEquals(Vibration.Status.FINISHED.getProtoEnumValue(), alarmMetrics.status);
+        assertEquals(Status.FINISHED.getProtoEnumValue(), alarmMetrics.status);
         assertFalse(alarmMetrics.endedBySameUid);
         assertEquals(-1, alarmMetrics.endedByUsage);
         assertEquals(VibrationAttributes.USAGE_TOUCH, alarmMetrics.interruptedUsage);
@@ -2742,12 +2874,12 @@
         VibrationStats.StatsInfo touchMetrics = argumentCaptor.getAllValues().get(0);
         assertEquals(UID, touchMetrics.uid);
         assertEquals(VibrationAttributes.USAGE_TOUCH, touchMetrics.usage);
-        assertEquals(Vibration.Status.IGNORED_FOR_POWER.getProtoEnumValue(), touchMetrics.status);
+        assertEquals(Status.IGNORED_FOR_POWER.getProtoEnumValue(), touchMetrics.status);
 
         VibrationStats.StatsInfo ringtoneMetrics = argumentCaptor.getAllValues().get(1);
         assertEquals(UID, ringtoneMetrics.uid);
         assertEquals(VibrationAttributes.USAGE_RINGTONE, ringtoneMetrics.usage);
-        assertEquals(Vibration.Status.IGNORED_FOR_SETTINGS.getProtoEnumValue(),
+        assertEquals(Status.IGNORED_FOR_SETTINGS.getProtoEnumValue(),
                 ringtoneMetrics.status);
 
         for (VibrationStats.StatsInfo metrics : argumentCaptor.getAllValues()) {
@@ -2814,7 +2946,7 @@
         assertEquals(FrameworkStatsLog.VIBRATION_REPORTED__VIBRATION_TYPE__SINGLE,
                 metrics.vibrationType);
         assertEquals(VibrationAttributes.USAGE_NOTIFICATION, metrics.usage);
-        assertEquals(Vibration.Status.FINISHED.getProtoEnumValue(), metrics.status);
+        assertEquals(Status.FINISHED.getProtoEnumValue(), metrics.status);
         assertTrue(metrics.totalDurationMillis >= 20);
 
         // vibratorOnMillis accumulates both vibrators, it's 20 for each constant.
diff --git a/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutManagerTests.java b/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutManagerTests.java
index d147325..0575d98 100644
--- a/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutManagerTests.java
+++ b/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutManagerTests.java
@@ -41,6 +41,8 @@
 import android.os.Handler;
 import android.os.Looper;
 import android.os.UserHandle;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.view.KeyEvent;
 import android.view.KeyboardShortcutGroup;
 import android.view.KeyboardShortcutInfo;
@@ -50,6 +52,8 @@
 import com.android.internal.R;
 
 import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
 import org.junit.Test;
 
 import java.util.Collections;
@@ -62,7 +66,13 @@
  */
 
 @SmallTest
+@EnableFlags(com.android.hardware.input.Flags.FLAG_MODIFIER_SHORTCUT_MANAGER_REFACTOR)
 public class ModifierShortcutManagerTests {
+
+    @ClassRule public static final SetFlagsRule.ClassRule SET_FLAGS_CLASS_RULE =
+            new SetFlagsRule.ClassRule();
+    @Rule public final SetFlagsRule mSetFlagsRule = SET_FLAGS_CLASS_RULE.createSetFlagsRule();
+
     private ModifierShortcutManager mModifierShortcutManager;
     private Handler mHandler;
     private Context mContext;
diff --git a/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutTests.java b/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutTests.java
index 71f90a2..43171f8 100644
--- a/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutTests.java
+++ b/services/tests/wmtests/src/com/android/server/policy/ModifierShortcutTests.java
@@ -44,26 +44,21 @@
 import android.content.ComponentName;
 import android.content.Intent;
 import android.os.RemoteException;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
-import android.platform.test.annotations.RequiresFlagsDisabled;
-import android.platform.test.annotations.RequiresFlagsEnabled;
-import android.platform.test.flag.junit.CheckFlagsRule;
-import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.util.SparseArray;
 
 import androidx.test.filters.SmallTest;
 
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 
 @Presubmit
 @SmallTest
+@EnableFlags(com.android.hardware.input.Flags.FLAG_MODIFIER_SHORTCUT_MANAGER_REFACTOR)
 public class ModifierShortcutTests extends ShortcutKeyTestBase {
 
-    @Rule
-    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
-
     private static final SparseArray<String> INTENT_SHORTCUTS =  new SparseArray<>();
     private static final SparseArray<String> ROLE_SHORTCUTS =  new SparseArray<>();
     static {
@@ -258,7 +253,7 @@
      * META+CTRL+BACKSPACE for taking a bugreport when the flag is enabled.
      */
     @Test
-    @RequiresFlagsEnabled(com.android.server.flags.Flags.FLAG_NEW_BUGREPORT_KEYBOARD_SHORTCUT)
+    @EnableFlags(com.android.server.flags.Flags.FLAG_NEW_BUGREPORT_KEYBOARD_SHORTCUT)
     public void testTakeBugReport_flagEnabled() throws RemoteException {
         sendKeyCombination(new int[]{KEYCODE_META_LEFT, KEYCODE_CTRL_LEFT, KEYCODE_DEL}, 0);
         mPhoneWindowManager.assertTakeBugreport(true);
@@ -268,7 +263,7 @@
      * META+CTRL+BACKSPACE for taking a bugreport does nothing when the flag is disabledd.
      */
     @Test
-    @RequiresFlagsDisabled(com.android.server.flags.Flags.FLAG_NEW_BUGREPORT_KEYBOARD_SHORTCUT)
+    @DisableFlags(com.android.server.flags.Flags.FLAG_NEW_BUGREPORT_KEYBOARD_SHORTCUT)
     public void testTakeBugReport_flagDisabled() throws RemoteException {
         sendKeyCombination(new int[]{KEYCODE_META_LEFT, KEYCODE_CTRL_LEFT, KEYCODE_DEL}, 0);
         mPhoneWindowManager.assertTakeBugreport(false);
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
index af4394a..0c1fbf3 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppTransitionControllerTest.java
@@ -25,11 +25,15 @@
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.view.WindowManager.TRANSIT_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_CLOSE;
 import static android.view.WindowManager.TRANSIT_OLD_ACTIVITY_OPEN;
 import static android.view.WindowManager.TRANSIT_OLD_DREAM_ACTIVITY_CLOSE;
 import static android.view.WindowManager.TRANSIT_OLD_DREAM_ACTIVITY_OPEN;
 import static android.view.WindowManager.TRANSIT_OLD_KEYGUARD_UNOCCLUDE;
 import static android.view.WindowManager.TRANSIT_OLD_TASK_CHANGE_WINDOWING_MODE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CHANGE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_CLOSE;
+import static android.view.WindowManager.TRANSIT_OLD_TASK_FRAGMENT_OPEN;
 import static android.view.WindowManager.TRANSIT_OLD_TASK_OPEN;
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_FRONT;
@@ -37,6 +41,7 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
+import static com.android.server.wm.WindowContainer.POSITION_BOTTOM;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -46,11 +51,14 @@
 import static org.junit.Assume.assumeFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 
 import android.annotation.Nullable;
+import android.graphics.Rect;
+import android.gui.DropInputMode;
 import android.os.Binder;
 import android.os.IBinder;
 import android.os.RemoteException;
@@ -833,6 +841,353 @@
     }
 
     @Test
+    public void testOverrideTaskFragmentAdapter_overrideWithEmbeddedActivity() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord activity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(activity);
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation run by the remote handler.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_noOverrideWithOnlyTaskFragmentFillingTask() {
+        final Task task = createTask(mDisplayContent);
+        final ActivityRecord closingActivity = createActivityRecord(task);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+
+        // Make sure the TaskFragment is not embedded.
+        assertFalse(taskFragment.isEmbeddedWithBoundsOverride());
+        final ActivityRecord openingActivity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(closingActivity);
+        prepareActivityForAppTransition(openingActivity);
+        final int uid = 12345;
+        closingActivity.info.applicationInfo.uid = uid;
+        openingActivity.info.applicationInfo.uid = uid;
+        task.effectiveUid = uid;
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity,
+                null /* changingTaskFragment */);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation is not run by the remote handler because the activity is filling the Task.
+        assertFalse(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_overrideWithTaskFragmentNotFillingTask() {
+        final Task task = createTask(mDisplayContent);
+        final ActivityRecord closingActivity = createActivityRecord(task);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+
+        // Make sure the TaskFragment is embedded.
+        taskFragment.setWindowingMode(WINDOWING_MODE_MULTI_WINDOW);
+        final Rect embeddedBounds = new Rect(task.getBounds());
+        embeddedBounds.right = embeddedBounds.left + embeddedBounds.width() / 2;
+        taskFragment.setBounds(embeddedBounds);
+        assertTrue(taskFragment.isEmbeddedWithBoundsOverride());
+        final ActivityRecord openingActivity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(closingActivity);
+        prepareActivityForAppTransition(openingActivity);
+        final int uid = 12345;
+        closingActivity.info.applicationInfo.uid = uid;
+        openingActivity.info.applicationInfo.uid = uid;
+        task.effectiveUid = uid;
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity,
+                null /* changingTaskFragment */);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation run by the remote handler.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_overrideWithNonEmbeddedActivity() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Closing non-embedded activity.
+        final ActivityRecord closingActivity = createActivityRecord(task);
+        prepareActivityForAppTransition(closingActivity);
+        // Opening TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord openingActivity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(openingActivity);
+        task.effectiveUid = openingActivity.getUid();
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation run by the remote handler.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_overrideEmbeddedActivityWithDiffUid() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Closing TaskFragment with embedded activity.
+        final TaskFragment taskFragment1 = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord closingActivity = taskFragment1.getTopMostActivity();
+        prepareActivityForAppTransition(closingActivity);
+        closingActivity.info.applicationInfo.uid = 12345;
+        // Opening TaskFragment with embedded activity with different UID.
+        final TaskFragment taskFragment2 = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord openingActivity = taskFragment2.getTopMostActivity();
+        prepareActivityForAppTransition(openingActivity);
+        openingActivity.info.applicationInfo.uid = 54321;
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment1);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation run by the remote handler.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_noOverrideWithTwoApps() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Closing activity in Task1.
+        final ActivityRecord closingActivity = createActivityRecord(mDisplayContent);
+        prepareActivityForAppTransition(closingActivity);
+        // Opening TaskFragment with embedded activity in Task2.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord openingActivity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(openingActivity);
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation not run by the remote handler.
+        assertFalse(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_noOverrideNonEmbeddedActivityWithDiffUid() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Closing TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord closingActivity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(closingActivity);
+        closingActivity.info.applicationInfo.uid = 12345;
+        task.effectiveUid = closingActivity.getUid();
+        // Opening non-embedded activity with different UID.
+        final ActivityRecord openingActivity = createActivityRecord(task);
+        prepareActivityForAppTransition(openingActivity);
+        openingActivity.info.applicationInfo.uid = 54321;
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(openingActivity, closingActivity, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation should not run by the remote handler when there are non-embedded activities of
+        // different UID.
+        assertFalse(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_noOverrideWithWallpaper() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with embedded activity.
+        final TaskFragment taskFragment = createTaskFragmentWithEmbeddedActivity(task, organizer);
+        final ActivityRecord activity = taskFragment.getTopMostActivity();
+        prepareActivityForAppTransition(activity);
+        // Set wallpaper as visible.
+        final WallpaperWindowToken wallpaperWindowToken = new WallpaperWindowToken(mWm,
+                mock(IBinder.class), true, mDisplayContent, true /* ownerCanManageAppTokens */);
+        spyOn(mDisplayContent.mWallpaperController);
+        doReturn(true).when(mDisplayContent.mWallpaperController).isWallpaperVisible();
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // Animation should not run by the remote handler when there is wallpaper in the transition.
+        assertFalse(remoteAnimationRunner.isAnimationStarted());
+    }
+
+    @Test
+    public void testOverrideTaskFragmentAdapter_inputProtectedForUntrustedAnimation() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with embedded activities, one is trusted embedded, and the other
+        // one is untrusted embedded.
+        final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(2)
+                .setOrganizer(organizer)
+                .build();
+        final ActivityRecord activity0 = taskFragment.getChildAt(0).asActivityRecord();
+        final ActivityRecord activity1 = taskFragment.getChildAt(1).asActivityRecord();
+        // Also create a non-embedded activity in the Task.
+        final ActivityRecord activity2 = new ActivityBuilder(mAtm).build();
+        task.addChild(activity2, POSITION_BOTTOM);
+        prepareActivityForAppTransition(activity0);
+        prepareActivityForAppTransition(activity1);
+        prepareActivityForAppTransition(activity2);
+        doReturn(false).when(taskFragment).isAllowedToEmbedActivityInTrustedMode(activity0);
+        doReturn(true).when(taskFragment).isAllowedToEmbedActivityInTrustedMode(activity1);
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(activity1, null /* closingActivity */, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // The animation will be animated remotely by client and all activities are input disabled
+        // for untrusted animation.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+        verify(activity0).setDropInputForAnimation(true);
+        verify(activity1).setDropInputForAnimation(true);
+        verify(activity2).setDropInputForAnimation(true);
+        verify(activity0).setDropInputMode(DropInputMode.ALL);
+        verify(activity1).setDropInputMode(DropInputMode.ALL);
+        verify(activity2).setDropInputMode(DropInputMode.ALL);
+
+        // Reset input after animation is finished.
+        clearInvocations(activity0);
+        clearInvocations(activity1);
+        clearInvocations(activity2);
+        remoteAnimationRunner.finishAnimation();
+
+        verify(activity0).setDropInputForAnimation(false);
+        verify(activity1).setDropInputForAnimation(false);
+        verify(activity2).setDropInputForAnimation(false);
+        verify(activity0).setDropInputMode(DropInputMode.OBSCURED);
+        verify(activity1).setDropInputMode(DropInputMode.NONE);
+        verify(activity2).setDropInputMode(DropInputMode.NONE);
+    }
+
+    /**
+     * Since we don't have any use case to rely on handling input during animation, disable it even
+     * if it is trusted embedding so that it could cover some edge-cases when a previously trusted
+     * host starts doing something bad.
+     */
+    @Test
+    public void testOverrideTaskFragmentAdapter_inputProtectedForTrustedAnimation() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with only trusted embedded activity
+        final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(1)
+                .setOrganizer(organizer)
+                .build();
+        final ActivityRecord activity = taskFragment.getChildAt(0).asActivityRecord();
+        prepareActivityForAppTransition(activity);
+        doReturn(true).when(taskFragment).isAllowedToEmbedActivityInTrustedMode(activity);
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // The animation will be animated remotely by client and all activities are input disabled
+        // for untrusted animation.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+        verify(activity).setDropInputForAnimation(true);
+        verify(activity).setDropInputMode(DropInputMode.ALL);
+
+        // Reset input after animation is finished.
+        clearInvocations(activity);
+        remoteAnimationRunner.finishAnimation();
+
+        verify(activity).setDropInputForAnimation(false);
+        verify(activity).setDropInputMode(DropInputMode.NONE);
+    }
+
+    /**
+     * We don't need to drop input for fully trusted embedding (system app, and embedding in the
+     * same app). This will allow users to do fast tapping.
+     */
+    @Test
+    public void testOverrideTaskFragmentAdapter_noInputProtectedForFullyTrustedAnimation() {
+        final Task task = createTask(mDisplayContent);
+        final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
+        final TestRemoteAnimationRunner remoteAnimationRunner = new TestRemoteAnimationRunner();
+        setupTaskFragmentRemoteAnimation(organizer, remoteAnimationRunner);
+
+        // Create a TaskFragment with only trusted embedded activity
+        final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(1)
+                .setOrganizer(organizer)
+                .build();
+        final ActivityRecord activity = taskFragment.getChildAt(0).asActivityRecord();
+        prepareActivityForAppTransition(activity);
+        final int uid = mAtm.mTaskFragmentOrganizerController.getTaskFragmentOrganizerUid(
+                getITaskFragmentOrganizer(organizer));
+        doReturn(true).when(task).isFullyTrustedEmbedding(uid);
+        spyOn(mDisplayContent.mAppTransition);
+
+        // Prepare and start transition.
+        prepareAndTriggerAppTransition(activity, null /* closingActivity */, taskFragment);
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+
+        // The animation will be animated remotely by client, but input should not be dropped for
+        // fully trusted.
+        assertTrue(remoteAnimationRunner.isAnimationStarted());
+        verify(activity, never()).setDropInputForAnimation(true);
+        verify(activity, never()).setDropInputMode(DropInputMode.ALL);
+    }
+
+    @Test
     public void testTransitionGoodToGoForTaskFragments() {
         final TaskFragmentOrganizer organizer = new TaskFragmentOrganizer(Runnable::run);
         final Task task = createTask(mDisplayContent);
@@ -898,6 +1253,22 @@
         verify(mDisplayContent.mAppTransition).goodToGo(anyInt(), any());
     }
 
+    /** Registers remote animation for the organizer. */
+    private void setupTaskFragmentRemoteAnimation(TaskFragmentOrganizer organizer,
+            TestRemoteAnimationRunner remoteAnimationRunner) {
+        final RemoteAnimationAdapter adapter = new RemoteAnimationAdapter(
+                remoteAnimationRunner, 10, 1);
+        final ITaskFragmentOrganizer iOrganizer = getITaskFragmentOrganizer(organizer);
+        final RemoteAnimationDefinition definition = new RemoteAnimationDefinition();
+        definition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CHANGE, adapter);
+        definition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_OPEN, adapter);
+        definition.addRemoteAnimation(TRANSIT_OLD_TASK_FRAGMENT_CLOSE, adapter);
+        definition.addRemoteAnimation(TRANSIT_OLD_ACTIVITY_OPEN, adapter);
+        definition.addRemoteAnimation(TRANSIT_OLD_ACTIVITY_CLOSE, adapter);
+        registerTaskFragmentOrganizer(iOrganizer);
+        mAtm.mTaskFragmentOrganizerController.registerRemoteAnimations(iOrganizer, definition);
+    }
+
     private static ITaskFragmentOrganizer getITaskFragmentOrganizer(
             TaskFragmentOrganizer organizer) {
         return ITaskFragmentOrganizer.Stub.asInterface(organizer.getOrganizerToken().asBinder());
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerExemptionTests.java b/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerExemptionTests.java
index 6e48818..3910904 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerExemptionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerExemptionTests.java
@@ -18,6 +18,7 @@
 
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED;
 import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS;
+import static android.app.ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE;
 
 import static com.android.server.wm.ActivityTaskManagerService.APP_SWITCH_ALLOW;
 import static com.android.server.wm.BackgroundActivityStartController.BAL_ALLOW_ALLOWLISTED_COMPONENT;
@@ -25,9 +26,11 @@
 import static com.android.server.wm.BackgroundActivityStartController.BAL_ALLOW_PERMISSION;
 import static com.android.server.wm.BackgroundActivityStartController.BAL_ALLOW_SAW_PERMISSION;
 import static com.android.server.wm.BackgroundActivityStartController.BAL_ALLOW_VISIBLE_WINDOW;
+import static com.android.server.wm.BackgroundActivityStartController.BAL_BLOCK;
 
 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.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -43,6 +46,9 @@
 import android.content.pm.PackageManagerInternal;
 import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.provider.DeviceConfig;
 import android.util.Pair;
 
@@ -52,6 +58,7 @@
 import com.android.modules.utils.testing.ExtendedMockitoRule;
 import com.android.server.am.PendingIntentRecord;
 import com.android.server.wm.BackgroundActivityStartController.BalVerdict;
+import com.android.window.flags.Flags;
 
 import org.junit.After;
 import org.junit.Before;
@@ -95,7 +102,9 @@
     @Rule
     public final ExtendedMockitoRule extendedMockitoRule =
             new ExtendedMockitoRule.Builder(this).setStrictness(Strictness.LENIENT).build();
-
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule =
+            DeviceFlagsValueProvider.createCheckFlagsRule();
     TestableBackgroundActivityStartController mController;
     @Mock
     ActivityMetricsLogger mActivityMetricsLogger;
@@ -186,7 +195,7 @@
         when(mAppOpsManager.checkOpNoThrow(
                 eq(AppOpsManager.OP_SYSTEM_EXEMPT_FROM_ACTIVITY_BG_START_RESTRICTION),
                 anyInt(), anyString())).thenReturn(AppOpsManager.MODE_DEFAULT);
-        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 BalVerdict.BLOCK);
     }
 
@@ -227,7 +236,7 @@
         // call
         BalVerdict callerVerdict = mController.checkBackgroundActivityStartAllowedByCaller(
                 balState);
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
 
         balState.setResultForCaller(callerVerdict);
@@ -295,7 +304,77 @@
                 checkedOptions);
 
         // call
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
+                balState);
+        balState.setResultForRealCaller(realCallerVerdict);
+
+        // assertions
+        assertWithMessage(balState.toString()).that(realCallerVerdict.getCode()).isEqualTo(
+                BAL_ALLOW_VISIBLE_WINDOW);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_BAL_ADDITIONAL_START_MODES)
+    public void testCaller_appHasVisibleWindowWithIfVisibleOptIn() {
+        int callingUid = REGULAR_UID_1;
+        int callingPid = REGULAR_PID_1;
+        final String callingPackage = REGULAR_PACKAGE_1;
+        int realCallingUid = REGULAR_UID_2;
+        int realCallingPid = REGULAR_PID_2;
+
+        // setup state
+        when(mService.hasActiveVisibleWindow(eq(callingUid))).thenReturn(true);
+        when(mService.getBalAppSwitchesState()).thenReturn(APP_SWITCH_ALLOW);
+
+        // prepare call
+        PendingIntentRecord originatingPendingIntent = mPendingIntentRecord;
+        BackgroundStartPrivileges forcedBalByPiSender = BackgroundStartPrivileges.NONE;
+        Intent intent = TEST_INTENT;
+        ActivityOptions checkedOptions = mCheckedOptions
+                .setPendingIntentCreatorBackgroundActivityStartMode(
+                        MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE);
+        BackgroundActivityStartController.BalState balState = mController.new BalState(callingUid,
+                callingPid, callingPackage, realCallingUid, realCallingPid, mCallerApp,
+                originatingPendingIntent, forcedBalByPiSender, mResultRecord, intent,
+                checkedOptions);
+
+        // call
+        BalVerdict callerVerdict = mController.checkBackgroundActivityStartAllowedByCaller(
+                balState);
+        balState.setResultForCaller(callerVerdict);
+
+        // assertions
+        assertWithMessage(balState.toString()).that(callerVerdict.getCode()).isEqualTo(
+                BAL_ALLOW_VISIBLE_WINDOW);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_BAL_ADDITIONAL_START_MODES)
+    public void testRealCaller_appHasVisibleWindowWithIfVisibleOptIn() {
+        int callingUid = REGULAR_UID_1;
+        int callingPid = REGULAR_PID_1;
+        final String callingPackage = REGULAR_PACKAGE_1;
+        int realCallingUid = REGULAR_UID_2;
+        int realCallingPid = REGULAR_PID_2;
+
+        // setup state
+        when(mService.hasActiveVisibleWindow(eq(realCallingUid))).thenReturn(true);
+        when(mService.getBalAppSwitchesState()).thenReturn(APP_SWITCH_ALLOW);
+
+        // prepare call
+        PendingIntentRecord originatingPendingIntent = mPendingIntentRecord;
+        BackgroundStartPrivileges forcedBalByPiSender = BackgroundStartPrivileges.NONE;
+        Intent intent = TEST_INTENT;
+        ActivityOptions checkedOptions = mCheckedOptions
+                .setPendingIntentCreatorBackgroundActivityStartMode(
+                        MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE);
+        BackgroundActivityStartController.BalState balState = mController.new BalState(callingUid,
+                callingPid, callingPackage, realCallingUid, realCallingPid, mCallerApp,
+                originatingPendingIntent, forcedBalByPiSender, mResultRecord, intent,
+                checkedOptions);
+
+        // call
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
         balState.setResultForRealCaller(realCallerVerdict);
 
@@ -320,7 +399,7 @@
         int realCallingPid = REGULAR_PID_2;
 
         // setup state
-        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 new BalVerdict(BAL_ALLOW_FOREGROUND, false, "allowed"));
         when(mService.getBalAppSwitchesState()).thenReturn(APP_SWITCH_ALLOW);
 
@@ -357,7 +436,7 @@
                 mService.getProcessController(eq(realCallingPid), eq(realCallingUid))).thenReturn(
                 mCallerApp);
         when(mService.getBalAppSwitchesState()).thenReturn(APP_SWITCH_ALLOW);
-        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 new BalVerdict(BAL_ALLOW_FOREGROUND, false, "allowed"));
 
         // prepare call
@@ -371,7 +450,7 @@
                 checkedOptions);
 
         // call
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
         balState.setResultForRealCaller(realCallerVerdict);
 
@@ -404,9 +483,9 @@
                 mService.getProcessController(eq(realCallingPid), eq(realCallingUid))).thenReturn(
                 mCallerApp);
         when(mService.getBalAppSwitchesState()).thenReturn(APP_SWITCH_ALLOW);
-        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 BalVerdict.BLOCK);
-        when(otherProcess.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        when(otherProcess.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 new BalVerdict(BAL_ALLOW_FOREGROUND, false, "allowed"));
 
         // prepare call
@@ -420,7 +499,7 @@
                 checkedOptions);
 
         // call
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
         balState.setResultForRealCaller(realCallerVerdict);
 
@@ -456,7 +535,7 @@
                 checkedOptions);
 
         // call
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
         balState.setResultForRealCaller(realCallerVerdict);
 
@@ -466,6 +545,45 @@
     }
 
     @Test
+    @RequiresFlagsEnabled(Flags.FLAG_BAL_ADDITIONAL_START_MODES)
+    public void testRealCaller_isCompanionAppWithOptInIfVisible() {
+        // The app has a service that is bound by a different, visible app. The app bound to the
+        // service must remain visible for the app in the background to start activities
+        // successfully.
+        int callingUid = REGULAR_UID_1;
+        int callingPid = REGULAR_PID_1;
+        final String callingPackage = REGULAR_PACKAGE_1;
+        int realCallingUid = REGULAR_UID_2;
+        int realCallingPid = REGULAR_PID_2;
+
+        // setup state
+        final int realCallingUserId = UserHandle.getUserId(realCallingUid);
+        when(mService.isAssociatedCompanionApp(eq(realCallingUserId),
+                eq(realCallingUid))).thenReturn(true);
+
+        // prepare call
+        PendingIntentRecord originatingPendingIntent = mPendingIntentRecord;
+        BackgroundStartPrivileges forcedBalByPiSender = BackgroundStartPrivileges.NONE;
+        Intent intent = TEST_INTENT;
+        ActivityOptions checkedOptions = mCheckedOptions
+                .setPendingIntentBackgroundActivityStartMode(
+                        MODE_BACKGROUND_ACTIVITY_START_ALLOW_IF_VISIBLE);
+        BackgroundActivityStartController.BalState balState = mController.new BalState(callingUid,
+                callingPid, callingPackage, realCallingUid, realCallingPid, null,
+                originatingPendingIntent, forcedBalByPiSender, mResultRecord, intent,
+                checkedOptions);
+
+        // call
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
+                balState);
+        balState.setResultForRealCaller(realCallerVerdict);
+
+        // assertions
+        assertWithMessage(balState.toString()).that(realCallerVerdict.getCode()).isEqualTo(
+                BAL_BLOCK);
+    }
+
+    @Test
     public void testCaller_balPermission() {
         int callingUid = REGULAR_UID_1;
         int callingPid = REGULAR_PID_1;
@@ -523,7 +641,7 @@
                 checkedOptions);
 
         // call
-        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedBySender(
+        BalVerdict realCallerVerdict = mController.checkBackgroundActivityStartAllowedByRealCaller(
                 balState);
         balState.setResultForRealCaller(realCallerVerdict);
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerTests.java
index e364264..6ec7895 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackgroundActivityStartControllerTests.java
@@ -24,6 +24,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.eq;
@@ -43,6 +44,7 @@
 import com.android.compatibility.common.util.DeviceConfigStateHelper;
 import com.android.server.am.PendingIntentRecord;
 import com.android.server.wm.BackgroundActivityStartController.BalVerdict;
+import com.android.server.wm.BackgroundLaunchProcessController.BalCheckConfiguration;
 
 import org.junit.After;
 import org.junit.Before;
@@ -167,9 +169,9 @@
         }
 
         @Override
-        BalVerdict checkBackgroundActivityStartAllowedBySender(BalState state) {
+        BalVerdict checkBackgroundActivityStartAllowedByRealCaller(BalState state) {
             return mRealCallerVerdict.orElseGet(
-                    () -> super.checkBackgroundActivityStartAllowedBySender(state));
+                    () -> super.checkBackgroundActivityStartAllowedByRealCaller(state));
         }
 
         public void setRealCallerVerdict(BalVerdict verdict) {
@@ -177,11 +179,12 @@
         }
 
         @Override
-        BalVerdict checkProcessAllowsBal(WindowProcessController app, BalState state) {
+        BalVerdict checkProcessAllowsBal(WindowProcessController app, BalState state,
+                BalCheckConfiguration checkConfiguration) {
             if (mProcessVerdicts.containsKey(app)) {
                 return mProcessVerdicts.get(app);
             }
-            return super.checkProcessAllowsBal(app, state);
+            return super.checkProcessAllowsBal(app, state, checkConfiguration);
         }
     }
 
@@ -209,7 +212,7 @@
         Mockito.when(mAppOpsManager.checkOpNoThrow(
                 eq(AppOpsManager.OP_SYSTEM_EXEMPT_FROM_ACTIVITY_BG_START_RESTRICTION),
                 anyInt(), anyString())).thenReturn(AppOpsManager.MODE_DEFAULT);
-        Mockito.when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt())).thenReturn(
+        Mockito.when(mCallerApp.areBackgroundActivityStartsAllowed(anyInt(), any())).thenReturn(
                 BalVerdict.BLOCK);
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/BackgroundLaunchProcessControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/BackgroundLaunchProcessControllerTests.java
index c9c7e92..27e147d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/BackgroundLaunchProcessControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/BackgroundLaunchProcessControllerTests.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static com.android.server.wm.ActivityTaskManagerService.ACTIVITY_BG_START_GRACE_PERIOD_MS;
 import static com.android.server.wm.ActivityTaskManagerService.APP_SWITCH_ALLOW;
 import static com.android.server.wm.ActivityTaskManagerService.APP_SWITCH_DISALLOW;
 import static com.android.server.wm.BackgroundActivityStartController.BAL_ALLOW_BOUND_BY_FOREGROUND;
@@ -95,7 +96,12 @@
     int mUid = 234;
     String mPackageName = "package.name";
     int mAppSwitchState = APP_SWITCH_DISALLOW;
-    boolean mIsCheckingForFgsStart = false;
+    BackgroundLaunchProcessController.BalCheckConfiguration mBalCheckConfiguration =
+            new BackgroundLaunchProcessController.BalCheckConfiguration(
+                    /* isCheckingForFgsStarts */ false,
+                    /* checkVisibility */ true,
+                    /* checkOtherExemptions */ true,
+                    ACTIVITY_BG_START_GRACE_PERIOD_MS);
     boolean mHasActivityInVisibleTask = false;
     boolean mHasBackgroundActivityStartPrivileges = false;
     long mLastStopAppSwitchesTime = 0L;
@@ -106,7 +112,7 @@
     public void testNothingAllows() {
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -118,7 +124,7 @@
         mHasBackgroundActivityStartPrivileges = true;
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -136,7 +142,7 @@
                 BackgroundStartPrivileges.ALLOW_BAL);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -154,7 +160,7 @@
                 BackgroundStartPrivileges.ALLOW_BAL);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -170,7 +176,7 @@
                 BackgroundStartPrivileges.ALLOW_BAL);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -186,7 +192,7 @@
                 BackgroundStartPrivileges.ALLOW_BAL);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -201,7 +207,7 @@
         mHasActiveVisibleWindow.add(999);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -216,7 +222,7 @@
         mHasActiveVisibleWindow.add(999);
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -229,7 +235,7 @@
         mHasActivityInVisibleTask = true;
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
@@ -245,7 +251,7 @@
         mLastActivityFinishTime = now - 100;
         BalVerdict balVerdict = mController.areBackgroundActivityStartsAllowed(
                 mPid, mUid, mPackageName,
-                mAppSwitchState, mIsCheckingForFgsStart,
+                mAppSwitchState, mBalCheckConfiguration,
                 mHasActivityInVisibleTask, mHasBackgroundActivityStartPrivileges,
                 mLastStopAppSwitchesTime, mLastActivityLaunchTime,
                 mLastActivityFinishTime);
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
index affaad6..14276ae 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentDeferredUpdateTests.java
@@ -16,7 +16,7 @@
 
 package com.android.server.wm;
 
-import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.LayoutParams.TYPE_BASE_APPLICATION;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.any;
@@ -277,7 +277,7 @@
         mDisplayContent.mDisplayUpdater.onDisplaySwitching(/* switching= */ true);
 
         mWmInternal.waitForAllWindowsDrawn(mScreenUnblocker,
-                /* timeout= */ Integer.MAX_VALUE, DEFAULT_DISPLAY);
+                /* timeout= */ Integer.MAX_VALUE, INVALID_DISPLAY);
         mWmInternal.waitForAllWindowsDrawn(mSecondaryScreenUnblocker,
                 /* timeout= */ Integer.MAX_VALUE, mSecondaryDisplayContent.getDisplayId());
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index f2ea1c9..eca4d21 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -1144,6 +1144,7 @@
 
     @Test
     public void testOrientationBehind() {
+        assertNull(mDisplayContent.getLastOrientationSource());
         final ActivityRecord prev = new ActivityBuilder(mAtm).setCreateTask(true)
                 .setScreenOrientation(getRotatedOrientation(mDisplayContent)).build();
         prev.setVisibleRequested(false);
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
index f1db713..957b5e0 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -371,8 +371,7 @@
         ensureTaskPlacement(fullscreenTask, firstActivity, secondActivity);
 
         // Move first activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(firstActivity,
-                null /* launchIntoPipHostActivity */, "initialMove");
+        mRootWindowContainer.moveActivityToPinnedRootTask(firstActivity, "initialMove");
 
         final TaskDisplayArea taskDisplayArea = fullscreenTask.getDisplayArea();
         Task pinnedRootTask = taskDisplayArea.getRootPinnedTask();
@@ -381,8 +380,7 @@
         ensureTaskPlacement(fullscreenTask, secondActivity);
 
         // Move second activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(secondActivity,
-                null /* launchIntoPipHostActivity */, "secondMove");
+        mRootWindowContainer.moveActivityToPinnedRootTask(secondActivity, "secondMove");
 
         // Need to get root tasks again as a new instance might have been created.
         pinnedRootTask = taskDisplayArea.getRootPinnedTask();
@@ -413,8 +411,7 @@
 
 
         // Move first activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(secondActivity,
-                null /* launchIntoPipHostActivity */, "initialMove");
+        mRootWindowContainer.moveActivityToPinnedRootTask(secondActivity, "initialMove");
 
         assertTrue(firstActivity.mRequestForceTransition);
     }
@@ -434,8 +431,7 @@
         transientActivity.setState(RESUMED, "test");
         transientActivity.getTask().moveToFront("test");
 
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity2,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity2, "test");
         assertEquals("Created PiP task must not change focus", transientActivity.getTask(),
                 mRootWindowContainer.getTopDisplayFocusedRootTask());
         final Task newPipTask = activity2.getTask();
@@ -460,8 +456,7 @@
         final Task task = activity.getTask();
 
         // Move activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity, "test");
 
         // Ensure a task has moved over.
         ensureTaskPlacement(task, activity);
@@ -499,8 +494,7 @@
         final Task task = activity.getTask();
 
         // Move activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity, "test");
 
         // Ensure a task has moved over.
         ensureTaskPlacement(task, activity);
@@ -524,8 +518,7 @@
         final ActivityRecord secondActivity = taskFragment.getBottomMostActivity();
 
         // Move first activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(firstActivity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(firstActivity, "test");
 
         final TaskDisplayArea taskDisplayArea = fullscreenTask.getDisplayArea();
         final Task pinnedRootTask = taskDisplayArea.getRootPinnedTask();
@@ -556,8 +549,7 @@
         final ActivityRecord topActivity = taskFragment.getTopMostActivity();
 
         // Move the top activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(topActivity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(topActivity, "test");
 
         final Task pinnedRootTask = task.getDisplayArea().getRootPinnedTask();
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimatorTest.java b/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimatorTest.java
index 9981a4d..9967cce 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimatorTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SurfaceAnimatorTest.java
@@ -36,7 +36,6 @@
 import android.view.SurfaceControl;
 import android.view.SurfaceControl.Builder;
 import android.view.SurfaceControl.Transaction;
-import android.view.SurfaceSession;
 
 import androidx.test.filters.SmallTest;
 
@@ -69,7 +68,6 @@
     @Mock AnimationAdapter mSpec2;
     @Mock Transaction mTransaction;
 
-    private SurfaceSession mSession = new SurfaceSession();
     private MyAnimatable mAnimatable;
     private MyAnimatable mAnimatable2;
     private DeferFinishAnimatable mDeferFinishAnimatable;
@@ -78,9 +76,9 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
-        mAnimatable = new MyAnimatable(mWm, mSession, mTransaction);
-        mAnimatable2 = new MyAnimatable(mWm, mSession, mTransaction);
-        mDeferFinishAnimatable = new DeferFinishAnimatable(mWm, mSession, mTransaction);
+        mAnimatable = new MyAnimatable(mWm, mTransaction);
+        mAnimatable2 = new MyAnimatable(mWm, mTransaction);
+        mDeferFinishAnimatable = new DeferFinishAnimatable(mWm, mTransaction);
     }
 
     @After
@@ -88,8 +86,6 @@
         mAnimatable = null;
         mAnimatable2 = null;
         mDeferFinishAnimatable = null;
-        mSession.kill();
-        mSession = null;
     }
 
     @Test
@@ -313,7 +309,6 @@
 
     private static class MyAnimatable implements Animatable {
 
-        private final SurfaceSession mSession;
         private final Transaction mTransaction;
         final SurfaceControl mParent;
         final SurfaceControl mSurface;
@@ -322,13 +317,12 @@
         boolean mFinishedCallbackCalled;
         @AnimationType int mFinishedAnimationType;
 
-        MyAnimatable(WindowManagerService wm, SurfaceSession session, Transaction transaction) {
-            mSession = session;
+        MyAnimatable(WindowManagerService wm, Transaction transaction) {
             mTransaction = transaction;
-            mParent = wm.makeSurfaceBuilder(mSession)
+            mParent = wm.makeSurfaceBuilder()
                     .setName("test surface parent")
                     .build();
-            mSurface = wm.makeSurfaceBuilder(mSession)
+            mSurface = wm.makeSurfaceBuilder()
                     .setName("test surface")
                     .build();
             mFinishedCallbackCalled = false;
@@ -360,7 +354,7 @@
 
         @Override
         public Builder makeAnimationLeash() {
-            return new Builder(mSession) {
+            return new Builder() {
 
                 @Override
                 public SurfaceControl build() {
@@ -406,9 +400,8 @@
 
         Runnable mEndDeferFinishCallback;
 
-        DeferFinishAnimatable(WindowManagerService wm, SurfaceSession session,
-                Transaction transaction) {
-            super(wm, session, transaction);
+        DeferFinishAnimatable(WindowManagerService wm, Transaction transaction) {
+            super(wm, transaction);
         }
 
         @Override
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 1e39f0b..ae722807 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -394,7 +394,7 @@
         mWmService = WindowManagerService.main(
                 mContext, mImService, false, wmPolicy, mAtmService,
                 testDisplayWindowSettingsProvider, StubTransaction::new,
-                (unused) -> new MockSurfaceControlBuilder());
+                MockSurfaceControlBuilder::new);
         spyOn(mWmService);
         spyOn(mWmService.mRoot);
         // Invoked during {@link ActivityStack} creation.
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
index d013053..0cd036f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentOrganizerControllerTest.java
@@ -90,6 +90,7 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
+import android.view.RemoteAnimationDefinition;
 import android.view.SurfaceControl;
 import android.window.IRemoteTransition;
 import android.window.ITaskFragmentOrganizer;
@@ -139,6 +140,7 @@
     private IBinder mFragmentToken;
     private WindowContainerTransaction mTransaction;
     private WindowContainerToken mFragmentWindowToken;
+    private RemoteAnimationDefinition mDefinition;
     private IBinder mErrorToken;
     private Rect mTaskFragBounds;
 
@@ -167,6 +169,7 @@
         mTransaction = new WindowContainerTransaction();
         mTransaction.setTaskFragmentOrganizer(mIOrganizer);
         mFragmentWindowToken = mTaskFragment.mRemoteToken.toWindowContainerToken();
+        mDefinition = new RemoteAnimationDefinition();
         mErrorToken = new Binder();
         final Rect displayBounds = mDisplayContent.getBounds();
         mTaskFragBounds = new Rect(displayBounds.left, displayBounds.top, displayBounds.centerX(),
@@ -576,6 +579,17 @@
     }
 
     @Test
+    public void testRegisterRemoteAnimations() {
+        mController.registerRemoteAnimations(mIOrganizer, mDefinition);
+
+        assertEquals(mDefinition, mController.getRemoteAnimationDefinition(mIOrganizer));
+
+        mController.unregisterRemoteAnimations(mIOrganizer);
+
+        assertNull(mController.getRemoteAnimationDefinition(mIOrganizer));
+    }
+
+    @Test
     public void testApplyTransaction_disallowRemoteTransitionForNonSystemOrganizer() {
         mTransaction.setRelativeBounds(mFragmentWindowToken, new Rect(0, 0, 100, 100));
         mTaskFragment.setTaskFragmentOrganizer(mOrganizerToken, 10 /* uid */,
@@ -2121,8 +2135,8 @@
     private static void setupMockParent(TaskFragment taskFragment, Task mockParent) {
         doReturn(mockParent).when(taskFragment).getTask();
         doReturn(new TaskFragmentParentInfo(
-                new Configuration(), DEFAULT_DISPLAY, true, true, null /* decorSurface */))
-                .when(mockParent).getTaskFragmentParentInfo();
+                new Configuration(), DEFAULT_DISPLAY, mockParent.mTaskId, true, true,
+                null /* decorSurface */)).when(mockParent).getTaskFragmentParentInfo();
 
         // Task needs to be visible
         mockParent.lastActiveTime = 100;
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
index 6be1af2..cc1805a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskFragmentTest.java
@@ -368,8 +368,7 @@
         assertNotEquals(TaskFragmentAnimationParams.DEFAULT, taskFragment0.getAnimationParams());
 
         // Move activity to pinned root task.
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity, "test");
 
         // Ensure taskFragment requested config is reset.
         assertEquals(taskFragment0, activity.getOrganizedTaskFragment());
@@ -399,8 +398,7 @@
         spyOn(mAtm.mTaskFragmentOrganizerController);
 
         // Move activity to pinned.
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity0,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity0, "test");
 
         // Ensure taskFragment requested config is reset.
         assertTrue(taskFragment0.mClearedTaskFragmentForPip);
@@ -434,8 +432,7 @@
                 .createActivityCount(1)
                 .build();
         final ActivityRecord activity = taskFragment.getTopMostActivity();
-        mRootWindowContainer.moveActivityToPinnedRootTask(activity,
-                null /* launchIntoPipHostActivity */, "test");
+        mRootWindowContainer.moveActivityToPinnedRootTask(activity, "test");
         spyOn(mAtm.mTaskFragmentOrganizerController);
         assertEquals(mIOrganizer, activity.mLastTaskFragmentOrganizerBeforePip);
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskTests.java b/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
index 55f74e9d..45082d2 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskTests.java
@@ -230,8 +230,7 @@
         final Task originalTask = activityMain.getTask();
         final ActivityRecord activityPip = new ActivityBuilder(mAtm).setTask(originalTask).build();
         activityPip.setState(RESUMED, "test");
-        mAtm.mRootWindowContainer.moveActivityToPinnedRootTask(activityPip,
-                null /* launchIntoPipHostActivity */, "test");
+        mAtm.mRootWindowContainer.moveActivityToPinnedRootTask(activityPip, "test");
         final Task pinnedActivityTask = activityPip.getTask();
 
         // Simulate pinnedActivityTask unintentionally added to recent during top activity resume.
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 56fca31..7320c0b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -1251,7 +1251,7 @@
         final Transition transition = app.mTransitionController.createTransition(TRANSIT_OPEN);
         app.mTransitionController.requestStartTransition(transition, app.getTask(),
                 null /* remoteTransition */, null /* displayChange */);
-        app.mTransitionController.collectExistenceChange(app.getTask());
+        transition.collectExistenceChange(app.getTask());
         mDisplayContent.setFixedRotationLaunchingAppUnchecked(app);
         final AsyncRotationController asyncRotationController =
                 mDisplayContent.getAsyncRotationController();
@@ -1416,7 +1416,8 @@
         activity1.setVisibleRequested(false);
         activity2.setVisibleRequested(true);
 
-        openTransition.finishTransition();
+        final ActionChain chain = ActionChain.testFinish(null);
+        openTransition.finishTransition(chain);
 
         // We finished the openTransition. Even though activity1 is visibleRequested=false, since
         // the closeTransition animation hasn't played yet, make sure that we didn't commit
@@ -1429,7 +1430,7 @@
         // normally.
         mWm.mSyncEngine.abort(closeTransition.getSyncId());
 
-        closeTransition.finishTransition();
+        closeTransition.finishTransition(chain);
 
         assertFalse(activity1.isVisible());
         assertTrue(activity2.isVisible());
@@ -1449,7 +1450,7 @@
         activity1.setState(ActivityRecord.State.INITIALIZING, "test");
         activity1.mLaunchTaskBehind = true;
         mWm.mSyncEngine.abort(noChangeTransition.getSyncId());
-        noChangeTransition.finishTransition();
+        noChangeTransition.finishTransition(chain);
         assertTrue(activity1.mLaunchTaskBehind);
     }
 
@@ -1468,7 +1469,7 @@
         // We didn't call abort on the transition itself, so it will still run onTransactionReady
         // normally.
         mWm.mSyncEngine.abort(transition1.getSyncId());
-        transition1.finishTransition();
+        transition1.finishTransition(ActionChain.testFinish(transition1));
 
         verify(transitionEndedListener).run();
 
@@ -1530,7 +1531,7 @@
 
         verify(taskSnapshotController, times(1)).recordSnapshot(eq(task2));
 
-        controller.finishTransition(openTransition);
+        controller.finishTransition(ActionChain.testFinish(openTransition));
 
         // We are now going to simulate closing task1 to return back to (open) task2.
         final Transition closeTransition = createTestTransition(TRANSIT_CLOSE, controller);
@@ -1595,7 +1596,7 @@
         doReturn(true).when(task1).isTranslucentForTransition();
         assertFalse(controller.canApplyDim(task1));
 
-        controller.finishTransition(closeTransition);
+        controller.finishTransition(ActionChain.testFinish(closeTransition));
         assertTrue(wasInFinishingTransition[0]);
         assertFalse(calledListenerOnOtherDisplay[0]);
         assertNull(controller.mFinishingTransition);
@@ -1651,7 +1652,7 @@
         // to avoid the latency to resume the current top, i.e. appB.
         assertTrue(controller.isTransientVisible(taskRecent));
         // The recent is paused after the transient transition is finished.
-        controller.finishTransition(transition);
+        controller.finishTransition(ActionChain.testFinish(transition));
         assertFalse(controller.isTransientVisible(taskRecent));
     }
 
@@ -2004,10 +2005,10 @@
     @DisableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_disableAnimOptionsPerChange() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeCommonAnimOptions("testPackage");
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         mTransition.overrideAnimationOptionsToInfoIfNecessary(mInfo);
@@ -2018,10 +2019,10 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_fromStyleAnimOptions() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeCommonAnimOptions("testPackage");
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         mTransition.overrideAnimationOptionsToInfoIfNecessary(mInfo);
@@ -2044,10 +2045,10 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_sceneAnimOptions() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeSceneTransitionAnimOptions();
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         mTransition.overrideAnimationOptionsToInfoIfNecessary(mInfo);
@@ -2070,10 +2071,10 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_crossProfileAnimOptions() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeCrossProfileAnimOptions();
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         final TransitionInfo.Change displayChange = mInfo.getChanges().get(0);
@@ -2098,13 +2099,13 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_customAnimOptions() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeCustomAnimOptions("testPackage", Resources.ID_NULL,
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         Color.GREEN, false /* overrideTaskTransition */);
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         mTransition.overrideAnimationOptionsToInfoIfNecessary(mInfo);
@@ -2131,7 +2132,7 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_haveTaskFragmentAnimParams() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
 
         final TaskFragment embeddedTf = mTransition.mTargets.get(2).mContainer.asTaskFragment();
         embeddedTf.setAnimationParams(new TaskFragmentAnimationParams.Builder()
@@ -2144,7 +2145,7 @@
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         Color.GREEN, false /* overrideTaskTransition */);
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         final TransitionInfo.Change displayChange = mInfo.getChanges().get(0);
@@ -2180,13 +2181,13 @@
     @EnableFlags(Flags.FLAG_MOVE_ANIMATION_OPTIONS_TO_CHANGE)
     @Test
     public void testOverrideAnimationOptionsToInfoIfNecessary_customAnimOptionsWithTaskOverride() {
-        initializeOverrideAnimationOptionsTest();
+        ActivityRecord r = initializeOverrideAnimationOptionsTest();
         TransitionInfo.AnimationOptions options = TransitionInfo.AnimationOptions
                 .makeCustomAnimOptions("testPackage", Resources.ID_NULL,
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         TransitionInfo.AnimationOptions.DEFAULT_ANIMATION_RESOURCES_ID,
                         Color.GREEN, true /* overrideTaskTransition */);
-        mTransition.setOverrideAnimation(options, null /* startCallback */,
+        mTransition.setOverrideAnimation(options, r, null /* startCallback */,
                 null /* finishCallback */);
 
         mTransition.overrideAnimationOptionsToInfoIfNecessary(mInfo);
@@ -2212,7 +2213,7 @@
                 options.getBackgroundColor(), activityChange.getBackgroundColor());
     }
 
-    private void initializeOverrideAnimationOptionsTest() {
+    private ActivityRecord initializeOverrideAnimationOptionsTest() {
         mTransition = createTestTransition(TRANSIT_OPEN);
 
         // Test set AnimationOptions for Activity and Task.
@@ -2240,6 +2241,7 @@
                 embeddedTf.getAnimationLeash()));
         mInfo.addChange(new TransitionInfo.Change(null /* container */,
                 nonEmbeddedActivity.getAnimationLeash()));
+        return nonEmbeddedActivity;
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
index 7652861..9602ae2 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WallpaperControllerTests.java
@@ -407,7 +407,7 @@
         final SurfaceControl.Transaction t = mock(SurfaceControl.Transaction.class);
         token.finishSync(t, token.getSyncGroup(), false /* cancel */);
         transit.onTransactionReady(transit.getSyncId(), t);
-        dc.mTransitionController.finishTransition(transit);
+        dc.mTransitionController.finishTransition(ActionChain.testFinish(transit));
         assertFalse(wallpaperWindow.isVisible());
         assertFalse(token.isVisible());
     }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
index 39640fb..916c237 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowManagerServiceTests.java
@@ -1189,7 +1189,7 @@
         final int displayId = mDisplayContent.mDisplayId;
         // Use real surface, so ViewportWindow's BlastBufferQueue can be created.
         final ArrayList<SurfaceControl> surfaceControls = new ArrayList<>();
-        mWm.mSurfaceControlFactory = s -> new SurfaceControl.Builder() {
+        mWm.mSurfaceControlFactory = () -> new SurfaceControl.Builder() {
             @Override
             public SurfaceControl build() {
                 final SurfaceControl sc = super.build();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index bcf4ebc..a215c0a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -2131,7 +2131,7 @@
         }
 
         public void finish() {
-            mController.finishTransition(mLastTransit);
+            mController.finishTransition(ActionChain.testFinish(mLastTransit));
         }
     }
 }
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
index 1d567b1..c45b99d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTracingPerfettoTest.java
@@ -39,8 +39,7 @@
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
+import org.mockito.Mockito;
 
 import perfetto.protos.PerfettoConfig.WindowManagerConfig.LogFrequency;
 
@@ -50,9 +49,7 @@
 @SmallTest
 @Presubmit
 public class WindowTracingPerfettoTest {
-    @Mock
     private WindowManagerService mWmMock;
-    @Mock
     private Choreographer mChoreographer;
     private WindowTracing mWindowTracing;
     private PerfettoTraceMonitor mTraceMonitor;
@@ -60,7 +57,10 @@
 
     @Before
     public void setUp() throws Exception {
-        MockitoAnnotations.initMocks(this);
+        mWmMock = Mockito.mock(WindowManagerService.class);
+        Mockito.doNothing().when(mWmMock).dumpDebugLocked(Mockito.any(), Mockito.anyInt());
+
+        mChoreographer = Mockito.mock(Choreographer.class);
 
         mWindowTracing = new WindowTracingPerfetto(mWmMock, mChoreographer,
                 new WindowManagerGlobalLock());
diff --git a/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java b/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
index 88ce3a6..4f60106 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ZOrderingTests.java
@@ -51,7 +51,6 @@
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 import android.view.SurfaceControl;
-import android.view.SurfaceSession;
 import android.window.ScreenCapture;
 
 import androidx.test.filters.SmallTest;
@@ -63,7 +62,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
-import java.util.function.Function;
+import java.util.function.Supplier;
 
 /**
  * Tests for the {@link DisplayContent#assignChildLayers(SurfaceControl.Transaction)} method.
@@ -126,8 +125,7 @@
         private LayerRecordingTransaction mTransaction;
         private SurfaceControl mPendingParent;
 
-        HierarchyRecorder(SurfaceSession s, LayerRecordingTransaction transaction) {
-            super(s);
+        HierarchyRecorder(LayerRecordingTransaction transaction) {
             mTransaction = transaction;
         }
 
@@ -146,8 +144,8 @@
         }
     }
 
-    private static class HierarchyRecordingBuilderFactory implements Function<SurfaceSession,
-            SurfaceControl.Builder> {
+    private static class HierarchyRecordingBuilderFactory
+            implements Supplier<SurfaceControl.Builder> {
         private LayerRecordingTransaction mTransaction;
 
         HierarchyRecordingBuilderFactory(LayerRecordingTransaction transaction) {
@@ -155,9 +153,8 @@
         }
 
         @Override
-        public SurfaceControl.Builder apply(SurfaceSession s) {
-            final LayerRecordingTransaction transaction = mTransaction;
-            return new HierarchyRecorder(s, transaction);
+        public SurfaceControl.Builder get() {
+            return new HierarchyRecorder(mTransaction);
         }
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/utils/DesktopModeFlagsUtilTest.java b/services/tests/wmtests/src/com/android/server/wm/utils/DesktopModeFlagsUtilTest.java
index 381e9e4..46b8e3a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/utils/DesktopModeFlagsUtilTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/utils/DesktopModeFlagsUtilTest.java
@@ -234,11 +234,11 @@
             FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
             FLAG_ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS
     })
-    public void isEnabled_dwFlagOn_overrideOff_featureFlagOn_returnsFalse() {
+    public void isEnabled_dwFlagOn_overrideOff_featureFlagOn_returnsTrue() {
         setOverride(OVERRIDE_OFF.getSetting());
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlagsUtil.DYNAMIC_INITIAL_BOUNDS.isEnabled(mContext)).isFalse();
+        assertThat(DesktopModeFlagsUtil.DYNAMIC_INITIAL_BOUNDS.isEnabled(mContext)).isTrue();
     }
 
     @Test
@@ -296,11 +296,11 @@
             FLAG_ENABLE_DESKTOP_WINDOWING_MODE,
             FLAG_ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS
     })
-    public void isEnabled_dwFlagOff_overrideOn_featureFlagOff_returnTrue() {
+    public void isEnabled_dwFlagOff_overrideOn_featureFlagOff_returnFalse() {
         setOverride(OVERRIDE_ON.getSetting());
 
         // Follow override if they exist, and is not equal to default toggle state (dw flag)
-        assertThat(DesktopModeFlagsUtil.DYNAMIC_INITIAL_BOUNDS.isEnabled(mContext)).isTrue();
+        assertThat(DesktopModeFlagsUtil.DYNAMIC_INITIAL_BOUNDS.isEnabled(mContext)).isFalse();
     }
 
     @Test
diff --git a/telecomm/java/android/telecom/CallControl.java b/telecomm/java/android/telecom/CallControl.java
index 808a575..1e1abf2 100644
--- a/telecomm/java/android/telecom/CallControl.java
+++ b/telecomm/java/android/telecom/CallControl.java
@@ -39,7 +39,10 @@
 
 /**
  * CallControl provides client side control of a call.  Each Call will get an individual CallControl
- * instance in which the client can alter the state of the associated call.
+ * instance in which the client can alter the state of the associated call.  Outgoing and incoming
+ * calls should move to active (via {@link CallControl#setActive(Executor, OutcomeReceiver)} or
+ * answered (via {@link CallControl#answer(int, Executor, OutcomeReceiver)} before 60 seconds.  If
+ * the new call is not moved to active or answered before 60 seconds, the call will be disconnected.
  *
  * <p>
  * Each method is Transactional meaning that it can succeed or fail. If a transaction succeeds,
diff --git a/telecomm/java/android/telecom/Connection.java b/telecomm/java/android/telecom/Connection.java
index 1df6cf7..ad7d987 100644
--- a/telecomm/java/android/telecom/Connection.java
+++ b/telecomm/java/android/telecom/Connection.java
@@ -2621,7 +2621,9 @@
     }
 
     /**
-     * Sets state to ringing (e.g., an inbound ringing connection).
+     * Sets state to ringing (e.g., an inbound ringing connection).  The Connection should not be
+     * in STATE_RINGING for more than 60 seconds. After 60 seconds, the Connection will
+     * be disconnected.
      */
     public final void setRinging() {
         checkImmutable();
@@ -2645,7 +2647,9 @@
     }
 
     /**
-     * Sets state to dialing (e.g., dialing an outbound connection).
+     * Sets state to dialing (e.g., dialing an outbound connection). The Connection should not be
+     * in STATE_DIALING for more than 60 seconds. After 60 seconds, the Connection will
+     * be disconnected.
      */
     public final void setDialing() {
         checkImmutable();
diff --git a/telephony/java/android/telephony/CarrierConfigManager.java b/telephony/java/android/telephony/CarrierConfigManager.java
index 1ba496d..13bd5eb 100644
--- a/telephony/java/android/telephony/CarrierConfigManager.java
+++ b/telephony/java/android/telephony/CarrierConfigManager.java
@@ -9994,6 +9994,17 @@
     public static final String KEY_SATELLITE_ESOS_SUPPORTED_BOOL = "satellite_esos_supported_bool";
 
     /**
+     * Indicate whether carrier roaming to satellite is using P2P SMS.
+     *
+     * This will need agreement with carriers before enabling this flag.
+     *
+     * The default value is false.
+     */
+    @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
+    public static final String KEY_SATELLITE_ROAMING_P2P_SMS_SUPPORTED_BOOL =
+            "satellite_roaming_p2p_sms_supported_bool";
+
+    /**
      * Defines the NIDD (Non-IP Data Delivery) APN to be used for carrier roaming to satellite
      * attachment. For more on NIDD, see 3GPP TS 29.542.
      * Note this config is the only source of truth regarding the definition of the APN.
@@ -10003,6 +10014,19 @@
     public static final String KEY_SATELLITE_NIDD_APN_NAME_STRING =
             "satellite_nidd_apn_name_string";
 
+    /**
+     * Default value {@code true}, meaning when an emergency call request comes in, if the device is
+     * in emergency satellite mode but hasn't sent the first satellite datagram, then exits
+     * satellite mode to allow the emergency call to go through.
+     *
+     * If {@code false}, the emergency call is always blocked if device is in emergency satellite
+     * mode. Note if device is NOT in emergency satellite mode, emergency call is always allowed.
+     *
+     * @hide
+     */
+    public static final String KEY_SATELLITE_ROAMING_TURN_OFF_SESSION_FOR_EMERGENCY_CALL_BOOL =
+            "satellite_roaming_turn_off_session_for_emergency_call_bool";
+
     /** @hide */
     @IntDef({
             CARRIER_ROAMING_NTN_CONNECT_AUTOMATIC,
@@ -10075,7 +10099,35 @@
      */
     @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
     public static final String KEY_SATELLITE_SCREEN_OFF_INACTIVITY_TIMEOUT_SEC_INT =
-            "satellite_screen_off_inactivity_timeout_duration_sec_int";
+            "satellite_screen_off_inactivity_timeout_sec_int";
+
+    /**
+     * An integer key holds the timeout duration in seconds used to determine whether to exit P2P
+     * SMS mode and start TN scanning.
+     *
+     * The timer is started when the device is not connected, user is not pointing to the
+     * satellite and no data transfer is happening.
+     * When the timer expires, the device will move to IDLE mode upon which TN scanning will start.
+     *
+     * The default value is 180 seconds.
+     */
+    @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
+    public static final String KEY_SATELLITE_P2P_SMS_INACTIVITY_TIMEOUT_SEC_INT =
+            "satellite_p2p_sms_inactivity_timeout_sec_int";
+
+    /**
+     * An integer key holds the timeout duration in seconds used to determine whether to exit ESOS
+     * mode and start TN scanning.
+     *
+     * The timer is started when the device is not connected, user is not pointing to the
+     * satellite and no data transfer is happening.
+     * When the timer expires, the device will move to IDLE mode upon which TN scanning will start.
+     *
+     * The default value is 600 seconds.
+     */
+    @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
+    public static final String KEY_SATELLITE_ESOS_INACTIVITY_TIMEOUT_SEC_INT =
+            "satellite_esos_inactivity_timeout_sec_int";
 
     /**
      * Indicating whether DUN APN should be disabled when the device is roaming. In that case,
@@ -11235,12 +11287,16 @@
         sDefaults.putInt(KEY_EMERGENCY_CALL_TO_SATELLITE_T911_HANDOVER_TIMEOUT_MILLIS_INT,
                 (int) TimeUnit.SECONDS.toMillis(30));
         sDefaults.putBoolean(KEY_SATELLITE_ESOS_SUPPORTED_BOOL, false);
+        sDefaults.putBoolean(KEY_SATELLITE_ROAMING_P2P_SMS_SUPPORTED_BOOL, false);
         sDefaults.putString(KEY_SATELLITE_NIDD_APN_NAME_STRING, "");
+        sDefaults.putBoolean(KEY_SATELLITE_ROAMING_TURN_OFF_SESSION_FOR_EMERGENCY_CALL_BOOL, true);
         sDefaults.putInt(KEY_CARRIER_ROAMING_NTN_CONNECT_TYPE_INT, 0);
         sDefaults.putInt(KEY_CARRIER_ROAMING_NTN_EMERGENCY_CALL_TO_SATELLITE_HANDOVER_TYPE_INT,
                 SatelliteManager.EMERGENCY_CALL_TO_SATELLITE_HANDOVER_TYPE_T911);
         sDefaults.putInt(KEY_CARRIER_SUPPORTED_SATELLITE_NOTIFICATION_HYSTERESIS_SEC_INT, 180);
         sDefaults.putInt(KEY_SATELLITE_SCREEN_OFF_INACTIVITY_TIMEOUT_SEC_INT, 30);
+        sDefaults.putInt(KEY_SATELLITE_P2P_SMS_INACTIVITY_TIMEOUT_SEC_INT, 180);
+        sDefaults.putInt(KEY_SATELLITE_ESOS_INACTIVITY_TIMEOUT_SEC_INT, 600);
         sDefaults.putString(KEY_DEFAULT_PREFERRED_APN_NAME_STRING, "");
         sDefaults.putBoolean(KEY_SUPPORTS_CALL_COMPOSER_BOOL, false);
         sDefaults.putBoolean(KEY_SUPPORTS_BUSINESS_CALL_COMPOSER_BOOL, false);
diff --git a/telephony/java/android/telephony/TelephonyManager.java b/telephony/java/android/telephony/TelephonyManager.java
index 7481daa..bc709ab 100644
--- a/telephony/java/android/telephony/TelephonyManager.java
+++ b/telephony/java/android/telephony/TelephonyManager.java
@@ -132,6 +132,7 @@
 import com.android.internal.telephony.PhoneConstants;
 import com.android.internal.telephony.RILConstants;
 import com.android.internal.telephony.flags.Flags;
+import com.android.internal.telephony.util.TelephonyUtils;
 import com.android.telephony.Rlog;
 
 import java.io.IOException;
@@ -9975,6 +9976,7 @@
             ALLOWED_NETWORK_TYPES_REASON_POWER,
             ALLOWED_NETWORK_TYPES_REASON_CARRIER,
             ALLOWED_NETWORK_TYPES_REASON_ENABLE_2G,
+            ALLOWED_NETWORK_TYPES_REASON_TEST,
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface AllowedNetworkTypesReason {
@@ -10013,6 +10015,14 @@
     public static final int ALLOWED_NETWORK_TYPES_REASON_ENABLE_2G = 3;
 
     /**
+     * To indicate allowed network type change is requested by Testing purpose, should be default to
+     * {@link #getAllNetworkTypesBitmask} when done testing.
+     *
+     * @hide
+     */
+    public static final int ALLOWED_NETWORK_TYPES_REASON_TEST = 4;
+
+    /**
      * Set the allowed network types of the device and provide the reason triggering the allowed
      * network change.
      * <p>Requires permission: {@link android.Manifest.permission#MODIFY_PHONE_STATE} or
@@ -10118,6 +10128,8 @@
             case TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_CARRIER:
             case TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_ENABLE_2G:
                 return true;
+            case TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_TEST:
+                return TelephonyUtils.IS_DEBUGGABLE;
         }
         return false;
     }
diff --git a/telephony/java/android/telephony/satellite/SatelliteManager.java b/telephony/java/android/telephony/satellite/SatelliteManager.java
index 6ef953c..83fc0dc 100644
--- a/telephony/java/android/telephony/satellite/SatelliteManager.java
+++ b/telephony/java/android/telephony/satellite/SatelliteManager.java
@@ -546,6 +546,16 @@
     public static final int EMERGENCY_CALL_TO_SATELLITE_HANDOVER_TYPE_T911 = 2;
 
     /**
+     * This intent will be broadcasted if there are any change to list of subscriber informations.
+     * This intent will be sent only to the app with component defined in
+     * config_satellite_carrier_roaming_esos_provisioned_class and package defined in
+     * config_satellite_gateway_service_package
+     * @hide
+     */
+    public static final String ACTION_SATELLITE_SUBSCRIBER_ID_LIST_CHANGED =
+            "android.telephony.action.ACTION_SATELLITE_SUBSCRIBER_ID_LIST_CHANGED";
+
+    /**
      * Request to enable or disable the satellite modem and demo mode.
      * If satellite modem and cellular modem cannot work concurrently,
      * then this will disable the cellular modem if satellite modem is enabled,
@@ -582,7 +592,7 @@
                                 () -> resultListener.accept(result)));
                     }
                 };
-                telephony.requestSatelliteEnabled(mSubId, attributes.isEnabled(),
+                telephony.requestSatelliteEnabled(attributes.isEnabled(),
                         attributes.isDemoMode(), attributes.isEmergencyMode(), errorCallback);
             } else {
                 Rlog.e(TAG, "requestEnabled() invalid telephony");
@@ -640,7 +650,7 @@
                         }
                     }
                 };
-                telephony.requestIsSatelliteEnabled(mSubId, receiver);
+                telephony.requestIsSatelliteEnabled(receiver);
             } else {
                 loge("requestIsEnabled() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -697,7 +707,7 @@
                         }
                     }
                 };
-                telephony.requestIsDemoModeEnabled(mSubId, receiver);
+                telephony.requestIsDemoModeEnabled(receiver);
             } else {
                 loge("requestIsDemoModeEnabled() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -754,7 +764,7 @@
                         }
                     }
                 };
-                telephony.requestIsEmergencyModeEnabled(mSubId, receiver);
+                telephony.requestIsEmergencyModeEnabled(receiver);
             } else {
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
                         new SatelliteException(SATELLITE_RESULT_ILLEGAL_STATE))));
@@ -811,7 +821,7 @@
                         }
                     }
                 };
-                telephony.requestIsSatelliteSupported(mSubId, receiver);
+                telephony.requestIsSatelliteSupported(receiver);
             } else {
                 loge("requestIsSupported() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -868,7 +878,7 @@
                         }
                     }
                 };
-                telephony.requestSatelliteCapabilities(mSubId, receiver);
+                telephony.requestSatelliteCapabilities(receiver);
             } else {
                 loge("requestCapabilities() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -1194,8 +1204,7 @@
                             }
                         };
                 sSatelliteTransmissionUpdateCallbackMap.put(callback, internalCallback);
-                telephony.startSatelliteTransmissionUpdates(mSubId, errorCallback,
-                        internalCallback);
+                telephony.startSatelliteTransmissionUpdates(errorCallback, internalCallback);
             } else {
                 loge("startTransmissionUpdates() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(
@@ -1245,8 +1254,7 @@
                                     () -> resultListener.accept(result)));
                         }
                     };
-                    telephony.stopSatelliteTransmissionUpdates(mSubId, errorCallback,
-                            internalCallback);
+                    telephony.stopSatelliteTransmissionUpdates(errorCallback, internalCallback);
                     // TODO: Notify SmsHandler that pointing UI stopped
                 } else {
                     loge("stopSatelliteTransmissionUpdates: No internal callback.");
@@ -1302,7 +1310,7 @@
                                 () -> resultListener.accept(result)));
                     }
                 };
-                cancelRemote = telephony.provisionSatelliteService(mSubId, token, provisionData,
+                cancelRemote = telephony.provisionSatelliteService(token, provisionData,
                         errorCallback);
             } else {
                 loge("provisionService() invalid telephony");
@@ -1354,7 +1362,7 @@
                                 () -> resultListener.accept(result)));
                     }
                 };
-                telephony.deprovisionSatelliteService(mSubId, token, errorCallback);
+                telephony.deprovisionSatelliteService(token, errorCallback);
             } else {
                 loge("deprovisionService() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(
@@ -1409,8 +1417,7 @@
                             }
                         };
                 sSatelliteProvisionStateCallbackMap.put(callback, internalCallback);
-                return telephony.registerForSatelliteProvisionStateChanged(
-                        mSubId, internalCallback);
+                return telephony.registerForSatelliteProvisionStateChanged(internalCallback);
             } else {
                 throw new IllegalStateException("telephony service is null.");
             }
@@ -1443,7 +1450,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForSatelliteProvisionStateChanged(mSubId, internalCallback);
+                    telephony.unregisterForSatelliteProvisionStateChanged(internalCallback);
                 } else {
                     loge("unregisterForProvisionStateChanged: No internal callback.");
                 }
@@ -1500,7 +1507,7 @@
                         }
                     }
                 };
-                telephony.requestIsSatelliteProvisioned(mSubId, receiver);
+                telephony.requestIsSatelliteProvisioned(receiver);
             } else {
                 loge("requestIsProvisioned() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -1550,7 +1557,7 @@
                     }
                 };
                 sSatelliteModemStateCallbackMap.put(callback, internalCallback);
-                return telephony.registerForSatelliteModemStateChanged(mSubId, internalCallback);
+                return telephony.registerForSatelliteModemStateChanged(internalCallback);
             } else {
                 throw new IllegalStateException("telephony service is null.");
             }
@@ -1583,7 +1590,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForModemStateChanged(mSubId, internalCallback);
+                    telephony.unregisterForModemStateChanged(internalCallback);
                 } else {
                     loge("unregisterForModemStateChanged: No internal callback.");
                 }
@@ -1646,7 +1653,7 @@
                             }
                         };
                 sSatelliteDatagramCallbackMap.put(callback, internalCallback);
-                return telephony.registerForIncomingDatagram(mSubId, internalCallback);
+                return telephony.registerForIncomingDatagram(internalCallback);
             } else {
                 throw new IllegalStateException("telephony service is null.");
             }
@@ -1678,7 +1685,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForIncomingDatagram(mSubId, internalCallback);
+                    telephony.unregisterForIncomingDatagram(internalCallback);
                 } else {
                     loge("unregisterForIncomingDatagram: No internal callback.");
                 }
@@ -1722,7 +1729,7 @@
                                 () -> resultListener.accept(result)));
                     }
                 };
-                telephony.pollPendingDatagrams(mSubId, internalCallback);
+                telephony.pollPendingDatagrams(internalCallback);
             } else {
                 loge("pollPendingDatagrams() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(
@@ -1780,7 +1787,7 @@
                                 () -> resultListener.accept(result)));
                     }
                 };
-                telephony.sendDatagram(mSubId, datagramType, datagram,
+                telephony.sendDatagram(datagramType, datagram,
                         needFullScreenPointingUI, internalCallback);
             } else {
                 loge("sendDatagram() invalid telephony");
@@ -1898,7 +1905,7 @@
                         }
                     }
                 };
-                telephony.requestTimeForNextSatelliteVisibility(mSubId, receiver);
+                telephony.requestTimeForNextSatelliteVisibility(receiver);
             } else {
                 loge("requestTimeForNextSatelliteVisibility() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -1929,7 +1936,7 @@
         try {
             ITelephony telephony = getITelephony();
             if (telephony != null) {
-                telephony.setDeviceAlignedWithSatellite(mSubId, isAligned);
+                telephony.setDeviceAlignedWithSatellite(isAligned);
             } else {
                 throw new IllegalStateException("telephony service is null.");
             }
@@ -2193,7 +2200,7 @@
                         }
                     }
                 };
-                telephony.requestNtnSignalStrength(mSubId, receiver);
+                telephony.requestNtnSignalStrength(receiver);
             } else {
                 loge("requestNtnSignalStrength() invalid telephony");
                 executor.execute(() -> Binder.withCleanCallingIdentity(() -> callback.onError(
@@ -2244,7 +2251,7 @@
                                                 ntnSignalStrength)));
                             }
                         };
-                telephony.registerForNtnSignalStrengthChanged(mSubId, internalCallback);
+                telephony.registerForNtnSignalStrengthChanged(internalCallback);
                 sNtnSignalStrengthCallbackMap.put(callback, internalCallback);
             } else {
                 throw new IllegalStateException("Telephony service is null.");
@@ -2284,7 +2291,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForNtnSignalStrengthChanged(mSubId, internalCallback);
+                    telephony.unregisterForNtnSignalStrengthChanged(internalCallback);
                 } else {
                     loge("unregisterForNtnSignalStrengthChanged: No internal callback.");
                     throw new IllegalArgumentException("callback is not valid");
@@ -2329,7 +2336,7 @@
                             }
                         };
                 sSatelliteCapabilitiesCallbackMap.put(callback, internalCallback);
-                return telephony.registerForCapabilitiesChanged(mSubId, internalCallback);
+                return telephony.registerForCapabilitiesChanged(internalCallback);
             } else {
                 throw new IllegalStateException("Telephony service is null.");
             }
@@ -2362,7 +2369,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForCapabilitiesChanged(mSubId, internalCallback);
+                    telephony.unregisterForCapabilitiesChanged(internalCallback);
                 } else {
                     loge("unregisterForCapabilitiesChanged: No internal callback.");
                 }
@@ -2438,7 +2445,7 @@
                         };
                 sSatelliteSupportedStateCallbackMap.put(callback, internalCallback);
                 return telephony.registerForSatelliteSupportedStateChanged(
-                        mSubId, internalCallback);
+                        internalCallback);
             } else {
                 throw new IllegalStateException("telephony service is null.");
             }
@@ -2473,7 +2480,7 @@
             ITelephony telephony = getITelephony();
             if (telephony != null) {
                 if (internalCallback != null) {
-                    telephony.unregisterForSatelliteSupportedStateChanged(mSubId, internalCallback);
+                    telephony.unregisterForSatelliteSupportedStateChanged(internalCallback);
                 } else {
                     loge("unregisterForSupportedStateChanged: No internal callback.");
                 }
diff --git a/telephony/java/com/android/internal/telephony/ITelephony.aidl b/telephony/java/com/android/internal/telephony/ITelephony.aidl
index 0c5f30f..e852e6b 100644
--- a/telephony/java/com/android/internal/telephony/ITelephony.aidl
+++ b/telephony/java/com/android/internal/telephony/ITelephony.aidl
@@ -2743,7 +2743,6 @@
     /**
      * Request to enable or disable the satellite modem.
      *
-     * @param subId The subId of the subscription to enable or disable the satellite modem for.
      * @param enableSatellite True to enable the satellite modem and false to disable.
      * @param enableDemoMode True if demo mode is enabled and false otherwise. When
      *                       disabling satellite, {@code enableDemoMode} is always considered as
@@ -2755,93 +2754,83 @@
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestSatelliteEnabled(int subId, boolean enableSatellite, boolean enableDemoMode,
+    void requestSatelliteEnabled(boolean enableSatellite, boolean enableDemoMode,
             boolean isEmergency, in IIntegerConsumer callback);
 
     /**
      * Request to get whether the satellite modem is enabled.
      *
-     * @param subId The subId of the subscription to request whether satellite is enabled for.
      * @param receiver Result receiver to get the error code of the request and whether the
      *                 satellite modem is enabled.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestIsSatelliteEnabled(int subId, in ResultReceiver receiver);
+    void requestIsSatelliteEnabled(in ResultReceiver receiver);
 
     /**
      * Request to get whether the satellite service demo mode is enabled.
      *
-     * @param subId The subId of the subscription to request whether the satellite demo mode is
-     *              enabled for.
      * @param receiver Result receiver to get the error code of the request and whether the
      *                 satellite demo mode is enabled.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestIsDemoModeEnabled(int subId, in ResultReceiver receiver);
+    void requestIsDemoModeEnabled(in ResultReceiver receiver);
 
     /**
      * Request to get whether the satellite service is enabled with emergency mode.
      *
-     * @param subId The subId of the subscription to request whether the satellite demo mode is
-     *              enabled for.
      * @param receiver Result receiver to get the error code of the request and whether the
      *                 satellite is enabled with emergency mode.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestIsEmergencyModeEnabled(int subId, in ResultReceiver receiver);
+    void requestIsEmergencyModeEnabled(in ResultReceiver receiver);
 
     /**
      * Request to get whether the satellite service is supported on the device.
      *
-     * @param subId The subId of the subscription to check whether satellite is supported for.
      * @param receiver Result receiver to get the error code of the request and whether the
      *                 satellite service is supported on the device.
      */
-    void requestIsSatelliteSupported(int subId, in ResultReceiver receiver);
+    void requestIsSatelliteSupported(in ResultReceiver receiver);
 
     /**
      * Request to get the capabilities of the satellite service.
      *
-     * @param subId The subId of the subscription to get the capabilities for.
      * @param receiver Result receiver to get the error code of the request and the requested
      *                 capabilities of the satellite service.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestSatelliteCapabilities(int subId, in ResultReceiver receiver);
+    void requestSatelliteCapabilities(in ResultReceiver receiver);
 
     /**
      * Start receiving satellite transmission updates.
      *
-     * @param subId The subId of the subscription to stop satellite transmission updates for.
      * @param resultCallback The callback to get the result of the request.
      * @param callback The callback to handle transmission updates.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void startSatelliteTransmissionUpdates(int subId, in IIntegerConsumer resultCallback,
+    void startSatelliteTransmissionUpdates(in IIntegerConsumer resultCallback,
             in ISatelliteTransmissionUpdateCallback callback);
 
     /**
      * Stop receiving satellite transmission updates.
      *
-     * @param subId The subId of the subscritpion to stop satellite transmission updates for.
      * @param resultCallback The callback to get the result of the request.
      * @param callback The callback that was passed to startSatelliteTransmissionUpdates.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void stopSatelliteTransmissionUpdates(int subId, in IIntegerConsumer resultCallback,
+    void stopSatelliteTransmissionUpdates(in IIntegerConsumer resultCallback,
             in ISatelliteTransmissionUpdateCallback callback);
 
     /**
      * Register the subscription with a satellite provider.
      * This is needed to register the subscription if the provider allows dynamic registration.
      *
-     * @param subId The subId of the subscription to be provisioned.
      * @param token The token to be used as a unique identifier for provisioning with satellite
      *              gateway.
      * @provisionData Data from the provisioning app that can be used by provisioning server
@@ -2851,7 +2840,7 @@
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    ICancellationSignal provisionSatelliteService(int subId, in String token,
+    ICancellationSignal provisionSatelliteService(in String token,
             in byte[] provisionData, in IIntegerConsumer callback);
 
     /**
@@ -2861,110 +2850,99 @@
      * {@link SatelliteCallback.SatelliteProvisionStateListener#onSatelliteProvisionStateChanged}
      * should report as deprovisioned.
      *
-     * @param subId The subId of the subscription to be deprovisioned.
      * @param token The token of the device/subscription to be deprovisioned.
      * @param callback The callback to get the result of the request.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void deprovisionSatelliteService(int subId, in String token, in IIntegerConsumer callback);
+    void deprovisionSatelliteService(in String token, in IIntegerConsumer callback);
 
     /**
      * Registers for provision state changed from satellite modem.
      *
-     * @param subId The subId of the subscription to register for provision state changed.
      * @param callback The callback to handle the satellite provision state changed event.
      *
      * @return The {@link SatelliteError} result of the operation.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    int registerForSatelliteProvisionStateChanged(int subId,
-            in ISatelliteProvisionStateCallback callback);
+    int registerForSatelliteProvisionStateChanged(in ISatelliteProvisionStateCallback callback);
 
     /**
      * Unregisters for provision state changed from satellite modem.
      * If callback was not registered before, the request will be ignored.
      *
-     * @param subId The subId of the subscription to unregister for provision state changed.
      * @param callback The callback that was passed to registerForSatelliteProvisionStateChanged.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForSatelliteProvisionStateChanged(int subId,
+    void unregisterForSatelliteProvisionStateChanged(
             in ISatelliteProvisionStateCallback callback);
 
     /**
      * Request to get whether the device is provisioned with a satellite provider.
      *
-     * @param subId The subId of the subscription to get whether the device is provisioned for.
      * @param receiver Result receiver to get the error code of the request and whether the
      *                 device is provisioned with a satellite provider.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestIsSatelliteProvisioned(int subId, in ResultReceiver receiver);
+    void requestIsSatelliteProvisioned(in ResultReceiver receiver);
 
     /**
      * Registers for modem state changed from satellite modem.
      *
-     * @param subId The subId of the subscription to register for satellite modem state changed.
      * @param callback The callback to handle the satellite modem state changed event.
      *
      * @return The {@link SatelliteError} result of the operation.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    int registerForSatelliteModemStateChanged(int subId, ISatelliteModemStateCallback callback);
+    int registerForSatelliteModemStateChanged(ISatelliteModemStateCallback callback);
 
     /**
      * Unregisters for modem state changed from satellite modem.
      * If callback was not registered before, the request will be ignored.
      *
-     * @param subId The subId of the subscription to unregister for satellite modem state changed.
      * @param callback The callback that was passed to registerForSatelliteStateChanged.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForModemStateChanged(int subId, ISatelliteModemStateCallback callback);
+    void unregisterForModemStateChanged(ISatelliteModemStateCallback callback);
 
    /**
      * Register to receive incoming datagrams over satellite.
      *
-     * @param subId The subId of the subscription to register for incoming satellite datagrams.
      * @param callback The callback to handle the incoming datagrams.
      *
      * @return The {@link SatelliteError} result of the operation.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    int registerForIncomingDatagram(int subId, ISatelliteDatagramCallback callback);
+    int registerForIncomingDatagram(ISatelliteDatagramCallback callback);
 
    /**
      * Unregister to stop receiving incoming datagrams over satellite.
      * If callback was not registered before, the request will be ignored.
      *
-     * @param subId The subId of the subscription to unregister for incoming satellite datagrams.
      * @param callback The callback that was passed to registerForIncomingDatagram.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForIncomingDatagram(int subId, ISatelliteDatagramCallback callback);
+    void unregisterForIncomingDatagram(ISatelliteDatagramCallback callback);
 
    /**
     * Poll pending satellite datagrams over satellite.
     *
-    * @param subId The subId of the subscription used for receiving datagrams.
     * @param callback The callback to get the result of the request.
     */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
                 + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void pollPendingDatagrams(int subId, IIntegerConsumer callback);
+    void pollPendingDatagrams(IIntegerConsumer callback);
 
    /**
     * Send datagram over satellite.
     *
-    * @param subId The subId of the subscription to send satellite datagrams for.
     * @param datagramType Type of datagram.
     * @param datagram Datagram to send over satellite.
     * @param needFullScreenPointingUI this is used to indicate pointingUI app to open in
@@ -2973,7 +2951,7 @@
     */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void sendDatagram(int subId, int datagramType, in SatelliteDatagram datagram,
+    void sendDatagram(int datagramType, in SatelliteDatagram datagram,
             in boolean needFullScreenPointingUI, IIntegerConsumer callback);
 
     /**
@@ -2991,13 +2969,12 @@
     /**
      * Request to get the time after which the satellite will be visible.
      *
-     * @param subId The subId to get the time after which the satellite will be visible for.
      * @param receiver Result receiver to get the error code of the request and the requested
      *                 time after which the satellite will be visible.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestTimeForNextSatelliteVisibility(int subId, in ResultReceiver receiver);
+    void requestTimeForNextSatelliteVisibility(in ResultReceiver receiver);
 
     /**
      * Inform whether the device is aligned with the satellite within in margin for demo mode.
@@ -3007,7 +2984,7 @@
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void setDeviceAlignedWithSatellite(int subId, boolean isAligned);
+    void setDeviceAlignedWithSatellite(boolean isAligned);
 
     /**
      * This API can be used by only CTS to update satellite vendor service package name.
@@ -3163,20 +3140,18 @@
     /**
      * Request to get the signal strength of the satellite connection.
      *
-     * @param subId The subId of the subscription to request for.
      * @param receiver Result receiver to get the error code of the request and the current signal
      * strength of the satellite connection.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void requestNtnSignalStrength(int subId, in ResultReceiver receiver);
+    void requestNtnSignalStrength(in ResultReceiver receiver);
 
     /**
      * Registers for NTN signal strength changed from satellite modem. If the registration operation
      * is not successful, a {@link SatelliteException} that contains {@link SatelliteResult} will be
      * thrown.
      *
-     * @param subId The subId of the subscription to request for.
      * @param callback The callback to handle the NTN signal strength changed event. If the
      * operation is successful, {@link NtnSignalStrengthCallback#onNtnSignalStrengthChanged(
      * NtnSignalStrength)} will return an instance of {@link NtnSignalStrength} with a value of
@@ -3185,30 +3160,27 @@
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void registerForNtnSignalStrengthChanged(int subId,
-            in INtnSignalStrengthCallback callback);
+    void registerForNtnSignalStrengthChanged(in INtnSignalStrengthCallback callback);
 
     /**
      * Unregisters for NTN signal strength changed from satellite modem.
      * If callback was not registered before, the request will be ignored.
      *
-     * @param subId The subId of the subscription to unregister for provision state changed.
      * @param callback The callback that was passed to
      * {@link #registerForNtnSignalStrengthChanged(Executor, NtnSignalStrengthCallback)}.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForNtnSignalStrengthChanged(int subId, in INtnSignalStrengthCallback callback);
+    void unregisterForNtnSignalStrengthChanged(in INtnSignalStrengthCallback callback);
 
     /**
      * Registers for satellite capabilities change event from the satellite service.
      *
-     * @param executor The executor on which the callback will be called.
      * @param callback The callback to handle the satellite capabilities changed event.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    int registerForCapabilitiesChanged(int subId, in ISatelliteCapabilitiesCallback callback);
+    int registerForCapabilitiesChanged(in ISatelliteCapabilitiesCallback callback);
 
     /**
      * Unregisters for satellite capabilities change event from the satellite service.
@@ -3219,8 +3191,7 @@
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForCapabilitiesChanged(int subId,
-            in ISatelliteCapabilitiesCallback callback);
+    void unregisterForCapabilitiesChanged(in ISatelliteCapabilitiesCallback callback);
 
     /**
      * This API can be used by only CTS to override the cached value for the device overlay config
@@ -3329,26 +3300,24 @@
     /**
      * Registers for supported state changed from satellite modem.
      *
-     * @param subId The subId of the subscription to register for supported state changed.
      * @param callback The callback to handle the satellite supported state changed event.
      *
      * @return The {@link SatelliteError} result of the operation.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    int registerForSatelliteSupportedStateChanged(int subId,
+    int registerForSatelliteSupportedStateChanged(
             in ISatelliteSupportedStateCallback callback);
 
     /**
      * Unregisters for supported state changed from satellite modem.
      * If callback was not registered before, the request will be ignored.
      *
-     * @param subId The subId of the subscription to unregister for supported state changed.
      * @param callback The callback that was passed to registerForSatelliteSupportedStateChanged.
      */
     @JavaPassthrough(annotation="@android.annotation.RequiresPermission("
             + "android.Manifest.permission.SATELLITE_COMMUNICATION)")
-    void unregisterForSatelliteSupportedStateChanged(int subId,
+    void unregisterForSatelliteSupportedStateChanged(
             in ISatelliteSupportedStateCallback callback);
 
     /**
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 379b45c..c3e1a1f 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
@@ -24,6 +24,7 @@
 import android.tools.flicker.legacy.LegacyFlickerTest
 import android.tools.flicker.legacy.LegacyFlickerTestFactory
 import android.tools.traces.parsers.toFlickerComponent
+import androidx.test.filters.FlakyTest
 import com.android.server.wm.flicker.activityembedding.ActivityEmbeddingTestBase
 import com.android.server.wm.flicker.helpers.ActivityEmbeddingAppHelper
 import com.android.server.wm.flicker.testapp.ActivityOptions
@@ -177,6 +178,13 @@
     @Ignore("Not applicable to this CUJ.")
     override fun visibleLayersShownMoreThanOneConsecutiveEntry() {}
 
+    @FlakyTest(bugId = 342596801)
+    override fun entireScreenCovered() = super.entireScreenCovered()
+
+    @FlakyTest(bugId = 342596801)
+    override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
+        super.visibleWindowsShownMoreThanOneConsecutiveEntry()
+
     companion object {
         /** {@inheritDoc} */
         private var startDisplayBounds = Rect()
diff --git a/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/OpenAppFromLockscreenViaIntentTest.kt b/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/OpenAppFromLockscreenViaIntentTest.kt
index adeba72..6e6a327 100644
--- a/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/OpenAppFromLockscreenViaIntentTest.kt
+++ b/tests/FlickerTests/AppLaunch/src/com/android/server/wm/flicker/launch/OpenAppFromLockscreenViaIntentTest.kt
@@ -196,7 +196,7 @@
         super.appLayerBecomesVisible()
     }
 
-    @Presubmit
+    @FlakyTest(bugId = 338296297)
     @Test
     override fun visibleWindowsShownMoreThanOneConsecutiveEntry() =
         super.visibleWindowsShownMoreThanOneConsecutiveEntry()
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt
index 753cb1f..3f6a0bf 100644
--- a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/DesktopModeAppHelper.kt
@@ -48,6 +48,13 @@
         RIGHT_BOTTOM
     }
 
+    enum class Edges {
+        LEFT,
+        RIGHT,
+        TOP,
+        BOTTOM
+    }
+
     /** Wait for an app moved to desktop to finish its transition. */
     private fun waitForAppToMoveToDesktop(wmHelper: WindowManagerStateHelper) {
         wmHelper
@@ -124,7 +131,8 @@
 
         val displayRect = getDisplayRect(wmHelper)
         val insets = getWindowInsets(
-            context, WindowInsets.Type.statusBars() or WindowInsets.Type.navigationBars())
+            context, WindowInsets.Type.statusBars() or WindowInsets.Type.navigationBars()
+        )
         displayRect.inset(insets)
 
         val expectedWidth = displayRect.width() / 2
@@ -187,6 +195,40 @@
         dragWindow(startX, startY, endX, endY, wmHelper, device)
     }
 
+    /** Resize a desktop app from its edges. */
+    fun edgeResize(
+        wmHelper: WindowManagerStateHelper,
+        motionEvent: MotionEventHelper,
+        edge: Edges
+    ) {
+        val windowRect = wmHelper.getWindowRegion(innerHelper).bounds
+        val (startX, startY) = getStartCoordinatesForEdgeResize(windowRect, edge)
+        val verticalChange = when (edge) {
+            Edges.LEFT -> 0
+            Edges.RIGHT -> 0
+            Edges.TOP -> -100
+            Edges.BOTTOM -> 100
+        }
+        val horizontalChange = when (edge) {
+            Edges.LEFT -> -100
+            Edges.RIGHT -> 100
+            Edges.TOP -> 0
+            Edges.BOTTOM -> 0
+        }
+
+        // The position we want to drag to
+        val endY = startY + verticalChange
+        val endX = startX + horizontalChange
+
+        motionEvent.actionDown(startX, startY)
+        motionEvent.actionMove(startX, startY, endX, endY, /* steps= */100)
+        motionEvent.actionUp(endX, endY)
+        wmHelper
+            .StateSyncBuilder()
+            .withAppTransitionIdle()
+            .waitForAndVerify()
+    }
+
     /** Drag a window from a source coordinate to a destination coordinate. */
     fun dragWindow(
         startX: Int, startY: Int,
@@ -237,6 +279,18 @@
         }
     }
 
+    private fun getStartCoordinatesForEdgeResize(
+        windowRect: Rect,
+        edge: Edges
+    ): Pair<Int, Int> {
+        return when (edge) {
+            Edges.LEFT -> Pair(windowRect.left, windowRect.bottom / 2)
+            Edges.RIGHT -> Pair(windowRect.right, windowRect.bottom / 2)
+            Edges.TOP -> Pair(windowRect.right / 2, windowRect.top)
+            Edges.BOTTOM -> Pair(windowRect.right / 2, windowRect.bottom)
+        }
+    }
+
     /** Exit desktop mode by dragging the app handle to the top drag zone. */
     fun exitDesktopWithDragToTopDragZone(
         wmHelper: WindowManagerStateHelper,
diff --git a/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/MotionEventHelper.kt b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/MotionEventHelper.kt
new file mode 100644
index 0000000..0835398
--- /dev/null
+++ b/tests/FlickerTests/test-apps/app-helpers/src/com/android/server/wm/flicker/helpers/MotionEventHelper.kt
@@ -0,0 +1,116 @@
+/*
+ * 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.wm.flicker.helpers
+
+import android.app.Instrumentation
+import android.os.SystemClock
+import android.view.ContentInfo.Source
+import android.view.InputDevice.SOURCE_MOUSE
+import android.view.InputDevice.SOURCE_STYLUS
+import android.view.MotionEvent
+import android.view.MotionEvent.ACTION_DOWN
+import android.view.MotionEvent.ACTION_MOVE
+import android.view.MotionEvent.ACTION_UP
+import android.view.MotionEvent.TOOL_TYPE_FINGER
+import android.view.MotionEvent.TOOL_TYPE_MOUSE
+import android.view.MotionEvent.TOOL_TYPE_STYLUS
+import android.view.MotionEvent.ToolType
+
+/**
+ * Helper class for injecting a custom motion event and performing some actions. This is used for
+ * instrumenting input injections like stylus, mouse and touchpad.
+ */
+class MotionEventHelper(
+    private val instr: Instrumentation,
+    private val inputMethod: InputMethod
+) {
+    enum class InputMethod(@ToolType val toolType: Int, @Source val source: Int) {
+        STYLUS(TOOL_TYPE_STYLUS, SOURCE_STYLUS),
+        MOUSE(TOOL_TYPE_MOUSE, SOURCE_MOUSE),
+        TOUCHPAD(TOOL_TYPE_FINGER, SOURCE_MOUSE)
+    }
+
+    fun actionDown(x: Int, y: Int) {
+        injectMotionEvent(ACTION_DOWN, x, y)
+    }
+
+    fun actionUp(x: Int, y: Int) {
+        injectMotionEvent(ACTION_UP, x, y)
+    }
+
+    fun actionMove(startX: Int, startY: Int, endX: Int, endY: Int, steps: Int) {
+        val incrementX = (endX - startX).toFloat() / (steps - 1)
+        val incrementY = (endY - startY).toFloat() / (steps - 1)
+
+        for (i in 0..steps) {
+            val time = SystemClock.uptimeMillis()
+            val x = startX + incrementX * i
+            val y = startY + incrementY * i
+
+            val moveEvent = getMotionEvent(time, time, ACTION_MOVE, x, y)
+            injectMotionEvent(moveEvent)
+        }
+    }
+
+    private fun injectMotionEvent(action: Int, x: Int, y: Int): MotionEvent {
+        val eventTime = SystemClock.uptimeMillis()
+        val event = getMotionEvent(eventTime, eventTime, action, x.toFloat(), y.toFloat())
+        injectMotionEvent(event)
+        return event
+    }
+
+    private fun injectMotionEvent(event: MotionEvent) {
+        instr.uiAutomation.injectInputEvent(event, true, false)
+    }
+
+    private fun getMotionEvent(
+        downTime: Long,
+        eventTime: Long,
+        action: Int,
+        x: Float,
+        y: Float,
+    ): MotionEvent {
+        val properties = MotionEvent.PointerProperties.createArray(1)
+        properties[0].toolType = inputMethod.toolType
+        properties[0].id = 1
+
+        val coords = MotionEvent.PointerCoords.createArray(1)
+        coords[0].x = x
+        coords[0].y = y
+        coords[0].pressure = 1f
+
+        val event =
+            MotionEvent.obtain(
+                downTime,
+                eventTime,
+                action,
+                /* pointerCount= */ 1,
+                properties,
+                coords,
+                /* metaState= */ 0,
+                /* buttonState= */ 0,
+                /* xPrecision = */ 1f,
+                /* yPrecision = */ 1f,
+                /* deviceId = */ 0,
+                /* edgeFlags = */ 0,
+                inputMethod.source,
+                /* flags = */ 0
+            )
+        event.displayId = 0
+        return event
+    }
+}
\ No newline at end of file
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 43fd57b..931e4f8 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
@@ -269,9 +269,23 @@
     /** Expand the PIP window back to full screen via intent and wait until the app is visible */
     fun exitPipToFullScreenViaIntent(wmHelper: WindowManagerStateHelper) = launchViaIntent(wmHelper)
 
-    fun changeAspectRatio() {
+    fun changeAspectRatio(wmHelper: WindowManagerStateHelper) {
         val intent = Intent("com.android.wm.shell.flicker.testapp.ASPECT_RATIO")
         context.sendBroadcast(intent)
+        // Wait on WMHelper on size change upon aspect ratio change
+        val windowRect = getWindowRect(wmHelper)
+        wmHelper
+            .StateSyncBuilder()
+            .add("pipAspectRatioChanged") {
+                val pipAppWindow =
+                    it.wmState.visibleWindows.firstOrNull { window ->
+                        this.windowMatchesAnyOf(window)
+                    }
+                        ?: return@add false
+                val pipRegion = pipAppWindow.frameRegion
+                return@add pipRegion != Region(windowRect)
+            }
+            .waitForAndVerify()
     }
 
     fun clickEnterPipButton(wmHelper: WindowManagerStateHelper) {
diff --git a/tests/Input/Android.bp b/tests/Input/Android.bp
index 06c2651..65398a2 100644
--- a/tests/Input/Android.bp
+++ b/tests/Input/Android.bp
@@ -40,7 +40,7 @@
         "frameworks-base-testutils",
         "hamcrest-library",
         "kotlin-test",
-        "mockito-target-minus-junit4",
+        "mockito-target-extended-minus-junit4",
         "platform-test-annotations",
         "platform-screenshot-diff-core",
         "services.core.unboosted",
diff --git a/tests/Input/assets/testPointerFillStyle.png b/tests/Input/assets/testPointerFillStyle.png
index b2354f8..297244f 100644
--- a/tests/Input/assets/testPointerFillStyle.png
+++ b/tests/Input/assets/testPointerFillStyle.png
Binary files differ
diff --git a/tests/Input/src/com/android/server/input/debug/TouchpadDebugViewTest.java b/tests/Input/src/com/android/server/input/debug/TouchpadDebugViewTest.java
new file mode 100644
index 0000000..ad0ef1b
--- /dev/null
+++ b/tests/Input/src/com/android/server/input/debug/TouchpadDebugViewTest.java
@@ -0,0 +1,292 @@
+/*
+ * Copyright 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.input.debug;
+
+import static android.view.InputDevice.SOURCE_TOUCHSCREEN;
+
+import static org.junit.Assert.assertEquals;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.graphics.Rect;
+import android.testing.TestableContext;
+import android.view.MotionEvent;
+import android.view.View;
+import android.view.ViewConfiguration;
+import android.view.WindowInsets;
+import android.view.WindowManager;
+import android.view.WindowMetrics;
+
+import androidx.test.platform.app.InstrumentationRegistry;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.cts.input.MotionEventBuilder;
+import com.android.cts.input.PointerBuilder;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Build/Install/Run:
+ * atest TouchpadDebugViewTest
+ */
+@RunWith(AndroidJUnit4.class)
+public class TouchpadDebugViewTest {
+    private static final int TOUCHPAD_DEVICE_ID = 6;
+
+    private TouchpadDebugView mTouchpadDebugView;
+    private WindowManager.LayoutParams mWindowLayoutParams;
+
+    @Mock
+    WindowManager mWindowManager;
+
+    Rect mWindowBounds;
+    WindowMetrics mWindowMetrics;
+    TestableContext mTestableContext;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        Context context = InstrumentationRegistry.getInstrumentation().getContext();
+        mTestableContext = new TestableContext(context);
+
+        mTestableContext.addMockSystemService(WindowManager.class, mWindowManager);
+
+        mWindowBounds = new Rect(0, 0, 2560, 1600);
+        mWindowMetrics = new WindowMetrics(mWindowBounds, new WindowInsets(mWindowBounds), 1.0f);
+
+        when(mWindowManager.getCurrentWindowMetrics()).thenReturn(mWindowMetrics);
+
+        mTouchpadDebugView = new TouchpadDebugView(mTestableContext, TOUCHPAD_DEVICE_ID);
+
+        mTouchpadDebugView.measure(
+                View.MeasureSpec.makeMeasureSpec(0, View.MeasureSpec.UNSPECIFIED),
+                View.MeasureSpec.makeMeasureSpec(0, View.MeasureSpec.UNSPECIFIED)
+        );
+
+        doAnswer(invocation -> {
+            mTouchpadDebugView.layout(0, 0, mTouchpadDebugView.getMeasuredWidth(),
+                    mTouchpadDebugView.getMeasuredHeight());
+            return null;
+        }).when(mWindowManager).addView(any(), any());
+
+        doAnswer(invocation -> {
+            mTouchpadDebugView.layout(0, 0, mTouchpadDebugView.getMeasuredWidth(),
+                    mTouchpadDebugView.getMeasuredHeight());
+            return null;
+        }).when(mWindowManager).updateViewLayout(any(), any());
+
+        mWindowLayoutParams = mTouchpadDebugView.getWindowLayoutParams();
+        mWindowLayoutParams.x = 20;
+        mWindowLayoutParams.y = 20;
+
+        mTouchpadDebugView.layout(0, 0, mTouchpadDebugView.getMeasuredWidth(),
+                mTouchpadDebugView.getMeasuredHeight());
+    }
+
+    @Test
+    public void testDragView() {
+        // Initial view position relative to screen.
+        int initialX = mWindowLayoutParams.x;
+        int initialY = mWindowLayoutParams.y;
+
+        float offsetX = ViewConfiguration.get(mTestableContext).getScaledTouchSlop() + 10;
+        float offsetY = ViewConfiguration.get(mTestableContext).getScaledTouchSlop() + 10;
+
+        // Simulate ACTION_DOWN event (initial touch).
+        MotionEvent actionDown = new MotionEventBuilder(MotionEvent.ACTION_DOWN, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(40f)
+                        .y(40f)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionDown);
+
+        verify(mWindowManager, times(0)).updateViewLayout(any(), any());
+
+        // Simulate ACTION_MOVE event (dragging to the right).
+        MotionEvent actionMove = new MotionEventBuilder(MotionEvent.ACTION_MOVE, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(40f + offsetX)
+                        .y(40f + offsetY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionMove);
+
+        ArgumentCaptor<WindowManager.LayoutParams> mWindowLayoutParamsCaptor =
+                ArgumentCaptor.forClass(WindowManager.LayoutParams.class);
+        verify(mWindowManager).updateViewLayout(any(), mWindowLayoutParamsCaptor.capture());
+
+        // Verify position after ACTION_MOVE
+        assertEquals(initialX + (long) offsetX, mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(initialY + (long) offsetY, mWindowLayoutParamsCaptor.getValue().y);
+
+        // Simulate ACTION_UP event (release touch).
+        MotionEvent actionUp = new MotionEventBuilder(MotionEvent.ACTION_UP, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(40f + offsetX)
+                        .y(40f + offsetY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionUp);
+
+        assertEquals(initialX + (long) offsetX, mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(initialY + (long) offsetY, mWindowLayoutParamsCaptor.getValue().y);
+    }
+
+    @Test
+    public void testDragViewOutOfBounds() {
+        int initialX = mWindowLayoutParams.x;
+        int initialY = mWindowLayoutParams.y;
+
+        MotionEvent actionDown = new MotionEventBuilder(MotionEvent.ACTION_DOWN, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX + 10f)
+                        .y(initialY + 10f)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionDown);
+
+        verify(mWindowManager, times(0)).updateViewLayout(any(), any());
+
+        // Simulate ACTION_MOVE event (dragging far to the right and bottom, beyond screen bounds)
+        MotionEvent actionMove = new MotionEventBuilder(MotionEvent.ACTION_MOVE, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(mWindowBounds.width() + mTouchpadDebugView.getWidth())
+                        .y(mWindowBounds.height() + mTouchpadDebugView.getHeight())
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionMove);
+
+        ArgumentCaptor<WindowManager.LayoutParams> mWindowLayoutParamsCaptor =
+                ArgumentCaptor.forClass(WindowManager.LayoutParams.class);
+        verify(mWindowManager).updateViewLayout(any(), mWindowLayoutParamsCaptor.capture());
+
+        // Verify the view has been clamped to the right and bottom edges of the screen
+        assertEquals(mWindowBounds.width() - mTouchpadDebugView.getWidth(),
+                mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(mWindowBounds.height() - mTouchpadDebugView.getHeight(),
+                mWindowLayoutParamsCaptor.getValue().y);
+
+        MotionEvent actionUp = new MotionEventBuilder(MotionEvent.ACTION_UP, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(mWindowBounds.width() + mTouchpadDebugView.getWidth())
+                        .y(mWindowBounds.height() + mTouchpadDebugView.getHeight())
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionUp);
+
+        // Verify the view has been clamped to the right and bottom edges of the screen
+        assertEquals(mWindowBounds.width() - mTouchpadDebugView.getWidth(),
+                mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(mWindowBounds.height() - mTouchpadDebugView.getHeight(),
+                mWindowLayoutParamsCaptor.getValue().y);
+    }
+
+    @Test
+    public void testSlopOffset() {
+        int initialX = mWindowLayoutParams.x;
+        int initialY = mWindowLayoutParams.y;
+
+        float offsetX = ViewConfiguration.get(mTestableContext).getScaledTouchSlop() / 2.0f;
+        float offsetY = -(ViewConfiguration.get(mTestableContext).getScaledTouchSlop() / 2.0f);
+
+        MotionEvent actionDown = new MotionEventBuilder(MotionEvent.ACTION_DOWN, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX)
+                        .y(initialY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionDown);
+
+        MotionEvent actionMove = new MotionEventBuilder(MotionEvent.ACTION_MOVE, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX + offsetX)
+                        .y(initialY + offsetY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionMove);
+
+        MotionEvent actionUp = new MotionEventBuilder(MotionEvent.ACTION_UP, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX)
+                        .y(initialY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionUp);
+
+        // In this case the updateViewLayout() method wouldn't be called because the drag
+        // distance hasn't exceeded the slop
+        verify(mWindowManager, times(0)).updateViewLayout(any(), any());
+    }
+
+    @Test
+    public void testViewReturnsToInitialPositionOnCancel() {
+        int initialX = mWindowLayoutParams.x;
+        int initialY = mWindowLayoutParams.y;
+
+        float offsetX = 50;
+        float offsetY = 50;
+
+        MotionEvent actionDown = new MotionEventBuilder(MotionEvent.ACTION_DOWN, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX)
+                        .y(initialY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionDown);
+
+        MotionEvent actionMove = new MotionEventBuilder(MotionEvent.ACTION_MOVE, SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX + offsetX)
+                        .y(initialY + offsetY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionMove);
+
+        ArgumentCaptor<WindowManager.LayoutParams> mWindowLayoutParamsCaptor =
+                ArgumentCaptor.forClass(WindowManager.LayoutParams.class);
+        verify(mWindowManager).updateViewLayout(any(), mWindowLayoutParamsCaptor.capture());
+
+        assertEquals(initialX + (long) offsetX, mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(initialY + (long) offsetY, mWindowLayoutParamsCaptor.getValue().y);
+
+        // Simulate ACTION_CANCEL event (canceling the touch event stream)
+        MotionEvent actionCancel = new MotionEventBuilder(MotionEvent.ACTION_CANCEL,
+                SOURCE_TOUCHSCREEN)
+                .pointer(new PointerBuilder(0, MotionEvent.TOOL_TYPE_FINGER)
+                        .x(initialX + offsetX)
+                        .y(initialY + offsetY)
+                )
+                .build();
+        mTouchpadDebugView.dispatchTouchEvent(actionCancel);
+
+        // Verify the view returns to its initial position
+        verify(mWindowManager, times(2)).updateViewLayout(any(),
+                mWindowLayoutParamsCaptor.capture());
+        assertEquals(initialX, mWindowLayoutParamsCaptor.getValue().x);
+        assertEquals(initialY, mWindowLayoutParamsCaptor.getValue().y);
+    }
+}
diff --git a/tools/aapt2/ResourceTable.cpp b/tools/aapt2/ResourceTable.cpp
index 7a4f40e..9751459 100644
--- a/tools/aapt2/ResourceTable.cpp
+++ b/tools/aapt2/ResourceTable.cpp
@@ -50,21 +50,21 @@
 
 template <typename T>
 bool less_than_struct_with_name(const std::unique_ptr<T>& lhs, StringPiece rhs) {
-  return lhs->name.compare(0, lhs->name.size(), rhs.data(), rhs.size()) < 0;
+  return lhs->name < rhs;
 }
 
 template <typename T>
 bool greater_than_struct_with_name(StringPiece lhs, const std::unique_ptr<T>& rhs) {
-  return rhs->name.compare(0, rhs->name.size(), lhs.data(), lhs.size()) > 0;
+  return rhs->name > lhs;
 }
 
 template <typename T>
 struct NameEqualRange {
   bool operator()(const std::unique_ptr<T>& lhs, StringPiece rhs) const {
-    return less_than_struct_with_name<T>(lhs, rhs);
+    return less_than_struct_with_name(lhs, rhs);
   }
   bool operator()(StringPiece lhs, const std::unique_ptr<T>& rhs) const {
-    return greater_than_struct_with_name<T>(lhs, rhs);
+    return greater_than_struct_with_name(lhs, rhs);
   }
 };
 
@@ -74,7 +74,7 @@
   if (lhs.id != rhs.second) {
     return lhs.id < rhs.second;
   }
-  return lhs.name.compare(0, lhs.name.size(), rhs.first.data(), rhs.first.size()) < 0;
+  return lhs.name < rhs.first;
 }
 
 template <typename T, typename Func, typename Elements>
@@ -90,14 +90,16 @@
   StringPiece product;
 };
 
-template <typename T>
-bool lt_config_key_ref(const T& lhs, const ConfigKey& rhs) {
-  int cmp = lhs->config.compare(*rhs.config);
-  if (cmp == 0) {
-    cmp = StringPiece(lhs->product).compare(rhs.product);
+struct lt_config_key_ref {
+  template <typename T>
+  bool operator()(const T& lhs, const ConfigKey& rhs) const noexcept {
+    int cmp = lhs->config.compare(*rhs.config);
+    if (cmp == 0) {
+      cmp = lhs->product.compare(rhs.product);
+    }
+    return cmp < 0;
   }
-  return cmp < 0;
-}
+};
 
 }  // namespace
 
@@ -159,10 +161,10 @@
 ResourceConfigValue* ResourceEntry::FindValue(const ConfigDescription& config,
                                               android::StringPiece product) {
   auto iter = std::lower_bound(values.begin(), values.end(), ConfigKey{&config, product},
-                               lt_config_key_ref<std::unique_ptr<ResourceConfigValue>>);
+                               lt_config_key_ref());
   if (iter != values.end()) {
     ResourceConfigValue* value = iter->get();
-    if (value->config == config && StringPiece(value->product) == product) {
+    if (value->config == config && value->product == product) {
       return value;
     }
   }
@@ -172,10 +174,10 @@
 const ResourceConfigValue* ResourceEntry::FindValue(const android::ConfigDescription& config,
                                                     android::StringPiece product) const {
   auto iter = std::lower_bound(values.begin(), values.end(), ConfigKey{&config, product},
-                               lt_config_key_ref<std::unique_ptr<ResourceConfigValue>>);
+                               lt_config_key_ref());
   if (iter != values.end()) {
     ResourceConfigValue* value = iter->get();
-    if (value->config == config && StringPiece(value->product) == product) {
+    if (value->config == config && value->product == product) {
       return value;
     }
   }
@@ -185,10 +187,10 @@
 ResourceConfigValue* ResourceEntry::FindOrCreateValue(const ConfigDescription& config,
                                                       StringPiece product) {
   auto iter = std::lower_bound(values.begin(), values.end(), ConfigKey{&config, product},
-                               lt_config_key_ref<std::unique_ptr<ResourceConfigValue>>);
+                               lt_config_key_ref());
   if (iter != values.end()) {
     ResourceConfigValue* value = iter->get();
-    if (value->config == config && StringPiece(value->product) == product) {
+    if (value->config == config && value->product == product) {
       return value;
     }
   }
@@ -199,36 +201,21 @@
 
 std::vector<ResourceConfigValue*> ResourceEntry::FindAllValues(const ConfigDescription& config) {
   std::vector<ResourceConfigValue*> results;
-
-  auto iter = values.begin();
+  auto iter =
+      std::lower_bound(values.begin(), values.end(), ConfigKey{&config, ""}, lt_config_key_ref());
   for (; iter != values.end(); ++iter) {
     ResourceConfigValue* value = iter->get();
-    if (value->config == config) {
-      results.push_back(value);
-      ++iter;
+    if (value->config != config) {
       break;
     }
-  }
-
-  for (; iter != values.end(); ++iter) {
-    ResourceConfigValue* value = iter->get();
-    if (value->config == config) {
-      results.push_back(value);
-    }
+    results.push_back(value);
   }
   return results;
 }
 
 bool ResourceEntry::HasDefaultValue() const {
-  const ConfigDescription& default_config = ConfigDescription::DefaultConfig();
-
   // The default config should be at the top of the list, since the list is sorted.
-  for (auto& config_value : values) {
-    if (config_value->config == default_config) {
-      return true;
-    }
-  }
-  return false;
+  return !values.empty() && values.front()->config == ConfigDescription::DefaultConfig();
 }
 
 ResourceTable::CollisionResult ResourceTable::ResolveFlagCollision(FlagStatus existing,
@@ -364,14 +351,14 @@
     if (found) {
       return &*it;
     }
-    return &*el.insert(it, std::forward<T>(value));
+    return &*el.insert(it, std::move(value));
   }
 };
 
 struct PackageViewComparer {
   bool operator()(const ResourceTablePackageView& lhs, const ResourceTablePackageView& rhs) {
     return less_than_struct_with_name_and_id<ResourceTablePackageView, uint8_t>(
-        lhs, std::make_pair(rhs.name, rhs.id));
+        lhs, std::tie(rhs.name, rhs.id));
   }
 };
 
@@ -384,7 +371,7 @@
 struct EntryViewComparer {
   bool operator()(const ResourceTableEntryView& lhs, const ResourceTableEntryView& rhs) {
     return less_than_struct_with_name_and_id<ResourceTableEntryView, uint16_t>(
-        lhs, std::make_pair(rhs.name, rhs.id));
+        lhs, std::tie(rhs.name, rhs.id));
   }
 };
 
@@ -429,10 +416,10 @@
 const ResourceConfigValue* ResourceTableEntryView::FindValue(const ConfigDescription& config,
                                                              android::StringPiece product) const {
   auto iter = std::lower_bound(values.begin(), values.end(), ConfigKey{&config, product},
-                               lt_config_key_ref<const ResourceConfigValue*>);
+                               lt_config_key_ref());
   if (iter != values.end()) {
     const ResourceConfigValue* value = *iter;
-    if (value->config == config && StringPiece(value->product) == product) {
+    if (value->config == config && value->product == product) {
       return value;
     }
   }
@@ -615,11 +602,15 @@
         result = ResolveValueCollision(config_value->value.get(), res.value.get());
       }
       switch (result) {
-        case CollisionResult::kKeepBoth:
+        case CollisionResult::kKeepBoth: {
           // Insert the value ignoring for duplicate configurations
-          entry->values.push_back(util::make_unique<ResourceConfigValue>(res.config, res.product));
-          entry->values.back()->value = std::move(res.value);
+          auto it = entry->values.insert(
+              std::lower_bound(entry->values.begin(), entry->values.end(),
+                               ConfigKey{&res.config, res.product}, lt_config_key_ref()),
+              util::make_unique<ResourceConfigValue>(res.config, res.product));
+          (*it)->value = std::move(res.value);
           break;
+        }
 
         case CollisionResult::kTakeNew:
           // Take the incoming value.
diff --git a/tools/hoststubgen/hoststubgen/Android.bp b/tools/hoststubgen/hoststubgen/Android.bp
index 682adbc..ea77b8d 100644
--- a/tools/hoststubgen/hoststubgen/Android.bp
+++ b/tools/hoststubgen/hoststubgen/Android.bp
@@ -118,7 +118,6 @@
 
 java_test_host {
     name: "hoststubgentest",
-    // main_class: "com.android.hoststubgen.Main",
     srcs: ["test/**/*.kt"],
     static_libs: [
         "hoststubgen",
@@ -143,8 +142,7 @@
     // "--policy-override-file $(location framework-policy-override.txt) " +
     "@$(location :hoststubgen-standard-options) " +
 
-    "--out-stub-jar $(location host_stub.jar) " +
-    "--out-impl-jar $(location host_impl.jar) " +
+    "--out-jar $(location host.jar) " +
 
     // "--keep-all-classes " + // Used it for an experiment. See KeepAllClassesFilter.
     "--gen-keep-all-file $(location hoststubgen_keep_all.txt) " +
@@ -159,10 +157,8 @@
     srcs: [
         ":hoststubgen-standard-options",
     ],
-    // Create two jar files.
     out: [
-        "host_stub.jar",
-        "host_impl.jar",
+        "host.jar",
 
         // Following files are created just as FYI.
         "hoststubgen_keep_all.txt",
diff --git a/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestStub.java b/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestStub.java
deleted file mode 100644
index cabdfe0..0000000
--- a/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestStub.java
+++ /dev/null
@@ -1,43 +0,0 @@
-/*
- * Copyright (C) 2023 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.hosttest.annotation;
-
-import static java.lang.annotation.ElementType.CONSTRUCTOR;
-import static java.lang.annotation.ElementType.FIELD;
-import static java.lang.annotation.ElementType.METHOD;
-import static java.lang.annotation.ElementType.TYPE;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
-
-/**
- * THIS ANNOTATION IS EXPERIMENTAL. REACH OUT TO g/ravenwood BEFORE USING IT, OR YOU HAVE ANY
- * QUESTIONS ABOUT IT.
- *
- * Mark a class, field or a method as "Stub", meaning tests can see the APIs.
- * When applied to a class, it will _not_ affect the visibility of its members. They need to be
- * individually marked.
- *
- * <p>In order to expose a class and all its members, use {@link HostSideTestWholeClassStub}
- * instead.
- *
- * @hide
- */
-@Target({TYPE, FIELD, METHOD, CONSTRUCTOR})
-@Retention(RetentionPolicy.CLASS)
-public @interface HostSideTestStub {
-}
diff --git a/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestWholeClassStub.java b/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestWholeClassStub.java
deleted file mode 100644
index 1824f6f..0000000
--- a/tools/hoststubgen/hoststubgen/annotations-src/android/hosttest/annotation/HostSideTestWholeClassStub.java
+++ /dev/null
@@ -1,35 +0,0 @@
-/*
- * Copyright (C) 2023 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.hosttest.annotation;
-
-import static java.lang.annotation.ElementType.TYPE;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
-
-/**
- * THIS ANNOTATION IS EXPERIMENTAL. REACH OUT TO g/ravenwood BEFORE USING IT, OR YOU HAVE ANY
- * QUESTIONS ABOUT IT.
- *
- * Same as {@link HostSideTestStub} but it'll change the visibility of all its members too.
- *
- * @hide
- */
-@Target({TYPE})
-@Retention(RetentionPolicy.CLASS)
-public @interface HostSideTestWholeClassStub {
-}
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInStub.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInStub.java
deleted file mode 100644
index 12b9875..0000000
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInStub.java
+++ /dev/null
@@ -1,38 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.hosthelper;
-
-import static java.lang.annotation.ElementType.CONSTRUCTOR;
-import static java.lang.annotation.ElementType.FIELD;
-import static java.lang.annotation.ElementType.METHOD;
-import static java.lang.annotation.ElementType.TYPE;
-
-import java.lang.annotation.Retention;
-import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
-
-/**
- * Annotation injected to all classes/methods/fields that are kept in the "stub" jar.
- *
- * All items in the stub jar are automatically kept in the impl jar as well, so
- * the items with this annotation will all have {@link HostStubGenKeptInImpl} too.
- */
-@Target({TYPE, METHOD, CONSTRUCTOR, FIELD})
-@Retention(RetentionPolicy.RUNTIME)
-public @interface HostStubGenKeptInStub {
-    String CLASS_INTERNAL_NAME = HostTestUtils.getInternalName(HostStubGenKeptInStub.class);
-    String CLASS_DESCRIPTOR = "L" + CLASS_INTERNAL_NAME + ";";
-}
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsIgnore.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsIgnore.java
index cb50404..b017103 100644
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsIgnore.java
+++ b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsIgnore.java
@@ -23,8 +23,6 @@
 
 /**
  * Annotation injected to all methods processed as "ignore".
- *
- * (This annotation is only added in the impl jar, but not the stub jar)
  */
 @Target({METHOD})
 @Retention(RetentionPolicy.RUNTIME)
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInImpl.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsKeep.java
similarity index 92%
rename from tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInImpl.java
rename to tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsKeep.java
index 2cc500f..18ef1ba 100644
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenKeptInImpl.java
+++ b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsKeep.java
@@ -25,11 +25,11 @@
 import java.lang.annotation.Target;
 
 /**
- * Annotation injected to all classes/methods/fields that are kept in the "impl" jar.
+ * Annotation injected to all classes/methods/fields that are kept in the processes jar.
  */
 @Target({TYPE, METHOD, CONSTRUCTOR, FIELD})
 @Retention(RetentionPolicy.RUNTIME)
-public @interface HostStubGenKeptInImpl {
-    String CLASS_INTERNAL_NAME = HostTestUtils.getInternalName(HostStubGenKeptInImpl.class);
+public @interface HostStubGenProcessedAsKeep {
+    String CLASS_INTERNAL_NAME = HostTestUtils.getInternalName(HostStubGenProcessedAsKeep.class);
     String CLASS_DESCRIPTOR = "L" + CLASS_INTERNAL_NAME + ";";
 }
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsSubstitute.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsSubstitute.java
index cfa4896..99e38c0 100644
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsSubstitute.java
+++ b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsSubstitute.java
@@ -26,8 +26,6 @@
 
 /**
  * Annotation injected to all methods that are processed as "substitute".
- *
- * (This annotation is only added in the impl jar, but not the stub jar)
  */
 @Target({TYPE, METHOD, CONSTRUCTOR, FIELD})
 @Retention(RetentionPolicy.RUNTIME)
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsThrow.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsThrow.java
index 0d2da11..4933cf8 100644
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsThrow.java
+++ b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostStubGenProcessedAsThrow.java
@@ -23,8 +23,6 @@
 
 /**
  * Annotation injected to all methods that are processed as "throw".
- *
- * (This annotation is only added in the impl jar, but not the stub jar)
  */
 @Target({METHOD})
 @Retention(RetentionPolicy.RUNTIME)
diff --git a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostTestUtils.java b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostTestUtils.java
index 60eb47ee..78fd8f7 100644
--- a/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostTestUtils.java
+++ b/tools/hoststubgen/hoststubgen/helper-runtime-src/com/android/hoststubgen/hosthelper/HostTestUtils.java
@@ -16,12 +16,8 @@
 package com.android.hoststubgen.hosthelper;
 
 import java.io.PrintStream;
-import java.lang.StackWalker.Option;
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
-import java.util.HashMap;
-
-import javax.annotation.concurrent.GuardedBy;
 
 /**
  * Utilities used in the host side test environment.
@@ -101,68 +97,6 @@
                 + methodName + methodDescriptor);
     }
 
-    private static final StackWalker sStackWalker =
-            StackWalker.getInstance(Option.RETAIN_CLASS_REFERENCE);
-
-    /**
-     * Return a {@link StackWalker} that supports {@link StackWalker#getCallerClass()}.
-     */
-    public static StackWalker getStackWalker() {
-        return sStackWalker;
-    }
-
-    /**
-     * Cache used by {@link #isClassAllowedToCallNonStubMethods}.
-     */
-    @GuardedBy("sAllowedClasses")
-    private static final HashMap<Class, Boolean> sAllowedClasses = new HashMap();
-
-    /**
-     * Return true if a given class is allowed to access non-stub methods -- that is, if the class
-     * is in the hoststubgen generated JARs. (not in the test jar.)
-     */
-    private static boolean isClassAllowedToCallNonStubMethods(Class<?> clazz) {
-        synchronized (sAllowedClasses) {
-            var cached = sAllowedClasses.get(clazz);
-            if (cached != null) {
-                return cached;
-            }
-        }
-        // All processed classes have this annotation.
-        var allowed = clazz.getAnnotation(HostStubGenKeptInImpl.class) != null;
-
-        // Java classes should be able to access any methods. (via callbacks, etc.)
-        if (!allowed) {
-            if (clazz.getPackageName().startsWith("java.")
-                    || clazz.getPackageName().startsWith("javax.")) {
-                allowed = true;
-            }
-        }
-        synchronized (sAllowedClasses) {
-            sAllowedClasses.put(clazz, allowed);
-        }
-        return allowed;
-    }
-
-    /**
-     * Called when non-stub methods are called. We do a host-unsupported method direct call check
-     * in here.
-     */
-    public static void onNonStubMethodCalled(
-            String methodClass,
-            String methodName,
-            String methodDescriptor,
-            Class<?> callerClass) {
-        if (SKIP_NON_STUB_METHOD_CHECK) {
-            return;
-        }
-        if (isClassAllowedToCallNonStubMethods(callerClass)) {
-            return; // Generated class is allowed to call framework class.
-        }
-        logPrintStream.println("! " + methodClass + "." + methodName + methodDescriptor
-                + " called by " + callerClass.getCanonicalName());
-    }
-
     /**
      * Called when any top level class (not nested classes) in the impl jar is loaded.
      *
diff --git a/tools/hoststubgen/hoststubgen/hoststubgen-standard-options.txt b/tools/hoststubgen/hoststubgen/hoststubgen-standard-options.txt
index c371b5d..e72c9a4 100644
--- a/tools/hoststubgen/hoststubgen/hoststubgen-standard-options.txt
+++ b/tools/hoststubgen/hoststubgen/hoststubgen-standard-options.txt
@@ -3,8 +3,6 @@
 --debug
 
 # Uncomment below lines to enable each feature.
---enable-non-stub-method-check
-# --no-non-stub-method-check
 
 #--default-method-call-hook
 #    com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -13,15 +11,10 @@
 
 # Standard annotations.
 # Note, each line is a single argument, so we need newlines after each `--xxx-annotation`.
---stub-annotation
-    android.hosttest.annotation.HostSideTestStub
 
 --keep-annotation
     android.hosttest.annotation.HostSideTestKeep
 
---stub-class-annotation
-    android.hosttest.annotation.HostSideTestWholeClassStub
-
 --keep-class-annotation
     android.hosttest.annotation.HostSideTestWholeClassKeep
 
diff --git a/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh b/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
index d97dd7c..5f0368a 100755
--- a/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
+++ b/tools/hoststubgen/hoststubgen/invoketest/hoststubgen-invoke-test.sh
@@ -43,9 +43,8 @@
 
 cleanup_temp
 
-JAR=hoststubgen-test-tiny-framework.jar
-STUB=$TEMP/stub.jar
-IMPL=$TEMP/impl.jar
+INJAR=hoststubgen-test-tiny-framework.jar
+OUTJAR=$TEMP/host.jar
 
 ANNOTATION_FILTER=$TEMP/annotation-filter.txt
 
@@ -81,27 +80,18 @@
     cat $ANNOTATION_FILTER
   fi
 
-  local stub_arg=""
-  local impl_arg=""
+  local out_arg=""
 
-  if [[ "$STUB" != "" ]] ; then
-    stub_arg="--out-stub-jar $STUB"
-  fi
-  if [[ "$IMPL" != "" ]] ; then
-    impl_arg="--out-impl-jar $IMPL"
+  if [[ "$OUTJAR" != "" ]] ; then
+    out_arg="--out-jar $OUTJAR"
   fi
 
   hoststubgen \
       --debug \
-      --in-jar $JAR \
-      $stub_arg \
-      $impl_arg \
-      --stub-annotation \
-          android.hosttest.annotation.HostSideTestStub \
+      --in-jar $INJAR \
+      $out_arg \
       --keep-annotation \
           android.hosttest.annotation.HostSideTestKeep \
-      --stub-class-annotation \
-          android.hosttest.annotation.HostSideTestWholeClassStub \
       --keep-class-annotation \
           android.hosttest.annotation.HostSideTestWholeClassKeep \
       --throw-annotation \
@@ -213,9 +203,9 @@
 "
 
 run_hoststubgen_for_failure "One specific class disallowed" \
-    "TinyFrameworkClassAnnotations is not allowed to have Ravenwood annotations" \
+    "TinyFrameworkAnnotations is not allowed to have Ravenwood annotations" \
     "
-!com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
+!com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
 * # All other classes allowed
 "
 
@@ -225,11 +215,7 @@
 * # All other classes allowed
 "
 
-STUB="" run_hoststubgen_for_success "No stub generation" ""
-
-IMPL="" run_hoststubgen_for_success "No impl generation" ""
-
-STUB="" IMPL="" run_hoststubgen_for_success "No stub, no impl generation" ""
+OUTJAR="" run_hoststubgen_for_success "No output generation" ""
 
 EXTRA_ARGS="--in-jar abc" run_hoststubgen_for_failure "Duplicate arg" \
     "Duplicate or conflicting argument found: --in-jar" \
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
index 7b08678..0f38fe7 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGen.kt
@@ -24,19 +24,13 @@
 import com.android.hoststubgen.filters.FilterPolicy
 import com.android.hoststubgen.filters.FilterRemapper
 import com.android.hoststubgen.filters.ImplicitOutputFilter
+import com.android.hoststubgen.filters.NativeFilter
 import com.android.hoststubgen.filters.OutputFilter
-import com.android.hoststubgen.filters.StubIntersectingFilter
 import com.android.hoststubgen.filters.createFilterFromTextPolicyFile
 import com.android.hoststubgen.filters.printAsTextPolicy
 import com.android.hoststubgen.utils.ClassFilter
 import com.android.hoststubgen.visitors.BaseAdapter
 import com.android.hoststubgen.visitors.PackageRedirectRemapper
-import org.objectweb.asm.ClassReader
-import org.objectweb.asm.ClassVisitor
-import org.objectweb.asm.ClassWriter
-import org.objectweb.asm.commons.ClassRemapper
-import org.objectweb.asm.commons.Remapper
-import org.objectweb.asm.util.CheckClassAdapter
 import java.io.BufferedInputStream
 import java.io.BufferedOutputStream
 import java.io.FileOutputStream
@@ -46,6 +40,12 @@
 import java.util.zip.ZipEntry
 import java.util.zip.ZipFile
 import java.util.zip.ZipOutputStream
+import org.objectweb.asm.ClassReader
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.ClassWriter
+import org.objectweb.asm.commons.ClassRemapper
+import org.objectweb.asm.commons.Remapper
+import org.objectweb.asm.util.CheckClassAdapter
 
 /**
  * Actual main class.
@@ -82,17 +82,16 @@
 
         // Transform the jar.
         convert(
-                options.inJar.get,
-                options.outStubJar.get,
-                options.outImplJar.get,
-                filter,
-                options.enableClassChecker.get,
-                allClasses,
-                errors,
-                stats,
-                filterRemapper,
-                options.numShards.get,
-                options.shard.get,
+            options.inJar.get,
+            options.outJar.get,
+            filter,
+            options.enableClassChecker.get,
+            allClasses,
+            errors,
+            stats,
+            filterRemapper,
+            options.numShards.get,
+            options.shard.get,
         )
 
         // Dump statistics, if specified.
@@ -117,10 +116,10 @@
      * jars, and "how". (e.g. with substitution?)
      */
     private fun buildFilter(
-            errors: HostStubGenErrors,
-            allClasses: ClassNodes,
-            options: HostStubGenOptions,
-            ): OutputFilter {
+        errors: HostStubGenErrors,
+        allClasses: ClassNodes,
+        options: HostStubGenOptions,
+    ): OutputFilter {
         // We build a "chain" of multiple filters here.
         //
         // The filters are build in from "inside", meaning the first filter created here is
@@ -134,6 +133,9 @@
         // The first filter is for the default policy from the command line options.
         var filter: OutputFilter = ConstantFilter(options.defaultPolicy.get, "default-by-options")
 
+        // Next, we build a filter that preserves all native methods by default
+        filter = NativeFilter(allClasses, filter)
+
         // Next, we need a filter that resolves "class-wide" policies.
         // This is used when a member (methods, fields, nested classes) don't get any polices
         // from upper filters. e.g. when a method has no annotations, then this filter will apply
@@ -159,9 +161,7 @@
         filter = AnnotationBasedFilter(
             errors,
             allClasses,
-            options.stubAnnotations,
             options.keepAnnotations,
-            options.stubClassAnnotations,
             options.keepClassAnnotations,
             options.throwAnnotations,
             options.removeAnnotations,
@@ -179,15 +179,6 @@
             filter = createFilterFromTextPolicyFile(it, allClasses, filter)
         }
 
-        // If `--intersect-stub-jar` is provided, load from these jar files too.
-        // We use this to restrict stub APIs to public/system/test APIs,
-        // by intersecting with a stub jar file created by metalava.
-        if (options.intersectStubJars.size > 0) {
-            val intersectingJars = loadIntersectingJars(options.intersectStubJars)
-
-            filter = StubIntersectingFilter(errors, intersectingJars, filter)
-        }
-
         // Apply the implicit filter.
         filter = ImplicitOutputFilter(errors, allClasses, filter)
 
@@ -195,34 +186,21 @@
     }
 
     /**
-     * Load jar files specified with "--intersect-stub-jar".
-     */
-    private fun loadIntersectingJars(filenames: Set<String>): Map<String, ClassNodes> {
-        val intersectingJars = mutableMapOf<String, ClassNodes>()
-
-        filenames.forEach { filename ->
-            intersectingJars[filename] = ClassNodes.loadClassStructures(filename)
-        }
-        return intersectingJars
-    }
-
-    /**
      * Convert a JAR file into "stub" and "impl" JAR files.
      */
     private fun convert(
-            inJar: String,
-            outStubJar: String?,
-            outImplJar: String?,
-            filter: OutputFilter,
-            enableChecker: Boolean,
-            classes: ClassNodes,
-            errors: HostStubGenErrors,
-            stats: HostStubGenStats,
-            remapper: Remapper?,
-            numShards: Int,
-            shard: Int,
-            ) {
-        log.i("Converting %s into [stub: %s, impl: %s] ...", inJar, outStubJar, outImplJar)
+        inJar: String,
+        outJar: String?,
+        filter: OutputFilter,
+        enableChecker: Boolean,
+        classes: ClassNodes,
+        errors: HostStubGenErrors,
+        stats: HostStubGenStats,
+        remapper: Remapper?,
+        numShards: Int,
+        shard: Int
+    ) {
+        log.i("Converting %s into %s ...", inJar, outJar)
         log.i("ASM CheckClassAdapter is %s", if (enableChecker) "enabled" else "disabled")
 
         log.iTime("Transforming jar") {
@@ -240,29 +218,26 @@
                     val shardStart = numItems * shard / numShards
                     val shardNextStart = numItems * (shard + 1) / numShards
 
-                    maybeWithZipOutputStream(outStubJar) { stubOutStream ->
-                        maybeWithZipOutputStream(outImplJar) { implOutStream ->
-                            val inEntries = inZip.entries()
-                            while (inEntries.hasMoreElements()) {
-                                val entry = inEntries.nextElement()
-                                val inShard = (shardStart <= itemIndex)
-                                        && (itemIndex < shardNextStart)
-                                itemIndex++
-                                if (!inShard) {
-                                    continue
-                                }
-                                convertSingleEntry(
-                                    inZip, entry, stubOutStream, implOutStream,
-                                    filter, packageRedirector, remapper,
-                                    enableChecker, classes, errors, stats
-                                )
-                                numItemsProcessed++
+                    maybeWithZipOutputStream(outJar) { outStream ->
+                        val inEntries = inZip.entries()
+                        while (inEntries.hasMoreElements()) {
+                            val entry = inEntries.nextElement()
+                            val inShard = (shardStart <= itemIndex)
+                                    && (itemIndex < shardNextStart)
+                            itemIndex++
+                            if (!inShard) {
+                                continue
                             }
-                            log.i("Converted all entries.")
+                            convertSingleEntry(
+                                inZip, entry, outStream, filter,
+                                packageRedirector, remapper, enableChecker,
+                                classes, errors, stats
+                            )
+                            numItemsProcessed++
                         }
+                        log.i("Converted all entries.")
                     }
-                    outStubJar?.let { log.i("Created stub: $it") }
-                    outImplJar?.let { log.i("Created impl: $it") }
+                    outJar?.let { log.i("Created: $it") }
                 }
             }
             log.i("%d / %d item(s) processed.", numItemsProcessed, numItems)
@@ -280,18 +255,17 @@
      * Convert a single ZIP entry, which may or may not be a class file.
      */
     private fun convertSingleEntry(
-            inZip: ZipFile,
-            entry: ZipEntry,
-            stubOutStream: ZipOutputStream?,
-            implOutStream: ZipOutputStream?,
-            filter: OutputFilter,
-            packageRedirector: PackageRedirectRemapper,
-            remapper: Remapper?,
-            enableChecker: Boolean,
-            classes: ClassNodes,
-            errors: HostStubGenErrors,
-            stats: HostStubGenStats,
-            ) {
+        inZip: ZipFile,
+        entry: ZipEntry,
+        outStream: ZipOutputStream?,
+        filter: OutputFilter,
+        packageRedirector: PackageRedirectRemapper,
+        remapper: Remapper?,
+        enableChecker: Boolean,
+        classes: ClassNodes,
+        errors: HostStubGenErrors,
+        stats: HostStubGenStats
+    ) {
         log.d("Entry: %s", entry.name)
         log.withIndent {
             val name = entry.name
@@ -303,8 +277,10 @@
 
             // If it's a class, convert it.
             if (name.endsWith(".class")) {
-                processSingleClass(inZip, entry, stubOutStream, implOutStream, filter,
-                        packageRedirector, remapper, enableChecker, classes, errors, stats)
+                processSingleClass(
+                    inZip, entry, outStream, filter, packageRedirector,
+                    remapper, enableChecker, classes, errors, stats
+                )
                 return
             }
 
@@ -312,17 +288,14 @@
 
             // - *.uau seems to contain hidden API information.
             // -  *_compat_config.xml is also about compat-framework.
-            if (name.endsWith(".uau") ||
-                    name.endsWith("_compat_config.xml")) {
+            if (name.endsWith(".uau") || name.endsWith("_compat_config.xml")) {
                 log.d("Not needed: %s", entry.name)
                 return
             }
 
             // Unknown type, we just copy it to both output zip files.
-            // TODO: We probably shouldn't do it for stub jar?
             log.v("Copying: %s", entry.name)
-            stubOutStream?.let { copyZipEntry(inZip, entry, it) }
-            implOutStream?.let { copyZipEntry(inZip, entry, it) }
+            outStream?.let { copyZipEntry(inZip, entry, it) }
         }
     }
 
@@ -330,10 +303,10 @@
      * Copy a single ZIP entry to the output.
      */
     private fun copyZipEntry(
-            inZip: ZipFile,
-            entry: ZipEntry,
-            out: ZipOutputStream,
-            ) {
+        inZip: ZipFile,
+        entry: ZipEntry,
+        out: ZipOutputStream,
+    ) {
         // TODO: It seems like copying entries this way is _very_ slow,
         // even with out.setLevel(0). Look for other ways to do it.
 
@@ -350,18 +323,17 @@
      * Convert a single class to "stub" and "impl".
      */
     private fun processSingleClass(
-            inZip: ZipFile,
-            entry: ZipEntry,
-            stubOutStream: ZipOutputStream?,
-            implOutStream: ZipOutputStream?,
-            filter: OutputFilter,
-            packageRedirector: PackageRedirectRemapper,
-            remapper: Remapper?,
-            enableChecker: Boolean,
-            classes: ClassNodes,
-            errors: HostStubGenErrors,
-            stats: HostStubGenStats,
-            ) {
+        inZip: ZipFile,
+        entry: ZipEntry,
+        outStream: ZipOutputStream?,
+        filter: OutputFilter,
+        packageRedirector: PackageRedirectRemapper,
+        remapper: Remapper?,
+        enableChecker: Boolean,
+        classes: ClassNodes,
+        errors: HostStubGenErrors,
+        stats: HostStubGenStats
+    ) {
         val classInternalName = entry.name.replaceFirst("\\.class$".toRegex(), "")
         val classPolicy = filter.getPolicyForClass(classInternalName)
         if (classPolicy.policy == FilterPolicy.Remove) {
@@ -373,33 +345,22 @@
         remapper?.mapType(classInternalName)?.let { remappedName ->
             if (remappedName != classInternalName) {
                 log.d("Renaming class file: %s -> %s", classInternalName, remappedName)
-                newName = remappedName + ".class"
+                newName = "$remappedName.class"
             }
         }
-        // Generate stub first.
-        if (stubOutStream != null && classPolicy.policy.needsInStub) {
-            log.v("Creating stub class: %s Policy: %s", classInternalName, classPolicy)
+
+        if (outStream != null) {
+            log.v("Creating class: %s Policy: %s", classInternalName, classPolicy)
             log.withIndent {
                 BufferedInputStream(inZip.getInputStream(entry)).use { bis ->
                     val newEntry = ZipEntry(newName)
-                    stubOutStream.putNextEntry(newEntry)
-                    convertClass(classInternalName, /*forImpl=*/false, bis,
-                            stubOutStream, filter, packageRedirector, remapper,
-                            enableChecker, classes, errors, null)
-                    stubOutStream.closeEntry()
-                }
-            }
-        }
-        if (implOutStream != null && classPolicy.policy.needsInImpl) {
-            log.v("Creating impl class: %s Policy: %s", classInternalName, classPolicy)
-            log.withIndent {
-                BufferedInputStream(inZip.getInputStream(entry)).use { bis ->
-                    val newEntry = ZipEntry(newName)
-                    implOutStream.putNextEntry(newEntry)
-                    convertClass(classInternalName, /*forImpl=*/true, bis,
-                            implOutStream, filter, packageRedirector, remapper,
-                            enableChecker, classes, errors, stats)
-                    implOutStream.closeEntry()
+                    outStream.putNextEntry(newEntry)
+                    convertClass(
+                        classInternalName, bis,
+                        outStream, filter, packageRedirector, remapper,
+                        enableChecker, classes, errors, stats
+                    )
+                    outStream.closeEntry()
                 }
             }
         }
@@ -409,18 +370,17 @@
      * Convert a single class to either "stub" or "impl".
      */
     private fun convertClass(
-            classInternalName: String,
-            forImpl: Boolean,
-            input: InputStream,
-            out: OutputStream,
-            filter: OutputFilter,
-            packageRedirector: PackageRedirectRemapper,
-            remapper: Remapper?,
-            enableChecker: Boolean,
-            classes: ClassNodes,
-            errors: HostStubGenErrors,
-            stats: HostStubGenStats?,
-            ) {
+        classInternalName: String,
+        input: InputStream,
+        out: OutputStream,
+        filter: OutputFilter,
+        packageRedirector: PackageRedirectRemapper,
+        remapper: Remapper?,
+        enableChecker: Boolean,
+        classes: ClassNodes,
+        errors: HostStubGenErrors,
+        stats: HostStubGenStats?
+    ) {
         val cr = ClassReader(input)
 
         // COMPUTE_FRAMES wouldn't be happy if code uses
@@ -439,14 +399,15 @@
         }
 
         val visitorOptions = BaseAdapter.Options(
-                enablePreTrace = options.enablePreTrace.get,
-                enablePostTrace = options.enablePostTrace.get,
-                enableNonStubMethodCallDetection = options.enableNonStubMethodCallDetection.get,
-                errors = errors,
-                stats = stats,
+            errors = errors,
+            stats = stats,
+            enablePreTrace = options.enablePreTrace.get,
+            enablePostTrace = options.enablePostTrace.get,
         )
-        outVisitor = BaseAdapter.getVisitor(classInternalName, classes, outVisitor, filter,
-                packageRedirector, remapper, forImpl, visitorOptions)
+        outVisitor = BaseAdapter.getVisitor(
+            classInternalName, classes, outVisitor, filter,
+            packageRedirector, visitorOptions
+        )
 
         cr.accept(outVisitor, ClassReader.EXPAND_FRAMES)
         val data = cw.toByteArray()
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
index f88b107..1cedcc3 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/HostStubGenOptions.kt
@@ -74,21 +74,16 @@
         /** Input jar file*/
         var inJar: SetOnce<String> = SetOnce(""),
 
-        /** Output stub jar file */
-        var outStubJar: SetOnce<String?> = SetOnce(null),
-
-        /** Output implementation jar file */
-        var outImplJar: SetOnce<String?> = SetOnce(null),
+        /** Output jar file */
+        var outJar: SetOnce<String?> = SetOnce(null),
 
         var inputJarDumpFile: SetOnce<String?> = SetOnce(null),
 
         var inputJarAsKeepAllFile: SetOnce<String?> = SetOnce(null),
 
-        var stubAnnotations: MutableSet<String> = mutableSetOf(),
         var keepAnnotations: MutableSet<String> = mutableSetOf(),
         var throwAnnotations: MutableSet<String> = mutableSetOf(),
         var removeAnnotations: MutableSet<String> = mutableSetOf(),
-        var stubClassAnnotations: MutableSet<String> = mutableSetOf(),
         var keepClassAnnotations: MutableSet<String> = mutableSetOf(),
 
         var substituteAnnotations: MutableSet<String> = mutableSetOf(),
@@ -103,8 +98,6 @@
         var defaultClassLoadHook: SetOnce<String?> = SetOnce(null),
         var defaultMethodCallHook: SetOnce<String?> = SetOnce(null),
 
-        var intersectStubJars: MutableSet<String> = mutableSetOf(),
-
         var policyOverrideFile: SetOnce<String?> = SetOnce(null),
 
         var defaultPolicy: SetOnce<FilterPolicy> = SetOnce(FilterPolicy.Remove),
@@ -115,8 +108,6 @@
         var enablePreTrace: SetOnce<Boolean> = SetOnce(false),
         var enablePostTrace: SetOnce<Boolean> = SetOnce(false),
 
-        var enableNonStubMethodCallDetection: SetOnce<Boolean> = SetOnce(false),
-
         var statsFile: SetOnce<String?> = SetOnce(null),
 
         var apiListFile: SetOnce<String?> = SetOnce(null),
@@ -150,10 +141,7 @@
             }
 
             while (true) {
-                val arg = ai.nextArgOptional()
-                if (arg == null) {
-                    break
-                }
+                val arg = ai.nextArgOptional() ?: break
 
                 // Define some shorthands...
                 fun nextArg(): String = ai.nextArgRequired(arg)
@@ -169,8 +157,9 @@
                         "-h", "--help" -> TODO("Help is not implemented yet")
 
                         "--in-jar" -> ret.inJar.set(nextArg()).ensureFileExists()
-                        "--out-stub-jar" -> ret.outStubJar.set(nextArg())
-                        "--out-impl-jar" -> ret.outImplJar.set(nextArg())
+                        // We support both arguments because some AOSP dependencies
+                        // still use the old argument
+                        "--out-jar", "--out-impl-jar" -> ret.outJar.set(nextArg())
 
                         "--policy-override-file" ->
                             ret.policyOverrideFile.set(nextArg())!!.ensureFileExists()
@@ -181,17 +170,10 @@
                         "--default-remove" -> ret.defaultPolicy.set(FilterPolicy.Remove)
                         "--default-throw" -> ret.defaultPolicy.set(FilterPolicy.Throw)
                         "--default-keep" -> ret.defaultPolicy.set(FilterPolicy.Keep)
-                        "--default-stub" -> ret.defaultPolicy.set(FilterPolicy.Stub)
-
-                        "--stub-annotation" ->
-                            ret.stubAnnotations.addUniqueAnnotationArg()
 
                         "--keep-annotation" ->
                             ret.keepAnnotations.addUniqueAnnotationArg()
 
-                        "--stub-class-annotation" ->
-                            ret.stubClassAnnotations.addUniqueAnnotationArg()
-
                         "--keep-class-annotation" ->
                             ret.keepClassAnnotations.addUniqueAnnotationArg()
 
@@ -225,9 +207,6 @@
                         "--default-method-call-hook" ->
                             ret.defaultMethodCallHook.set(nextArg())
 
-                        "--intersect-stub-jar" ->
-                            ret.intersectStubJars += nextArg().ensureFileExists()
-
                         "--gen-keep-all-file" ->
                             ret.inputJarAsKeepAllFile.set(nextArg())
 
@@ -241,12 +220,6 @@
                         "--enable-post-trace" -> ret.enablePostTrace.set(true)
                         "--no-post-trace" -> ret.enablePostTrace.set(false)
 
-                        "--enable-non-stub-method-check" ->
-                            ret.enableNonStubMethodCallDetection.set(true)
-
-                        "--no-non-stub-method-check" ->
-                            ret.enableNonStubMethodCallDetection.set(false)
-
                         "--gen-input-dump-file" -> ret.inputJarDumpFile.set(nextArg())
 
                         "--stats-file" -> ret.statsFile.set(nextArg())
@@ -273,9 +246,8 @@
             if (!ret.inJar.isSet) {
                 throw ArgumentsException("Required option missing: --in-jar")
             }
-            if (!ret.outStubJar.isSet && !ret.outImplJar.isSet) {
-                log.w("Neither --out-stub-jar nor --out-impl-jar is set." +
-                        " $executableName will not generate jar files.")
+            if (!ret.outJar.isSet) {
+                log.w("--out-jar is not set. $executableName will not generate jar files.")
             }
             if (ret.numShards.isSet != ret.shard.isSet) {
                 throw ArgumentsException("--num-shards and --shard-index must be used together")
@@ -287,11 +259,6 @@
                 }
             }
 
-            if (ret.enableNonStubMethodCallDetection.get) {
-                log.w("--enable-non-stub-method-check is not fully implemented yet." +
-                    " See the todo in doesMethodNeedNonStubCallCheck().")
-            }
-
             return ret
         }
     }
@@ -300,32 +267,27 @@
         return """
             HostStubGenOptions{
               inJar='$inJar',
-              outStubJar='$outStubJar',
-              outImplJar='$outImplJar',
+              outJar='$outJar',
               inputJarDumpFile=$inputJarDumpFile,
               inputJarAsKeepAllFile=$inputJarAsKeepAllFile,
-              stubAnnotations=$stubAnnotations,
               keepAnnotations=$keepAnnotations,
               throwAnnotations=$throwAnnotations,
               removeAnnotations=$removeAnnotations,
-              stubClassAnnotations=$stubClassAnnotations,
               keepClassAnnotations=$keepClassAnnotations,
               substituteAnnotations=$substituteAnnotations,
               nativeSubstituteAnnotations=$nativeSubstituteAnnotations,
               classLoadHookAnnotations=$classLoadHookAnnotations,
               keepStaticInitializerAnnotations=$keepStaticInitializerAnnotations,
               packageRedirects=$packageRedirects,
-              $annotationAllowedClassesFile=$annotationAllowedClassesFile,
+              annotationAllowedClassesFile=$annotationAllowedClassesFile,
               defaultClassLoadHook=$defaultClassLoadHook,
               defaultMethodCallHook=$defaultMethodCallHook,
-              intersectStubJars=$intersectStubJars,
               policyOverrideFile=$policyOverrideFile,
               defaultPolicy=$defaultPolicy,
               cleanUpOnError=$cleanUpOnError,
               enableClassChecker=$enableClassChecker,
               enablePreTrace=$enablePreTrace,
               enablePostTrace=$enablePostTrace,
-              enableNonStubMethodCallDetection=$enableNonStubMethodCallDetection,
               statsFile=$statsFile,
               apiListFile=$apiListFile,
               numShards=$numShards,
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/asm/AsmUtils.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/asm/AsmUtils.kt
index 6cf2143..7197e0e 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/asm/AsmUtils.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/asm/AsmUtils.kt
@@ -202,18 +202,6 @@
 }
 
 /**
- * Take a class name. If it's a nested class, then return the name of its direct outer class name.
- * Otherwise, return null.
- */
-fun getDirectOuterClassName(className: String): String? {
-    val pos = className.lastIndexOf('$')
-    if (pos < 0) {
-        return null
-    }
-    return className.substring(0, pos)
-}
-
-/**
  * Write bytecode to push all the method arguments to the stack.
  * The number of arguments and their type are taken from [methodDescriptor].
  */
@@ -339,6 +327,10 @@
     return (this.access and Opcodes.ACC_PUBLIC) != 0
 }
 
+fun MethodNode.isNative(): Boolean {
+    return (this.access and Opcodes.ACC_NATIVE) != 0
+}
+
 fun MethodNode.isSpecial(): Boolean {
     return CTOR_NAME == this.name || CLASS_INITIALIZER_NAME == this.name
 }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/dumper/ApiDumper.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/dumper/ApiDumper.kt
index aaefee4..5e4e70f 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/dumper/ApiDumper.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/dumper/ApiDumper.kt
@@ -44,7 +44,7 @@
         val descriptor: String,
     )
 
-    val javaStandardApiPolicy = FilterPolicy.Stub.withReason("Java standard API")
+    private val javaStandardApiPolicy = FilterPolicy.Keep.withReason("Java standard API")
 
     private val shownMethods = mutableSetOf<MethodKey>()
 
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/AnnotationBasedFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/AnnotationBasedFilter.kt
index 248121c..38a41b2 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/AnnotationBasedFilter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/AnnotationBasedFilter.kt
@@ -17,7 +17,6 @@
 
 import com.android.hoststubgen.ClassParseException
 import com.android.hoststubgen.HostStubGenErrors
-import com.android.hoststubgen.HostStubGenInternalException
 import com.android.hoststubgen.InvalidAnnotationException
 import com.android.hoststubgen.addNonNullElement
 import com.android.hoststubgen.asm.CLASS_INITIALIZER_DESC
@@ -35,258 +34,125 @@
 
 // TODO: Detect invalid cases, such as...
 // - Class's visibility is lower than the members'.
-// - HostSideTestSubstituteWith is set, but it doesn't have @Stub or @Keep
 
 /**
  * [OutputFilter] using Java annotations.
  */
 class AnnotationBasedFilter(
-        private val errors: HostStubGenErrors,
-        private val classes: ClassNodes,
-        stubAnnotations_: Set<String>,
-        keepAnnotations_: Set<String>,
-        stubClassAnnotations_: Set<String>,
-        keepClassAnnotations_: Set<String>,
-        throwAnnotations_: Set<String>,
-        removeAnnotations_: Set<String>,
-        substituteAnnotations_: Set<String>,
-        nativeSubstituteAnnotations_: Set<String>,
-        classLoadHookAnnotations_: Set<String>,
-        keepStaticInitializerAnnotations_: Set<String>,
-        private val annotationAllowedClassesFilter: ClassFilter,
-        fallback: OutputFilter,
+    private val errors: HostStubGenErrors,
+    private val classes: ClassNodes,
+    keepAnnotations_: Set<String>,
+    keepClassAnnotations_: Set<String>,
+    throwAnnotations_: Set<String>,
+    removeAnnotations_: Set<String>,
+    substituteAnnotations_: Set<String>,
+    nativeSubstituteAnnotations_: Set<String>,
+    classLoadHookAnnotations_: Set<String>,
+    keepStaticInitializerAnnotations_: Set<String>,
+    private val annotationAllowedClassesFilter: ClassFilter,
+    fallback: OutputFilter,
 ) : DelegatingFilter(fallback) {
-    private var stubAnnotations = convertToInternalNames(stubAnnotations_)
-    private var keepAnnotations = convertToInternalNames(keepAnnotations_)
-    private var stubClassAnnotations = convertToInternalNames(stubClassAnnotations_)
-    private var keepClassAnnotations = convertToInternalNames(keepClassAnnotations_)
-    private var throwAnnotations = convertToInternalNames(throwAnnotations_)
-    private var removeAnnotations = convertToInternalNames(removeAnnotations_)
-    private var substituteAnnotations = convertToInternalNames(substituteAnnotations_)
-    private var nativeSubstituteAnnotations = convertToInternalNames(nativeSubstituteAnnotations_)
-    private var classLoadHookAnnotations = convertToInternalNames(classLoadHookAnnotations_)
-    private var keepStaticInitializerAnnotations =
-            convertToInternalNames(keepStaticInitializerAnnotations_)
+    private val keepAnnotations = convertToInternalNames(keepAnnotations_)
+    private val keepClassAnnotations = convertToInternalNames(keepClassAnnotations_)
+    private val throwAnnotations = convertToInternalNames(throwAnnotations_)
+    private val removeAnnotations = convertToInternalNames(removeAnnotations_)
+    private val substituteAnnotations = convertToInternalNames(substituteAnnotations_)
+    private val nativeSubstituteAnnotations = convertToInternalNames(nativeSubstituteAnnotations_)
+    private val classLoadHookAnnotations = convertToInternalNames(classLoadHookAnnotations_)
+    private val keepStaticInitializerAnnotations =
+        convertToInternalNames(keepStaticInitializerAnnotations_)
 
     /** Annotations that control API visibility. */
-    private var visibilityAnnotations: Set<String> = convertToInternalNames(
-        stubAnnotations_ +
-        keepAnnotations_ +
-        stubClassAnnotations_ +
-        keepClassAnnotations_ +
-        throwAnnotations_ +
-        removeAnnotations_)
+    private val visibilityAnnotations = keepAnnotations +
+            keepClassAnnotations +
+            throwAnnotations +
+            removeAnnotations +
+            substituteAnnotations
+
+    /** All the annotations we use. */
+    private val allAnnotations = visibilityAnnotations +
+            nativeSubstituteAnnotations +
+            classLoadHookAnnotations +
+            keepStaticInitializerAnnotations
 
     /**
      * All the annotations we use. Note, this one is in a [convertToJvmNames] format unlike
      * other ones, because of how it's used.
      */
-    private var allAnnotations: Set<String> = convertToJvmNames(
-        stubAnnotations_ +
-                keepAnnotations_ +
-                stubClassAnnotations_ +
+    private val allAnnotationClasses: Set<String> = convertToJvmNames(
+        keepAnnotations_ +
                 keepClassAnnotations_ +
                 throwAnnotations_ +
                 removeAnnotations_ +
                 substituteAnnotations_ +
                 nativeSubstituteAnnotations_ +
-                classLoadHookAnnotations_)
+                classLoadHookAnnotations_ +
+                keepStaticInitializerAnnotations_
+    )
 
-    private val substitutionHelper = SubstitutionHelper()
+    private val policyCache = mutableMapOf<String, ClassAnnotations>()
 
-    private val reasonAnnotation = "annotation"
-    private val reasonClassAnnotation = "class-annotation"
-
-    /**
-     * Throw if an item has more than one visibility annotations.
-     *
-     * name1 - 4 are only used in exception messages. We take them as separate strings
-     * to avoid unnecessary string concatenations.
-     */
-    private fun detectInvalidAnnotations(
-        visibles: List<AnnotationNode>?,
-        invisibles: List<AnnotationNode>?,
-        type: String,
-        name1: String,
-        name2: String,
-        name3: String,
-    ) {
-        var count = 0
-        for (an in visibles ?: emptyList()) {
-            if (visibilityAnnotations.contains(an.desc)) {
-                count++
-            }
-        }
-        for (an in invisibles ?: emptyList()) {
-            if (visibilityAnnotations.contains(an.desc)) {
-                count++
-            }
-        }
-        if (count > 1) {
-            val description = if (name2 == "" && name3 == "") {
-                "$type $name1"
-            } else {
-                "$type $name1.$name2$name3"
-            }
-            throw InvalidAnnotationException(
-                "Found more than one visibility annotations on $description")
+    private val AnnotationNode.policy: FilterPolicyWithReason? get() {
+        return when (desc) {
+            in keepAnnotations -> FilterPolicy.Keep.withReason(REASON_ANNOTATION)
+            in keepClassAnnotations -> FilterPolicy.KeepClass.withReason(REASON_CLASS_ANNOTATION)
+            in substituteAnnotations -> FilterPolicy.Substitute.withReason(REASON_ANNOTATION)
+            in throwAnnotations -> FilterPolicy.Throw.withReason(REASON_ANNOTATION)
+            in removeAnnotations -> FilterPolicy.Remove.withReason(REASON_ANNOTATION)
+            else -> null
         }
     }
 
-    fun findAnyAnnotation(
-            className: String,
-            anyAnnotations: Set<String>,
-            visibleAnnotations: List<AnnotationNode>?,
-            invisibleAnnotations: List<AnnotationNode>?,
-    ): AnnotationNode? {
-        val ret = findAnyAnnotation(anyAnnotations, visibleAnnotations, invisibleAnnotations)
-
-        if (ret != null) {
-            if (!annotationAllowedClassesFilter.matches(className)) {
-                throw InvalidAnnotationException(
-                        "Class ${className.toHumanReadableClassName()} is not allowed to have " +
-                                "Ravenwood annotations. Contact g/ravenwood for more details.")
-            }
-        }
-
-        return ret
-    }
-
-    /**
-     * Find a visibility annotation.
-     *
-     * name1 - 4 are only used in exception messages.
-     */
-    private fun findAnnotation(
-            className: String,
-            visibles: List<AnnotationNode>?,
-            invisibles: List<AnnotationNode>?,
-            type: String,
-            name1: String,
-            name2: String = "",
-            name3: String = "",
-    ): FilterPolicyWithReason? {
-        detectInvalidAnnotations(visibles, invisibles, type, name1, name2, name3)
-
-        findAnyAnnotation(className, stubAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.Stub.withReason(reasonAnnotation)
-        }
-        findAnyAnnotation(className, stubClassAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.StubClass.withReason(reasonClassAnnotation)
-        }
-        findAnyAnnotation(className, keepAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.Keep.withReason(reasonAnnotation)
-        }
-        findAnyAnnotation(className, keepClassAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.KeepClass.withReason(reasonClassAnnotation)
-        }
-        findAnyAnnotation(className, throwAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.Throw.withReason(reasonAnnotation)
-        }
-        findAnyAnnotation(className, removeAnnotations, visibles, invisibles)?.let {
-            return FilterPolicy.Remove.withReason(reasonAnnotation)
-        }
-
-        return null
+    private fun getAnnotationPolicy(cn: ClassNode): ClassAnnotations {
+        return policyCache.getOrPut(cn.name) { ClassAnnotations(cn) }
     }
 
     override fun getPolicyForClass(className: String): FilterPolicyWithReason {
-        val cn = classes.getClass(className)
-
-        findAnnotation(
-            cn.name,
-            cn.visibleAnnotations,
-            cn.invisibleAnnotations,
-            "class",
-            className)?.let {
-            return it
-        }
-
         // If it's any of the annotations, then always keep it.
-        if (allAnnotations.contains(className)) {
+        if (allAnnotationClasses.contains(className)) {
             return FilterPolicy.KeepClass.withReason("HostStubGen Annotation")
         }
 
-        return super.getPolicyForClass(className)
+        val cn = classes.getClass(className)
+        return getAnnotationPolicy(cn).classPolicy ?: super.getPolicyForClass(className)
     }
 
-    override fun getPolicyForField(
-            className: String,
-            fieldName: String
-    ): FilterPolicyWithReason {
+    override fun getPolicyForField(className: String, fieldName: String): FilterPolicyWithReason {
         val cn = classes.getClass(className)
-
-        cn.fields?.firstOrNull { it.name == fieldName }?.let {fn ->
-            findAnnotation(
-                cn.name,
-                fn.visibleAnnotations,
-                fn.invisibleAnnotations,
-                "field",
-                className,
-                fieldName
-                )?.let { policy ->
-                // If the item has an annotation, then use it.
-                return policy
-            }
-        }
-        return super.getPolicyForField(className, fieldName)
+        return getAnnotationPolicy(cn).fieldPolicies[fieldName]
+            ?: super.getPolicyForField(className, fieldName)
     }
 
     override fun getPolicyForMethod(
-            className: String,
-            methodName: String,
-            descriptor: String
+        className: String,
+        methodName: String,
+        descriptor: String
     ): FilterPolicyWithReason {
         val cn = classes.getClass(className)
 
         if (methodName == CLASS_INITIALIZER_NAME && descriptor == CLASS_INITIALIZER_DESC) {
-            findAnyAnnotation(cn.name, keepStaticInitializerAnnotations,
-                    cn.visibleAnnotations, cn.invisibleAnnotations)?.let {
-                return FilterPolicy.Keep.withReason(reasonAnnotation)
+            if (cn.findAnyAnnotation(keepStaticInitializerAnnotations) != null) {
+                return FilterPolicy.Keep.withReason(REASON_ANNOTATION)
             }
         }
 
-        cn.methods?.firstOrNull { it.name == methodName && it.desc == descriptor }?.let { mn ->
-            // @SubstituteWith is going to complicate the policy here, so we ask helper
-            // what to do.
-            substitutionHelper.getPolicyFromSubstitution(cn, mn.name, mn.desc)?.let {
-                return it
-            }
-
-            // If there's no substitution, then we check the annotation.
-            findAnnotation(
-                cn.name,
-                mn.visibleAnnotations,
-                mn.invisibleAnnotations,
-                "method",
-                className,
-                methodName,
-                descriptor
-            )?.let { policy ->
-                return policy
-            }
-        }
-        return super.getPolicyForMethod(className, methodName, descriptor)
+        return getAnnotationPolicy(cn).methodPolicies[MethodKey(methodName, descriptor)]
+            ?: super.getPolicyForMethod(className, methodName, descriptor)
     }
 
     override fun getRenameTo(
-            className: String,
-            methodName: String,
-            descriptor: String
+        className: String,
+        methodName: String,
+        descriptor: String
     ): String? {
         val cn = classes.getClass(className)
-
-        // If the method has a "substitute with" annotation, then return its "value" parameter.
-        cn.methods?.firstOrNull { it.name == methodName && it.desc == descriptor }?.let { mn ->
-            return substitutionHelper.getRenameTo(cn, mn.name, mn.desc)
-        }
-        return null
+        return getAnnotationPolicy(cn).renamedMethods[MethodKey(methodName, descriptor)]
+            ?: super.getRenameTo(className, methodName, descriptor)
     }
 
     override fun getNativeSubstitutionClass(className: String): String? {
         classes.getClass(className).let { cn ->
-            findAnyAnnotation(nativeSubstituteAnnotations,
-                    cn.visibleAnnotations, cn.invisibleAnnotations)?.let { an ->
+            cn.findAnyAnnotation(nativeSubstituteAnnotations)?.let { an ->
                 return getAnnotationField(an, "value")?.toJvmClassName()
             }
         }
@@ -295,8 +161,7 @@
 
     override fun getClassLoadHooks(className: String): List<String> {
         val e = classes.getClass(className).let { cn ->
-            findAnyAnnotation(classLoadHookAnnotations,
-                cn.visibleAnnotations, cn.invisibleAnnotations)?.let { an ->
+            cn.findAnyAnnotation(classLoadHookAnnotations)?.let { an ->
                 getAnnotationField(an, "value")?.toHumanReadableMethodName()
             }
         }
@@ -306,102 +171,130 @@
     private data class MethodKey(val name: String, val desc: String)
 
     /**
-     * In order to handle substitution, we need to build a reverse mapping of substitution
-     * methods.
+     * Every time we see a class, we scan all its methods for substitution attributes,
+     * and compute (implicit) policies caused by them.
      *
-     * This class automatically builds such a map internally that the above methods can
-     * take advantage of.
+     * For example, for the following methods:
+     *
+     *   @Substitute(suffix = "_host")
+     *   private void foo() {
+     *      // This isn't supported on the host side.
+     *   }
+     *   private void foo_host() {
+     *      // Host side implementation
+     *   }
+     *
+     * We internally handle them as:
+     *
+     *   foo() -> Substitute
+     *   foo_host() -> Stub, and then rename it to foo().
      */
-    private inner class SubstitutionHelper {
-        private var currentClass: ClassNode? = null
+    private inner class ClassAnnotations(cn: ClassNode) {
 
-        private var policiesFromSubstitution = mutableMapOf<MethodKey, FilterPolicyWithReason>()
-        private var substituteToMethods = mutableMapOf<MethodKey, String>()
+        val classPolicy: FilterPolicyWithReason?
+        val fieldPolicies = mutableMapOf<String, FilterPolicyWithReason>()
+        val methodPolicies = mutableMapOf<MethodKey, FilterPolicyWithReason>()
+        val renamedMethods = mutableMapOf<MethodKey, String>()
 
-        fun getPolicyFromSubstitution(cn: ClassNode, methodName: String, descriptor: String):
-                FilterPolicyWithReason? {
-            setClass(cn)
-            return policiesFromSubstitution[MethodKey(methodName, descriptor)]
-        }
+        init {
+            val allowAnnotation = annotationAllowedClassesFilter.matches(cn.name)
+            detectInvalidAnnotations(
+                cn.name, allowAnnotation,
+                cn.visibleAnnotations, cn.invisibleAnnotations,
+                "class", cn.name
+            )
+            classPolicy = cn.findAnyAnnotation(visibilityAnnotations)?.policy
 
-        fun getRenameTo(cn: ClassNode, methodName: String, descriptor: String): String? {
-            setClass(cn)
-            return substituteToMethods[MethodKey(methodName, descriptor)]
+            for (fn in cn.fields ?: emptyList()) {
+                detectInvalidAnnotations(
+                    cn.name, allowAnnotation,
+                    fn.visibleAnnotations, fn.invisibleAnnotations,
+                    "field", cn.name, fn.name
+                )
+                fn.findAnyAnnotation(visibilityAnnotations)?.policy?.let {
+                    fieldPolicies[fn.name] = it
+                }
+            }
+
+            for (mn in cn.methods ?: emptyList()) {
+                detectInvalidAnnotations(
+                    cn.name, allowAnnotation,
+                    mn.visibleAnnotations, mn.invisibleAnnotations,
+                    "method", cn.name, mn.name, mn.desc
+                )
+
+                val an = mn.findAnyAnnotation(visibilityAnnotations) ?: continue
+                val policy = an.policy ?: continue
+                methodPolicies[MethodKey(mn.name, mn.desc)] = policy
+
+                if (policy.policy != FilterPolicy.Substitute) continue
+
+                // Handle substitution
+                val suffix = getAnnotationField(an, "suffix", false) ?: "\$ravenwood"
+                val replacement = mn.name + suffix
+
+                if (replacement == mn.name) {
+                    errors.onErrorFound("@SubstituteWith require a different name")
+                } else {
+                    // The replacement method has to be renamed
+                    methodPolicies[MethodKey(replacement, mn.desc)] =
+                        FilterPolicy.Keep.withReason(REASON_ANNOTATION)
+                    renamedMethods[MethodKey(replacement, mn.desc)] = mn.name
+
+                    log.v("Substitution found: %s%s -> %s", replacement, mn.desc, mn.name)
+                }
+            }
         }
 
         /**
-         * Every time we see a different class, we scan all its methods for substitution attributes,
-         * and compute (implicit) policies caused by them.
+         * Throw if an item has more than one visibility annotations, or the class is not allowed
          *
-         * For example, for the following methods:
-         *
-         *   @Stub
-         *   @Substitute(suffix = "_host")
-         *   private void foo() {
-         *      // This isn't supported on the host side.
-         *   }
-         *   private void foo_host() {
-         *      // Host side implementation
-         *   }
-         *
-         * We internally handle them as:
-         *
-         *   foo() -> Remove
-         *   foo_host() -> Stub, and then rename it to foo().
+         * name1 - 4 are only used in exception messages. We take them as separate strings
+         * to avoid unnecessary string concatenations.
          */
-        private fun setClass(cn: ClassNode) {
-            if (currentClass == cn) {
-                return
-            }
-            // If the class is changing, we'll rebuild the internal structure.
-            currentClass = cn
-
-            policiesFromSubstitution.clear()
-            substituteToMethods.clear()
-
-            for (mn in cn.methods ?: emptyList()) {
-                findAnyAnnotation(substituteAnnotations,
-                        mn.visibleAnnotations,
-                        mn.invisibleAnnotations)?.let { an ->
-
-                    // Find the policy for this method.
-                    val policy = outermostFilter.getPolicyForMethod(cn.name, mn.name, mn.desc)
-                            .policy.resolveClassWidePolicy()
-                    // Make sure it's either Stub or Keep.
-                    if (!(policy.needsInStub || policy.needsInImpl)) {
-                        // TODO: Use the real annotation names in the message
-                        errors.onErrorFound("@SubstituteWith must have either @Stub or @Keep")
-                        return@let
-                    }
-                    if (!policy.isUsableWithMethods) {
-                        throw HostStubGenInternalException("Policy $policy shouldn't show up here")
-                    }
-
-                    val suffix = getAnnotationField(an, "suffix", false) ?: "\$ravenwood"
-                    val renameFrom = mn.name + suffix
-                    val renameTo = mn.name
-
-                    if (renameFrom == renameTo) {
-                        errors.onErrorFound("@SubstituteWith have a different name")
-                        return@let
-                    }
-
-                    // This mn has "SubstituteWith". This means,
-                    // 1. Re move the "rename-to" method, so add it to substitutedMethods.
-                    policiesFromSubstitution[MethodKey(renameTo, mn.desc)] =
-                            FilterPolicy.Remove.withReason("substitute-to")
-
-                    // If the policy is "stub", use "stub".
-                    // Otherwise, it must be "keep" or "throw", but there's no point in using
-                    // "throw", so let's use "keep".
-                    val newPolicy = if (policy.needsInStub) policy else FilterPolicy.Keep
-                    // 2. We also keep the from-to in the map.
-                    policiesFromSubstitution[MethodKey(renameFrom, mn.desc)] =
-                            newPolicy.withReason("substitute-from")
-                    substituteToMethods[MethodKey(renameFrom, mn.desc)] = renameTo
-
-                    log.v("Substitution found: %s%s -> %s", renameFrom, mn.desc, renameTo)
+        private fun detectInvalidAnnotations(
+            className: String,
+            allowAnnotation: Boolean,
+            visibles: List<AnnotationNode>?,
+            invisibles: List<AnnotationNode>?,
+            type: String,
+            name1: String,
+            name2: String = "",
+            name3: String = "",
+        ) {
+            var count = 0
+            var visibleCount = 0
+            for (an in visibles ?: emptyList()) {
+                if (visibilityAnnotations.contains(an.desc)) {
+                    visibleCount++
                 }
+                if (allAnnotations.contains(an.desc)) {
+                    count++
+                }
+            }
+            for (an in invisibles ?: emptyList()) {
+                if (visibilityAnnotations.contains(an.desc)) {
+                    visibleCount++
+                }
+                if (allAnnotations.contains(an.desc)) {
+                    count++
+                }
+            }
+            if (count > 0 && !allowAnnotation) {
+                throw InvalidAnnotationException(
+                    "Class ${className.toHumanReadableClassName()} is not allowed to have " +
+                            "Ravenwood annotations. Contact g/ravenwood for more details."
+                )
+            }
+            if (visibleCount > 1) {
+                val description = if (name2 == "" && name3 == "") {
+                    "$type $name1"
+                } else {
+                    "$type $name1.$name2$name3"
+                }
+                throw InvalidAnnotationException(
+                    "Found more than one visibility annotations on $description"
+                )
             }
         }
     }
@@ -409,8 +302,11 @@
     /**
      * Return the (String) value of 'value' parameter from an annotation.
      */
-    private fun getAnnotationField(an: AnnotationNode, name: String,
-                                   required: Boolean = true): String? {
+    private fun getAnnotationField(
+        an: AnnotationNode,
+        name: String,
+        required: Boolean = true
+    ): String? {
         try {
             val suffix = findAnnotationValueAsString(an, name)
             if (suffix == null && required) {
@@ -424,6 +320,9 @@
     }
 
     companion object {
+        private const val REASON_ANNOTATION = "annotation"
+        private const val REASON_CLASS_ANNOTATION = "class-annotation"
+
         /**
          * Convert from human-readable type names (e.g. "com.android.TypeName") to the internal type
          * names (e.g. "Lcom/android/TypeName).
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ClassWidePolicyPropagatingFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ClassWidePolicyPropagatingFilter.kt
index 47790b1..8ee3a94 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ClassWidePolicyPropagatingFilter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ClassWidePolicyPropagatingFilter.kt
@@ -16,7 +16,7 @@
 package com.android.hoststubgen.filters
 
 import com.android.hoststubgen.asm.ClassNodes
-import com.android.hoststubgen.asm.getDirectOuterClassName
+import com.android.hoststubgen.asm.isNative
 
 /**
  * This is used as the second last fallback filter. This filter propagates the class-wide policy
@@ -24,74 +24,80 @@
  */
 class ClassWidePolicyPropagatingFilter(
     private val classes: ClassNodes,
-    fallback: OutputFilter,
-    ) : DelegatingFilter(fallback) {
+    fallback: OutputFilter
+) : DelegatingFilter(fallback) {
 
-    private fun getClassWidePolicy(className: String, resolve: Boolean): FilterPolicyWithReason? {
+    /**
+     * We don't use ClassNode.outerClass, because it gives as the top-level
+     * outer class (A$B$C -> A), not the direct outer class (A$B$C -> A$B).
+     *
+     * Sometimes a class name includes `$`, but is not as a nested class name separator
+     * (e.g. a class name like `MyClass$$`). In this case, `MyClass$` is not actually a class.
+     *
+     * So before getting the class policy on a nonexistent class, which may cause an
+     * incorrect result, we make sure the class actually exists.
+     */
+    private fun getDirectOuterClass(className: String): String? {
         var currentClass = className
-
-
-        // If the class name is `a.b.c.A$B$C`, then we try to get the class wide policy
-        // from a.b.c.A$B$C, then a.b.c.A$B, and then a.b.c.A.
         while (true) {
-            // Sometimes a class name has a `$` in it but not as a nest class name separator --
-            // e.g. class name like "MyClass$$". In this case, `MyClass$` may not actually be
-            // a class name.
-            // So before getting the class policy on a nonexistent class, which may cause an
-            // incorrect result, we make sure if className actually exists.
-            if (classes.hasClass(className)) {
-                outermostFilter.getPolicyForClass(className).let { policy ->
-                    if (policy.policy.isClassWidePolicy) {
-                        val p = if (resolve) {
-                            policy.policy.resolveClassWidePolicy()
-                        } else {
-                            policy.policy
-                        }
-
-                        return p.withReason(policy.reason)
-                            .wrapReason("class-wide in $currentClass")
-                    }
-                    // If the class's policy is remove, then remove it.
-                    if (policy.policy == FilterPolicy.Remove) {
-                        return FilterPolicy.Remove.withReason("class-wide in $currentClass")
-                    }
-                }
-            }
-
-            // Next, look at the outer class...
-            val outer = getDirectOuterClassName(currentClass)
-            if (outer == null) {
+            val pos = currentClass.lastIndexOf('$')
+            if (pos < 0) {
                 return null
             }
-            currentClass = outer
+            currentClass = currentClass.substring(0, pos)
+            if (classes.hasClass(currentClass)) {
+                return currentClass
+            }
         }
     }
 
+    private fun getClassWidePolicy(className: String, resolve: Boolean): FilterPolicyWithReason? {
+        outermostFilter.getPolicyForClass(className).let { policy ->
+            if (policy.policy == FilterPolicy.KeepClass) {
+                val p = if (resolve) {
+                    policy.policy.resolveClassWidePolicy()
+                } else {
+                    policy.policy
+                }
+
+                return p.withReason(policy.reason)
+                    .wrapReason("class-wide in $className")
+            }
+            // If the class's policy is remove, then remove it.
+            if (policy.policy == FilterPolicy.Remove) {
+                return FilterPolicy.Remove.withReason("class-wide in $className")
+            }
+        }
+        return null
+    }
+
     override fun getPolicyForClass(className: String): FilterPolicyWithReason {
-        // If it's a nested class, use the outer class's policy.
-        getDirectOuterClassName(className)?.let { outerName ->
-            getClassWidePolicy(outerName, resolve = false)?.let { policy ->
-                return policy
-            }
-        }
-
-        return super.getPolicyForClass(className)
+        // If the class name is `a.b.c.A$B$C`, then we try to get the class wide policy
+        // from a.b.c.A$B$C, then a.b.c.A$B, and then a.b.c.A, recursively
+        return getDirectOuterClass(className)?.let { getClassWidePolicy(it, resolve = false) }
+            ?: super.getPolicyForClass(className)
     }
 
-    override fun getPolicyForField(
-            className: String,
-            fieldName: String
-    ): FilterPolicyWithReason {
+    override fun getPolicyForField(className: String, fieldName: String): FilterPolicyWithReason {
         return getClassWidePolicy(className, resolve = true)
                 ?: super.getPolicyForField(className, fieldName)
     }
 
     override fun getPolicyForMethod(
-            className: String,
-            methodName: String,
-            descriptor: String
+        className: String,
+        methodName: String,
+        descriptor: String
     ): FilterPolicyWithReason {
-        return getClassWidePolicy(className, resolve = true)
-                ?: super.getPolicyForMethod(className, methodName, descriptor)
+        return outermostFilter.getNativeSubstitutionClass(className)?.let {
+            // First check native substitution
+            classes.findMethod(className, methodName, descriptor)?.let { mn ->
+                if (mn.isNative()) {
+                    FilterPolicy.NativeSubstitute.withReason("class-wide in $className")
+                } else {
+                    null
+                }
+            }
+        } ?: getClassWidePolicy(className, resolve = true)
+        ?: super.getPolicyForMethod(className, methodName, descriptor)
     }
 }
\ No newline at end of file
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ConstantFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ConstantFilter.kt
index 678e6ea..be3c59c 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ConstantFilter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ConstantFilter.kt
@@ -26,23 +26,17 @@
  * @param policy the policy. Cannot be a "substitute" policy.
  */
 class ConstantFilter(
-        policy: FilterPolicy,
-        val reason: String
+    policy: FilterPolicy,
+    private val reason: String
 ) : OutputFilter() {
-    val classPolicy: FilterPolicy
-    val fieldPolicy: FilterPolicy
-    val methodPolicy: FilterPolicy
+
+    private val classPolicy: FilterPolicy
+    private val fieldPolicy: FilterPolicy
+    private val methodPolicy: FilterPolicy
 
     init {
-        if (policy.isSubstitute) {
-            throw HostStubGenInternalException(
-                    "ConstantFilter doesn't allow substitution policies.")
-        }
-        if (policy.isClassWidePolicy) {
-            // We prevent it, because there's no point in using class-wide policies because
-            // all members get othe same policy too anyway.
-            throw HostStubGenInternalException(
-                    "ConstantFilter doesn't allow class-wide policies.")
+        if (!policy.isUsableWithDefault) {
+            throw HostStubGenInternalException("ConstantFilter doesn't support $policy.")
         }
         methodPolicy = policy
 
@@ -63,10 +57,10 @@
     }
 
     override fun getPolicyForMethod(
-            className: String,
-            methodName: String,
-            descriptor: String,
-            ): FilterPolicyWithReason {
+        className: String,
+        methodName: String,
+        descriptor: String,
+    ): FilterPolicyWithReason {
         return methodPolicy.withReason(reason)
     }
 }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicy.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicy.kt
index f839444..ab03874 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicy.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicy.kt
@@ -17,37 +17,25 @@
 
 enum class FilterPolicy {
     /**
-     * Keep the item in the stub jar file, so tests can use it.
-     */
-    Stub,
-
-    /**
-     * Keep the item in the impl jar file, but not in the stub file. Tests cannot use it directly,
-     * but indirectly.
+     * Keep the item in the jar file.
      */
     Keep,
 
     /**
-     * Only used for types. Keep the class in the stub, and also all its members.
-     * But each member can have another annotations to override it.
-     */
-    StubClass,
-
-    /**
-     * Only used for types. Keep the class in the impl, not in the stub, and also all its members.
-     * But each member can have another annotations to override it.
+     * Only usable with classes. Keep the class in the jar, and also all its members.
+     * Each member can have another policy to override it.
      */
     KeepClass,
 
     /**
-     * Same as [Stub], but replace it with a "substitution" method. Only usable with methods.
+     * Only usable with methods. Replace a method with a "substitution" method.
      */
-    SubstituteAndStub,
+    Substitute,
 
     /**
-     * Same as [Keep], but replace it with a "substitution" method. Only usable with methods.
+     * Only usable with methods. Replace a native method with a "substitution" method,
      */
-    SubstituteAndKeep,
+    NativeSubstitute,
 
     /**
      * Only usable with methods. The item will be kept in the impl jar file, but when called,
@@ -57,7 +45,7 @@
 
     /**
      * Only usable with methods. The item will be kept in the impl jar file, but when called,
-     * it'll no-op.  Currently only supported for methods returning `void`.
+     * it'll no-op.
      */
     Ignore,
 
@@ -66,20 +54,19 @@
      */
     Remove;
 
-    val isSubstitute: Boolean
-        get() = this == SubstituteAndStub || this == SubstituteAndKeep
-
-    val needsInStub: Boolean
-        get() = this == Stub || this == StubClass || this == SubstituteAndStub || this == Ignore
-
-    val needsInImpl: Boolean
-        get() = this != Remove
+    val needsInOutput: Boolean
+        get() {
+            return when (this) {
+                Remove -> false
+                else -> true
+            }
+        }
 
     /** Returns whether a policy can be used with classes */
     val isUsableWithClasses: Boolean
         get() {
             return when (this) {
-                Stub, StubClass, Keep, KeepClass, Remove -> true
+                Keep, KeepClass, Remove -> true
                 else -> false
             }
         }
@@ -88,7 +75,7 @@
     val isUsableWithFields: Boolean
         get() {
             return when (this) {
-                Stub, Keep, Remove -> true
+                Keep, Remove -> true
                 else -> false
             }
         }
@@ -97,16 +84,16 @@
     val isUsableWithMethods: Boolean
         get() {
             return when (this) {
-                StubClass, KeepClass -> false
+                KeepClass -> false
                 else -> true
             }
         }
 
-    /** Returns whether a policy is a class-wide one. */
-    val isClassWidePolicy: Boolean
+    /** Returns whether a policy can be used as default policy. */
+    val isUsableWithDefault: Boolean
         get() {
             return when (this) {
-                StubClass, KeepClass -> true
+                Keep, Throw, Remove -> true
                 else -> false
             }
         }
@@ -116,25 +103,24 @@
         get() {
             return when (this) {
                 // TODO: handle native method with no substitution as being unsupported
-                Stub, StubClass, Keep, KeepClass, SubstituteAndStub, SubstituteAndKeep -> true
+                Keep, KeepClass, Substitute, NativeSubstitute -> true
                 else -> false
             }
         }
 
-    fun getSubstitutionBasePolicy(): FilterPolicy {
-        return when (this) {
-            SubstituteAndKeep -> Keep
-            SubstituteAndStub -> Stub
-            else -> this
+    val isMethodRewriteBody: Boolean
+        get() {
+            return when (this) {
+                NativeSubstitute, Throw, Ignore -> true
+                else -> false
+            }
         }
-    }
 
     /**
-     * Convert {Stub,Keep}Class to the corresponding Stub or Keep.
+     * Convert KeepClass to Keep, or return itself.
      */
     fun resolveClassWidePolicy(): FilterPolicy {
         return when (this) {
-            StubClass -> Stub
             KeepClass -> Keep
             else -> this
         }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicyWithReason.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicyWithReason.kt
index eb03f66..b10165b 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicyWithReason.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/FilterPolicyWithReason.kt
@@ -30,36 +30,6 @@
         return FilterPolicyWithReason(policy, "$reason [inner-reason: ${this.reason}]")
     }
 
-    /**
-     * If the visibility is lower than "Keep" (meaning if it's "remove"),
-     * then return a new [FilterPolicy] with "Keep".
-     * Otherwise, return itself
-     */
-    fun promoteToKeep(promotionReason: String): FilterPolicyWithReason {
-        if (policy.needsInImpl) {
-            return this
-        }
-        val newPolicy = if (policy.isClassWidePolicy) FilterPolicy.KeepClass else FilterPolicy.Keep
-
-        return FilterPolicyWithReason(newPolicy,
-                "$promotionReason [original remove reason: ${this.reason}]")
-    }
-
-    /**
-     * If the visibility is above "Keep" (meaning if it's "stub"),
-     * then return a new [FilterPolicy] with "Keep".
-     * Otherwise, return itself
-     */
-    fun demoteToKeep(promotionReason: String): FilterPolicyWithReason {
-        if (!policy.needsInStub) {
-            return this
-        }
-        val newPolicy = if (policy.isClassWidePolicy) FilterPolicy.KeepClass else FilterPolicy.Keep
-
-        return FilterPolicyWithReason(newPolicy,
-                "$promotionReason [original stub reason: ${this.reason}]")
-    }
-
     override fun toString(): String {
         return "[$policy - reason: $reason]"
     }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ImplicitOutputFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ImplicitOutputFilter.kt
index 5a26fc6..474da6d 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ImplicitOutputFilter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/ImplicitOutputFilter.kt
@@ -25,7 +25,6 @@
 import com.android.hoststubgen.asm.isAutoGeneratedEnumMember
 import com.android.hoststubgen.asm.isEnum
 import com.android.hoststubgen.asm.isSynthetic
-import com.android.hoststubgen.asm.isVisibilityPrivateOrPackagePrivate
 import com.android.hoststubgen.log
 import org.objectweb.asm.tree.ClassNode
 
@@ -53,7 +52,7 @@
             }
             // If the outer class needs to be in impl, it should be in impl too.
             val outerPolicy = outermostFilter.getPolicyForClass(cn.outerClass)
-            if (outerPolicy.policy.needsInImpl) {
+            if (outerPolicy.policy.needsInOutput) {
                 return FilterPolicy.KeepClass.withReason("anonymous-inner-class")
             }
         }
@@ -79,19 +78,6 @@
         val fallback = super.getPolicyForMethod(className, methodName, descriptor)
         val classPolicy = outermostFilter.getPolicyForClass(className)
 
-        // If the class is in the stub, then we need to put the private constructor in the stub too,
-        // to prevent the class from getting instantiated.
-        if (classPolicy.policy.needsInStub &&
-                !fallback.policy.needsInStub &&
-                (methodName == "<init>") && // Constructor?
-                (descriptor == "()V")) { // Has zero parameters?
-            classes.findMethod(className, methodName, descriptor)?.let { mn ->
-                if (isVisibilityPrivateOrPackagePrivate(mn.access)) {
-                    return FilterPolicy.Stub.withReason("private constructor in stub class")
-                }
-            }
-        }
-
         val cn = classes.getClass(className)
 
         // If we throw from the static initializer, the class would be useless, so we convert it
@@ -107,7 +93,7 @@
         }
 
         log.d("Class ${cn.name} Class policy: $classPolicy")
-        if (classPolicy.policy.needsInImpl) {
+        if (classPolicy.policy.needsInOutput) {
             // Do it only when the class needs to be kept...
 
             // Member policy should be "keep" or "stub".
@@ -152,7 +138,7 @@
         val classPolicy = outermostFilter.getPolicyForClass(className)
 
         log.d("Class ${cn.name} Class policy: $classPolicy")
-        if (classPolicy.policy.needsInImpl) {
+        if (classPolicy.policy.needsInOutput) {
             // Do it only when the class needs to be kept...
 
             // Member policy should be "keep" or "stub".
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/NativeFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/NativeFilter.kt
new file mode 100644
index 0000000..bd71931
--- /dev/null
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/NativeFilter.kt
@@ -0,0 +1,40 @@
+/*
+ * 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.hoststubgen.filters
+
+import com.android.hoststubgen.asm.ClassNodes
+import com.android.hoststubgen.asm.isNative
+
+class NativeFilter(
+    private val classes: ClassNodes,
+    fallback: OutputFilter
+) : DelegatingFilter(fallback) {
+    override fun getPolicyForMethod(
+        className: String,
+        methodName: String,
+        descriptor: String,
+    ): FilterPolicyWithReason {
+        return classes.findMethod(className, methodName, descriptor)?.let { mn ->
+            // For native methods that weren't handled by outer filters,
+            // we keep it so that native method registration will not crash.
+            if (mn.isNative()) {
+                FilterPolicy.Keep.withReason("native-preserve")
+            } else {
+                null
+            }
+        } ?: super.getPolicyForMethod(className, methodName, descriptor)
+    }
+}
\ No newline at end of file
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/StubIntersectingFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/StubIntersectingFilter.kt
deleted file mode 100644
index f92a027..0000000
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/StubIntersectingFilter.kt
+++ /dev/null
@@ -1,91 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.filters
-
-import com.android.hoststubgen.HostStubGenErrors
-import com.android.hoststubgen.asm.ClassNodes
-
-private const val REASON = "demoted, not in intersect jars"
-
-/**
- * An [OutputFilter] that will restrict what to put in stub to only what shows up in "intersecting
- * jar" files.
- *
- * For example, if the Android public API stub jar is provided, then the HostStubGen's output
- * stub will be restricted to public APIs.
- */
-class StubIntersectingFilter(
-        private val errors: HostStubGenErrors,
-        /**
-         * If a class / field / method is not in any of these jars, then we will not put it in
-         * stub.
-         */
-        private val intersectingJars: Map<String, ClassNodes>,
-        fallback: OutputFilter,
-) : DelegatingFilter(fallback) {
-    private inline fun exists(predicate: (ClassNodes) -> Boolean): Boolean {
-        intersectingJars.forEach { entry ->
-            if (predicate(entry.value)) {
-                return true
-            }
-        }
-        return false
-    }
-
-    /**
-     * If [origPolicy] is less than "Stub", then return it as-is.
-     *
-     * Otherwise, call [inStubChecker] to see if the API is in any of [intersectingJars].
-     * If yes, then return [origPolicy] as-is. Otherwise, demote to "Keep".
-     */
-    private fun intersectWithStub(
-            origPolicy: FilterPolicyWithReason,
-            inStubChecker: () -> Boolean,
-    ): FilterPolicyWithReason {
-        if (origPolicy.policy.needsInStub) {
-            // Only check the stub jars, when the class is supposed to be in stub otherwise.
-            if (!inStubChecker()) {
-                return origPolicy.demoteToKeep(REASON)
-            }
-        }
-        return origPolicy
-    }
-
-    override fun getPolicyForClass(className: String): FilterPolicyWithReason {
-        return intersectWithStub(super.getPolicyForClass(className)) {
-            exists { classes -> classes.findClass(className) != null }
-        }
-    }
-
-    override fun getPolicyForField(
-            className: String,
-            fieldName: String
-    ): FilterPolicyWithReason {
-        return intersectWithStub(super.getPolicyForField(className, fieldName)) {
-            exists { classes -> classes.findField(className, fieldName) != null }
-        }
-    }
-
-    override fun getPolicyForMethod(
-            className: String,
-            methodName: String,
-            descriptor: String
-    ): FilterPolicyWithReason {
-        return intersectWithStub(super.getPolicyForMethod(className, methodName, descriptor)) {
-            exists { classes -> classes.findMethod(className, methodName, descriptor) != null }
-        }
-    }
-}
\ No newline at end of file
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/TextFileFilterPolicyParser.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/TextFileFilterPolicyParser.kt
index 53bcf10..14fd82b 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/TextFileFilterPolicyParser.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/filters/TextFileFilterPolicyParser.kt
@@ -240,7 +240,7 @@
 
                             imf.setPolicyForMethod(className, name, signature,
                                     policy.withReason(FILTER_REASON))
-                            if (policy.isSubstitute) {
+                            if (policy == FilterPolicy.Substitute) {
                                 val fromName = fields[3].substring(1)
 
                                 if (fromName == name) {
@@ -248,10 +248,9 @@
                                             "Substitution must have a different name")
                                 }
 
-                                // Set the policy  for the "from" method.
+                                // Set the policy for the "from" method.
                                 imf.setPolicyForMethod(className, fromName, signature,
-                                        policy.getSubstitutionBasePolicy()
-                                                .withReason(FILTER_REASON))
+                                    FilterPolicy.Keep.withReason(FILTER_REASON))
 
                                 val classAndMethod = splitWithLastPeriod(fromName)
                                 if (classAndMethod != null) {
@@ -346,18 +345,14 @@
 
 private fun parsePolicy(s: String): FilterPolicy {
     return when (s.lowercase()) {
-        "s", "stub" -> FilterPolicy.Stub
         "k", "keep" -> FilterPolicy.Keep
         "t", "throw" -> FilterPolicy.Throw
         "r", "remove" -> FilterPolicy.Remove
-        "sc", "stubclass" -> FilterPolicy.StubClass
         "kc", "keepclass" -> FilterPolicy.KeepClass
         "i", "ignore" -> FilterPolicy.Ignore
         else -> {
             if (s.startsWith("@")) {
-                FilterPolicy.SubstituteAndStub
-            } else if (s.startsWith("%")) {
-                FilterPolicy.SubstituteAndKeep
+                FilterPolicy.Substitute
             } else {
                 throw ParseException("Invalid policy \"$s\"")
             }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/utils/ClassFilter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/utils/ClassFilter.kt
index 01a7ab3..7440b94 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/utils/ClassFilter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/utils/ClassFilter.kt
@@ -24,19 +24,19 @@
 /**
  * General purpose filter for class names.
  */
-class ClassFilter private constructor (
-        val defaultResult: Boolean,
+class ClassFilter private constructor(
+    private val defaultResult: Boolean,
 ) {
-    private data class FilterElement(
-            val allowed: Boolean,
-            val internalName: String,
-            val isPrefix: Boolean,
+    private class FilterElement(
+        val allowed: Boolean,
+        val internalName: String,
+        val isPrefix: Boolean,
     ) {
         fun matches(classInternalName: String): Boolean {
-            if (isPrefix) {
-                return classInternalName.startsWith(internalName)
+            return if (isPrefix) {
+                classInternalName.startsWith(internalName)
             } else {
-                return classInternalName == internalName
+                classInternalName == internalName
             }
         }
     }
@@ -54,15 +54,16 @@
             return it
         }
 
-        var result = defaultResult
-        run outer@{
-            elements.forEach { e ->
-                if (e.matches(classInternalName)) {
-                    result = e.allowed
-                    return@outer // break equivalent.
-                }
+        val testClasses = sequence {
+            // Yield itself and its outer class(es) one by one
+            var idx = classInternalName.length
+            while (idx > 0) {
+                yield(classInternalName.substring(0, idx))
+                idx = classInternalName.lastIndexOf('$', idx - 1)
             }
         }
+
+        val result = elements.find { testClasses.any(it::matches) }?.allowed ?: defaultResult
         cache[classInternalName] = result
 
         return result
@@ -87,9 +88,9 @@
 
         /** Build a filter from a string (for unit tests). */
         fun buildFromString(
-                filterString: String,
-                defaultResult: Boolean,
-                filenameForErrorMessage: String
+            filterString: String,
+            defaultResult: Boolean,
+            filenameForErrorMessage: String
         ): ClassFilter {
             val ret = ClassFilter(defaultResult)
 
@@ -119,17 +120,20 @@
 
                 // Handle wildcard -- e.g. "package.name.*"
                 if (line.endsWith(".*")) {
-                    ret.elements.add(FilterElement(
-                            allow, line.substring(0, line.length - 2).toJvmClassName(), true))
+                    ret.elements.add(
+                        FilterElement(
+                            allow, line.substring(0, line.length - 2).toJvmClassName(), true
+                        )
+                    )
                     return@forEach
                 }
 
                 // Any other uses of "*" would be an error.
                 if (line.contains('*')) {
                     throw ParseException(
-                            "Wildcard (*) can only show up as the last element",
-                            filenameForErrorMessage,
-                            lineNo
+                        "Wildcard (*) can only show up as the last element",
+                        filenameForErrorMessage,
+                        lineNo
                     )
                 }
                 ret.elements.add(FilterElement(allow, line.toJvmClassName(), false))
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BaseAdapter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BaseAdapter.kt
index bad0449..41ba928 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BaseAdapter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BaseAdapter.kt
@@ -26,67 +26,47 @@
 import com.android.hoststubgen.filters.FilterPolicy
 import com.android.hoststubgen.filters.FilterPolicyWithReason
 import com.android.hoststubgen.filters.OutputFilter
-import com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-import com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
+import com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 import com.android.hoststubgen.log
+import java.io.PrintWriter
 import org.objectweb.asm.ClassVisitor
 import org.objectweb.asm.FieldVisitor
 import org.objectweb.asm.MethodVisitor
 import org.objectweb.asm.Opcodes
 import org.objectweb.asm.commons.ClassRemapper
-import org.objectweb.asm.commons.Remapper
 import org.objectweb.asm.util.TraceClassVisitor
-import java.io.PrintWriter
 
-val OPCODE_VERSION = Opcodes.ASM9
+const val OPCODE_VERSION = Opcodes.ASM9
 
-abstract class BaseAdapter (
-        protected val classes: ClassNodes,
-        nextVisitor: ClassVisitor,
-        protected val filter: OutputFilter,
-        protected val options: Options,
+abstract class BaseAdapter(
+    protected val classes: ClassNodes,
+    nextVisitor: ClassVisitor,
+    protected val filter: OutputFilter,
+    protected val options: Options,
 ) : ClassVisitor(OPCODE_VERSION, nextVisitor) {
 
     /**
      * Options to control the behavior.
      */
-    data class Options (
-            val errors: HostStubGenErrors,
-            val stats: HostStubGenStats?,
-            val enablePreTrace: Boolean,
-            val enablePostTrace: Boolean,
-            val enableNonStubMethodCallDetection: Boolean,
-            )
+    data class Options(
+        val errors: HostStubGenErrors,
+        val stats: HostStubGenStats?,
+        val enablePreTrace: Boolean,
+        val enablePostTrace: Boolean
+    )
 
     protected lateinit var currentPackageName: String
     protected lateinit var currentClassName: String
     protected var nativeSubstitutionClass: String? = null
     protected lateinit var classPolicy: FilterPolicyWithReason
 
-    /**
-     * Return whether an item with a given policy should be included in the output.
-     */
-    protected abstract fun shouldEmit(policy: FilterPolicy): Boolean
-
-    /**
-     * Inject [HostStubGenKeptInStub] and [HostStubGenKeptInImpl] as needed to an item.
-     */
-    protected fun injectInStubAndKeepAnnotations(policy: FilterPolicy, v: UnifiedVisitor) {
-        if (policy.needsInStub) {
-            v.visitAnnotation(HostStubGenKeptInStub.CLASS_DESCRIPTOR, true)
-        }
-        if (policy.needsInImpl) {
-            v.visitAnnotation(HostStubGenKeptInImpl.CLASS_DESCRIPTOR, true)
-        }
-    }
-
     override fun visit(
-            version: Int,
-            access: Int,
-            name: String,
-            signature: String?,
-            superName: String?,
-            interfaces: Array<String>,
+        version: Int,
+        access: Int,
+        name: String,
+        signature: String?,
+        superName: String?,
+        interfaces: Array<String>,
     ) {
         super.visit(version, access, name, signature, superName, interfaces)
         currentClassName = name
@@ -103,21 +83,25 @@
                 .toJvmClassName()
             log.d("  NativeSubstitutionClass: $fullClassName")
             if (classes.findClass(fullClassName) == null) {
-                log.w("Native substitution class $fullClassName not found. Class must be " +
-                        "available at runtime.")
+                log.w(
+                    "Native substitution class $fullClassName not found. Class must be " +
+                            "available at runtime."
+                )
             } else {
                 // If the class exists, it must have a KeepClass policy.
                 if (filter.getPolicyForClass(fullClassName).policy != FilterPolicy.KeepClass) {
                     // TODO: Use real annotation name.
                     options.errors.onErrorFound(
-                            "Native substitution class $fullClassName should have @Keep.")
+                        "Native substitution class $fullClassName should have @Keep."
+                    )
                 }
             }
 
             nativeSubstitutionClass = fullClassName
         }
+
         // Inject annotations to generated classes.
-        injectInStubAndKeepAnnotations(classPolicy.policy, UnifiedVisitor.on(this))
+        UnifiedVisitor.on(this).visitAnnotation(HostStubGenProcessedAsKeep.CLASS_DESCRIPTOR, true)
     }
 
     override fun visitEnd() {
@@ -141,11 +125,11 @@
     }
 
     override fun visitField(
-            access: Int,
-            name: String,
-            descriptor: String,
-            signature: String?,
-            value: Any?,
+        access: Int,
+        name: String,
+        descriptor: String,
+        signature: String?,
+        value: Any?,
     ): FieldVisitor? {
         if (skipMemberModificationNestCount > 0) {
             return super.visitField(access, name, descriptor, signature, value)
@@ -154,7 +138,7 @@
         log.d("visitField: %s %s [%x] Policy: %s", name, descriptor, access, policy)
 
         log.withIndent {
-            if (!shouldEmit(policy.policy)) {
+            if (policy.policy == FilterPolicy.Remove) {
                 log.d("Removing %s %s", name, policy)
                 return null
             }
@@ -162,18 +146,19 @@
             log.v("Emitting field: %s %s %s", name, descriptor, policy)
             val ret = super.visitField(access, name, descriptor, signature, value)
 
-            injectInStubAndKeepAnnotations(policy.policy, UnifiedVisitor.on(ret))
+            UnifiedVisitor.on(ret)
+                .visitAnnotation(HostStubGenProcessedAsKeep.CLASS_DESCRIPTOR, true)
 
             return ret
         }
     }
 
     override fun visitMethod(
-            access: Int,
-            name: String,
-            descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
+        access: Int,
+        name: String,
+        descriptor: String,
+        signature: String?,
+        exceptions: Array<String>?,
     ): MethodVisitor? {
         if (skipMemberModificationNestCount > 0) {
             return super.visitMethod(access, name, descriptor, signature, exceptions)
@@ -187,11 +172,11 @@
             // Instead of this method, we rename the substitute-to method with the original
             // name, in the "Maybe rename the method" part below.
             val policy = filter.getPolicyForMethod(currentClassName, name, descriptor)
-            if (policy.policy.isSubstitute) {
+            if (policy.policy == FilterPolicy.Substitute) {
                 log.d("Skipping %s%s %s", name, descriptor, policy)
                 return null
             }
-            if (!shouldEmit(p.policy)) {
+            if (p.policy == FilterPolicy.Remove) {
                 log.d("Removing %s%s %s", name, descriptor, policy)
                 return null
             }
@@ -209,13 +194,16 @@
                 // `name` is the name of the method we're currently visiting, so it's usually a
                 // "...$ravewnwood" name.
                 newAccess = checkSubstitutionMethodCompatibility(
-                        classes, currentClassName, newName, name, descriptor, options.errors)
+                    classes, currentClassName, newName, name, descriptor, options.errors
+                )
                 if (newAccess == NOT_COMPATIBLE) {
                     return null
                 }
 
-                log.v("Emitting %s.%s%s as %s %s", currentClassName, name, descriptor,
-                        newName, policy)
+                log.v(
+                    "Emitting %s.%s%s as %s %s", currentClassName, name, descriptor,
+                    newName, policy
+                )
             } else {
                 log.v("Emitting method: %s%s %s", name, descriptor, policy)
                 newName = name
@@ -225,14 +213,17 @@
             // But note, we only use it when calling the super's method,
             // but not for visitMethodInner(), because when subclass wants to change access,
             // it can do so inside visitMethodInner().
-            newAccess = updateAccessFlags(newAccess, name, descriptor)
+            newAccess = updateAccessFlags(newAccess, name, descriptor, policy.policy)
 
-            val ret = visitMethodInner(access, newName, descriptor, signature, exceptions, policy,
+            val ret = visitMethodInner(
+                access, newName, descriptor, signature, exceptions, policy,
                 renameTo != null,
-                super.visitMethod(newAccess, newName, descriptor, signature, exceptions))
+                super.visitMethod(newAccess, newName, descriptor, signature, exceptions)
+            )
 
             ret?.let {
-                injectInStubAndKeepAnnotations(policy.policy, UnifiedVisitor.on(ret))
+                UnifiedVisitor.on(ret)
+                    .visitAnnotation(HostStubGenProcessedAsKeep.CLASS_DESCRIPTOR, true)
             }
 
             return ret
@@ -240,9 +231,10 @@
     }
 
     open fun updateAccessFlags(
-            access: Int,
-            name: String,
-            descriptor: String,
+        access: Int,
+        name: String,
+        descriptor: String,
+        policy: FilterPolicy,
     ): Int {
         return access
     }
@@ -256,7 +248,7 @@
         policy: FilterPolicyWithReason,
         substituted: Boolean,
         superVisitor: MethodVisitor?,
-        ): MethodVisitor?
+    ): MethodVisitor?
 
     companion object {
         fun getVisitor(
@@ -265,8 +257,6 @@
             nextVisitor: ClassVisitor,
             filter: OutputFilter,
             packageRedirector: PackageRedirectRemapper,
-            remapper: Remapper?,
-            forImpl: Boolean,
             options: Options,
         ): ClassVisitor {
             var next = nextVisitor
@@ -289,23 +279,20 @@
                 if (!packageRedirector.isTarget(classInternalName)) {
                     next = ClassRemapper(next, packageRedirector)
                 } else {
-                    log.v("Class $classInternalName is a redirect-from class, not applying" +
-                            " --package-redirect")
+                    log.v(
+                        "Class $classInternalName is a redirect-from class, not applying" +
+                                " --package-redirect"
+                    )
                 }
             }
 
-            var ret: ClassVisitor
-            if (forImpl) {
-                ret = ImplGeneratingAdapter(classes, next, filter, options)
-            } else {
-                ret = StubGeneratingAdapter(classes, next, filter, options)
-            }
+            next = ImplGeneratingAdapter(classes, next, filter, options)
 
             // Inject TraceClassVisitor for debugging.
             if (options.enablePreTrace) {
-                ret = TraceClassVisitor(ret, verbosePrinter)
+                next = TraceClassVisitor(next, verbosePrinter)
             }
-            return ret
+            return next
         }
     }
 }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BodyReplacingMethodVisitor.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BodyReplacingMethodVisitor.kt
index 8250412..55d0c0e 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BodyReplacingMethodVisitor.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/BodyReplacingMethodVisitor.kt
@@ -20,38 +20,23 @@
 import org.objectweb.asm.Handle
 import org.objectweb.asm.Label
 import org.objectweb.asm.MethodVisitor
-import org.objectweb.asm.Opcodes
 import org.objectweb.asm.TypePath
 
 /**
- * A method visitor that removes everything from method body.
+ * A method visitor that creates or replaces a method body.
  *
- * To inject a method body, override [visitCode] and create the opcodes there.
+ * Override [emitNewCode] to build the method body.
  */
 abstract class BodyReplacingMethodVisitor(
-    access: Int,
-    name: String,
-    descriptor: String,
-    signature: String?,
-    exceptions: Array<String>?,
-    next: MethodVisitor?,
+    private val createBody: Boolean,
+    next: MethodVisitor?
 ) : MethodVisitor(OPCODE_VERSION, next) {
-    val isVoid: Boolean
-    val isStatic: Boolean
-
-    init {
-        isVoid = descriptor.endsWith(")V")
-        isStatic = access and Opcodes.ACC_STATIC != 0
-    }
 
     // Following methods are for things that we need to keep.
     // Since they're all calling the super method, we can just remove them, but we keep them
     // just to clarify what we're keeping.
 
-    final override fun visitParameter(
-            name: String?,
-            access: Int
-    ) {
+    final override fun visitParameter(name: String?, access: Int) {
         super.visitParameter(name, access)
     }
 
@@ -59,10 +44,7 @@
         return super.visitAnnotationDefault()
     }
 
-    final override fun visitAnnotation(
-            descriptor: String?,
-            visible: Boolean
-    ): AnnotationVisitor? {
+    final override fun visitAnnotation(descriptor: String?, visible: Boolean): AnnotationVisitor? {
         return super.visitAnnotation(descriptor, visible)
     }
 
@@ -75,17 +57,14 @@
         return super.visitTypeAnnotation(typeRef, typePath, descriptor, visible)
     }
 
-    final override fun visitAnnotableParameterCount(
-            parameterCount: Int,
-            visible: Boolean
-    ) {
+    final override fun visitAnnotableParameterCount(parameterCount: Int, visible: Boolean) {
         super.visitAnnotableParameterCount(parameterCount, visible)
     }
 
     final override fun visitParameterAnnotation(
-            parameter: Int,
-            descriptor: String?,
-            visible: Boolean
+        parameter: Int,
+        descriptor: String?,
+        visible: Boolean
     ): AnnotationVisitor? {
         return super.visitParameterAnnotation(parameter, descriptor, visible)
     }
@@ -94,10 +73,6 @@
         super.visitAttribute(attribute)
     }
 
-    override fun visitEnd() {
-        super.visitEnd()
-    }
-
     /**
      * Control when to emit the code. We use this to ignore all visitXxx method calls caused by
      * the original method, so we'll remove all the original code.
@@ -108,9 +83,18 @@
      * (See also https://asm.ow2.io/asm4-guide.pdf section 3.2.1 about the MethovVisitor
      * call order.)
      */
-    var emitCode = false
+    private var emitCode = false
+
+    /**
+     * This value will be set as true when [visitCode] is called. In [visitEnd], if this value
+     * is still false, this means that the original method does not have a body.
+     *
+     * We want to forcefully inject a method body in [visitEnd] if [createBody] is true.
+     */
+    private var visitedCode = false
 
     final override fun visitCode() {
+        visitedCode = true
         super.visitCode()
 
         try {
@@ -122,15 +106,19 @@
         }
     }
 
+    final override fun visitEnd() {
+        if (!visitedCode && createBody) {
+            visitCode()
+        }
+        super.visitEnd()
+    }
+
     /**
      * Subclass must implement it and emit code, and call [visitMaxs] at the end.
      */
     abstract fun emitNewCode()
 
-    final override fun visitMaxs(
-            maxStack: Int,
-            maxLocals: Int
-    ) {
+    final override fun visitMaxs(maxStack: Int, maxLocals: Int) {
         if (emitCode) {
             super.visitMaxs(maxStack, maxLocals)
         }
@@ -140,11 +128,11 @@
     // emit any of them, so they are all no-op.
 
     final override fun visitFrame(
-            type: Int,
-            numLocal: Int,
-            local: Array<out Any>?,
-            numStack: Int,
-            stack: Array<out Any>?
+        type: Int,
+        numLocal: Int,
+        local: Array<out Any>?,
+        numStack: Int,
+        stack: Array<out Any>?
     ) {
         if (emitCode) {
             super.visitFrame(type, numLocal, local, numStack, stack)
@@ -157,38 +145,29 @@
         }
     }
 
-    final override fun visitIntInsn(
-            opcode: Int,
-            operand: Int
-    ) {
+    final override fun visitIntInsn(opcode: Int, operand: Int) {
         if (emitCode) {
             super.visitIntInsn(opcode, operand)
         }
     }
 
-    final override fun visitVarInsn(
-            opcode: Int,
-            varIndex: Int
-    ) {
+    final override fun visitVarInsn(opcode: Int, varIndex: Int) {
         if (emitCode) {
             super.visitVarInsn(opcode, varIndex)
         }
     }
 
-    final override fun visitTypeInsn(
-            opcode: Int,
-            type: String?
-    ) {
+    final override fun visitTypeInsn(opcode: Int, type: String?) {
         if (emitCode) {
             super.visitTypeInsn(opcode, type)
         }
     }
 
     final override fun visitFieldInsn(
-            opcode: Int,
-            owner: String?,
-            name: String?,
-            descriptor: String?
+        opcode: Int,
+        owner: String?,
+        name: String?,
+        descriptor: String?
     ) {
         if (emitCode) {
             super.visitFieldInsn(opcode, owner, name, descriptor)
@@ -196,11 +175,11 @@
     }
 
     final override fun visitMethodInsn(
-            opcode: Int,
-            owner: String?,
-            name: String?,
-            descriptor: String?,
-            isInterface: Boolean
+        opcode: Int,
+        owner: String?,
+        name: String?,
+        descriptor: String?,
+        isInterface: Boolean
     ) {
         if (emitCode) {
             super.visitMethodInsn(opcode, owner, name, descriptor, isInterface)
@@ -208,21 +187,20 @@
     }
 
     final override fun visitInvokeDynamicInsn(
-            name: String?,
-            descriptor: String?,
-            bootstrapMethodHandle: Handle?,
-            vararg bootstrapMethodArguments: Any?
+        name: String?,
+        descriptor: String?,
+        bootstrapMethodHandle: Handle?,
+        vararg bootstrapMethodArguments: Any?
     ) {
         if (emitCode) {
-            super.visitInvokeDynamicInsn(name, descriptor, bootstrapMethodHandle,
-                    *bootstrapMethodArguments)
+            super.visitInvokeDynamicInsn(
+                name, descriptor, bootstrapMethodHandle,
+                *bootstrapMethodArguments
+            )
         }
     }
 
-    final override fun visitJumpInsn(
-            opcode: Int,
-            label: Label?
-    ) {
+    final override fun visitJumpInsn(opcode: Int, label: Label?) {
         if (emitCode) {
             super.visitJumpInsn(opcode, label)
         }
@@ -240,20 +218,17 @@
         }
     }
 
-    final override fun visitIincInsn(
-            varIndex: Int,
-            increment: Int
-    ) {
+    final override fun visitIincInsn(varIndex: Int, increment: Int) {
         if (emitCode) {
             super.visitIincInsn(varIndex, increment)
         }
     }
 
     final override fun visitTableSwitchInsn(
-            min: Int,
-            max: Int,
-            dflt: Label?,
-            vararg labels: Label?
+        min: Int,
+        max: Int,
+        dflt: Label?,
+        vararg labels: Label?
     ) {
         if (emitCode) {
             super.visitTableSwitchInsn(min, max, dflt, *labels)
@@ -261,29 +236,26 @@
     }
 
     final override fun visitLookupSwitchInsn(
-            dflt: Label?,
-            keys: IntArray?,
-            labels: Array<out Label>?
+        dflt: Label?,
+        keys: IntArray?,
+        labels: Array<out Label>?
     ) {
         if (emitCode) {
             super.visitLookupSwitchInsn(dflt, keys, labels)
         }
     }
 
-    final override fun visitMultiANewArrayInsn(
-            descriptor: String?,
-            numDimensions: Int
-    ) {
+    final override fun visitMultiANewArrayInsn(descriptor: String?, numDimensions: Int) {
         if (emitCode) {
             super.visitMultiANewArrayInsn(descriptor, numDimensions)
         }
     }
 
     final override fun visitInsnAnnotation(
-            typeRef: Int,
-            typePath: TypePath?,
-            descriptor: String?,
-            visible: Boolean
+        typeRef: Int,
+        typePath: TypePath?,
+        descriptor: String?,
+        visible: Boolean
     ): AnnotationVisitor? {
         if (emitCode) {
             return super.visitInsnAnnotation(typeRef, typePath, descriptor, visible)
@@ -292,10 +264,10 @@
     }
 
     final override fun visitTryCatchBlock(
-            start: Label?,
-            end: Label?,
-            handler: Label?,
-            type: String?
+        start: Label?,
+        end: Label?,
+        handler: Label?,
+        type: String?
     ) {
         if (emitCode) {
             super.visitTryCatchBlock(start, end, handler, type)
@@ -303,10 +275,10 @@
     }
 
     final override fun visitTryCatchAnnotation(
-            typeRef: Int,
-            typePath: TypePath?,
-            descriptor: String?,
-            visible: Boolean
+        typeRef: Int,
+        typePath: TypePath?,
+        descriptor: String?,
+        visible: Boolean
     ): AnnotationVisitor? {
         if (emitCode) {
             return super.visitTryCatchAnnotation(typeRef, typePath, descriptor, visible)
@@ -315,12 +287,12 @@
     }
 
     final override fun visitLocalVariable(
-            name: String?,
-            descriptor: String?,
-            signature: String?,
-            start: Label?,
-            end: Label?,
-            index: Int
+        name: String?,
+        descriptor: String?,
+        signature: String?,
+        start: Label?,
+        end: Label?,
+        index: Int
     ) {
         if (emitCode) {
             super.visitLocalVariable(name, descriptor, signature, start, end, index)
@@ -328,25 +300,23 @@
     }
 
     final override fun visitLocalVariableAnnotation(
-            typeRef: Int,
-            typePath: TypePath?,
-            start: Array<out Label>?,
-            end: Array<out Label>?,
-            index: IntArray?,
-            descriptor: String?,
-            visible: Boolean
+        typeRef: Int,
+        typePath: TypePath?,
+        start: Array<out Label>?,
+        end: Array<out Label>?,
+        index: IntArray?,
+        descriptor: String?,
+        visible: Boolean
     ): AnnotationVisitor? {
         if (emitCode) {
             return super.visitLocalVariableAnnotation(
-                    typeRef, typePath, start, end, index, descriptor, visible)
+                typeRef, typePath, start, end, index, descriptor, visible
+            )
         }
         return null
     }
 
-    final override fun visitLineNumber(
-            line: Int,
-            start: Label?
-    ) {
+    final override fun visitLineNumber(line: Int, start: Label?) {
         if (emitCode) {
             super.visitLineNumber(line, start)
         }
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/ImplGeneratingAdapter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/ImplGeneratingAdapter.kt
index 3d2e142..057d653 100644
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/ImplGeneratingAdapter.kt
+++ b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/ImplGeneratingAdapter.kt
@@ -18,7 +18,6 @@
 import com.android.hoststubgen.asm.CLASS_INITIALIZER_DESC
 import com.android.hoststubgen.asm.CLASS_INITIALIZER_NAME
 import com.android.hoststubgen.asm.ClassNodes
-import com.android.hoststubgen.asm.isVisibilityPrivateOrPackagePrivate
 import com.android.hoststubgen.asm.prependArgTypeToMethodDescriptor
 import com.android.hoststubgen.asm.writeByteCodeToPushArguments
 import com.android.hoststubgen.asm.writeByteCodeToReturn
@@ -42,16 +41,12 @@
  * An adapter that generates the "impl" class file from an input class file.
  */
 class ImplGeneratingAdapter(
-        classes: ClassNodes,
-        nextVisitor: ClassVisitor,
-        filter: OutputFilter,
-        options: Options,
+    classes: ClassNodes,
+    nextVisitor: ClassVisitor,
+    filter: OutputFilter,
+    options: Options,
 ) : BaseAdapter(classes, nextVisitor, filter, options) {
 
-    override fun shouldEmit(policy: FilterPolicy): Boolean {
-        return policy.needsInImpl
-    }
-
     private var classLoadHooks: List<String> = emptyList()
 
     override fun visit(
@@ -107,14 +102,14 @@
     private fun writeClassLoadHookCalls(mv: MethodVisitor) {
         classLoadHooks.forEach { classLoadHook ->
             // First argument: the class type.
-            mv.visitLdcInsn(Type.getType("L" + currentClassName + ";"))
+            mv.visitLdcInsn(Type.getType("L$currentClassName;"))
 
             // Second argument: method name
             mv.visitLdcInsn(classLoadHook)
 
             // Call HostTestUtils.onClassLoaded().
             mv.visitMethodInsn(
-                Opcodes.INVOKESTATIC,
+                INVOKESTATIC,
                 HostTestUtils.CLASS_INTERNAL_NAME,
                 "onClassLoaded",
                 "(Ljava/lang/Class;Ljava/lang/String;)V",
@@ -124,69 +119,49 @@
     }
 
     override fun updateAccessFlags(
-            access: Int,
-            name: String,
-            descriptor: String,
+        access: Int,
+        name: String,
+        descriptor: String,
+        policy: FilterPolicy,
     ): Int {
-        if ((access and Opcodes.ACC_NATIVE) != 0 && nativeSubstitutionClass != null) {
+        if (policy.isMethodRewriteBody) {
+            // If we are rewriting the entire method body, we need
+            // to convert native methods to non-native
             return access and Opcodes.ACC_NATIVE.inv()
         }
         return access
     }
 
     override fun visitMethodInner(
-            access: Int,
-            name: String,
-            descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            policy: FilterPolicyWithReason,
-            substituted: Boolean,
-            superVisitor: MethodVisitor?,
+        access: Int,
+        name: String,
+        descriptor: String,
+        signature: String?,
+        exceptions: Array<String>?,
+        policy: FilterPolicyWithReason,
+        substituted: Boolean,
+        superVisitor: MethodVisitor?,
     ): MethodVisitor? {
-        // Inject method log, if needed.
         var innerVisitor = superVisitor
 
         //  If method logging is enabled, inject call to the logging method.
         val methodCallHooks = filter.getMethodCallHooks(currentClassName, name, descriptor)
         if (methodCallHooks.isNotEmpty()) {
             innerVisitor = MethodCallHookInjectingAdapter(
-                access,
                 name,
                 descriptor,
-                signature,
-                exceptions,
-                innerVisitor,
                 methodCallHooks,
-                )
+                innerVisitor,
+            )
         }
 
         // If this class already has a class initializer and a class load hook is needed, then
         // we inject code.
         if (classLoadHooks.isNotEmpty() &&
             name == CLASS_INITIALIZER_NAME &&
-            descriptor == CLASS_INITIALIZER_DESC) {
-            innerVisitor = ClassLoadHookInjectingMethodAdapter(
-                access,
-                name,
-                descriptor,
-                signature,
-                exceptions,
-                innerVisitor,
-            )
-        }
-
-        // If non-stub method call detection is enabled, then inject a call to the checker.
-        if (options.enableNonStubMethodCallDetection && doesMethodNeedNonStubCallCheck(
-                access, name, descriptor, policy) ) {
-            innerVisitor = NonStubMethodCallDetectingAdapter(
-                    access,
-                    name,
-                    descriptor,
-                    signature,
-                    exceptions,
-                    innerVisitor,
-            )
+            descriptor == CLASS_INITIALIZER_DESC
+        ) {
+            innerVisitor = ClassLoadHookInjectingMethodAdapter(innerVisitor)
         }
 
         fun MethodVisitor.withAnnotation(descriptor: String): MethodVisitor {
@@ -195,34 +170,31 @@
         }
 
         log.withIndent {
-            var willThrow = false
-            if (policy.policy == FilterPolicy.Throw) {
-                log.v("Making method throw...")
-                willThrow = true
-                innerVisitor = ThrowingMethodAdapter(
-                    access, name, descriptor, signature, exceptions, innerVisitor)
-                    .withAnnotation(HostStubGenProcessedAsThrow.CLASS_DESCRIPTOR)
+            // When we encounter native methods, we want to forcefully
+            // inject a method body. Also see [updateAccessFlags].
+            val forceCreateBody = (access and Opcodes.ACC_NATIVE) != 0
+            when (policy.policy) {
+                FilterPolicy.Throw -> {
+                    log.v("Making method throw...")
+                    return ThrowingMethodAdapter(forceCreateBody, innerVisitor)
+                        .withAnnotation(HostStubGenProcessedAsThrow.CLASS_DESCRIPTOR)
+                }
+                FilterPolicy.Ignore -> {
+                    log.v("Making method ignored...")
+                    return IgnoreMethodAdapter(descriptor, forceCreateBody, innerVisitor)
+                        .withAnnotation(HostStubGenProcessedAsIgnore.CLASS_DESCRIPTOR)
+                }
+                FilterPolicy.NativeSubstitute -> {
+                    log.v("Rewriting native method...")
+                    return NativeSubstitutingMethodAdapter(access, name, descriptor, innerVisitor)
+                        .withAnnotation(HostStubGenProcessedAsSubstitute.CLASS_DESCRIPTOR)
+                }
+                else -> {}
             }
-            if ((access and Opcodes.ACC_NATIVE) != 0 && nativeSubstitutionClass != null) {
-                log.v("Rewriting native method...")
-                return NativeSubstitutingMethodAdapter(
-                        access, name, descriptor, signature, exceptions, innerVisitor)
-                    .withAnnotation(HostStubGenProcessedAsSubstitute.CLASS_DESCRIPTOR)
-            }
-            if (willThrow) {
-                return innerVisitor
-            }
+        }
 
-            if (policy.policy == FilterPolicy.Ignore) {
-                log.v("Making method ignored...")
-                return IgnoreMethodAdapter(
-                    access, name, descriptor, signature, exceptions, innerVisitor)
-                    .withAnnotation(HostStubGenProcessedAsIgnore.CLASS_DESCRIPTOR)
-            }
-            if (filter.hasAnyMethodCallReplace()) {
-                innerVisitor = MethodCallReplacingAdapter(
-                    access, name, descriptor, signature, exceptions, innerVisitor)
-            }
+        if (filter.hasAnyMethodCallReplace()) {
+            innerVisitor = MethodCallReplacingAdapter(name, innerVisitor)
         }
         if (substituted) {
             innerVisitor?.withAnnotation(HostStubGenProcessedAsSubstitute.CLASS_DESCRIPTOR)
@@ -231,53 +203,32 @@
         return innerVisitor
     }
 
-    fun doesMethodNeedNonStubCallCheck(
-            access: Int,
-            name: String,
-            descriptor: String,
-            policy: FilterPolicyWithReason,
-    ): Boolean {
-        // If a method is in the stub, then no need to check.
-        if (policy.policy.needsInStub) {
-            return false
-        }
-        // If a method is private or package-private, no need to check.
-        // Technically test code can use framework package name, so it's a bit too lenient.
-        if (isVisibilityPrivateOrPackagePrivate(access)) {
-            return false
-        }
-        // TODO: If the method overrides a method that's accessible by tests, then we shouldn't
-        // do the check. (e.g. overrides a stub method or java standard method.)
-
-        return true
-    }
-
     /**
      * A method adapter that replaces the method body with a HostTestUtils.onThrowMethodCalled()
      * call.
      */
     private inner class ThrowingMethodAdapter(
-            access: Int,
-            val name: String,
-            descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?
-    ) : BodyReplacingMethodVisitor(access, name, descriptor, signature, exceptions, next) {
+        createBody: Boolean,
+        next: MethodVisitor?
+    ) : BodyReplacingMethodVisitor(createBody, next) {
         override fun emitNewCode() {
-            visitMethodInsn(Opcodes.INVOKESTATIC,
-                    HostTestUtils.CLASS_INTERNAL_NAME,
-                    "onThrowMethodCalled",
-                    "()V",
-                    false)
+            visitMethodInsn(
+                INVOKESTATIC,
+                HostTestUtils.CLASS_INTERNAL_NAME,
+                "onThrowMethodCalled",
+                "()V",
+                false
+            )
 
             // We still need a RETURN opcode for the return type.
             // For now, let's just inject a `throw`.
             visitTypeInsn(Opcodes.NEW, "java/lang/RuntimeException")
             visitInsn(Opcodes.DUP)
             visitLdcInsn("Unreachable")
-            visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/RuntimeException",
-                    "<init>", "(Ljava/lang/String;)V", false)
+            visitMethodInsn(
+                Opcodes.INVOKESPECIAL, "java/lang/RuntimeException",
+                "<init>", "(Ljava/lang/String;)V", false
+            )
             visitInsn(Opcodes.ATHROW)
 
             // visitMaxs(3, if (isStatic) 0 else 1)
@@ -289,13 +240,10 @@
      * A method adapter that replaces the method body with a no-op return.
      */
     private inner class IgnoreMethodAdapter(
-            access: Int,
-            name: String,
-            val descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?
-    ) : BodyReplacingMethodVisitor(access, name, descriptor, signature, exceptions, next) {
+        val descriptor: String,
+        createBody: Boolean,
+        next: MethodVisitor?
+    ) : BodyReplacingMethodVisitor(createBody, next) {
         override fun emitNewCode() {
             when (Type.getReturnType(descriptor)) {
                 Type.VOID_TYPE -> visitInsn(Opcodes.RETURN)
@@ -326,30 +274,24 @@
     }
 
     /**
-     * A method adapter that replaces a native method call with a call to the "native substitution"
-     * class.
+     * A method adapter that rewrite a native method body with a
+     * call to a method in the "native substitution" class.
      */
     private inner class NativeSubstitutingMethodAdapter(
-            val access: Int,
-            private val name: String,
-            private val descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?
-    ) : MethodVisitor(OPCODE_VERSION, next) {
-        override fun visitCode() {
-            throw RuntimeException("NativeSubstitutingMethodVisitor should be called on " +
-                    " native method, where visitCode() shouldn't be called.")
-        }
+        access: Int,
+        private val name: String,
+        private val descriptor: String,
+        next: MethodVisitor?
+    ) : BodyReplacingMethodVisitor(true, next) {
 
-        override fun visitEnd() {
-            super.visitCode()
+        private val isStatic = (access and Opcodes.ACC_STATIC) != 0
 
+        override fun emitNewCode() {
             var targetDescriptor = descriptor
             var argOffset = 0
 
             // For non-static native method, we need to tweak it a bit.
-            if ((access and Opcodes.ACC_STATIC) == 0) {
+            if (!isStatic) {
                 // Push `this` as the first argument.
                 this.visitVarInsn(Opcodes.ALOAD, 0)
 
@@ -366,16 +308,17 @@
 
             writeByteCodeToPushArguments(descriptor, this, argOffset)
 
-            visitMethodInsn(Opcodes.INVOKESTATIC,
-                    nativeSubstitutionClass,
-                    name,
-                    targetDescriptor,
-                    false)
+            visitMethodInsn(
+                INVOKESTATIC,
+                nativeSubstitutionClass,
+                name,
+                targetDescriptor,
+                false
+            )
 
             writeByteCodeToReturn(descriptor, this)
 
             visitMaxs(99, 0) // We let ASM figure them out.
-            super.visitEnd()
         }
     }
 
@@ -386,25 +329,22 @@
      * `this(...)`. The logging code will be injected *before* such calls.
      */
     private inner class MethodCallHookInjectingAdapter(
-            access: Int,
-            val name: String,
-            val descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?,
-            val hooks: List<String>,
+        val name: String,
+        val descriptor: String,
+        val hooks: List<String>,
+        next: MethodVisitor?,
     ) : MethodVisitor(OPCODE_VERSION, next) {
         override fun visitCode() {
             super.visitCode()
 
             hooks.forEach { hook ->
-                mv.visitLdcInsn(Type.getType("L" + currentClassName + ";"))
+                mv.visitLdcInsn(Type.getType("L$currentClassName;"))
                 visitLdcInsn(name)
                 visitLdcInsn(descriptor)
                 visitLdcInsn(hook)
 
                 visitMethodInsn(
-                    Opcodes.INVOKESTATIC,
+                    INVOKESTATIC,
                     HostTestUtils.CLASS_INTERNAL_NAME,
                     "callMethodCallHook",
                     "(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V",
@@ -418,11 +358,6 @@
      * Inject a class load hook call.
      */
     private inner class ClassLoadHookInjectingMethodAdapter(
-        access: Int,
-        val name: String,
-        val descriptor: String,
-        signature: String?,
-        exceptions: Array<String>?,
         next: MethodVisitor?
     ) : MethodVisitor(OPCODE_VERSION, next) {
         override fun visitCode() {
@@ -432,53 +367,8 @@
         }
     }
 
-    /**
-     * A method adapter that detects calls to non-stub methods.
-     */
-    private inner class NonStubMethodCallDetectingAdapter(
-            access: Int,
-            val name: String,
-            val descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?
-    ) : MethodVisitor(OPCODE_VERSION, next) {
-        override fun visitCode() {
-            super.visitCode()
-
-            // First three arguments to HostTestUtils.onNonStubMethodCalled().
-            visitLdcInsn(currentClassName)
-            visitLdcInsn(name)
-            visitLdcInsn(descriptor)
-
-            // Call: HostTestUtils.getStackWalker().getCallerClass().
-            // This push the caller Class in the stack.
-            visitMethodInsn(Opcodes.INVOKESTATIC,
-                    HostTestUtils.CLASS_INTERNAL_NAME,
-                    "getStackWalker",
-                    "()Ljava/lang/StackWalker;",
-                    false)
-            visitMethodInsn(Opcodes.INVOKEVIRTUAL,
-                    "java/lang/StackWalker",
-                    "getCallerClass",
-                    "()Ljava/lang/Class;",
-                    false)
-
-            // Then call onNonStubMethodCalled().
-            visitMethodInsn(Opcodes.INVOKESTATIC,
-                    HostTestUtils.CLASS_INTERNAL_NAME,
-                    "onNonStubMethodCalled",
-                    "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V",
-                    false)
-        }
-    }
-
     private inner class MethodCallReplacingAdapter(
-        access: Int,
         val callerMethodName: String,
-        val descriptor: String,
-        signature: String?,
-        exceptions: Array<String>?,
         next: MethodVisitor?,
     ) : MethodVisitor(OPCODE_VERSION, next) {
         override fun visitMethodInsn(
@@ -497,7 +387,8 @@
                 }
             }
             val to = filter.getMethodCallReplaceTo(
-                currentClassName, callerMethodName, owner!!, name!!, descriptor!!)
+                currentClassName, callerMethodName, owner!!, name!!, descriptor!!
+            )
 
             if (to == null
                 // Don't replace if the target is the callsite.
diff --git a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/StubGeneratingAdapter.kt b/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/StubGeneratingAdapter.kt
deleted file mode 100644
index fc20f28..0000000
--- a/tools/hoststubgen/hoststubgen/src/com/android/hoststubgen/visitors/StubGeneratingAdapter.kt
+++ /dev/null
@@ -1,86 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.visitors
-
-import com.android.hoststubgen.asm.ClassNodes
-import com.android.hoststubgen.filters.FilterPolicy
-import com.android.hoststubgen.filters.FilterPolicyWithReason
-import com.android.hoststubgen.filters.OutputFilter
-import com.android.hoststubgen.log
-import org.objectweb.asm.ClassVisitor
-import org.objectweb.asm.MethodVisitor
-import org.objectweb.asm.Opcodes
-
-/**
- * An adapter that generates the "impl" class file from an input class file.
- */
-class StubGeneratingAdapter(
-        classes: ClassNodes,
-        nextVisitor: ClassVisitor,
-        filter: OutputFilter,
-        options: Options,
-) : BaseAdapter(classes, nextVisitor, filter, options) {
-
-    override fun shouldEmit(policy: FilterPolicy): Boolean {
-        return policy.needsInStub
-    }
-
-    override fun visitMethodInner(
-            access: Int,
-            name: String,
-            descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            policy: FilterPolicyWithReason,
-            substituted: Boolean,
-            superVisitor: MethodVisitor?,
-    ): MethodVisitor? {
-        return StubMethodVisitor(access, name, descriptor, signature, exceptions, superVisitor)
-    }
-
-    private inner class StubMethodVisitor(
-            access: Int,
-            val name: String,
-            descriptor: String,
-            signature: String?,
-            exceptions: Array<String>?,
-            next: MethodVisitor?
-    ) : BodyReplacingMethodVisitor(access, name, descriptor, signature, exceptions, next) {
-        override fun emitNewCode() {
-            log.d("  Generating stub method for $currentClassName.$name")
-
-            // Inject the following code:
-            //   throw new RuntimeException("Stub!");
-
-            /*
-                NEW java/lang/RuntimeException
-                DUP
-                LDC "not supported on host side"
-                INVOKESPECIAL java/lang/RuntimeException.<init> (Ljava/lang/String;)V
-                ATHROW
-                MAXSTACK = 3
-                MAXLOCALS = 2 <- 1 for this, 1 for return value.
-             */
-            visitTypeInsn(Opcodes.NEW, "java/lang/RuntimeException")
-            visitInsn(Opcodes.DUP)
-            visitLdcInsn("Stub!")
-            visitMethodInsn(Opcodes.INVOKESPECIAL, "java/lang/RuntimeException",
-                    "<init>", "(Ljava/lang/String;)V", false)
-            visitInsn(Opcodes.ATHROW)
-            visitMaxs(0, 0) // We let ASM figure them out.
-        }
-    }
-}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/Android.bp b/tools/hoststubgen/hoststubgen/test-tiny-framework/Android.bp
index e7873d6..ba2c869 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/Android.bp
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/Android.bp
@@ -21,7 +21,7 @@
 
 // Create stub/impl jars from "hoststubgen-test-tiny-framework", using the following 3 rules.
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host",
+    name: "hoststubgen-test-tiny-framework-host-base",
     defaults: ["hoststubgen-command-defaults"],
     cmd: hoststubgen_common_options +
         "--in-jar $(location :hoststubgen-test-tiny-framework) " +
@@ -35,25 +35,13 @@
 }
 
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-stub",
+    name: "hoststubgen-test-tiny-framework-host",
     cmd: "cp $(in) $(out)",
     srcs: [
-        ":hoststubgen-test-tiny-framework-host{host_stub.jar}",
+        ":hoststubgen-test-tiny-framework-host-base{host.jar}",
     ],
     out: [
-        "host_stub.jar",
-    ],
-    visibility: ["//visibility:private"],
-}
-
-java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-impl",
-    cmd: "cp $(in) $(out)",
-    srcs: [
-        ":hoststubgen-test-tiny-framework-host{host_impl.jar}",
-    ],
-    out: [
-        "host_impl.jar",
+        "host.jar",
     ],
     visibility: ["//visibility:private"],
 }
@@ -61,7 +49,7 @@
 // Same as "hoststubgen-test-tiny-framework-host", but with more options, to test more hoststubgen
 // features.
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-ext",
+    name: "hoststubgen-test-tiny-framework-host-ext-base",
     defaults: ["hoststubgen-command-defaults"],
     cmd: hoststubgen_common_options +
         "--in-jar $(location :hoststubgen-test-tiny-framework) " +
@@ -79,37 +67,25 @@
 }
 
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-ext-stub",
+    name: "hoststubgen-test-tiny-framework-host-ext",
     cmd: "cp $(in) $(out)",
     srcs: [
-        ":hoststubgen-test-tiny-framework-host-ext{host_stub.jar}",
+        ":hoststubgen-test-tiny-framework-host-ext-base{host.jar}",
     ],
     out: [
-        "host_stub.jar",
-    ],
-    visibility: ["//visibility:private"],
-}
-
-java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-ext-impl",
-    cmd: "cp $(in) $(out)",
-    srcs: [
-        ":hoststubgen-test-tiny-framework-host-ext{host_impl.jar}",
-    ],
-    out: [
-        "host_impl.jar",
+        "host.jar",
     ],
     visibility: ["//visibility:private"],
 }
 
 // Compile the test jar, using 2 rules.
-// 1. Build the test against the stub.
+// 1. Build the test against the original framework.
 java_library_host {
     name: "hoststubgen-test-tiny-test-lib",
     srcs: ["tiny-test/src/**/*.java"],
 
     libs: [
-        "hoststubgen-test-tiny-framework-host-stub",
+        "hoststubgen-test-tiny-framework",
     ],
     static_libs: [
         "junit",
@@ -129,7 +105,7 @@
     static_libs: [
         "hoststubgen-test-tiny-test-lib",
         "hoststubgen-helper-runtime",
-        "hoststubgen-test-tiny-framework-host-impl",
+        "hoststubgen-test-tiny-framework-host",
     ],
     test_suites: ["general-tests"],
 }
@@ -149,49 +125,25 @@
 }
 
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-stub-dump",
+    name: "hoststubgen-test-tiny-framework-host-dump",
     defaults: ["hoststubgen-jar-dump-defaults"],
     srcs: [
-        ":hoststubgen-test-tiny-framework-host-stub",
+        ":hoststubgen-test-tiny-framework-host",
     ],
     out: [
-        "02-hoststubgen-test-tiny-framework-host-stub-dump.txt",
+        "03-hoststubgen-test-tiny-framework-host-dump.txt",
     ],
     visibility: ["//visibility:private"],
 }
 
 java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-impl-dump",
+    name: "hoststubgen-test-tiny-framework-host-ext-dump",
     defaults: ["hoststubgen-jar-dump-defaults"],
     srcs: [
-        ":hoststubgen-test-tiny-framework-host-impl",
+        ":hoststubgen-test-tiny-framework-host-ext",
     ],
     out: [
-        "03-hoststubgen-test-tiny-framework-host-impl-dump.txt",
-    ],
-    visibility: ["//visibility:private"],
-}
-
-java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-ext-stub-dump",
-    defaults: ["hoststubgen-jar-dump-defaults"],
-    srcs: [
-        ":hoststubgen-test-tiny-framework-host-ext-stub",
-    ],
-    out: [
-        "12-hoststubgen-test-tiny-framework-host-ext-stub-dump.txt",
-    ],
-    visibility: ["//visibility:private"],
-}
-
-java_genrule_host {
-    name: "hoststubgen-test-tiny-framework-host-ext-impl-dump",
-    defaults: ["hoststubgen-jar-dump-defaults"],
-    srcs: [
-        ":hoststubgen-test-tiny-framework-host-ext-impl",
-    ],
-    out: [
-        "13-hoststubgen-test-tiny-framework-host-ext-impl-dump.txt",
+        "13-hoststubgen-test-tiny-framework-host-ext-dump.txt",
     ],
     visibility: ["//visibility:private"],
 }
@@ -206,11 +158,9 @@
         "golden-output/*.txt",
     ],
     java_data: [
-        "hoststubgen-test-tiny-framework-host-stub-dump",
-        "hoststubgen-test-tiny-framework-host-impl-dump",
         "hoststubgen-test-tiny-framework-orig-dump",
-        "hoststubgen-test-tiny-framework-host-ext-stub-dump",
-        "hoststubgen-test-tiny-framework-host-ext-impl-dump",
+        "hoststubgen-test-tiny-framework-host-dump",
+        "hoststubgen-test-tiny-framework-host-ext-dump",
     ],
     test_suites: ["general-tests"],
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/annotation-allowed-classes-tiny-framework.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/annotation-allowed-classes-tiny-framework.txt
index bd9e85e..de4cb0c 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/annotation-allowed-classes-tiny-framework.txt
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/annotation-allowed-classes-tiny-framework.txt
@@ -6,10 +6,10 @@
 
 
 # To allow a specific class to use annotations:
-# com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
+# com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
 
 # To disallow a specific class to use annotations:
-# !com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
+# !com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
 
 # To allow a specific package to use annotations:
 # com.android.hoststubgen.test.*
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/01-hoststubgen-test-tiny-framework-orig-dump.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/01-hoststubgen-test-tiny-framework-orig-dump.txt
index c2f593c..5fde14f 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/01-hoststubgen-test-tiny-framework-orig-dump.txt
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/01-hoststubgen-test-tiny-framework-orig-dump.txt
@@ -104,26 +104,6 @@
     java.lang.annotation.Retention(
       value=Ljava/lang/annotation/RetentionPolicy;.CLASS
     )
-## Class: android/hosttest/annotation/HostSideTestStub.class
-  Compiled from "HostSideTestStub.java"
-public interface android.hosttest.annotation.HostSideTestStub extends java.lang.annotation.Annotation
-  minor version: 0
-  major version: 61
-  flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "HostSideTestStub.java"
-RuntimeVisibleAnnotations:
-  x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
-    java.lang.annotation.Target(
-      value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
-    )
-  x: #x(#x=e#x.#x)
-    java.lang.annotation.Retention(
-      value=Ljava/lang/annotation/RetentionPolicy;.CLASS
-    )
 ## Class: android/hosttest/annotation/HostSideTestSubstitute.class
   Compiled from "HostSideTestSubstitute.java"
 public interface android.hosttest.annotation.HostSideTestSubstitute extends java.lang.annotation.Annotation
@@ -187,26 +167,6 @@
     java.lang.annotation.Retention(
       value=Ljava/lang/annotation/RetentionPolicy;.CLASS
     )
-## Class: android/hosttest/annotation/HostSideTestWholeClassStub.class
-  Compiled from "HostSideTestWholeClassStub.java"
-public interface android.hosttest.annotation.HostSideTestWholeClassStub extends java.lang.annotation.Annotation
-  minor version: 0
-  major version: 61
-  flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestWholeClassStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "HostSideTestWholeClassStub.java"
-RuntimeVisibleAnnotations:
-  x: #x(#x=[e#x.#x])
-    java.lang.annotation.Target(
-      value=[Ljava/lang/annotation/ElementType;.TYPE]
-    )
-  x: #x(#x=e#x.#x)
-    java.lang.annotation.Retention(
-      value=Ljava/lang/annotation/RetentionPolicy;.CLASS
-    )
 ## Class: android/hosttest/annotation/tests/HostSideTestSuppress.class
   Compiled from "HostSideTestSuppress.java"
 public interface android.hosttest.annotation.tests.HostSideTestSuppress extends java.lang.annotation.Annotation
@@ -394,120 +354,15 @@
   com/android/hoststubgen/test/tinyframework/R$Nested
 InnerClasses:
   public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl
-  minor version: 0
-  major version: 61
-  flags: (0x0020) ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 3
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                  // Method java/lang/Object."<init>":()V
-         x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl;
-
-  public static int getOneKeep();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: iconst_1
-         x: ireturn
-      LineNumberTable:
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestKeep
-
-  public static int getOneStub();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: iconst_1
-         x: ireturn
-      LineNumberTable:
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-SourceFile: "TinyFrameworkCallerCheck.java"
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.class
+  Compiled from "TinyFrameworkAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
   minor version: 0
   major version: 61
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                         // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                  // Method java/lang/Object."<init>":()V
-         x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck;
-
-  public static int getOne_withCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: invokestatic  #x                  // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneKeep:()I
-         x: ireturn
-      LineNumberTable:
-
-  public static int getOne_noCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneStub:()I
-         x: ireturn
-      LineNumberTable:
-}
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.class
-  Compiled from "TinyFrameworkClassAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 10, attributes: 2
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
+  interfaces: 0, fields: 2, methods: 8, attributes: 2
   public int keep;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
@@ -519,7 +374,7 @@
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
 
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations();
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations();
     descriptor: ()V
     flags: (0x0001) ACC_PUBLIC
     Code:
@@ -528,42 +383,21 @@
          x: invokespecial #x                  // Method java/lang/Object."<init>":()V
          x: aload_0
          x: iconst_1
-         x: putfield      #x                  // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
+         x: putfield      #x                  // Field keep:I
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public int addOne(int);
     descriptor: (I)I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=2, locals=2, args_size=2
-         x: aload_0
-         x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-            0       6     1 value   I
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
          x: iload_1
          x: iconst_1
          x: iadd
@@ -571,7 +405,7 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
             0       4     1 value   I
     RuntimeInvisibleAnnotations:
       x: #x()
@@ -589,7 +423,7 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
             0       8     1   foo   Ljava/lang/String;
     RuntimeInvisibleAnnotations:
       x: #x()
@@ -608,11 +442,9 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
             0      10     1 value   I
     RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
       x: #x(#x=s#x)
         android.hosttest.annotation.HostSideTestSubstitute(
           suffix="_host"
@@ -630,15 +462,13 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
             0       4     1 value   I
 
   public static native int nativeAddThree(int);
     descriptor: (I)I
     flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
     RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
       x: #x(#x=s#x)
         android.hosttest.annotation.HostSideTestSubstitute(
           suffix="_host"
@@ -668,212 +498,19 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
 }
-SourceFile: "TinyFrameworkClassAnnotations.java"
+SourceFile: "TinyFrameworkAnnotations.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestClassLoadHook(
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.class
-  Compiled from "TinyFrameworkClassClassWideAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 10, attributes: 2
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-
-  public int remove;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                  // Method java/lang/Object."<init>":()V
-         x: aload_0
-         x: iconst_1
-         x: putfield      #x                  // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: aload_0
-         x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       6     1 value   I
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: iload_1
-         x: iconst_1
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       4     1 value   I
-
-  public void toBeRemoved(java.lang.String);
-    descriptor: (Ljava/lang/String;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":()V
-         x: athrow
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       8     1   foo   Ljava/lang/String;
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String not supported on host side
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0      10     1 value   I
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-      x: #x(#x=s#x)
-        android.hosttest.annotation.HostSideTestSubstitute(
-          suffix="_host"
-        )
-
-  public int addTwo_host(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: iload_1
-         x: iconst_2
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       4     1 value   I
-
-  public static native int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-      x: #x(#x=s#x)
-        android.hosttest.annotation.HostSideTestSubstitute(
-          suffix="_host"
-        )
-
-  public static int nativeAddThree_host(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=2, locals=1, args_size=1
-         x: iload_0
-         x: iconst_3
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0 value   I
-
-  public java.lang.String unsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: ldc           #x                 // String This value shouldn\'t be seen on the host side.
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-}
-SourceFile: "TinyFrameworkClassClassWideAnnotations.java"
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.class
   Compiled from "TinyFrameworkClassLoadHook.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook
@@ -935,7 +572,131 @@
 SourceFile: "TinyFrameworkClassLoadHook.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.class
+  Compiled from "TinyFrameworkClassWideAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 2, methods: 6, attributes: 2
+  public int keep;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+
+  public int remove;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestRemove
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=1, args_size=1
+         x: aload_0
+         x: invokespecial #x                  // Method java/lang/Object."<init>":()V
+         x: aload_0
+         x: iconst_1
+         x: putfield      #x                  // Field keep:I
+         x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+
+  public int addOne(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=2, args_size=2
+         x: iload_1
+         x: iconst_1
+         x: iadd
+         x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0       4     1 value   I
+
+  public int addTwo(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=3, locals=2, args_size=2
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: ldc           #x                 // String not supported on host side
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+         x: athrow
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0      10     1 value   I
+    RuntimeInvisibleAnnotations:
+      x: #x(#x=s#x)
+        android.hosttest.annotation.HostSideTestSubstitute(
+          suffix="_host"
+        )
+
+  public int addTwo_host(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=2, args_size=2
+         x: iload_1
+         x: iconst_2
+         x: iadd
+         x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0       4     1 value   I
+
+  public void toBeRemoved(java.lang.String);
+    descriptor: (Ljava/lang/String;)V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=2, args_size=2
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":()V
+         x: athrow
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0       8     1   foo   Ljava/lang/String;
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestRemove
+
+  public java.lang.String unsupportedMethod();
+    descriptor: ()Ljava/lang/String;
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=1, locals=1, args_size=1
+         x: ldc           #x                 // String This value shouldn\'t be seen on the host side.
+         x: areturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestThrow
+}
+SourceFile: "TinyFrameworkClassWideAnnotations.java"
+RuntimeInvisibleAnnotations:
+  x: #x()
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.class
   Compiled from "TinyFrameworkClassWithInitializerDefault.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault
@@ -950,14 +711,14 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault();
     descriptor: ()V
@@ -989,7 +750,7 @@
 SourceFile: "TinyFrameworkClassWithInitializerDefault.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.class
   Compiled from "TinyFrameworkClassWithInitializerStub.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub
@@ -1004,14 +765,14 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub();
     descriptor: ()V
@@ -1047,7 +808,7 @@
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.class
@@ -1064,21 +825,21 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex GREEN;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex BLUE;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private final java.lang.String mLongName;
     descriptor: Ljava/lang/String;
@@ -1158,7 +919,7 @@
     Signature: #x                          // (Ljava/lang/String;Ljava/lang/String;)V
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.lang.String getLongName();
     descriptor: ()Ljava/lang/String;
@@ -1174,7 +935,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.lang.String getShortName();
     descriptor: ()Ljava/lang/String;
@@ -1190,7 +951,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1252,7 +1013,7 @@
 SourceFile: "TinyFrameworkEnumComplex.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.class
   Compiled from "TinyFrameworkEnumSimple.java"
 public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple>
@@ -1267,14 +1028,14 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple DOG;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $VALUES;
     descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
@@ -1373,7 +1134,7 @@
 SourceFile: "TinyFrameworkEnumSimple.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.class
   Compiled from "TinyFrameworkExceptionTester.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester
@@ -1427,7 +1188,7 @@
 SourceFile: "TinyFrameworkExceptionTester.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.class
   Compiled from "TinyFrameworkForTextPolicy.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy
@@ -1436,15 +1197,11 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 19, attributes: 1
+  interfaces: 0, fields: 2, methods: 17, attributes: 1
   public int stub;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
 
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-
   public int remove;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
@@ -1459,35 +1216,17 @@
          x: aload_0
          x: iconst_1
          x: putfield      #x                  // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
 
   public int addOne(int);
     descriptor: (I)I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=2, locals=2, args_size=2
-         x: aload_0
-         x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-            0       6     1 value   I
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
          x: iload_1
          x: iconst_1
          x: iadd
@@ -1699,19 +1438,6 @@
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
             0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
 }
 SourceFile: "TinyFrameworkForTextPolicy.java"
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested.class
@@ -1729,7 +1455,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -1737,7 +1463,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested();
     descriptor: ()V
@@ -1756,7 +1482,7 @@
             0      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -1772,7 +1498,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -1785,7 +1511,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -1840,7 +1566,7 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
@@ -1883,7 +1609,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -1891,7 +1617,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas();
     descriptor: ()V
@@ -1910,7 +1636,7 @@
             0      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -1926,7 +1652,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -1939,7 +1665,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -1994,7 +1720,7 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 NestMembers:
@@ -2179,7 +1905,7 @@
 SourceFile: "TinyFrameworkMethodCallReplace.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
 BootstrapMethods:
@@ -2199,7 +1925,7 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                         // com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 11, attributes: 2
+  interfaces: 0, fields: 1, methods: 12, attributes: 2
   int value;
     descriptor: I
     flags: (0x0000)
@@ -2294,6 +2020,13 @@
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
 
+  public static native void nativeStillKeep();
+    descriptor: ()V
+    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestKeep
+
   public static void nativeStillNotSupported_should_be_like_this();
     descriptor: ()V
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
@@ -2312,7 +2045,7 @@
 SourceFile: "TinyFrameworkNative.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestNativeSubstitutionClass(
       value="TinyFrameworkNative_host"
@@ -2684,7 +2417,7 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 1, attributes: 4
+  interfaces: 0, fields: 2, methods: 1, attributes: 3
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
@@ -2717,9 +2450,6 @@
       <no name>                      final mandated
 }
 SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 InnerClasses:
   public #x= #x of #x;                  // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
@@ -2778,6 +2508,40 @@
 InnerClasses:
   public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass.class
+  Compiled from "TinyFrameworkNestedClasses.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 1, methods: 1, attributes: 3
+  public int value;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=1, args_size=1
+         x: aload_0
+         x: invokespecial #x                  // Method java/lang/Object."<init>":()V
+         x: aload_0
+         x: bipush        8
+         x: putfield      #x                  // Field value:I
+        x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0      11     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass;
+}
+SourceFile: "TinyFrameworkNestedClasses.java"
+NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+InnerClasses:
+  public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass
@@ -2786,7 +2550,7 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 4
+  interfaces: 0, fields: 1, methods: 2, attributes: 3
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
@@ -2820,13 +2584,11 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
 }
 SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 InnerClasses:
   public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$SubClass extends com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass
@@ -2937,11 +2699,12 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
+  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
@@ -2957,6 +2720,7 @@
   public static #x= #x of #x;          // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public #x= #x of #x;                 // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;          // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.class
   Compiled from "TinyFrameworkPackageRedirect.java"
@@ -2999,7 +2763,7 @@
 SourceFile: "TinyFrameworkPackageRedirect.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.class
   Compiled from "TinyFrameworkRenamedClassCaller.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller
@@ -3041,7 +2805,7 @@
 SourceFile: "TinyFrameworkRenamedClassCaller.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.class
   Compiled from "TinyFrameworkToBeRenamed.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed
@@ -3088,7 +2852,7 @@
 SourceFile: "TinyFrameworkToBeRenamed.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/packagetest/A.class
   Compiled from "A.java"
 public class com.android.hoststubgen.test.tinyframework.packagetest.A
@@ -3834,4 +3598,4 @@
 SourceFile: "UnsupportedClass.java"
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/02-hoststubgen-test-tiny-framework-host-stub-dump.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/02-hoststubgen-test-tiny-framework-host-stub-dump.txt
deleted file mode 100644
index 1b83d24..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/02-hoststubgen-test-tiny-framework-host-stub-dump.txt
+++ /dev/null
@@ -1,2829 +0,0 @@
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy.class
-  Compiled from "IPretendingAidl.java"
-public class com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub$Proxy
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub$Proxy();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub.class
-  Compiled from "IPretendingAidl.java"
-public class com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int addOne(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl.class
-  Compiled from "IPretendingAidl.java"
-public interface com.android.hoststubgen.test.tinyframework.IPretendingAidl
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 4
-}
-InnerClasses:
-  public static #x= #x of #x;            // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-  com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
-## Class: com/android/hoststubgen/test/tinyframework/R$Nested.class
-  Compiled from "R.java"
-public class com.android.hoststubgen.test.tinyframework.R$Nested
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/R$Nested
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 4
-  public static int[] ARRAY;
-    descriptor: [I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.R$Nested();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
-SourceFile: "R.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/R
-## Class: com/android/hoststubgen/test/tinyframework/R.class
-  Compiled from "R.java"
-public class com.android.hoststubgen.test.tinyframework.R
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/R
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 1, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.R();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
-SourceFile: "R.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/R$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl
-  minor version: 0
-  major version: 61
-  flags: (0x0020) ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOneStub();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-InnerClasses:
-  private static #x= #x of #x;           // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 5
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_withCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_noCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.class
-  Compiled from "TinyFrameworkClassAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 5, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-SourceFile: "TinyFrameworkClassAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestClassLoadHook(
-      value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
-    )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.class
-  Compiled from "TinyFrameworkClassClassWideAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 8, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int remove;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeRemoved(java.lang.String);
-    descriptor: (Ljava/lang/String;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String unsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassClassWideAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.class
-  Compiled from "TinyFrameworkClassLoadHook.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 3, attributes: 3
-  public static final java.util.Set<java.lang.Class<?>> sLoadedClasses;
-    descriptor: Ljava/util/Set;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/Set<Ljava/lang/Class<*>;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void onClassLoaded(java.lang.Class<?>);
-    descriptor: (Ljava/lang/Class;)V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // (Ljava/lang/Class<*>;)V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassLoadHook.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.class
-  Compiled from "TinyFrameworkClassWithInitializerDefault.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 0, attributes: 3
-  public static boolean sInitialized;
-    descriptor: Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.lang.Object sObject;
-    descriptor: Ljava/lang/Object;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-}
-SourceFile: "TinyFrameworkClassWithInitializerDefault.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.class
-  Compiled from "TinyFrameworkClassWithInitializerStub.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 0, attributes: 3
-  public static boolean sInitialized;
-    descriptor: Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.lang.Object sObject;
-    descriptor: Ljava/lang/Object;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-}
-SourceFile: "TinyFrameworkClassWithInitializerStub.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestClassLoadHook(
-      value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
-    )
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.class
-  Compiled from "TinyFrameworkEnumComplex.java"
-public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex>
-  minor version: 0
-  major version: 61
-  flags: (0x4031) ACC_PUBLIC, ACC_FINAL, ACC_SUPER, ACC_ENUM
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex
-  super_class: #x                         // java/lang/Enum
-  interfaces: 0, fields: 4, methods: 7, attributes: 4
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex RED;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex GREEN;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex BLUE;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $VALUES;
-    descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex valueOf(java.lang.String);
-    descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      mandated
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex(java.lang.String, java.lang.String);
-    descriptor: (Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;)V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=5, args_size=5
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // (Ljava/lang/String;Ljava/lang/String;)V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-    MethodParameters:
-      Name                           Flags
-      <no name>                      synthetic
-      <no name>                      synthetic
-      <no name>
-      <no name>
-
-  public java.lang.String getLongName();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.lang.String getShortName();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;>;
-SourceFile: "TinyFrameworkEnumComplex.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.class
-  Compiled from "TinyFrameworkEnumSimple.java"
-public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple>
-  minor version: 0
-  major version: 61
-  flags: (0x4031) ACC_PUBLIC, ACC_FINAL, ACC_SUPER, ACC_ENUM
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple
-  super_class: #x                         // java/lang/Enum
-  interfaces: 0, fields: 3, methods: 5, attributes: 4
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple CAT;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple DOG;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $VALUES;
-    descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple valueOf(java.lang.String);
-    descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      mandated
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple();
-    descriptor: (Ljava/lang/String;I)V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=3, args_size=3
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      synthetic
-      <no name>                      synthetic
-
-  private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;>;
-SourceFile: "TinyFrameworkEnumSimple.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.class
-  Compiled from "TinyFrameworkExceptionTester.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int testException();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkExceptionTester.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.class
-  Compiled from "TinyFrameworkForTextPolicy.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 14, attributes: 2
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String toBeIgnoredObj();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeIgnoredV();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public boolean toBeIgnoredZ();
-    descriptor: ()Z
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public byte toBeIgnoredB();
-    descriptor: ()B
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public char toBeIgnoredC();
-    descriptor: ()C
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public short toBeIgnoredS();
-    descriptor: ()S
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int toBeIgnoredI();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public float toBeIgnoredF();
-    descriptor: ()F
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public double toBeIgnoredD();
-    descriptor: ()D
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkForTextPolicy.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested.class
-  Compiled from "TinyFrameworkLambdas.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 7, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static java.lang.Integer lambda$getSupplier_static$3();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$getSupplier$2();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$static$1();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$new$0();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkLambdas.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.class
-  Compiled from "TinyFrameworkLambdas.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 7, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static java.lang.Integer lambda$getSupplier_static$3();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$getSupplier$2();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$static$1();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$new$0();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkLambdas.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo.class
-  Compiled from "TinyFrameworkMethodCallReplace.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace$ReplaceTo
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace$ReplaceTo();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void startThread(java.lang.Thread);
-    descriptor: (Ljava/lang/Thread;)V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int add(int, int);
-    descriptor: (II)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-SourceFile: "TinyFrameworkMethodCallReplace.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.class
-  Compiled from "TinyFrameworkMethodCallReplace.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 4, attributes: 5
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static boolean nonStaticMethodCallReplaceTester() throws java.lang.Exception;
-    descriptor: ()Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Exceptions:
-      throws java.lang.Exception
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int staticMethodCallReplaceTester();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static void lambda$nonStaticMethodCallReplaceTester$0(java.util.concurrent.atomic.AtomicBoolean);
-    descriptor: (Ljava/util/concurrent/atomic/AtomicBoolean;)V
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkMethodCallReplace.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.class
-  Compiled from "TinyFrameworkNative.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNative
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 10, attributes: 3
-  int value;
-    descriptor: I
-    flags: (0x0000)
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNative();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native int nativeAddTwo(int);
-    descriptor: (I)I
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddTwo_should_be_like_this(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native long nativeLongPlus(long, long);
-    descriptor: (JJ)J
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static long nativeLongPlus_should_be_like_this(long, long);
-    descriptor: (JJ)J
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=4, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void setValue(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public native int nativeNonStaticAddToValue(int);
-    descriptor: (I)I
-    flags: (0x0101) ACC_PUBLIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int nativeNonStaticAddToValue_should_be_like_this(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void nativeStillNotSupported_should_be_like_this();
-    descriptor: ()V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native byte nativeBytePlus(byte, byte);
-    descriptor: (BB)B
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkNative.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestNativeSubstitutionClass(
-      value="TinyFrameworkNative_host"
-    )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 1, attributes: 4
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$InnerClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 1, attributes: 5
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  final com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses this$0;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
-    flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$InnerClass(com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses);
-    descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      final mandated
-}
-InnerClasses:
-  public #x= #x of #x;                   // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 5
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$SubClass extends com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  interfaces: 0, fields: 0, methods: 1, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$SubClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;            // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 4, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  public static #x= #x of #x;           // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public #x= #x of #x;                  // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.class
-  Compiled from "TinyFrameworkPackageRedirect.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkPackageRedirect
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkPackageRedirect();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int foo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkPackageRedirect.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.class
-  Compiled from "TinyFrameworkRenamedClassCaller.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int foo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkRenamedClassCaller.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/packagetest/A.class
-  Compiled from "A.java"
-public class com.android.hoststubgen.test.tinyframework.packagetest.A
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/packagetest/A
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "A.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/packagetest/sub/A.class
-  Compiled from "A.java"
-public class com.android.hoststubgen.test.tinyframework.packagetest.sub.A
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/packagetest/sub/A
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "A.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C1.class
-  Compiled from "C1.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C1
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C1
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C1.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C2.class
-  Compiled from "C2.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C2
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/subclasstest/C1
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C2.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C3.class
-  Compiled from "C3.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C3
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/subclasstest/C2
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C3.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CA.class
-  Compiled from "CA.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.CA
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/CA
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "CA.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CB.class
-  Compiled from "CB.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.CB
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/CB
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "CB.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I1.class
-  Compiled from "I1.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I1
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I1
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I1.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I2.class
-  Compiled from "I2.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I2 extends com.android.hoststubgen.test.tinyframework.subclasstest.I1
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I2
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I2.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I3.class
-  Compiled from "I3.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I3 extends com.android.hoststubgen.test.tinyframework.subclasstest.I2
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I3
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I3.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IA.class
-  Compiled from "IA.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.IA
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/IA
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "IA.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IB.class
-  Compiled from "IB.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.IB
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/IB
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "IB.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/unsupported/UnsupportedClass.class
-  Compiled from "UnsupportedClass.java"
-public class com.unsupported.UnsupportedClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/unsupported/UnsupportedClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.unsupported.UnsupportedClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int getValue();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "UnsupportedClass.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.class
-  Compiled from "TinyFrameworkToBeRenamed.java"
-public class rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 3
-  private final int mValue;
-    descriptor: I
-    flags: (0x0012) ACC_PRIVATE, ACC_FINAL
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int getValue();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkToBeRenamed.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-impl-dump.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-dump.txt
similarity index 67%
rename from tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-impl-dump.txt
rename to tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-dump.txt
index d23b450..e41d46d 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-impl-dump.txt
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/03-hoststubgen-test-tiny-framework-host-dump.txt
@@ -12,12 +12,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestClassLoadHook.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -39,7 +39,7 @@
 SourceFile: "HostSideTestKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -62,12 +62,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestNativeSubstitutionClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -89,7 +89,7 @@
 SourceFile: "HostSideTestRemove.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -98,20 +98,20 @@
     java.lang.annotation.Retention(
       value=Ljava/lang/annotation/RetentionPolicy;.CLASS
     )
-## Class: android/hosttest/annotation/HostSideTestStub.class
-  Compiled from "HostSideTestStub.java"
-public interface android.hosttest.annotation.HostSideTestStub extends java.lang.annotation.Annotation
+## Class: android/hosttest/annotation/HostSideTestStaticInitializerKeep.class
+  Compiled from "HostSideTestStaticInitializerKeep.java"
+public interface android.hosttest.annotation.HostSideTestStaticInitializerKeep extends java.lang.annotation.Annotation
   minor version: 0
   major version: 61
   flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestStub
+  this_class: #x                          // android/hosttest/annotation/HostSideTestStaticInitializerKeep
   super_class: #x                         // java/lang/Object
   interfaces: 1, fields: 0, methods: 0, attributes: 2
 }
-SourceFile: "HostSideTestStub.java"
+SourceFile: "HostSideTestStaticInitializerKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -134,12 +134,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestSubstitute.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.METHOD]
@@ -161,7 +161,7 @@
 SourceFile: "HostSideTestThrow.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -183,29 +183,7 @@
 SourceFile: "HostSideTestWholeClassKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-  x: #x(#x=[e#x.#x])
-    java.lang.annotation.Target(
-      value=[Ljava/lang/annotation/ElementType;.TYPE]
-    )
-  x: #x(#x=e#x.#x)
-    java.lang.annotation.Retention(
-      value=Ljava/lang/annotation/RetentionPolicy;.CLASS
-    )
-## Class: android/hosttest/annotation/HostSideTestWholeClassStub.class
-  Compiled from "HostSideTestWholeClassStub.java"
-public interface android.hosttest.annotation.HostSideTestWholeClassStub extends java.lang.annotation.Annotation
-  minor version: 0
-  major version: 61
-  flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestWholeClassStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "HostSideTestWholeClassStub.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -237,9 +215,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int addTwo(int);
     descriptor: (I)I
@@ -256,9 +232,7 @@
             0       4     0     a   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
@@ -266,9 +240,7 @@
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
 ## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub.class
   Compiled from "IPretendingAidl.java"
@@ -293,9 +265,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int addOne(int);
     descriptor: (I)I
@@ -312,19 +282,15 @@
             0       4     0     a   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
+  public static #x= #x of #x;             // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
   public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
 ## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl.class
   Compiled from "IPretendingAidl.java"
@@ -342,9 +308,7 @@
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
   com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
@@ -362,9 +326,7 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.R$Nested();
     descriptor: ()V
@@ -380,9 +342,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/R$Nested;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -400,18 +360,14 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
+  public static #x= #x of #x;             // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
 SourceFile: "R.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/R
 ## Class: com/android/hoststubgen/test/tinyframework/R.class
   Compiled from "R.java"
@@ -436,189 +392,31 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/R;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
+  public static #x= #x of #x;             // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
 SourceFile: "R.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/R$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl
-  minor version: 0
-  major version: 61
-  flags: (0x0020) ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 4
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-         x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOneKeep();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-         x: ldc           #x                 // String getOneKeep
-         x: ldc           #x                 // String ()I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_1
-        x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestKeep
-
-  public static int getOneStub();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: iconst_1
-         x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-InnerClasses:
-  private static #x= #x of #x;           // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.class
+  Compiled from "TinyFrameworkAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
   minor version: 0
   major version: 61
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 5
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-         x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_withCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneKeep:()I
-         x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_noCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=1, locals=0, args_size=0
-         x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneStub:()I
-         x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.class
-  Compiled from "TinyFrameworkClassAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 8, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
+  interfaces: 0, fields: 1, methods: 6, attributes: 3
   public int keep;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -628,12 +426,12 @@
     flags: (0x000a) ACC_PRIVATE, ACC_STATIC
     Code:
       stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
+         x: ldc           #x                  // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations();
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations();
     descriptor: ()V
     flags: (0x0001) ACC_PUBLIC
     Code:
@@ -642,70 +440,36 @@
          x: invokespecial #x                 // Method java/lang/Object."<init>":()V
          x: aload_0
          x: iconst_1
-         x: putfield      #x                 // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
+         x: putfield      #x                 // Field keep:I
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public int addOne(int);
     descriptor: (I)I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=2, locals=2, args_size=2
-         x: aload_0
          x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
+         x: iconst_1
+         x: iadd
          x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-            0       6     1 value   I
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
+            0       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String addOneInner
-         x: ldc           #x                 // String (I)I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iload_1
-        x: iconst_1
-        x: iadd
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           15       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-           15       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -722,15 +486,13 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
             0       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddThree(int);
     descriptor: (I)I
@@ -749,277 +511,39 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String unsupportedMethod();
     descriptor: ()Ljava/lang/String;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String unsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
-        x: new           #x                 // class java/lang/RuntimeException
-        x: dup
-        x: ldc           #x                 // String Unreachable
-        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+      stack=3, locals=1, args_size=1
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: ldc           #x                 // String Unreachable
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
         x: athrow
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
 }
-SourceFile: "TinyFrameworkClassAnnotations.java"
+SourceFile: "TinyFrameworkAnnotations.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestClassLoadHook(
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.class
-  Compiled from "TinyFrameworkClassClassWideAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 8, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int remove;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=1, args_size=1
-         x: aload_0
-         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-         x: aload_0
-         x: iconst_1
-         x: putfield      #x                 // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: aload_0
-         x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       6     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: iload_1
-         x: iconst_1
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeRemoved(java.lang.String);
-    descriptor: (Ljava/lang/String;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":()V
-         x: athrow
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       8     1   foo   Ljava/lang/String;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=2, locals=2, args_size=2
-         x: iload_1
-         x: iconst_2
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-            0       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=2, locals=1, args_size=1
-         x: iload_0
-         x: iconst_3
-         x: iadd
-         x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       4     0 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String unsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: ldc           #x                 // String This value shouldn\'t be seen on the host side.
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassClassWideAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.class
   Compiled from "TinyFrameworkClassLoadHook.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook
@@ -1035,9 +559,7 @@
     Signature: #x                          // Ljava/util/Set<Ljava/lang/Class<*>;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook();
     descriptor: ()V
@@ -1053,9 +575,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void onClassLoaded(java.lang.Class<?>);
     descriptor: (Ljava/lang/Class;)V
@@ -1077,9 +597,7 @@
     Signature: #x                          // (Ljava/lang/Class<*>;)V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -1094,19 +612,115 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkClassLoadHook.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.class
+  Compiled from "TinyFrameworkClassWideAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 1, methods: 4, attributes: 3
+  public int keep;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=1, args_size=1
+         x: aload_0
+         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+         x: aload_0
+         x: iconst_1
+         x: putfield      #x                 // Field keep:I
+         x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public int addOne(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=2, args_size=2
+         x: iload_1
+         x: iconst_1
+         x: iadd
+         x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0       4     1 value   I
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public int addTwo(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=2, args_size=2
+         x: iload_1
+         x: iconst_2
+         x: iadd
+         x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+            0       4     1 value   I
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public java.lang.String unsupportedMethod();
+    descriptor: ()Ljava/lang/String;
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=3, locals=1, args_size=1
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: ldc           #x                 // String Unreachable
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+        x: athrow
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestThrow
+}
+SourceFile: "TinyFrameworkClassWideAnnotations.java"
+RuntimeVisibleAnnotations:
+  x: #x()
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+RuntimeInvisibleAnnotations:
+  x: #x()
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.class
   Compiled from "TinyFrameworkClassWithInitializerDefault.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault
@@ -1121,35 +735,29 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
 }
 SourceFile: "TinyFrameworkClassWithInitializerDefault.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.class
   Compiled from "TinyFrameworkClassWithInitializerStub.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub
@@ -1164,24 +772,20 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   static {};
     descriptor: ()V
@@ -1201,21 +805,19 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkClassWithInitializerStub.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestClassLoadHook(
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.class
@@ -1232,43 +834,37 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex GREEN;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex BLUE;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private final java.lang.String mLongName;
     descriptor: Ljava/lang/String;
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -1278,7 +874,7 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -1288,9 +884,7 @@
     flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1304,9 +898,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex valueOf(java.lang.String);
     descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1324,9 +916,7 @@
             0      10     0  name   Ljava/lang/String;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      mandated
@@ -1356,12 +946,10 @@
     Signature: #x                          // (Ljava/lang/String;Ljava/lang/String;)V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
     MethodParameters:
       Name                           Flags
       <no name>                      synthetic
@@ -1383,12 +971,10 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.lang.String getShortName();
     descriptor: ()Ljava/lang/String;
@@ -1404,12 +990,10 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1434,9 +1018,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -1473,20 +1055,16 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;>;
 SourceFile: "TinyFrameworkEnumComplex.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.class
   Compiled from "TinyFrameworkEnumSimple.java"
 public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple>
@@ -1501,33 +1079,27 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple DOG;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $VALUES;
     descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
     flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
@@ -1541,9 +1113,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple valueOf(java.lang.String);
     descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
@@ -1561,9 +1131,7 @@
             0      10     0  name   Ljava/lang/String;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      mandated
@@ -1585,9 +1153,7 @@
     Signature: #x                          // ()V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      synthetic
@@ -1612,9 +1178,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -1639,20 +1203,16 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;>;
 SourceFile: "TinyFrameworkEnumSimple.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.class
   Compiled from "TinyFrameworkExceptionTester.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester
@@ -1676,9 +1236,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int testException();
     descriptor: ()I
@@ -1709,19 +1267,15 @@
            11      11     0     e   Ljava/lang/Exception;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkExceptionTester.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.class
   Compiled from "TinyFrameworkForTextPolicy.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy
@@ -1730,22 +1284,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 17, attributes: 2
+  interfaces: 0, fields: 1, methods: 15, attributes: 2
   public int stub;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -1753,7 +1298,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                 // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -1767,63 +1312,32 @@
          x: aload_0
          x: iconst_1
          x: putfield      #x                 // Field stub:I
-         x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
+            0      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int addOne(int);
     descriptor: (I)I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=2, locals=2, args_size=2
-         x: aload_0
          x: iload_1
-         x: invokevirtual #x                 // Method addOneInner:(I)I
+         x: iconst_1
+         x: iadd
          x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-            0       6     1 value   I
+            0       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
+            0       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                 // String addOneInner
-         x: ldc           #x                 // String (I)I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iload_1
-        x: iconst_1
-        x: iadd
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           15       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-           15       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String toBeIgnoredObj();
     descriptor: ()Ljava/lang/String;
@@ -1836,9 +1350,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public void toBeIgnoredV();
     descriptor: ()V
@@ -1850,9 +1362,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public boolean toBeIgnoredZ();
     descriptor: ()Z
@@ -1865,9 +1375,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public byte toBeIgnoredB();
     descriptor: ()B
@@ -1880,9 +1388,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public char toBeIgnoredC();
     descriptor: ()C
@@ -1895,9 +1401,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public short toBeIgnoredS();
     descriptor: ()S
@@ -1910,9 +1414,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int toBeIgnoredI();
     descriptor: ()I
@@ -1925,9 +1427,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public float toBeIgnoredF();
     descriptor: ()F
@@ -1940,9 +1440,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public double toBeIgnoredD();
     descriptor: ()D
@@ -1955,9 +1453,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int addTwo(int);
     descriptor: (I)I
@@ -1977,9 +1473,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddThree(int);
     descriptor: (I)I
@@ -1998,57 +1492,29 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String unsupportedMethod();
     descriptor: ()Ljava/lang/String;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                 // String unsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
-        x: new           #x                 // class java/lang/RuntimeException
-        x: dup
-        x: ldc           #x                 // String Unreachable
-        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+      stack=3, locals=1, args_size=1
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: ldc           #x                 // String Unreachable
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
         x: athrow
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=1, locals=1, args_size=1
-         x: aload_0
-         x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-         x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkForTextPolicy.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested.class
   Compiled from "TinyFrameworkLambdas.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested
@@ -2064,12 +1530,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -2077,12 +1541,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested();
     descriptor: ()V
@@ -2101,12 +1563,10 @@
             0      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2122,12 +1582,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2140,12 +1598,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -2158,9 +1614,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$getSupplier$2();
     descriptor: ()Ljava/lang/Integer;
@@ -2173,9 +1627,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$static$1();
     descriptor: ()Ljava/lang/Integer;
@@ -2188,9 +1640,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$new$0();
     descriptor: ()Ljava/lang/Integer;
@@ -2203,9 +1653,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -2218,7 +1666,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
@@ -2226,12 +1674,10 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 BootstrapMethods:
@@ -2271,12 +1717,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -2284,12 +1728,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas();
     descriptor: ()V
@@ -2308,12 +1750,10 @@
             0      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2329,12 +1769,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2347,12 +1785,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -2365,9 +1801,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$getSupplier$2();
     descriptor: ()Ljava/lang/Integer;
@@ -2380,9 +1814,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$static$1();
     descriptor: ()Ljava/lang/Integer;
@@ -2395,9 +1827,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$new$0();
     descriptor: ()Ljava/lang/Integer;
@@ -2410,9 +1840,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -2425,7 +1853,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
@@ -2433,12 +1861,10 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 BootstrapMethods:
@@ -2487,9 +1913,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void startThread(java.lang.Thread);
     descriptor: (Ljava/lang/Thread;)V
@@ -2508,9 +1932,7 @@
             0      10     0 thread   Ljava/lang/Thread;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int add(int, int);
     descriptor: (II)I
@@ -2528,18 +1950,14 @@
             0       4     1     b   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
+  public static #x= #x of #x;             // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
 SourceFile: "TinyFrameworkMethodCallReplace.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.class
   Compiled from "TinyFrameworkMethodCallReplace.java"
@@ -2564,9 +1982,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static boolean nonStaticMethodCallReplaceTester() throws java.lang.Exception;
     descriptor: ()Z
@@ -2600,9 +2016,7 @@
       throws java.lang.Exception
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int staticMethodCallReplaceTester();
     descriptor: ()I
@@ -2616,9 +2030,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static void lambda$nonStaticMethodCallReplaceTester$0(java.util.concurrent.atomic.AtomicBoolean);
     descriptor: (Ljava/util/concurrent/atomic/AtomicBoolean;)V
@@ -2636,22 +2048,18 @@
             0      11     0    ab   Ljava/util/concurrent/atomic/AtomicBoolean;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;           // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
+  public static #x= #x of #x;            // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
   public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
 SourceFile: "TinyFrameworkMethodCallReplace.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 BootstrapMethods:
   x: #x REF_invokeStatic java/lang/invoke/LambdaMetafactory.metafactory:(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;Ljava/lang/invoke/MethodType;Ljava/lang/invoke/MethodHandle;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;
     Method arguments:
@@ -2668,15 +2076,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 11, attributes: 3
+  interfaces: 0, fields: 1, methods: 12, attributes: 3
   int value;
     descriptor: I
     flags: (0x0000)
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNative();
     descriptor: ()V
@@ -2692,9 +2098,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo(int);
     descriptor: (I)I
@@ -2708,9 +2112,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo_should_be_like_this(int);
     descriptor: (I)I
@@ -2726,9 +2128,7 @@
             0       5     0   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus(long, long);
     descriptor: (JJ)J
@@ -2743,9 +2143,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus_should_be_like_this(long, long);
     descriptor: (JJ)J
@@ -2763,9 +2161,7 @@
             0       6     2  arg2   J
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public void setValue(int);
     descriptor: (I)V
@@ -2783,9 +2179,7 @@
             0       6     1     v   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int nativeNonStaticAddToValue(int);
     descriptor: (I)I
@@ -2800,9 +2194,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int nativeNonStaticAddToValue_should_be_like_this(int);
     descriptor: (I)I
@@ -2820,38 +2212,38 @@
             0       6     1   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void nativeStillNotSupported();
     descriptor: ()V
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-         x: ldc           #x                 // String nativeStillNotSupported
-         x: ldc           #x                 // String ()V
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
-        x: new           #x                 // class java/lang/RuntimeException
-        x: dup
-        x: ldc           #x                 // String Unreachable
-        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+      stack=3, locals=0, args_size=0
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+         x: new           #x                 // class java/lang/RuntimeException
+         x: dup
+         x: ldc           #x                 // String Unreachable
+         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
         x: athrow
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
 
+  public static native void nativeStillKeep();
+    descriptor: ()V
+    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestKeep
+
   public static void nativeStillNotSupported_should_be_like_this();
     descriptor: ()V
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
@@ -2864,9 +2256,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static byte nativeBytePlus(byte, byte);
     descriptor: (BB)B
@@ -2881,19 +2271,15 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkNative.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestNativeSubstitutionClass(
       value="TinyFrameworkNative_host"
@@ -2911,125 +2297,95 @@
     descriptor: ()V
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-         x: ldc           #x                 // String <init>
-         x: ldc           #x                 // String ()V
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-        x: return
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo(int);
     descriptor: (I)I
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-         x: ldc           #x                 // String nativeAddTwo
-         x: ldc           #x                 // String (I)I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iload_0
-        x: iconst_2
-        x: iadd
-        x: ireturn
+      stack=2, locals=1, args_size=1
+         x: iload_0
+         x: iconst_2
+         x: iadd
+         x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       4     0   arg   I
+            0       4     0   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus(long, long);
     descriptor: (JJ)J
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
       stack=4, locals=4, args_size=2
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-         x: ldc           #x                 // String nativeLongPlus
-         x: ldc           #x                 // String (JJ)J
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: lload_0
-        x: lload_2
-        x: ladd
-        x: lreturn
+         x: lload_0
+         x: lload_2
+         x: ladd
+         x: lreturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       4     0  arg1   J
-           15       4     2  arg2   J
+            0       4     0  arg1   J
+            0       4     2  arg2   J
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeNonStaticAddToValue(com.android.hoststubgen.test.tinyframework.TinyFrameworkNative, int);
     descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;I)I
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-         x: ldc           #x                 // String nativeNonStaticAddToValue
-         x: ldc           #x                 // String (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;I)I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: getfield      #x                 // Field com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.value:I
-        x: iload_1
-        x: iadd
-        x: ireturn
+      stack=2, locals=2, args_size=2
+         x: aload_0
+         x: getfield      #x                 // Field com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.value:I
+         x: iload_1
+         x: iadd
+         x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       7     0 source   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
-           15       7     1   arg   I
+            0       7     0 source   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
+            0       7     1   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static byte nativeBytePlus(byte, byte);
     descriptor: (BB)B
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-         x: ldc           #x                 // String nativeBytePlus
-         x: ldc           #x                 // String (BB)B
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iload_0
-        x: iload_1
-        x: iadd
-        x: i2b
-        x: ireturn
+      stack=2, locals=2, args_size=2
+         x: iload_0
+         x: iload_1
+         x: iadd
+         x: i2b
+         x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  arg1   B
-           15       5     1  arg2   B
+            0       5     0  arg1   B
+            0       5     1  arg2   B
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkNative_host.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
     android.hosttest.annotation.HostSideTestWholeClassKeep
@@ -3047,7 +2403,7 @@
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$1(com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses);
     descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;)V
@@ -3067,7 +2423,7 @@
             0      10     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
@@ -3076,45 +2432,33 @@
     descriptor: ()Ljava/lang/Integer;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Integer;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_1
-        x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: iconst_1
+         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
     flags: (0x1041) ACC_PUBLIC, ACC_BRIDGE, ACC_SYNTHETIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Object;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
@@ -3123,7 +2467,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3148,51 +2492,39 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Integer;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_2
-        x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: iconst_2
+         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
     flags: (0x1041) ACC_PUBLIC, ACC_BRIDGE, ACC_SYNTHETIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Object;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
@@ -3201,7 +2533,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3217,7 +2549,7 @@
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$3(com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses);
     descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;)V
@@ -3237,7 +2569,7 @@
             0      10     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
@@ -3246,45 +2578,33 @@
     descriptor: ()Ljava/lang/Integer;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Integer;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_3
-        x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: iconst_3
+         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
     flags: (0x1041) ACC_PUBLIC, ACC_BRIDGE, ACC_SYNTHETIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Object;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
@@ -3293,7 +2613,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3318,51 +2638,39 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Integer;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_4
-        x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: iconst_4
+         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
     flags: (0x1041) ACC_PUBLIC, ACC_BRIDGE, ACC_SYNTHETIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Object;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
@@ -3371,7 +2679,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3387,9 +2695,7 @@
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass(int);
     descriptor: (I)V
@@ -3409,18 +2715,14 @@
             0      10     1     x   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;             // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3430,24 +2732,20 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 1, attributes: 5
+  interfaces: 0, fields: 2, methods: 1, attributes: 4
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   final com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses this$0;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$InnerClass(com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses);
     descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;)V
@@ -3470,24 +2768,17 @@
             0      15     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
 }
 InnerClasses:
-  public #x= #x of #x;                   // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public #x= #x of #x;                    // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3512,51 +2803,39 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Integer;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: bipush        7
-        x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: bipush        7
+         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
+            0       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
     flags: (0x1041) ACC_PUBLIC, ACC_BRIDGE, ACC_SYNTHETIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-         x: ldc           #x                 // String get
-         x: ldc           #x                 // String ()Ljava/lang/Object;
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
-        x: areturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
+         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
+            0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
@@ -3566,7 +2845,50 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass.class
+  Compiled from "TinyFrameworkNestedClasses.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 1, methods: 1, attributes: 4
+  public int value;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=2, locals=1, args_size=1
+         x: aload_0
+         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+         x: aload_0
+         x: bipush        8
+         x: putfield      #x                 // Field value:I
+        x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+            0      11     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass;
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+}
+InnerClasses:
+  public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
+SourceFile: "TinyFrameworkNestedClasses.java"
+RuntimeVisibleAnnotations:
+  x: #x()
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3576,15 +2898,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 5
+  interfaces: 0, fields: 1, methods: 2, attributes: 4
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass();
     descriptor: ()V
@@ -3603,9 +2923,7 @@
             0      11     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -3620,22 +2938,16 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;             // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3662,19 +2974,15 @@
             0       6     1     x   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;             // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;            // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3691,9 +2999,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -3701,9 +3007,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses();
     descriptor: ()V
@@ -3725,9 +3029,7 @@
             0      17     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -3746,9 +3048,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -3763,9 +3063,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -3780,33 +3078,31 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  public static #x= #x of #x;           // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;            // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public #x= #x of #x;                  // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;          // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
+  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
@@ -3836,9 +3132,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int foo(int);
     descriptor: (I)I
@@ -3857,19 +3151,15 @@
             0      12     0 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkPackageRedirect.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.class
   Compiled from "TinyFrameworkRenamedClassCaller.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller
@@ -3893,9 +3183,7 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int foo(int);
     descriptor: (I)I
@@ -3914,19 +3202,15 @@
             0      12     0 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkRenamedClassCaller.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/packagetest/A.class
   Compiled from "A.java"
 public class com.android.hoststubgen.test.tinyframework.packagetest.A
@@ -3940,9 +3224,7 @@
 SourceFile: "A.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/packagetest/sub/A.class
   Compiled from "A.java"
 public class com.android.hoststubgen.test.tinyframework.packagetest.sub.A
@@ -3956,9 +3238,7 @@
 SourceFile: "A.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C1.class
   Compiled from "C1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -3972,9 +3252,7 @@
 SourceFile: "C1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C2.class
   Compiled from "C2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -3988,9 +3266,7 @@
 SourceFile: "C2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C3.class
   Compiled from "C3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
@@ -4004,9 +3280,7 @@
 SourceFile: "C3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CA.class
   Compiled from "CA.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.CA
@@ -4020,9 +3294,7 @@
 SourceFile: "CA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CB.class
   Compiled from "CB.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.CB
@@ -4036,9 +3308,7 @@
 SourceFile: "CB.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C1.class
   Compiled from "Class_C1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C1 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -4052,7 +3322,7 @@
 SourceFile: "Class_C1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C2.class
   Compiled from "Class_C2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
@@ -4066,7 +3336,7 @@
 SourceFile: "Class_C2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C3.class
   Compiled from "Class_C3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C3
@@ -4080,7 +3350,7 @@
 SourceFile: "Class_C3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I1.class
   Compiled from "Class_I1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I1 implements com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -4094,7 +3364,7 @@
 SourceFile: "Class_I1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I1_IA.class
   Compiled from "Class_I1_IA.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I1_IA implements com.android.hoststubgen.test.tinyframework.subclasstest.I1,com.android.hoststubgen.test.tinyframework.subclasstest.IA
@@ -4108,7 +3378,7 @@
 SourceFile: "Class_I1_IA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I2.class
   Compiled from "Class_I2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I2 implements com.android.hoststubgen.test.tinyframework.subclasstest.I2
@@ -4122,7 +3392,7 @@
 SourceFile: "Class_I2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I3.class
   Compiled from "Class_I3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I3 implements com.android.hoststubgen.test.tinyframework.subclasstest.I3
@@ -4136,7 +3406,7 @@
 SourceFile: "Class_I3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I1.class
   Compiled from "I1.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -4150,9 +3420,7 @@
 SourceFile: "I1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I2.class
   Compiled from "I2.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I2 extends com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -4166,9 +3434,7 @@
 SourceFile: "I2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I3.class
   Compiled from "I3.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I3 extends com.android.hoststubgen.test.tinyframework.subclasstest.I2
@@ -4182,9 +3448,7 @@
 SourceFile: "I3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IA.class
   Compiled from "IA.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.IA
@@ -4198,9 +3462,7 @@
 SourceFile: "IA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IB.class
   Compiled from "IB.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.IB
@@ -4214,9 +3476,7 @@
 SourceFile: "IB.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/supported/UnsupportedClass.class
   Compiled from "UnsupportedClass.java"
 public class com.supported.UnsupportedClass
@@ -4231,60 +3491,48 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.supported.UnsupportedClass(int);
     descriptor: (I)V
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                 // String com/supported/UnsupportedClass
-         x: ldc           #x                 // String <init>
-         x: ldc           #x                 // String (I)V
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-        x: aload_0
-        x: iload_1
-        x: putfield      #x                 // Field mValue:I
-        x: return
+      stack=2, locals=2, args_size=2
+         x: aload_0
+         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+         x: aload_0
+         x: iload_1
+         x: putfield      #x                 // Field mValue:I
+         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15      10     0  this   Lcom/supported/UnsupportedClass;
-           15      10     1 value   I
+            0      10     0  this   Lcom/supported/UnsupportedClass;
+            0      10     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
     flags: (0x0001) ACC_PUBLIC
     Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                 // String com/supported/UnsupportedClass
-         x: ldc           #x                 // String getValue
-         x: ldc           #x                 // String ()I
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-         x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: aload_0
-        x: getfield      #x                 // Field mValue:I
-        x: ireturn
+      stack=1, locals=1, args_size=1
+         x: aload_0
+         x: getfield      #x                 // Field mValue:I
+         x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           15       5     0  this   Lcom/supported/UnsupportedClass;
+            0       5     0  this   Lcom/supported/UnsupportedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "UnsupportedClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
     android.hosttest.annotation.HostSideTestWholeClassKeep
@@ -4316,9 +3564,7 @@
             0      14     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
@@ -4336,19 +3582,15 @@
             0      10     0  this   Lcom/unsupported/UnsupportedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "UnsupportedClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.class
   Compiled from "TinyFrameworkToBeRenamed.java"
 public class rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed
@@ -4363,9 +3605,7 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed(int);
     descriptor: (I)V
@@ -4385,9 +3625,7 @@
             0      10     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
@@ -4403,16 +3641,12 @@
             0       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkToBeRenamed.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/12-hoststubgen-test-tiny-framework-host-ext-stub-dump.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/12-hoststubgen-test-tiny-framework-host-ext-stub-dump.txt
deleted file mode 100644
index 1b83d24..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/12-hoststubgen-test-tiny-framework-host-ext-stub-dump.txt
+++ /dev/null
@@ -1,2829 +0,0 @@
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy.class
-  Compiled from "IPretendingAidl.java"
-public class com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub$Proxy
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub$Proxy();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub.class
-  Compiled from "IPretendingAidl.java"
-public class com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.IPretendingAidl$Stub();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int addOne(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl.class
-  Compiled from "IPretendingAidl.java"
-public interface com.android.hoststubgen.test.tinyframework.IPretendingAidl
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 4
-}
-InnerClasses:
-  public static #x= #x of #x;            // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-  public static #x= #x of #x;           // Proxy=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy of class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-SourceFile: "IPretendingAidl.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-  com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
-## Class: com/android/hoststubgen/test/tinyframework/R$Nested.class
-  Compiled from "R.java"
-public class com.android.hoststubgen.test.tinyframework.R$Nested
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/R$Nested
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 4
-  public static int[] ARRAY;
-    descriptor: [I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.R$Nested();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
-SourceFile: "R.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/R
-## Class: com/android/hoststubgen/test/tinyframework/R.class
-  Compiled from "R.java"
-public class com.android.hoststubgen.test.tinyframework.R
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/R
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 1, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.R();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
-SourceFile: "R.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/R$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl
-  minor version: 0
-  major version: 61
-  flags: (0x0020) ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 4
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOneStub();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-InnerClasses:
-  private static #x= #x of #x;           // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 5
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_withCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_noCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.class
-  Compiled from "TinyFrameworkClassAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 5, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-SourceFile: "TinyFrameworkClassAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestClassLoadHook(
-      value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
-    )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.class
-  Compiled from "TinyFrameworkClassClassWideAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 8, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int remove;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeRemoved(java.lang.String);
-    descriptor: (Ljava/lang/String;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String unsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassClassWideAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.class
-  Compiled from "TinyFrameworkClassLoadHook.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 3, attributes: 3
-  public static final java.util.Set<java.lang.Class<?>> sLoadedClasses;
-    descriptor: Ljava/util/Set;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/Set<Ljava/lang/Class<*>;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void onClassLoaded(java.lang.Class<?>);
-    descriptor: (Ljava/lang/Class;)V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // (Ljava/lang/Class<*>;)V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassLoadHook.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.class
-  Compiled from "TinyFrameworkClassWithInitializerDefault.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 0, attributes: 3
-  public static boolean sInitialized;
-    descriptor: Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.lang.Object sObject;
-    descriptor: Ljava/lang/Object;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-}
-SourceFile: "TinyFrameworkClassWithInitializerDefault.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.class
-  Compiled from "TinyFrameworkClassWithInitializerStub.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 0, attributes: 3
-  public static boolean sInitialized;
-    descriptor: Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.lang.Object sObject;
-    descriptor: Ljava/lang/Object;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-}
-SourceFile: "TinyFrameworkClassWithInitializerStub.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestClassLoadHook(
-      value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
-    )
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.class
-  Compiled from "TinyFrameworkEnumComplex.java"
-public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex>
-  minor version: 0
-  major version: 61
-  flags: (0x4031) ACC_PUBLIC, ACC_FINAL, ACC_SUPER, ACC_ENUM
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex
-  super_class: #x                         // java/lang/Enum
-  interfaces: 0, fields: 4, methods: 7, attributes: 4
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex RED;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex GREEN;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex BLUE;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $VALUES;
-    descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex valueOf(java.lang.String);
-    descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      mandated
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex(java.lang.String, java.lang.String);
-    descriptor: (Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;)V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=5, args_size=5
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // (Ljava/lang/String;Ljava/lang/String;)V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-    MethodParameters:
-      Name                           Flags
-      <no name>                      synthetic
-      <no name>                      synthetic
-      <no name>
-      <no name>
-
-  public java.lang.String getLongName();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.lang.String getShortName();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;>;
-SourceFile: "TinyFrameworkEnumComplex.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.class
-  Compiled from "TinyFrameworkEnumSimple.java"
-public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple>
-  minor version: 0
-  major version: 61
-  flags: (0x4031) ACC_PUBLIC, ACC_FINAL, ACC_SUPER, ACC_ENUM
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple
-  super_class: #x                         // java/lang/Enum
-  interfaces: 0, fields: 3, methods: 5, attributes: 4
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple CAT;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple DOG;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $VALUES;
-    descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple valueOf(java.lang.String);
-    descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      mandated
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple();
-    descriptor: (Ljava/lang/String;I)V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=3, locals=3, args_size=3
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()V
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      synthetic
-      <no name>                      synthetic
-
-  private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $values();
-    descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;>;
-SourceFile: "TinyFrameworkEnumSimple.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.class
-  Compiled from "TinyFrameworkExceptionTester.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int testException();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkExceptionTester.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.class
-  Compiled from "TinyFrameworkForTextPolicy.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 14, attributes: 2
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String toBeIgnoredObj();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeIgnoredV();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public boolean toBeIgnoredZ();
-    descriptor: ()Z
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public byte toBeIgnoredB();
-    descriptor: ()B
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public char toBeIgnoredC();
-    descriptor: ()C
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public short toBeIgnoredS();
-    descriptor: ()S
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int toBeIgnoredI();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public float toBeIgnoredF();
-    descriptor: ()F
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public double toBeIgnoredD();
-    descriptor: ()D
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkForTextPolicy.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested.class
-  Compiled from "TinyFrameworkLambdas.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 7, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static java.lang.Integer lambda$getSupplier_static$3();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$getSupplier$2();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$static$1();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$new$0();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkLambdas.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.class
-  Compiled from "TinyFrameworkLambdas.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 7, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  private static java.lang.Integer lambda$getSupplier_static$3();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$getSupplier$2();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$static$1();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static java.lang.Integer lambda$new$0();
-    descriptor: ()Ljava/lang/Integer;
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkLambdas.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestStub
-  x: #x()
-    android.hosttest.annotation.HostSideTestStaticInitializerKeep
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo.class
-  Compiled from "TinyFrameworkMethodCallReplace.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace$ReplaceTo
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 3, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace$ReplaceTo();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void startThread(java.lang.Thread);
-    descriptor: (Ljava/lang/Thread;)V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int add(int, int);
-    descriptor: (II)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-SourceFile: "TinyFrameworkMethodCallReplace.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.class
-  Compiled from "TinyFrameworkMethodCallReplace.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 4, attributes: 5
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static boolean nonStaticMethodCallReplaceTester() throws java.lang.Exception;
-    descriptor: ()Z
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Exceptions:
-      throws java.lang.Exception
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int staticMethodCallReplaceTester();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static void lambda$nonStaticMethodCallReplaceTester$0(java.util.concurrent.atomic.AtomicBoolean);
-    descriptor: (Ljava/util/concurrent/atomic/AtomicBoolean;)V
-    flags: (0x100a) ACC_PRIVATE, ACC_STATIC, ACC_SYNTHETIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;           // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-  public static final #x= #x of #x;    // Lookup=class java/lang/invoke/MethodHandles$Lookup of class java/lang/invoke/MethodHandles
-SourceFile: "TinyFrameworkMethodCallReplace.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.class
-  Compiled from "TinyFrameworkNative.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNative
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 10, attributes: 3
-  int value;
-    descriptor: I
-    flags: (0x0000)
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNative();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native int nativeAddTwo(int);
-    descriptor: (I)I
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddTwo_should_be_like_this(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native long nativeLongPlus(long, long);
-    descriptor: (JJ)J
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static long nativeLongPlus_should_be_like_this(long, long);
-    descriptor: (JJ)J
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=4, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void setValue(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public native int nativeNonStaticAddToValue(int);
-    descriptor: (I)I
-    flags: (0x0101) ACC_PUBLIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int nativeNonStaticAddToValue_should_be_like_this(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static void nativeStillNotSupported_should_be_like_this();
-    descriptor: ()V
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static native byte nativeBytePlus(byte, byte);
-    descriptor: (BB)B
-    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkNative.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-  x: #x(#x=s#x)
-    android.hosttest.annotation.HostSideTestNativeSubstitutionClass(
-      value="TinyFrameworkNative_host"
-    )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 1, attributes: 4
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$InnerClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 1, attributes: 5
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  final com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses this$0;
-    descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
-    flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$InnerClass(com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses);
-    descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    MethodParameters:
-      Name                           Flags
-      <no name>                      final mandated
-}
-InnerClasses:
-  public #x= #x of #x;                   // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 5
-  public int value;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$SubClass extends com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$BaseClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  interfaces: 0, fields: 0, methods: 1, attributes: 4
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$SubClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  public static #x= #x of #x;            // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;            // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.class
-  Compiled from "TinyFrameworkNestedClasses.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 4, attributes: 5
-  public final java.util.function.Supplier<java.lang.Integer> mSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0011) ACC_PUBLIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
-    descriptor: Ljava/util/function/Supplier;
-    flags: (0x0019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL
-    Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.util.function.Supplier<java.lang.Integer> getSupplier();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
-    descriptor: ()Ljava/util/function/Supplier;
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  static {};
-    descriptor: ()V
-    flags: (0x0008) ACC_STATIC
-    Code:
-      stack=3, locals=0, args_size=0
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  public static #x= #x of #x;           // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  public #x= #x of #x;                  // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
-  #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-SourceFile: "TinyFrameworkNestedClasses.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.class
-  Compiled from "TinyFrameworkPackageRedirect.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkPackageRedirect
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkPackageRedirect();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int foo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkPackageRedirect.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.class
-  Compiled from "TinyFrameworkRenamedClassCaller.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int foo(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkRenamedClassCaller.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: com/android/hoststubgen/test/tinyframework/packagetest/A.class
-  Compiled from "A.java"
-public class com.android.hoststubgen.test.tinyframework.packagetest.A
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/packagetest/A
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "A.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/packagetest/sub/A.class
-  Compiled from "A.java"
-public class com.android.hoststubgen.test.tinyframework.packagetest.sub.A
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/packagetest/sub/A
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "A.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C1.class
-  Compiled from "C1.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C1
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C1
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C1.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C2.class
-  Compiled from "C2.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C2
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/subclasstest/C1
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C2.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C3.class
-  Compiled from "C3.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/C3
-  super_class: #x                         // com/android/hoststubgen/test/tinyframework/subclasstest/C2
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "C3.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CA.class
-  Compiled from "CA.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.CA
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/CA
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "CA.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CB.class
-  Compiled from "CB.java"
-public class com.android.hoststubgen.test.tinyframework.subclasstest.CB
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/CB
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "CB.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I1.class
-  Compiled from "I1.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I1
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I1
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I1.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I2.class
-  Compiled from "I2.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I2 extends com.android.hoststubgen.test.tinyframework.subclasstest.I1
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I2
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I2.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I3.class
-  Compiled from "I3.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.I3 extends com.android.hoststubgen.test.tinyframework.subclasstest.I2
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/I3
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "I3.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IA.class
-  Compiled from "IA.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.IA
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/IA
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "IA.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IB.class
-  Compiled from "IB.java"
-public interface com.android.hoststubgen.test.tinyframework.subclasstest.IB
-  minor version: 0
-  major version: 61
-  flags: (0x0601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/subclasstest/IB
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 0, attributes: 2
-}
-SourceFile: "IB.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-## Class: com/unsupported/UnsupportedClass.class
-  Compiled from "UnsupportedClass.java"
-public class com.unsupported.UnsupportedClass
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/unsupported/UnsupportedClass
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 2, attributes: 3
-  public com.unsupported.UnsupportedClass(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int getValue();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "UnsupportedClass.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-## Class: rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.class
-  Compiled from "TinyFrameworkToBeRenamed.java"
-public class rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 2, attributes: 3
-  private final int mValue;
-    descriptor: I
-    flags: (0x0012) ACC_PRIVATE, ACC_FINAL
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed(int);
-    descriptor: (I)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=2, args_size=2
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int getValue();
-    descriptor: ()I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=3, locals=1, args_size=1
-         x: new           #x                 // class java/lang/RuntimeException
-         x: dup
-         x: ldc           #x                 // String Stub!
-         x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
-         x: athrow
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkToBeRenamed.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-impl-dump.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-dump.txt
similarity index 75%
rename from tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-impl-dump.txt
rename to tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-dump.txt
index d12a23d..2ca723b 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-impl-dump.txt
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/golden-output/13-hoststubgen-test-tiny-framework-host-ext-dump.txt
@@ -22,12 +22,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestClassLoadHook.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -58,7 +58,7 @@
 SourceFile: "HostSideTestKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -91,12 +91,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestNativeSubstitutionClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -127,7 +127,7 @@
 SourceFile: "HostSideTestRemove.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -136,13 +136,13 @@
     java.lang.annotation.Retention(
       value=Ljava/lang/annotation/RetentionPolicy;.CLASS
     )
-## Class: android/hosttest/annotation/HostSideTestStub.class
-  Compiled from "HostSideTestStub.java"
-public interface android.hosttest.annotation.HostSideTestStub extends java.lang.annotation.Annotation
+## Class: android/hosttest/annotation/HostSideTestStaticInitializerKeep.class
+  Compiled from "HostSideTestStaticInitializerKeep.java"
+public interface android.hosttest.annotation.HostSideTestStaticInitializerKeep extends java.lang.annotation.Annotation
   minor version: 0
   major version: 61
   flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestStub
+  this_class: #x                          // android/hosttest/annotation/HostSideTestStaticInitializerKeep
   super_class: #x                         // java/lang/Object
   interfaces: 1, fields: 0, methods: 1, attributes: 2
   private static {};
@@ -150,15 +150,15 @@
     flags: (0x000a) ACC_PRIVATE, ACC_STATIC
     Code:
       stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class android/hosttest/annotation/HostSideTestStub
+         x: ldc           #x                  // class android/hosttest/annotation/HostSideTestStaticInitializerKeep
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
-SourceFile: "HostSideTestStub.java"
+SourceFile: "HostSideTestStaticInitializerKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x,e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE,Ljava/lang/annotation/ElementType;.FIELD,Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -191,12 +191,12 @@
     flags: (0x0401) ACC_PUBLIC, ACC_ABSTRACT
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "HostSideTestSubstitute.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.METHOD]
@@ -227,7 +227,7 @@
 SourceFile: "HostSideTestThrow.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x,e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.METHOD,Ljava/lang/annotation/ElementType;.CONSTRUCTOR]
@@ -258,38 +258,7 @@
 SourceFile: "HostSideTestWholeClassKeep.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-  x: #x(#x=[e#x.#x])
-    java.lang.annotation.Target(
-      value=[Ljava/lang/annotation/ElementType;.TYPE]
-    )
-  x: #x(#x=e#x.#x)
-    java.lang.annotation.Retention(
-      value=Ljava/lang/annotation/RetentionPolicy;.CLASS
-    )
-## Class: android/hosttest/annotation/HostSideTestWholeClassStub.class
-  Compiled from "HostSideTestWholeClassStub.java"
-public interface android.hosttest.annotation.HostSideTestWholeClassStub extends java.lang.annotation.Annotation
-  minor version: 0
-  major version: 61
-  flags: (0x2601) ACC_PUBLIC, ACC_INTERFACE, ACC_ABSTRACT, ACC_ANNOTATION
-  this_class: #x                          // android/hosttest/annotation/HostSideTestWholeClassStub
-  super_class: #x                         // java/lang/Object
-  interfaces: 1, fields: 0, methods: 1, attributes: 2
-  private static {};
-    descriptor: ()V
-    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
-    Code:
-      stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class android/hosttest/annotation/HostSideTestWholeClassStub
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
-         x: return
-}
-SourceFile: "HostSideTestWholeClassStub.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
   x: #x(#x=[e#x.#x])
     java.lang.annotation.Target(
       value=[Ljava/lang/annotation/ElementType;.TYPE]
@@ -313,7 +282,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -336,9 +305,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int addTwo(int);
     descriptor: (I)I
@@ -360,9 +327,7 @@
            11       4     0     a   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;          // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
@@ -370,9 +335,7 @@
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
 ## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub.class
   Compiled from "IPretendingAidl.java"
@@ -389,7 +352,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -412,9 +375,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int addOne(int);
     descriptor: (I)I
@@ -436,9 +397,7 @@
            11       4     0     a   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // Stub=class com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub of class com/android/hoststubgen/test/tinyframework/IPretendingAidl
@@ -446,9 +405,7 @@
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/IPretendingAidl
 ## Class: com/android/hoststubgen/test/tinyframework/IPretendingAidl.class
   Compiled from "IPretendingAidl.java"
@@ -465,7 +422,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/IPretendingAidl
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
@@ -475,9 +432,7 @@
 SourceFile: "IPretendingAidl.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub
   com/android/hoststubgen/test/tinyframework/IPretendingAidl$Stub$Proxy
@@ -495,9 +450,7 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.R$Nested();
     descriptor: ()V
@@ -518,9 +471,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/R$Nested;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -546,18 +497,14 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
-  public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
+  public static #x= #x of #x;             // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
 SourceFile: "R.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/R
 ## Class: com/android/hoststubgen/test/tinyframework/R.class
   Compiled from "R.java"
@@ -574,7 +521,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/R
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -597,239 +544,31 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/R;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/R$Nested of class com/android/hoststubgen/test/tinyframework/R
 SourceFile: "R.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/R$Nested
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl
-  minor version: 0
-  major version: 61
-  flags: (0x0020) ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 4, attributes: 4
-  private static {};
-    descriptor: ()V
-    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
-    Code:
-      stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
-         x: return
-
-  private com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck$Impl();
-    descriptor: ()V
-    flags: (0x0002) ACC_PRIVATE
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-         x: ldc           #x                 // String <init>
-         x: ldc           #x                 // String ()V
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-        x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOneKeep();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-         x: ldc           #x                 // String getOneKeep
-         x: ldc           #x                 // String ()I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-        x: ldc           #x                 // String getOneKeep
-        x: ldc           #x                 // String ()I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: iconst_1
-        x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestKeep
-
-  public static int getOneStub();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-         x: ldc           #x                 // String getOneStub
-         x: ldc           #x                 // String ()I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: iconst_1
-        x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-}
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.class
-  Compiled from "TinyFrameworkCallerCheck.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.class
+  Compiled from "TinyFrameworkAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations
   minor version: 0
   major version: 61
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 0, methods: 4, attributes: 5
-  private static {};
-    descriptor: ()V
-    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
-    Code:
-      stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
-         x: return
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkCallerCheck();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-         x: ldc           #x                 // String <init>
-         x: ldc           #x                 // String ()V
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-        x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_withCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-         x: ldc           #x                 // String getOne_withCheck
-         x: ldc           #x                 // String ()I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneKeep:()I
-        x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int getOne_noCheck();
-    descriptor: ()I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-         x: ldc           #x                 // String getOne_noCheck
-         x: ldc           #x                 // String ()I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl.getOneStub:()I
-        x: ireturn
-      LineNumberTable:
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-InnerClasses:
-  private static #x= #x of #x;          // Impl=class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl of class com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck
-SourceFile: "TinyFrameworkCallerCheck.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
-NestMembers:
-  com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck$Impl
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.class
-  Compiled from "TinyFrameworkClassAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 8, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
+  interfaces: 0, fields: 1, methods: 6, attributes: 3
   public int keep;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -839,20 +578,20 @@
     flags: (0x000a) ACC_PRIVATE, ACC_STATIC
     Code:
       stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
         x: return
 
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassAnnotations();
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkAnnotations();
     descriptor: ()V
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String <init>
          x: ldc           #x                 // String ()V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -861,68 +600,29 @@
         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
         x: aload_0
         x: iconst_1
-        x: putfield      #x                 // Field stub:I
-        x: aload_0
-        x: iconst_2
         x: putfield      #x                 // Field keep:I
         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           11      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+           11      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public int addOne(int);
     descriptor: (I)I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String addOne
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: iload_1
-        x: invokevirtual #x                 // Method addOneInner:(I)I
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-           11       6     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String addOneInner
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-        x: ldc           #x                 // String addOneInner
-        x: ldc           #x                 // String (I)I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iload_1
         x: iconst_1
         x: iadd
@@ -930,11 +630,11 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-           26       4     1 value   I
+           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
+           11       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -944,7 +644,7 @@
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String addTwo
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -956,22 +656,20 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
+           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations;
            11       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddThree(int);
     descriptor: (I)I
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
       stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String nativeAddThree
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -988,26 +686,18 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String unsupportedMethod();
     descriptor: ()Ljava/lang/String;
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations
          x: ldc           #x                 // String unsupportedMethod
          x: ldc           #x                 // String ()Ljava/lang/String;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-        x: ldc           #x                 // String unsupportedMethod
-        x: ldc           #x                 // String ()Ljava/lang/String;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
         x: new           #x                 // class java/lang/RuntimeException
         x: dup
@@ -1018,307 +708,22 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations
-         x: ldc           #x                 // String visibleButUsesUnsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-        x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-    RuntimeInvisibleAnnotations:
-      x: #x()
-        android.hosttest.annotation.HostSideTestStub
 }
-SourceFile: "TinyFrameworkClassAnnotations.java"
+SourceFile: "TinyFrameworkAnnotations.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestClassLoadHook(
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
-## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.class
-  Compiled from "TinyFrameworkClassClassWideAnnotations.java"
-public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations
-  minor version: 0
-  major version: 61
-  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
-  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-  super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 3, methods: 9, attributes: 3
-  public int stub;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int remove;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  private static {};
-    descriptor: ()V
-    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
-    Code:
-      stack=2, locals=0, args_size=0
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
-         x: return
-
-  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassClassWideAnnotations();
-    descriptor: ()V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String <init>
-         x: ldc           #x                 // String ()V
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
-        x: aload_0
-        x: iconst_1
-        x: putfield      #x                 // Field stub:I
-        x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
-        x: return
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOne(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String addOne
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: iload_1
-        x: invokevirtual #x                 // Method addOneInner:(I)I
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-           11       6     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String addOneInner
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: iload_1
-        x: iconst_1
-        x: iadd
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-           11       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public void toBeRemoved(java.lang.String);
-    descriptor: (Ljava/lang/String;)V
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String toBeRemoved
-         x: ldc           #x                 // String (Ljava/lang/String;)V
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: new           #x                 // class java/lang/RuntimeException
-        x: dup
-        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":()V
-        x: athrow
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       8     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-           11       8     1   foo   Ljava/lang/String;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addTwo(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String addTwo
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: iload_1
-        x: iconst_2
-        x: iadd
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-           11       4     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public static int nativeAddThree(int);
-    descriptor: (I)I
-    flags: (0x0009) ACC_PUBLIC, ACC_STATIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String nativeAddThree
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: iload_0
-        x: iconst_3
-        x: iadd
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       4     0 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String unsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String unsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String This value shouldn\'t be seen on the host side.
-        x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       3     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations
-         x: ldc           #x                 // String visibleButUsesUnsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokevirtual #x                 // Method unsupportedMethod:()Ljava/lang/String;
-        x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-}
-SourceFile: "TinyFrameworkClassClassWideAnnotations.java"
-RuntimeVisibleAnnotations:
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.class
   Compiled from "TinyFrameworkClassLoadHook.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook
@@ -1334,9 +739,7 @@
     Signature: #x                          // Ljava/util/Set<Ljava/lang/Class<*>;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook();
     descriptor: ()V
@@ -1357,9 +760,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void onClassLoaded(java.lang.Class<?>);
     descriptor: (Ljava/lang/Class;)V
@@ -1386,9 +787,7 @@
     Signature: #x                          // (Ljava/lang/Class<*>;)V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -1411,19 +810,145 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkClassLoadHook.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.class
+  Compiled from "TinyFrameworkClassWideAnnotations.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 1, methods: 5, attributes: 3
+  public int keep;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  private static {};
+    descriptor: ()V
+    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
+    Code:
+      stack=2, locals=0, args_size=0
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
+         x: return
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWideAnnotations();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=4, locals=1, args_size=1
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+         x: ldc           #x                 // String <init>
+         x: ldc           #x                 // String ()V
+         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
+        x: aload_0
+        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+        x: aload_0
+        x: iconst_1
+        x: putfield      #x                 // Field keep:I
+        x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+           11      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public int addOne(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=4, locals=2, args_size=2
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+         x: ldc           #x                 // String addOne
+         x: ldc           #x                 // String (I)I
+         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
+        x: iload_1
+        x: iconst_1
+        x: iadd
+        x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+           11       4     1 value   I
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public int addTwo(int);
+    descriptor: (I)I
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=4, locals=2, args_size=2
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+         x: ldc           #x                 // String addTwo
+         x: ldc           #x                 // String (I)I
+         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
+        x: iload_1
+        x: iconst_2
+        x: iadd
+        x: ireturn
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations;
+           11       4     1 value   I
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  public java.lang.String unsupportedMethod();
+    descriptor: ()Ljava/lang/String;
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=4, locals=1, args_size=1
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations
+         x: ldc           #x                 // String unsupportedMethod
+         x: ldc           #x                 // String ()Ljava/lang/String;
+         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
+        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+        x: new           #x                 // class java/lang/RuntimeException
+        x: dup
+        x: ldc           #x                 // String Unreachable
+        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+        x: athrow
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestThrow
+}
+SourceFile: "TinyFrameworkClassWideAnnotations.java"
+RuntimeVisibleAnnotations:
+  x: #x()
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+RuntimeInvisibleAnnotations:
+  x: #x()
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.class
   Compiled from "TinyFrameworkClassWithInitializerDefault.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerDefault
@@ -1438,35 +963,29 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
 }
 SourceFile: "TinyFrameworkClassWithInitializerDefault.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.class
   Compiled from "TinyFrameworkClassWithInitializerStub.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkClassWithInitializerStub
@@ -1481,24 +1000,20 @@
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.lang.Object sObject;
     descriptor: Ljava/lang/Object;
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   static {};
     descriptor: ()V
@@ -1526,21 +1041,19 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkClassWithInitializerStub.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestClassLoadHook(
       value="com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded"
     )
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.class
@@ -1557,43 +1070,37 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex GREEN;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex BLUE;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private final java.lang.String mLongName;
     descriptor: Ljava/lang/String;
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -1603,7 +1110,7 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestKeep
@@ -1613,9 +1120,7 @@
     flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1634,9 +1139,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex valueOf(java.lang.String);
     descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1659,9 +1162,7 @@
            11      10     0  name   Ljava/lang/String;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      mandated
@@ -1696,12 +1197,10 @@
     Signature: #x                          // (Ljava/lang/String;Ljava/lang/String;)V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
     MethodParameters:
       Name                           Flags
       <no name>                      synthetic
@@ -1728,12 +1227,10 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.lang.String getShortName();
     descriptor: ()Ljava/lang/String;
@@ -1754,12 +1251,10 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumComplex[] $values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1789,9 +1284,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -1826,7 +1319,7 @@
         x: dup
         x: ldc           #x                 // String BLUE
         x: iconst_2
-        x: ldc           #x                // String Blue
+        x: ldc           #x                 // String Blue
         x: ldc           #x                // String B
         x: invokespecial #x                 // Method "<init>":(Ljava/lang/String;ILjava/lang/String;Ljava/lang/String;)V
         x: putstatic     #x                 // Field BLUE:Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;
@@ -1836,20 +1329,16 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex;>;
 SourceFile: "TinyFrameworkEnumComplex.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.class
   Compiled from "TinyFrameworkEnumSimple.java"
 public final class com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple extends java.lang.Enum<com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple>
@@ -1864,33 +1353,27 @@
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple DOG;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
     flags: (0x4019) ACC_PUBLIC, ACC_STATIC, ACC_FINAL, ACC_ENUM
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static final com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] $VALUES;
     descriptor: [Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
     flags: (0x101a) ACC_PRIVATE, ACC_STATIC, ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple[] values();
     descriptor: ()[Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
@@ -1909,9 +1392,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static com.android.hoststubgen.test.tinyframework.TinyFrameworkEnumSimple valueOf(java.lang.String);
     descriptor: (Ljava/lang/String;)Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;
@@ -1934,9 +1415,7 @@
            11      10     0  name   Ljava/lang/String;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      mandated
@@ -1963,9 +1442,7 @@
     Signature: #x                          // ()V
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      synthetic
@@ -1995,9 +1472,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -2030,20 +1505,16 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 Signature: #x                           // Ljava/lang/Enum<Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple;>;
 SourceFile: "TinyFrameworkEnumSimple.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.class
   Compiled from "TinyFrameworkExceptionTester.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkExceptionTester
@@ -2059,7 +1530,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -2082,9 +1553,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int testException();
     descriptor: ()I
@@ -2120,19 +1589,15 @@
            22      11     0     e   Ljava/lang/Exception;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkExceptionTester.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.class
   Compiled from "TinyFrameworkForTextPolicy.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkForTextPolicy
@@ -2141,22 +1606,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 17, attributes: 2
+  interfaces: 0, fields: 1, methods: 15, attributes: 2
   public int stub;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int keep;
-    descriptor: I
-    flags: (0x0001) ACC_PUBLIC
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -2164,7 +1620,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
          x: ldc           #x                 // String com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
@@ -2186,19 +1642,14 @@
         x: aload_0
         x: iconst_1
         x: putfield      #x                 // Field stub:I
-        x: aload_0
-        x: iconst_2
-        x: putfield      #x                 // Field keep:I
         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           11      15     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
+           11      10     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int addOne(int);
     descriptor: (I)I
@@ -2210,37 +1661,6 @@
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: iload_1
-        x: invokevirtual #x                 // Method addOneInner:(I)I
-        x: ireturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-           11       6     1 value   I
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public int addOneInner(int);
-    descriptor: (I)I
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=2, args_size=2
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                 // String addOneInner
-         x: ldc           #x                 // String (I)I
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-        x: ldc           #x                 // String addOneInner
-        x: ldc           #x                 // String (I)I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iload_1
         x: iconst_1
         x: iadd
@@ -2248,11 +1668,11 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-           26       4     1 value   I
+           11       4     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
+           11       4     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String toBeIgnoredObj();
     descriptor: ()Ljava/lang/String;
@@ -2270,9 +1690,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public void toBeIgnoredV();
     descriptor: ()V
@@ -2289,9 +1707,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public boolean toBeIgnoredZ();
     descriptor: ()Z
@@ -2309,9 +1725,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public byte toBeIgnoredB();
     descriptor: ()B
@@ -2329,9 +1743,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public char toBeIgnoredC();
     descriptor: ()C
@@ -2349,9 +1761,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public short toBeIgnoredS();
     descriptor: ()S
@@ -2369,9 +1779,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int toBeIgnoredI();
     descriptor: ()I
@@ -2389,9 +1797,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public float toBeIgnoredF();
     descriptor: ()F
@@ -2409,9 +1815,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public double toBeIgnoredD();
     descriptor: ()D
@@ -2429,9 +1833,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsIgnore
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int addTwo(int);
     descriptor: (I)I
@@ -2439,7 +1841,7 @@
     Code:
       stack=4, locals=2, args_size=2
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                // String addTwo
+         x: ldc           #x                 // String addTwo
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
@@ -2456,9 +1858,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddThree(int);
     descriptor: (I)I
@@ -2466,7 +1866,7 @@
     Code:
       stack=4, locals=1, args_size=1
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                // String nativeAddThree
+         x: ldc           #x                 // String nativeAddThree
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
@@ -2482,9 +1882,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.String unsupportedMethod();
     descriptor: ()Ljava/lang/String;
@@ -2492,57 +1890,26 @@
     Code:
       stack=4, locals=1, args_size=1
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                // String unsupportedMethod
+         x: ldc           #x                 // String unsupportedMethod
          x: ldc           #x                 // String ()Ljava/lang/String;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-        x: ldc           #x                // String unsupportedMethod
-        x: ldc           #x                 // String ()Ljava/lang/String;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
-        x: invokestatic  #x                // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
-        x: new           #x                // class java/lang/RuntimeException
+        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
+        x: new           #x                 // class java/lang/RuntimeException
         x: dup
-        x: ldc           #x                // String Unreachable
-        x: invokespecial #x                // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+        x: ldc           #x                 // String Unreachable
+        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
         x: athrow
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-
-  public java.lang.String visibleButUsesUnsupportedMethod();
-    descriptor: ()Ljava/lang/String;
-    flags: (0x0001) ACC_PUBLIC
-    Code:
-      stack=4, locals=1, args_size=1
-         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy
-         x: ldc           #x                // String visibleButUsesUnsupportedMethod
-         x: ldc           #x                 // String ()Ljava/lang/String;
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
-         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: aload_0
-        x: invokevirtual #x                // Method unsupportedMethod:()Ljava/lang/String;
-        x: areturn
-      LineNumberTable:
-      LocalVariableTable:
-        Start  Length  Slot  Name   Signature
-           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy;
-    RuntimeVisibleAnnotations:
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkForTextPolicy.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested.class
   Compiled from "TinyFrameworkLambdas.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested
@@ -2558,12 +1925,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -2571,12 +1936,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas$Nested();
     descriptor: ()V
@@ -2600,12 +1963,10 @@
            11      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2626,12 +1987,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2649,12 +2008,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -2672,9 +2029,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$getSupplier$2();
     descriptor: ()Ljava/lang/Integer;
@@ -2692,9 +2047,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$static$1();
     descriptor: ()Ljava/lang/Integer;
@@ -2712,9 +2065,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$new$0();
     descriptor: ()Ljava/lang/Integer;
@@ -2732,9 +2083,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -2750,12 +2099,12 @@
         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
         x: invokedynamic #x,  0             // InvokeDynamic #x:get:()Ljava/util/function/Supplier;
-        x: putstatic     #x                // Field sSupplier:Ljava/util/function/Supplier;
+        x: putstatic     #x                 // Field sSupplier:Ljava/util/function/Supplier;
         x: return
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;            // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
@@ -2763,12 +2112,10 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 BootstrapMethods:
@@ -2808,12 +2155,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -2821,12 +2166,10 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkLambdas();
     descriptor: ()V
@@ -2850,12 +2193,10 @@
            11      14     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2876,12 +2217,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -2899,12 +2238,10 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
-        android.hosttest.annotation.HostSideTestStub
+        android.hosttest.annotation.HostSideTestKeep
 
   private static java.lang.Integer lambda$getSupplier_static$3();
     descriptor: ()Ljava/lang/Integer;
@@ -2922,9 +2259,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$getSupplier$2();
     descriptor: ()Ljava/lang/Integer;
@@ -2942,9 +2277,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$static$1();
     descriptor: ()Ljava/lang/Integer;
@@ -2962,9 +2295,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static java.lang.Integer lambda$new$0();
     descriptor: ()Ljava/lang/Integer;
@@ -2982,9 +2313,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -3000,12 +2329,12 @@
         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
         x: invokedynamic #x,  0             // InvokeDynamic #x:get:()Ljava/util/function/Supplier;
-        x: putstatic     #x                // Field sSupplier:Ljava/util/function/Supplier;
+        x: putstatic     #x                 // Field sSupplier:Ljava/util/function/Supplier;
         x: return
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // Nested=class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas$Nested of class com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas
@@ -3013,12 +2342,10 @@
 SourceFile: "TinyFrameworkLambdas.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestStub
+    android.hosttest.annotation.HostSideTestKeep
   x: #x()
     android.hosttest.annotation.HostSideTestStaticInitializerKeep
 BootstrapMethods:
@@ -3059,7 +2386,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -3082,9 +2409,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void startThread(java.lang.Thread);
     descriptor: (Ljava/lang/Thread;)V
@@ -3108,9 +2433,7 @@
            11      10     0 thread   Ljava/lang/Thread;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int add(int, int);
     descriptor: (II)I
@@ -3133,18 +2456,14 @@
            11       4     1     b   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
 SourceFile: "TinyFrameworkMethodCallReplace.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.class
   Compiled from "TinyFrameworkMethodCallReplace.java"
@@ -3161,7 +2480,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -3184,9 +2503,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static boolean nonStaticMethodCallReplaceTester() throws java.lang.Exception;
     descriptor: ()Z
@@ -3225,9 +2542,7 @@
       throws java.lang.Exception
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int staticMethodCallReplaceTester();
     descriptor: ()I
@@ -3246,9 +2561,7 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static void lambda$nonStaticMethodCallReplaceTester$0(java.util.concurrent.atomic.AtomicBoolean);
     descriptor: (Ljava/util/concurrent/atomic/AtomicBoolean;)V
@@ -3261,7 +2574,7 @@
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
         x: aload_0
-        x: invokestatic  #x                // Method java/lang/Thread.currentThread:()Ljava/lang/Thread;
+        x: invokestatic  #x                 // Method java/lang/Thread.currentThread:()Ljava/lang/Thread;
         x: invokevirtual #x                // Method java/lang/Thread.isDaemon:()Z
         x: invokevirtual #x                // Method java/util/concurrent/atomic/AtomicBoolean.set:(Z)V
         x: return
@@ -3271,9 +2584,7 @@
            11      11     0    ab   Ljava/util/concurrent/atomic/AtomicBoolean;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // ReplaceTo=class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace$ReplaceTo of class com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace
@@ -3281,12 +2592,10 @@
 SourceFile: "TinyFrameworkMethodCallReplace.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 BootstrapMethods:
   x: #x REF_invokeStatic java/lang/invoke/LambdaMetafactory.metafactory:(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;Ljava/lang/invoke/MethodType;Ljava/lang/invoke/MethodHandle;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;
     Method arguments:
@@ -3303,15 +2612,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 12, attributes: 3
+  interfaces: 0, fields: 1, methods: 13, attributes: 3
   int value;
     descriptor: I
     flags: (0x0000)
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -3319,7 +2626,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -3342,9 +2649,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo(int);
     descriptor: (I)I
@@ -3363,9 +2668,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo_should_be_like_this(int);
     descriptor: (I)I
@@ -3386,9 +2689,7 @@
            11       5     0   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus(long, long);
     descriptor: (JJ)J
@@ -3408,9 +2709,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus_should_be_like_this(long, long);
     descriptor: (JJ)J
@@ -3433,9 +2732,7 @@
            11       6     2  arg2   J
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public void setValue(int);
     descriptor: (I)V
@@ -3458,9 +2755,7 @@
            11       6     1     v   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int nativeNonStaticAddToValue(int);
     descriptor: (I)I
@@ -3480,9 +2775,7 @@
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int nativeNonStaticAddToValue_should_be_like_this(int);
     descriptor: (I)I
@@ -3505,9 +2798,7 @@
            11       6     1   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static void nativeStillNotSupported();
     descriptor: ()V
@@ -3519,49 +2810,49 @@
          x: ldc           #x                 // String ()V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-        x: ldc           #x                 // String nativeStillNotSupported
-        x: ldc           #x                 // String ()V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onThrowMethodCalled:()V
         x: new           #x                 // class java/lang/RuntimeException
         x: dup
         x: ldc           #x                 // String Unreachable
-        x: invokespecial #x                // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
+        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":(Ljava/lang/String;)V
         x: athrow
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsThrow
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     RuntimeInvisibleAnnotations:
       x: #x()
         android.hosttest.annotation.HostSideTestThrow
 
+  public static native void nativeStillKeep();
+    descriptor: ()V
+    flags: (0x0109) ACC_PUBLIC, ACC_STATIC, ACC_NATIVE
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+    RuntimeInvisibleAnnotations:
+      x: #x()
+        android.hosttest.annotation.HostSideTestKeep
+
   public static void nativeStillNotSupported_should_be_like_this();
     descriptor: ()V
     flags: (0x0009) ACC_PUBLIC, ACC_STATIC
     Code:
       stack=4, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-         x: ldc           #x                // String nativeStillNotSupported_should_be_like_this
+         x: ldc           #x                 // String nativeStillNotSupported_should_be_like_this
          x: ldc           #x                 // String ()V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
         x: new           #x                 // class java/lang/RuntimeException
         x: dup
-        x: invokespecial #x                // Method java/lang/RuntimeException."<init>":()V
+        x: invokespecial #x                 // Method java/lang/RuntimeException."<init>":()V
         x: athrow
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static byte nativeBytePlus(byte, byte);
     descriptor: (BB)B
@@ -3569,31 +2860,27 @@
     Code:
       stack=4, locals=2, args_size=2
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNative
-         x: ldc           #x                // String nativeBytePlus
-         x: ldc           #x                // String (BB)B
+         x: ldc           #x                 // String nativeBytePlus
+         x: ldc           #x                 // String (BB)B
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
         x: iload_0
         x: iload_1
-        x: invokestatic  #x                // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host.nativeBytePlus:(BB)B
+        x: invokestatic  #x                 // Method com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host.nativeBytePlus:(BB)B
         x: ireturn
     RuntimeVisibleAnnotations:
       x: #x()
         com.android.hoststubgen.hosthelper.HostStubGenProcessedAsSubstitute
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkNative.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
   x: #x(#x=s#x)
     android.hosttest.annotation.HostSideTestNativeSubstitutionClass(
       value="TinyFrameworkNative_host"
@@ -3627,22 +2914,16 @@
          x: ldc           #x                 // String ()V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-        x: ldc           #x                 // String <init>
-        x: ldc           #x                 // String ()V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
         x: return
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeAddTwo(int);
     descriptor: (I)I
@@ -3654,12 +2935,6 @@
          x: ldc           #x                 // String (I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-        x: ldc           #x                 // String nativeAddTwo
-        x: ldc           #x                 // String (I)I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iload_0
         x: iconst_2
         x: iadd
@@ -3667,10 +2942,10 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       4     0   arg   I
+           11       4     0   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static long nativeLongPlus(long, long);
     descriptor: (JJ)J
@@ -3682,12 +2957,6 @@
          x: ldc           #x                 // String (JJ)J
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-        x: ldc           #x                 // String nativeLongPlus
-        x: ldc           #x                 // String (JJ)J
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: lload_0
         x: lload_2
         x: ladd
@@ -3695,11 +2964,11 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       4     0  arg1   J
-           26       4     2  arg2   J
+           11       4     0  arg1   J
+           11       4     2  arg2   J
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int nativeNonStaticAddToValue(com.android.hoststubgen.test.tinyframework.TinyFrameworkNative, int);
     descriptor: (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;I)I
@@ -3711,12 +2980,6 @@
          x: ldc           #x                 // String (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;I)I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-        x: ldc           #x                 // String nativeNonStaticAddToValue
-        x: ldc           #x                 // String (Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;I)I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: getfield      #x                 // Field com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.value:I
         x: iload_1
@@ -3725,11 +2988,11 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       7     0 source   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
-           26       7     1   arg   I
+           11       7     0 source   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNative;
+           11       7     1   arg   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static byte nativeBytePlus(byte, byte);
     descriptor: (BB)B
@@ -3741,12 +3004,6 @@
          x: ldc           #x                 // String (BB)B
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNative_host
-        x: ldc           #x                 // String nativeBytePlus
-        x: ldc           #x                 // String (BB)B
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iload_0
         x: iload_1
         x: iadd
@@ -3755,16 +3012,16 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  arg1   B
-           26       5     1  arg2   B
+           11       5     0  arg1   B
+           11       5     1  arg2   B
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkNative_host.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
     android.hosttest.annotation.HostSideTestWholeClassKeep
@@ -3782,7 +3039,7 @@
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -3817,7 +3074,7 @@
            11      10     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
@@ -3832,22 +3089,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Integer;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Integer;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iconst_1
         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
@@ -3859,22 +3110,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Object;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Object;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
@@ -3883,7 +3128,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -3923,7 +3168,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
@@ -3935,22 +3180,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Integer;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Integer;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iconst_2
         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
@@ -3962,22 +3201,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Object;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Object;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
@@ -3986,7 +3219,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4002,7 +3235,7 @@
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -4037,7 +3270,7 @@
            11      10     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
@@ -4052,22 +3285,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Integer;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Integer;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iconst_3
         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
@@ -4079,22 +3306,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Object;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Object;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
@@ -4103,7 +3324,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4143,7 +3364,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
@@ -4155,22 +3376,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Integer;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Integer;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: iconst_4
         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
@@ -4182,22 +3397,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Object;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Object;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                     // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
@@ -4206,7 +3415,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4222,9 +3431,7 @@
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -4232,7 +3439,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4259,18 +3466,14 @@
            11      10     1     x   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4280,24 +3483,20 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 2, methods: 2, attributes: 5
+  interfaces: 0, fields: 2, methods: 2, attributes: 4
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   final com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses this$0;
     descriptor: Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     flags: (0x1010) ACC_FINAL, ACC_SYNTHETIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -4305,7 +3504,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4335,9 +3534,7 @@
            11      15     1 this$0   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
     MethodParameters:
       Name                           Flags
       <no name>                      final mandated
@@ -4347,12 +3544,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4392,7 +3584,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Integer get();
     descriptor: ()Ljava/lang/Integer;
@@ -4404,22 +3596,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Integer;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Integer;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: bipush        7
         x: invokestatic  #x                 // Method java/lang/Integer.valueOf:(I)Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
+           11       6     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.lang.Object get();
     descriptor: ()Ljava/lang/Object;
@@ -4431,22 +3617,16 @@
          x: ldc           #x                 // String ()Ljava/lang/Object;
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
-        x: ldc           #x                 // String get
-        x: ldc           #x                 // String ()Ljava/lang/Object;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokevirtual #x                 // Method get:()Ljava/lang/Integer;
         x: areturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
+           11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
@@ -4456,7 +3636,65 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass.class
+  Compiled from "TinyFrameworkNestedClasses.java"
+public class com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  minor version: 0
+  major version: 61
+  flags: (0x0021) ACC_PUBLIC, ACC_SUPER
+  this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+  super_class: #x                         // java/lang/Object
+  interfaces: 0, fields: 1, methods: 2, attributes: 4
+  public int value;
+    descriptor: I
+    flags: (0x0001) ACC_PUBLIC
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+
+  private static {};
+    descriptor: ()V
+    flags: (0x000a) ACC_PRIVATE, ACC_STATIC
+    Code:
+      stack=2, locals=0, args_size=0
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
+         x: return
+
+  public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass();
+    descriptor: ()V
+    flags: (0x0001) ACC_PUBLIC
+    Code:
+      stack=4, locals=1, args_size=1
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
+         x: ldc           #x                 // String <init>
+         x: ldc           #x                 // String ()V
+         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
+         x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
+        x: aload_0
+        x: invokespecial #x                 // Method java/lang/Object."<init>":()V
+        x: aload_0
+        x: bipush        8
+        x: putfield      #x                 // Field value:I
+        x: return
+      LineNumberTable:
+      LocalVariableTable:
+        Start  Length  Slot  Name   Signature
+           11      11     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass;
+    RuntimeVisibleAnnotations:
+      x: #x()
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
+}
+InnerClasses:
+  public static #x= #x of #x;          // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
+SourceFile: "TinyFrameworkNestedClasses.java"
+RuntimeVisibleAnnotations:
+  x: #x()
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4466,15 +3704,13 @@
   flags: (0x0021) ACC_PUBLIC, ACC_SUPER
   this_class: #x                          // com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   super_class: #x                         // java/lang/Object
-  interfaces: 0, fields: 1, methods: 3, attributes: 5
+  interfaces: 0, fields: 1, methods: 3, attributes: 4
   public int value;
     descriptor: I
     flags: (0x0001) ACC_PUBLIC
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -4482,7 +3718,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4508,9 +3744,7 @@
            11      11     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -4530,22 +3764,16 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
+  public static #x= #x of #x;           // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
-RuntimeInvisibleAnnotations:
-  x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4562,7 +3790,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4587,9 +3815,7 @@
            11       6     1     x   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
@@ -4597,9 +3823,7 @@
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 NestHost: class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.class
   Compiled from "TinyFrameworkNestedClasses.java"
@@ -4616,9 +3840,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static final java.util.function.Supplier<java.lang.Integer> sSupplier;
     descriptor: Ljava/util/function/Supplier;
@@ -4626,9 +3848,7 @@
     Signature: #x                          // Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses();
     descriptor: ()V
@@ -4655,9 +3875,7 @@
            11      17     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public java.util.function.Supplier<java.lang.Integer> getSupplier();
     descriptor: ()Ljava/util/function/Supplier;
@@ -4681,9 +3899,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static java.util.function.Supplier<java.lang.Integer> getSupplier_static();
     descriptor: ()Ljava/util/function/Supplier;
@@ -4703,9 +3919,7 @@
     Signature: #x                          // ()Ljava/util/function/Supplier<Ljava/lang/Integer;>;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   static {};
     descriptor: ()V
@@ -4728,33 +3942,31 @@
       LineNumberTable:
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 InnerClasses:
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$1
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$3
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$2
-  public static #x= #x of #x;           // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;            // SubClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;           // BaseClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public static #x= #x of #x;           // StaticNestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
   public #x= #x of #x;                  // InnerClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses
+  public static #x= #x of #x;          // Double$NestedClass=class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass of class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
   #x;                                    // class com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
 SourceFile: "TinyFrameworkNestedClasses.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 NestMembers:
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$SubClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$BaseClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass
+  com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$Double$NestedClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$StaticNestedClass$1
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$InnerClass
   com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses$4
@@ -4776,7 +3988,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4799,9 +4011,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int foo(int);
     descriptor: (I)I
@@ -4825,19 +4035,15 @@
            11      12     0 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkPackageRedirect.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.class
   Compiled from "TinyFrameworkRenamedClassCaller.java"
 public class com.android.hoststubgen.test.tinyframework.TinyFrameworkRenamedClassCaller
@@ -4853,7 +4059,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -4876,9 +4082,7 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public static int foo(int);
     descriptor: (I)I
@@ -4902,19 +4106,15 @@
            11      12     0 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkRenamedClassCaller.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: com/android/hoststubgen/test/tinyframework/packagetest/A.class
   Compiled from "A.java"
 public class com.android.hoststubgen.test.tinyframework.packagetest.A
@@ -4930,16 +4130,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/packagetest/A
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "A.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/packagetest/sub/A.class
   Compiled from "A.java"
 public class com.android.hoststubgen.test.tinyframework.packagetest.sub.A
@@ -4955,16 +4153,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/packagetest/sub/A
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "A.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C1.class
   Compiled from "C1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -4980,16 +4176,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/C1
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "C1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C2.class
   Compiled from "C2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -5005,16 +4199,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/C2
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "C2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/C3.class
   Compiled from "C3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
@@ -5030,16 +4222,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/C3
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "C3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CA.class
   Compiled from "CA.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.CA
@@ -5055,16 +4245,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/CA
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "CA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/CB.class
   Compiled from "CB.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.CB
@@ -5080,16 +4268,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/CB
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "CB.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C1.class
   Compiled from "Class_C1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C1 extends com.android.hoststubgen.test.tinyframework.subclasstest.C1
@@ -5112,7 +4298,7 @@
 SourceFile: "Class_C1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C2.class
   Compiled from "Class_C2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C2 extends com.android.hoststubgen.test.tinyframework.subclasstest.C2
@@ -5135,7 +4321,7 @@
 SourceFile: "Class_C2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_C3.class
   Compiled from "Class_C3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_C3 extends com.android.hoststubgen.test.tinyframework.subclasstest.C3
@@ -5158,7 +4344,7 @@
 SourceFile: "Class_C3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I1.class
   Compiled from "Class_I1.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I1 implements com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -5181,7 +4367,7 @@
 SourceFile: "Class_I1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I1_IA.class
   Compiled from "Class_I1_IA.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I1_IA implements com.android.hoststubgen.test.tinyframework.subclasstest.I1,com.android.hoststubgen.test.tinyframework.subclasstest.IA
@@ -5204,7 +4390,7 @@
 SourceFile: "Class_I1_IA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I2.class
   Compiled from "Class_I2.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I2 implements com.android.hoststubgen.test.tinyframework.subclasstest.I2
@@ -5227,7 +4413,7 @@
 SourceFile: "Class_I2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/Class_I3.class
   Compiled from "Class_I3.java"
 public class com.android.hoststubgen.test.tinyframework.subclasstest.Class_I3 implements com.android.hoststubgen.test.tinyframework.subclasstest.I3
@@ -5250,7 +4436,7 @@
 SourceFile: "Class_I3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I1.class
   Compiled from "I1.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -5266,16 +4452,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/I1
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "I1.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I2.class
   Compiled from "I2.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I2 extends com.android.hoststubgen.test.tinyframework.subclasstest.I1
@@ -5298,9 +4482,7 @@
 SourceFile: "I2.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/I3.class
   Compiled from "I3.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.I3 extends com.android.hoststubgen.test.tinyframework.subclasstest.I2
@@ -5323,9 +4505,7 @@
 SourceFile: "I3.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IA.class
   Compiled from "IA.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.IA
@@ -5341,16 +4521,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/IA
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "IA.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/android/hoststubgen/test/tinyframework/subclasstest/IB.class
   Compiled from "IB.java"
 public interface com.android.hoststubgen.test.tinyframework.subclasstest.IB
@@ -5366,16 +4544,14 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/subclasstest/IB
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 }
 SourceFile: "IB.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 ## Class: com/supported/UnsupportedClass.class
   Compiled from "UnsupportedClass.java"
 public class com.supported.UnsupportedClass
@@ -5390,7 +4566,7 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
@@ -5412,12 +4588,6 @@
          x: ldc           #x                 // String (I)V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/supported/UnsupportedClass
-        x: ldc           #x                 // String <init>
-        x: ldc           #x                 // String (I)V
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: invokespecial #x                 // Method java/lang/Object."<init>":()V
         x: aload_0
@@ -5427,11 +4597,11 @@
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26      10     0  this   Lcom/supported/UnsupportedClass;
-           26      10     1 value   I
+           11      10     0  this   Lcom/supported/UnsupportedClass;
+           11      10     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
@@ -5443,27 +4613,21 @@
          x: ldc           #x                 // String ()I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.callMethodCallHook:(Ljava/lang/Class;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)V
-        x: ldc           #x                 // String com/supported/UnsupportedClass
-        x: ldc           #x                 // String getValue
-        x: ldc           #x                 // String ()I
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.getStackWalker:()Ljava/lang/StackWalker;
-        x: invokevirtual #x                 // Method java/lang/StackWalker.getCallerClass:()Ljava/lang/Class;
-        x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onNonStubMethodCalled:(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/Class;)V
         x: aload_0
         x: getfield      #x                 // Field mValue:I
         x: ireturn
       LineNumberTable:
       LocalVariableTable:
         Start  Length  Slot  Name   Signature
-           26       5     0  this   Lcom/supported/UnsupportedClass;
+           11       5     0  this   Lcom/supported/UnsupportedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "UnsupportedClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
     android.hosttest.annotation.HostSideTestWholeClassKeep
@@ -5482,7 +4646,7 @@
     Code:
       stack=2, locals=0, args_size=0
          x: ldc           #x                  // class com/unsupported/UnsupportedClass
-         x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
+         x: ldc           #x                  // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
 
@@ -5510,9 +4674,7 @@
            11      14     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
@@ -5535,19 +4697,15 @@
            11      10     0  this   Lcom/unsupported/UnsupportedClass;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "UnsupportedClass.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
 ## Class: rename_prefix/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.class
   Compiled from "TinyFrameworkToBeRenamed.java"
 public class rename_prefix.com.android.hoststubgen.test.tinyframework.TinyFrameworkToBeRenamed
@@ -5562,16 +4720,14 @@
     flags: (0x0012) ACC_PRIVATE, ACC_FINAL
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   private static {};
     descriptor: ()V
     flags: (0x000a) ACC_PRIVATE, ACC_STATIC
     Code:
       stack=2, locals=0, args_size=0
-         x: ldc           #x                 // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logClassLoaded
          x: invokestatic  #x                 // Method com/android/hoststubgen/hosthelper/HostTestUtils.onClassLoaded:(Ljava/lang/Class;Ljava/lang/String;)V
          x: return
@@ -5581,7 +4737,7 @@
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=2, args_size=2
-         x: ldc           #x                 // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
          x: ldc           #x                 // String <init>
          x: ldc           #x                 // String (I)V
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -5599,16 +4755,14 @@
            11      10     1 value   I
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 
   public int getValue();
     descriptor: ()I
     flags: (0x0001) ACC_PUBLIC
     Code:
       stack=4, locals=1, args_size=1
-         x: ldc           #x                 // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
+         x: ldc           #x                  // class com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed
          x: ldc           #x                 // String getValue
          x: ldc           #x                 // String ()I
          x: ldc           #x                 // String com.android.hoststubgen.hosthelper.HostTestUtils.logMethodCall
@@ -5622,16 +4776,12 @@
            11       5     0  this   Lcom/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed;
     RuntimeVisibleAnnotations:
       x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-      x: #x()
-        com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+        com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 }
 SourceFile: "TinyFrameworkToBeRenamed.java"
 RuntimeVisibleAnnotations:
   x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInStub
-  x: #x()
-    com.android.hoststubgen.hosthelper.HostStubGenKeptInImpl
+    com.android.hoststubgen.hosthelper.HostStubGenProcessedAsKeep
 RuntimeInvisibleAnnotations:
   x: #x()
-    android.hosttest.annotation.HostSideTestWholeClassStub
+    android.hosttest.annotation.HostSideTestWholeClassKeep
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/policy-override-tiny-framework.txt b/tools/hoststubgen/hoststubgen/test-tiny-framework/policy-override-tiny-framework.txt
index f064433..3c138d2 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/policy-override-tiny-framework.txt
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/policy-override-tiny-framework.txt
@@ -1,9 +1,8 @@
-class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy	stub
-  field stub	stub
-  field keep	keep
+class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy	keep
+  field stub	keep
   # field remove	remove # Implicitly remove
-  method <init>	()V	            stub
-  method addOne	(I)I	        stub
+  method <init>	()V	            keep
+  method addOne	(I)I	        keep
   method addOneInner	(I)I	keep
   method toBeRemoved	(Ljava/lang/String;)V	remove
   method addTwo	(I)I	        @addTwo_host
@@ -11,7 +10,7 @@
   method nativeAddThree	(I)I	@addThree_host
   # method addThree_host	(I)I	# used as a substitute
   method unsupportedMethod	()Ljava/lang/String;	throw
-  method visibleButUsesUnsupportedMethod	()Ljava/lang/String;	stub
+  method visibleButUsesUnsupportedMethod	()Ljava/lang/String;	keep
   method toBeIgnoredObj	()Ljava/lang/String;	ignore
   method toBeIgnoredV	()V	ignore
   method toBeIgnoredZ	()Z	ignore
@@ -27,22 +26,22 @@
 class com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy	~com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded
 
 # Heuristics rule: Stub all the AIDL classes.
-class :aidl stubclass
+class :aidl keepclass
 
 # Heuristics rule: Stub all the R classes.
-class :r stubclass
+class :r keepclass
 
 # Default is "remove", so let's put all the base classes / interfaces in the stub first.
-class com.android.hoststubgen.test.tinyframework.subclasstest.C1 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.C2 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.C3 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.CA stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.CB stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.I1 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.I2 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.I3 stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.IA stub
-class com.android.hoststubgen.test.tinyframework.subclasstest.IB stub
+class com.android.hoststubgen.test.tinyframework.subclasstest.C1 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.C2 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.C3 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.CA keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.CB keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.I1 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.I2 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.I3 keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.IA keep
+class com.android.hoststubgen.test.tinyframework.subclasstest.IB keep
 
 # Then define inheritance based policies.
 class *com.android.hoststubgen.test.tinyframework.subclasstest.C1 keep
@@ -52,15 +51,15 @@
 class *com.android.hoststubgen.test.tinyframework.subclasstest.IA remove
 
 # Test package directive
-package com.android.hoststubgen.test.tinyframework.packagetest stub
+package com.android.hoststubgen.test.tinyframework.packagetest keep
 class com.android.hoststubgen.test.tinyframework.packagetest.B remove
 class com.android.hoststubgen.test.tinyframework.packagetest.sub.B remove
 # The following rules are the same as above
-# class com.android.hoststubgen.test.tinyframework.packagetest.A stub
-# class com.android.hoststubgen.test.tinyframework.packagetest.sub.A stub
+# class com.android.hoststubgen.test.tinyframework.packagetest.A keep
+# class com.android.hoststubgen.test.tinyframework.packagetest.sub.A keep
 
 # Used to test method call replacement.
-class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace stubclass
+class com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace keepclass
   method originalAdd (II)I @com.android.hoststubgen.test.tinyframework.TinyFrameworkMethodCallReplace$ReplaceTo.add
 
 # Used to test method call replacement.
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/run-test-manually.sh b/tools/hoststubgen/hoststubgen/test-tiny-framework/run-test-manually.sh
index 872bbf8..80ebf3a 100755
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/run-test-manually.sh
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/run-test-manually.sh
@@ -43,8 +43,7 @@
 
 tiny_framework_classes=$out/tiny-framework/classes/
 tiny_framework_jar=$out/tiny-framework.jar
-tiny_framework_host_stub_jar=$out/tiny-framework_host_stub.jar
-tiny_framework_host_impl_jar=$out/tiny-framework_host_impl.jar
+tiny_framework_host_jar=$out/tiny-framework_host.jar
 
 tiny_test_classes=$out/tiny-test/classes/
 tiny_test_jar=$out/tiny-test.jar
@@ -87,8 +86,7 @@
 run $HOSTSTUBGEN \
     @../hoststubgen-standard-options.txt \
     --in-jar $tiny_framework_jar \
-    --out-stub-jar $tiny_framework_host_stub_jar \
-    --out-impl-jar $tiny_framework_host_impl_jar \
+    --out-jar $tiny_framework_host_jar \
     --policy-override-file policy-override-tiny-framework.txt \
     --gen-keep-all-file out/tiny-framework_keep_all.txt \
     --gen-input-dump-file out/tiny-framework_dump.txt \
@@ -97,14 +95,14 @@
     $HOSTSTUBGEN_OPTS
 
 # Extract the jar files, so we can look into them.
-extract $tiny_framework_host_stub_jar $tiny_framework_host_impl_jar
+extract $tiny_framework_host_jar
 
 # Build the test
 echo "# Building tiny-test..."
 run $JAVAC \
     -cp $( \
         join : \
-        $tiny_framework_host_stub_jar \
+        $tiny_framework_jar \
         "${test_compile_classpaths[@]}" \
         ) \
     -d $tiny_test_classes \
@@ -124,7 +122,7 @@
     -cp $( \
         join : \
         $tiny_test_jar \
-        $tiny_framework_host_impl_jar \
+        $tiny_framework_host_jar \
         "${test_compile_classpaths[@]}" \
         "${test_runtime_classpaths[@]}" \
         ) \
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.java
similarity index 80%
rename from tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.java
rename to tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.java
index 6d8a48a..ed0fa26 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassAnnotations.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotations.java
@@ -18,39 +18,30 @@
 import android.hosttest.annotation.HostSideTestClassLoadHook;
 import android.hosttest.annotation.HostSideTestKeep;
 import android.hosttest.annotation.HostSideTestRemove;
-import android.hosttest.annotation.HostSideTestStub;
 import android.hosttest.annotation.HostSideTestSubstitute;
 import android.hosttest.annotation.HostSideTestThrow;
 
 /**
  * Test without class-wide annotations.
  */
-@HostSideTestStub
+@HostSideTestKeep
 @HostSideTestClassLoadHook(
         "com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded")
-public class TinyFrameworkClassAnnotations {
-    @HostSideTestStub
-    public TinyFrameworkClassAnnotations() {
+public class TinyFrameworkAnnotations {
+    @HostSideTestKeep
+    public TinyFrameworkAnnotations() {
     }
 
-    @HostSideTestStub
-    public int stub = 1;
-
     @HostSideTestKeep
-    public int keep = 2;
+    public int keep = 1;
 
     // Members will be deleted by default.
     // Deleted fields cannot have an initial value, because otherwise .ctor will fail to set it at
     // runtime.
     public int remove;
 
-    @HostSideTestStub
-    public int addOne(int value) {
-        return addOneInner(value);
-    }
-
     @HostSideTestKeep
-    public int addOneInner(int value) {
+    public int addOne(int value) {
         return value + 1;
     }
 
@@ -59,7 +50,6 @@
         throw new RuntimeException();
     }
 
-    @HostSideTestStub
     @HostSideTestSubstitute(suffix = "_host")
     public int addTwo(int value) {
         throw new RuntimeException("not supported on host side");
@@ -69,7 +59,6 @@
         return value + 2;
     }
 
-    @HostSideTestStub
     @HostSideTestSubstitute(suffix = "_host")
     public static native int nativeAddThree(int value);
 
@@ -82,9 +71,4 @@
     public String unsupportedMethod() {
         return "This value shouldn't be seen on the host side.";
     }
-
-    @HostSideTestStub
-    public String visibleButUsesUnsupportedMethod() {
-        return unsupportedMethod();
-    }
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.java
deleted file mode 100644
index f530207..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkCallerCheck.java
+++ /dev/null
@@ -1,57 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
-
-import android.hosttest.annotation.HostSideTestKeep;
-import android.hosttest.annotation.HostSideTestStub;
-import android.hosttest.annotation.HostSideTestWholeClassStub;
-
-/**
- * Used by the benchmark.
- */
-@HostSideTestWholeClassStub
-public class TinyFrameworkCallerCheck {
-
-    /**
-     * This method uses an inner method (which has the caller check).
-     *
-     * Benchmark result: 768ns
-     */
-    public static int getOne_withCheck() {
-        return Impl.getOneKeep();
-    }
-
-    /**
-     * This method doesn't have any caller check.
-     *
-     * Benchmark result: 2ns
-     */
-    public static int getOne_noCheck() {
-        return Impl.getOneStub();
-    }
-
-    private static class Impl {
-        @HostSideTestKeep
-        public static int getOneKeep() {
-            return 1;
-        }
-
-        @HostSideTestStub
-        public static int getOneStub() {
-            return 1;
-        }
-    }
-}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.java
deleted file mode 100644
index 145b65a..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassClassWideAnnotations.java
+++ /dev/null
@@ -1,74 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
-
-import android.hosttest.annotation.HostSideTestStub;
-import android.hosttest.annotation.HostSideTestSubstitute;
-import android.hosttest.annotation.HostSideTestWholeClassStub;
-
-@HostSideTestWholeClassStub
-public class TinyFrameworkClassClassWideAnnotations {
-    public TinyFrameworkClassClassWideAnnotations() {
-    }
-
-    public int stub = 1;
-
-    public int keep = 2;
-
-    // Cannot have an initial value, because otherwise .ctor will fail to set it at runtime.
-    public int remove;
-
-    // @Stub
-    public int addOne(int value) {
-        return addOneInner(value);
-    }
-
-    // @Keep
-    public int addOneInner(int value) {
-        return value + 1;
-    }
-
-    // @Remove
-    public void toBeRemoved(String foo) {
-        throw new RuntimeException();
-    }
-
-    @HostSideTestStub
-    @HostSideTestSubstitute(suffix = "_host")
-    public int addTwo(int value) {
-        throw new RuntimeException("not supported on host side");
-    }
-
-    public int addTwo_host(int value) {
-        return value + 2;
-    }
-
-    @HostSideTestStub
-    @HostSideTestSubstitute(suffix = "_host")
-    public static native int nativeAddThree(int value);
-
-    public static int nativeAddThree_host(int value) {
-        return value + 3;
-    }
-
-    public String unsupportedMethod() {
-        return "This value shouldn't be seen on the host side.";
-    }
-
-    public String visibleButUsesUnsupportedMethod() {
-        return unsupportedMethod();
-    }
-}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.java
index 98fc634..f734790 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassLoadHook.java
@@ -15,12 +15,12 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
 import java.util.HashSet;
 import java.util.Set;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkClassLoadHook {
     private TinyFrameworkClassLoadHook() {
     }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.java
new file mode 100644
index 0000000..e83163e
--- /dev/null
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotations.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
+
+import android.hosttest.annotation.HostSideTestRemove;
+import android.hosttest.annotation.HostSideTestSubstitute;
+import android.hosttest.annotation.HostSideTestThrow;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
+
+@HostSideTestWholeClassKeep
+public class TinyFrameworkClassWideAnnotations {
+    public TinyFrameworkClassWideAnnotations() {
+    }
+
+    public int keep = 1;
+
+    @HostSideTestRemove
+    public int remove;
+
+    public int addOne(int value) {
+        return value + 1;
+    }
+
+    @HostSideTestSubstitute(suffix = "_host")
+    public int addTwo(int value) {
+        throw new RuntimeException("not supported on host side");
+    }
+
+    public int addTwo_host(int value) {
+        return value + 2;
+    }
+
+    @HostSideTestRemove
+    public void toBeRemoved(String foo) {
+        throw new RuntimeException();
+    }
+
+    @HostSideTestThrow
+    public String unsupportedMethod() {
+        return "This value shouldn't be seen on the host side.";
+    }
+}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.java
index 8324ed9..3df21d9 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerDefault.java
@@ -15,18 +15,16 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestClassLoadHook;
-import android.hosttest.annotation.HostSideTestStub;
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestKeep;
 
-@HostSideTestStub
+@HostSideTestKeep
 public class TinyFrameworkClassWithInitializerDefault {
     static {
         sInitialized = true;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public static boolean sInitialized;
-    @HostSideTestStub
+    @HostSideTestKeep
     public static Object sObject = new Object();
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.java
index ea1ad93..cc665de 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithInitializerStub.java
@@ -16,20 +16,20 @@
 package com.android.hoststubgen.test.tinyframework;
 
 import android.hosttest.annotation.HostSideTestClassLoadHook;
+import android.hosttest.annotation.HostSideTestKeep;
 import android.hosttest.annotation.HostSideTestStaticInitializerKeep;
-import android.hosttest.annotation.HostSideTestStub;
 
 @HostSideTestClassLoadHook(
         "com.android.hoststubgen.test.tinyframework.TinyFrameworkClassLoadHook.onClassLoaded")
-@HostSideTestStub
+@HostSideTestKeep
 @HostSideTestStaticInitializerKeep
 public class TinyFrameworkClassWithInitializerStub {
     static {
         sInitialized = true;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public static boolean sInitialized;
-    @HostSideTestStub
+    @HostSideTestKeep
     public static Object sObject = new Object();
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.java
index 51f4818..f833ad8 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumComplex.java
@@ -16,15 +16,14 @@
 package com.android.hoststubgen.test.tinyframework;
 
 import android.hosttest.annotation.HostSideTestKeep;
-import android.hosttest.annotation.HostSideTestStub;
 
-@HostSideTestStub
+@HostSideTestKeep
 public enum TinyFrameworkEnumComplex {
-    @HostSideTestStub
+    @HostSideTestKeep
     RED("Red", "R"),
-    @HostSideTestStub
+    @HostSideTestKeep
     GREEN("Green", "G"),
-    @HostSideTestStub
+    @HostSideTestKeep
     BLUE("Blue", "B");
 
     @HostSideTestKeep
@@ -33,18 +32,18 @@
     @HostSideTestKeep
     private final String mShortName;
 
-    @HostSideTestStub
+    @HostSideTestKeep
     TinyFrameworkEnumComplex(String longName, String shortName) {
         mLongName = longName;
         mShortName = shortName;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public String getLongName() {
         return mLongName;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public String getShortName() {
         return mShortName;
     }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.java
index f440d86..c023169 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkEnumSimple.java
@@ -15,12 +15,12 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestStub;
+import android.hosttest.annotation.HostSideTestKeep;
 
-@HostSideTestStub
+@HostSideTestKeep
 public enum TinyFrameworkEnumSimple {
-    @HostSideTestStub
+    @HostSideTestKeep
     CAT,
-    @HostSideTestStub
+    @HostSideTestKeep
     DOG,
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.java
index 909d3b4..f7cae7d 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkExceptionTester.java
@@ -15,9 +15,9 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkExceptionTester {
     public static int testException() {
         try {
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.java
index 1977c90..ec1efba 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkForTextPolicy.java
@@ -24,17 +24,11 @@
 
     public int stub = 1;
 
-    public int keep = 2;
-
     // Removed fields cannot have an initial value, because otherwise .ctor will fail to set it at
     // runtime.
     public int remove;
 
     public int addOne(int value) {
-        return addOneInner(value);
-    }
-
-    public int addOneInner(int value) {
         return value + 1;
     }
 
@@ -95,8 +89,4 @@
     public String unsupportedMethod() {
         return "This value shouldn't be seen on the host side.";
     }
-
-    public String visibleButUsesUnsupportedMethod() {
-        return unsupportedMethod();
-    }
 }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.java
index 0d1203b..1ca653e 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkLambdas.java
@@ -15,8 +15,8 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
+import android.hosttest.annotation.HostSideTestKeep;
 import android.hosttest.annotation.HostSideTestStaticInitializerKeep;
-import android.hosttest.annotation.HostSideTestStub;
 
 import java.util.function.Supplier;
 
@@ -28,48 +28,48 @@
  *
  * Implicit filter should take care of them.
  */
-@HostSideTestStub
+@HostSideTestKeep
 @HostSideTestStaticInitializerKeep
 public class TinyFrameworkLambdas {
-    @HostSideTestStub
+    @HostSideTestKeep
     public TinyFrameworkLambdas() {
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public final Supplier<Integer> mSupplier = () -> 1;
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public static final Supplier<Integer> sSupplier = () -> 2;
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public Supplier<Integer> getSupplier() {
         return () -> 3;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     public static Supplier<Integer> getSupplier_static() {
         return () -> 4;
     }
 
-    @HostSideTestStub
+    @HostSideTestKeep
     @HostSideTestStaticInitializerKeep
     public static class Nested {
-        @HostSideTestStub
+        @HostSideTestKeep
         public Nested() {
         }
 
-        @HostSideTestStub
+        @HostSideTestKeep
         public final Supplier<Integer> mSupplier = () -> 5;
 
-        @HostSideTestStub
+        @HostSideTestKeep
         public static final Supplier<Integer> sSupplier = () -> 6;
 
-        @HostSideTestStub
+        @HostSideTestKeep
         public Supplier<Integer> getSupplier() {
             return () -> 7;
         }
 
-        @HostSideTestStub
+        @HostSideTestKeep
         public static Supplier<Integer> getSupplier_static() {
             return () -> 8;
         }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.java
index 1ff3744..57c69a3 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkMethodCallReplace.java
@@ -15,11 +15,11 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkMethodCallReplace {
     //  This method should return true.
     public static boolean nonStaticMethodCallReplaceTester() throws Exception {
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.java
index 09ee183..73b5e2f 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNative.java
@@ -15,11 +15,12 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
+import android.hosttest.annotation.HostSideTestKeep;
 import android.hosttest.annotation.HostSideTestNativeSubstitutionClass;
 import android.hosttest.annotation.HostSideTestThrow;
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 @HostSideTestNativeSubstitutionClass("TinyFrameworkNative_host")
 public class TinyFrameworkNative {
     public static native int nativeAddTwo(int arg);
@@ -49,6 +50,9 @@
     @HostSideTestThrow
     public static native void nativeStillNotSupported();
 
+    @HostSideTestKeep
+    public static native void nativeStillKeep();
+
     public static void nativeStillNotSupported_should_be_like_this() {
         throw new RuntimeException();
     }
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.java
index e1c48bb..c1ea2ee 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkNestedClasses.java
@@ -15,11 +15,11 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
 import java.util.function.Supplier;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkNestedClasses {
     public final Supplier<Integer> mSupplier = new Supplier<Integer>() {
         @Override
@@ -34,6 +34,7 @@
             return 2;
         }
     };
+
     public Supplier<Integer> getSupplier() {
         return new Supplier<Integer>() {
             @Override
@@ -52,12 +53,10 @@
         };
     }
 
-    @HostSideTestWholeClassStub
     public class InnerClass {
         public int value = 5;
     }
 
-    @HostSideTestWholeClassStub
     public static class StaticNestedClass {
         public int value = 6;
 
@@ -70,6 +69,10 @@
                 }
             };
         }
+
+        public static class Double$NestedClass {
+            public int value = 8;
+        }
     }
 
     public static class BaseClass {
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.java
index a82be54..941fcff 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkPackageRedirect.java
@@ -15,9 +15,9 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkPackageRedirect {
     /**
      * A method that uses "unsupported" class. HostStubGen will redirect them to the "supported"
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.java
index 31a164a..707bc0e 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkRenamedClassCaller.java
@@ -15,9 +15,9 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkRenamedClassCaller {
     /** Calls the class that'll be renamed. */
     public static int foo(int value) {
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.java
index 1430bcb..8319ced 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkToBeRenamed.java
@@ -15,12 +15,12 @@
  */
 package com.android.hoststubgen.test.tinyframework;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
 /**
  * This class will be renamed by the "rename" directive in the policy file.
  */
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class TinyFrameworkToBeRenamed {
     private final int mValue;
 
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/unsupported/UnsupportedClass.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/unsupported/UnsupportedClass.java
index 0409b02..92f41ac 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/unsupported/UnsupportedClass.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-framework/src/com/unsupported/UnsupportedClass.java
@@ -15,10 +15,10 @@
  */
 package com.unsupported;
 
-import android.hosttest.annotation.HostSideTestWholeClassStub;
+import android.hosttest.annotation.HostSideTestWholeClassKeep;
 
 // Used for testing --package-redirect.
-@HostSideTestWholeClassStub
+@HostSideTestWholeClassKeep
 public class UnsupportedClass {
     public UnsupportedClass(int value) {
         throw new RuntimeException("This class is not supported");
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotationsTest.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotationsTest.java
new file mode 100644
index 0000000..1ae0493
--- /dev/null
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkAnnotationsTest.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+public class TinyFrameworkAnnotationsTest {
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+
+    @Test
+    public void testSimple() {
+        TinyFrameworkAnnotations tfc = new TinyFrameworkAnnotations();
+        assertThat(tfc.addOne(1)).isEqualTo(2);
+        assertThat(tfc.keep).isEqualTo(1);
+    }
+
+    @Test
+    public void testRemove() {
+        TinyFrameworkAnnotations tfc = new TinyFrameworkAnnotations();
+        assertThrows(NoSuchMethodError.class, () -> tfc.toBeRemoved("abc"));
+        assertThrows(NoSuchFieldError.class, () -> tfc.remove = 1);
+    }
+
+    @Test
+    public void testSubstitute() {
+        TinyFrameworkAnnotations tfc = new TinyFrameworkAnnotations();
+        assertThat(tfc.addTwo(1)).isEqualTo(3);
+    }
+
+    @Test
+    public void testSubstituteNative() {
+        TinyFrameworkAnnotations tfc = new TinyFrameworkAnnotations();
+        assertThat(tfc.nativeAddThree(1)).isEqualTo(4);
+    }
+
+    @Test
+    public void testUnsupportedMethod() {
+        TinyFrameworkAnnotations tfc = new TinyFrameworkAnnotations();
+
+        thrown.expect(RuntimeException.class);
+        thrown.expectMessage("not yet supported");
+        tfc.unsupportedMethod();
+    }
+}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkBenchmark.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkBenchmark.java
deleted file mode 100644
index d57735b..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkBenchmark.java
+++ /dev/null
@@ -1,71 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
-
-import org.junit.Test;
-
-import java.text.DecimalFormat;
-
-/**
- * Contains simple micro-benchmarks.
- */
-@LargeTest
-public class TinyFrameworkBenchmark {
-    private static final int MINIMAL_ITERATION = 1000;
-    private static final int MEASURE_SECONDS = 1;
-
-    private static final DecimalFormat sFormatter = new DecimalFormat("#,###");
-
-    private void doBenchmark(String name, Runnable r) {
-        // Worm up
-        for (int i = 0; i < MINIMAL_ITERATION; i++) {
-            r.run();
-        }
-
-        // Start measuring.
-        final long start = System.nanoTime();
-        final long end = start + MEASURE_SECONDS * 1_000_000_000L;
-
-        double iteration = 0;
-        while (System.nanoTime() <= end) {
-            for (int i = 0; i < MINIMAL_ITERATION; i++) {
-                r.run();
-            }
-            iteration += MINIMAL_ITERATION;
-        }
-
-        final long realEnd = System.nanoTime();
-
-        System.out.println(String.format("%s\t%s", name,
-                sFormatter.format((((double) realEnd - start)) / iteration)));
-    }
-
-    /**
-     * Micro-benchmark for a method without a non-stub caller check.
-     */
-    @Test
-    public void benchNoCallerCheck() {
-        doBenchmark("No caller check", TinyFrameworkCallerCheck::getOne_noCheck);
-    }
-
-    /**
-     * Micro-benchmark for a method with a non-stub caller check.
-     */
-    @Test
-    public void benchWithCallerCheck() {
-        doBenchmark("With caller check", TinyFrameworkCallerCheck::getOne_withCheck);
-    }
-}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassTest.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassTest.java
index 1692c6e89..14229a0 100644
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassTest.java
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassTest.java
@@ -17,10 +17,9 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.fail;
+import static org.junit.Assert.assertThrows;
 
 import com.android.hoststubgen.test.tinyframework.R.Nested;
-import com.android.hoststubgen.test.tinyframework.TinyFrameworkNestedClasses.SubClass;
 
 import org.junit.Rule;
 import org.junit.Test;
@@ -41,16 +40,12 @@
         assertThat(tfc.stub).isEqualTo(1);
     }
 
-//    @Test
-//    public void testDoesntCompile() {
-//        TinyFrameworkClass tfc = new TinyFrameworkClass();
-//
-//        tfc.addOneInner(1); // Shouldn't compile.
-//        tfc.toBeRemoved("abc"); // Shouldn't compile.
-//        tfc.unsupportedMethod(); // Shouldn't compile.
-//        int a = tfc.keep; // Shouldn't compile
-//        int b = tfc.remove; // Shouldn't compile
-//    }
+    @Test
+    public void testRemove() {
+        TinyFrameworkForTextPolicy tfc = new TinyFrameworkForTextPolicy();
+        assertThrows(NoSuchMethodError.class, () -> tfc.toBeRemoved("abc"));
+        assertThrows(NoSuchFieldError.class, () -> tfc.remove = 1);
+    }
 
     @Test
     public void testIgnore() {
@@ -79,48 +74,12 @@
     }
 
     @Test
-    public void testVisibleButUsesUnsupportedMethod() {
+    public void testUnsupportedMethod() {
         TinyFrameworkForTextPolicy tfc = new TinyFrameworkForTextPolicy();
 
         thrown.expect(RuntimeException.class);
         thrown.expectMessage("not yet supported");
-        tfc.visibleButUsesUnsupportedMethod();
-    }
-
-    @Test
-    public void testNestedClass1() {
-        assertThat(new TinyFrameworkNestedClasses().mSupplier.get()).isEqualTo(1);
-    }
-
-    @Test
-    public void testNestedClass2() {
-        assertThat(TinyFrameworkNestedClasses.sSupplier.get()).isEqualTo(2);
-    }
-
-    @Test
-    public void testNestedClass3() {
-        assertThat(new TinyFrameworkNestedClasses().getSupplier().get()).isEqualTo(3);
-    }
-
-    @Test
-    public void testNestedClass4() {
-        assertThat(TinyFrameworkNestedClasses.getSupplier_static().get()).isEqualTo(4);
-    }
-
-    @Test
-    public void testNestedClass5() {
-        assertThat((new TinyFrameworkNestedClasses()).new InnerClass().value).isEqualTo(5);
-    }
-
-    @Test
-    public void testNestedClass6() {
-        assertThat(new TinyFrameworkNestedClasses.StaticNestedClass().value).isEqualTo(6);
-    }
-
-    @Test
-    public void testNestedClass7() {
-        assertThat(TinyFrameworkNestedClasses.StaticNestedClass.getSupplier_static().get())
-                .isEqualTo(7);
+        tfc.unsupportedMethod();
     }
 
     @Test
@@ -186,28 +145,22 @@
     }
 
     @Test
-    public void testSubstituteNativeWithThrow() throws Exception {
-        // We can't use TinyFrameworkNative.nativeStillNotSupported() directly in this class,
-        // because @Throw implies @Keep (not @Stub), and we currently compile this test
-        // against the stub jar (so it won't contain @Throw methods).
-        //
-        // But the method exists at runtime, so we can use reflections to call it.
-        //
-        // In the real Ravenwood environment, we don't use HostStubGen's stub jar at all,
-        // so it's not a problem.
+    public void testSubstituteNativeWithThrow() {
+        thrown.expect(RuntimeException.class);
+        thrown.expectMessage("not yet supported");
 
-        final var clazz = TinyFrameworkNative.class;
-        final var method = clazz.getMethod("nativeStillNotSupported");
+        TinyFrameworkNative.nativeStillNotSupported();
+    }
 
-        try {
-            method.invoke(null);
+    @Test
+    public void testSubstituteNativeWithKeep() {
+        // We don't want to complicate the test by setting up JNI,
+        // so to test out whether the native method is preserved, we
+        // check whether calling it will throw UnsatisfiedLinkError,
+        // which would only happen on native methods.
+        thrown.expect(UnsatisfiedLinkError.class);
 
-            fail("java.lang.reflect.InvocationTargetException expected");
-
-        } catch (java.lang.reflect.InvocationTargetException e) {
-            var inner = e.getCause();
-            assertThat(inner.getMessage()).contains("not yet supported");
-        }
+        TinyFrameworkNative.nativeStillKeep();
     }
 
     @Test
@@ -216,12 +169,6 @@
         thrown.expectMessage("Outer exception");
 
         TinyFrameworkExceptionTester.testException();
-
-    }
-
-    @Test
-    public void testMethodCallBeforeSuperCall() {
-        assertThat(new SubClass(3).value).isEqualTo(3);
     }
 
     @Test
@@ -231,7 +178,7 @@
         // Having this line before assertThat() will ensure these class are already loaded.
         var classes = new Class[]{
                 TinyFrameworkClassWithInitializerStub.class,
-                TinyFrameworkClassAnnotations.class,
+                TinyFrameworkAnnotations.class,
                 TinyFrameworkForTextPolicy.class,
         };
 
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotationsTest.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotationsTest.java
new file mode 100644
index 0000000..34c98e9
--- /dev/null
+++ b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWideAnnotationsTest.java
@@ -0,0 +1,105 @@
+/*
+ * 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.hoststubgen.test.tinyframework;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.junit.Assert.assertThrows;
+
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+
+public class TinyFrameworkClassWideAnnotationsTest {
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+
+    @Test
+    public void testSimple() {
+        var tfc = new TinyFrameworkClassWideAnnotations();
+        assertThat(tfc.addOne(1)).isEqualTo(2);
+        assertThat(tfc.keep).isEqualTo(1);
+    }
+
+    @Test
+    public void testRemove() {
+        var tfc = new TinyFrameworkClassWideAnnotations();
+        assertThrows(NoSuchMethodError.class, () -> tfc.toBeRemoved("abc"));
+        assertThrows(NoSuchFieldError.class, () -> tfc.remove = 1);
+    }
+
+    @Test
+    public void testSubstitute() {
+        var tfc = new TinyFrameworkClassWideAnnotations();
+        assertThat(tfc.addTwo(1)).isEqualTo(3);
+    }
+
+    @Test
+    public void testUnsupportedMethod() {
+        var tfc = new TinyFrameworkClassWideAnnotations();
+
+        thrown.expect(RuntimeException.class);
+        thrown.expectMessage("not yet supported");
+        tfc.unsupportedMethod();
+    }
+
+    @Test
+    public void testMethodCallBeforeSuperCall() {
+        assertThat(new TinyFrameworkNestedClasses.SubClass(3).value).isEqualTo(3);
+    }
+
+    @Test
+    public void testNestedClass1() {
+        assertThat(new TinyFrameworkNestedClasses().mSupplier.get()).isEqualTo(1);
+    }
+
+    @Test
+    public void testNestedClass2() {
+        assertThat(TinyFrameworkNestedClasses.sSupplier.get()).isEqualTo(2);
+    }
+
+    @Test
+    public void testNestedClass3() {
+        assertThat(new TinyFrameworkNestedClasses().getSupplier().get()).isEqualTo(3);
+    }
+
+    @Test
+    public void testNestedClass4() {
+        assertThat(TinyFrameworkNestedClasses.getSupplier_static().get()).isEqualTo(4);
+    }
+
+    @Test
+    public void testNestedClass5() {
+        assertThat((new TinyFrameworkNestedClasses()).new InnerClass().value).isEqualTo(5);
+    }
+
+    @Test
+    public void testNestedClass6() {
+        assertThat(new TinyFrameworkNestedClasses.StaticNestedClass().value).isEqualTo(6);
+    }
+
+    @Test
+    public void testNestedClass7() {
+        assertThat(TinyFrameworkNestedClasses.StaticNestedClass.getSupplier_static().get())
+                .isEqualTo(7);
+    }
+
+    @Test
+    public void testNestedClass8() {
+        assertThat(new TinyFrameworkNestedClasses.StaticNestedClass.Double$NestedClass().value)
+                .isEqualTo(8);
+    }
+}
diff --git a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithAnnotTest.java b/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithAnnotTest.java
deleted file mode 100644
index 288c716..0000000
--- a/tools/hoststubgen/hoststubgen/test-tiny-framework/tiny-test/src/com/android/hoststubgen/test/tinyframework/TinyFrameworkClassWithAnnotTest.java
+++ /dev/null
@@ -1,66 +0,0 @@
-/*
- * Copyright (C) 2023 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.hoststubgen.test.tinyframework;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.rules.ExpectedException;
-
-public class TinyFrameworkClassWithAnnotTest {
-    @Rule
-    public ExpectedException thrown = ExpectedException.none();
-
-    @Test
-    public void testSimple() {
-        TinyFrameworkClassAnnotations tfc = new TinyFrameworkClassAnnotations();
-        assertThat(tfc.addOne(1)).isEqualTo(2);
-        assertThat(tfc.stub).isEqualTo(1);
-    }
-
-//    @Test
-//    public void testDoesntCompile() {
-//        TinyFrameworkClassWithAnnot tfc = new TinyFrameworkClassWithAnnot();
-//
-//        tfc.addOneInner(1); // Shouldn't compile.
-//        tfc.toBeRemoved("abc"); // Shouldn't compile.
-//        tfc.unsupportedMethod(); // Shouldn't compile.
-//        int a = tfc.keep; // Shouldn't compile
-//        int b = tfc.remove; // Shouldn't compile
-//    }
-
-    @Test
-    public void testSubstitute() {
-        TinyFrameworkClassAnnotations tfc = new TinyFrameworkClassAnnotations();
-        assertThat(tfc.addTwo(1)).isEqualTo(3);
-    }
-
-    @Test
-    public void testSubstituteNative() {
-        TinyFrameworkClassAnnotations tfc = new TinyFrameworkClassAnnotations();
-        assertThat(tfc.nativeAddThree(1)).isEqualTo(4);
-    }
-
-    @Test
-    public void testVisibleButUsesUnsupportedMethod() {
-        TinyFrameworkClassAnnotations tfc = new TinyFrameworkClassAnnotations();
-
-        thrown.expect(RuntimeException.class);
-        thrown.expectMessage("not yet supported");
-        tfc.visibleButUsesUnsupportedMethod();
-    }
-}
diff --git a/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/asm/AsmUtilsTest.kt b/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/asm/AsmUtilsTest.kt
index 6b46c84..5b2795c 100644
--- a/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/asm/AsmUtilsTest.kt
+++ b/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/asm/AsmUtilsTest.kt
@@ -23,18 +23,6 @@
 import org.objectweb.asm.Opcodes.ACC_STATIC
 
 class AsmUtilsTest {
-    private fun checkGetDirectOuterClassName(input: String, expected: String?) {
-        assertThat(getDirectOuterClassName(input)).isEqualTo(expected)
-    }
-
-    @Test
-    fun testGetDirectOuterClassName() {
-        checkGetDirectOuterClassName("a", null)
-        checkGetDirectOuterClassName("a\$x", "a")
-        checkGetDirectOuterClassName("a.b.c\$x", "a.b.c")
-        checkGetDirectOuterClassName("a.b.c\$x\$y", "a.b.c\$x")
-    }
-
     @Test
     fun testVisibility() {
         fun test(access: Int, expected: Visibility) {
diff --git a/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/utils/ClassFilterTest.kt b/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/utils/ClassFilterTest.kt
index f651514..85b6e80 100644
--- a/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/utils/ClassFilterTest.kt
+++ b/tools/hoststubgen/hoststubgen/test/com/android/hoststubgen/utils/ClassFilterTest.kt
@@ -72,6 +72,18 @@
     }
 
     @Test
+    fun testNestedClass() {
+        val f = ClassFilter.buildFromString("a.b.c\nm.n.o\$p\n", false, "X")
+        assertThat(f.matches("a/b/c")).isEqualTo(true)
+        assertThat(f.matches("a/b/c\$d")).isEqualTo(true)
+        assertThat(f.matches("a/b/c\$d\$e")).isEqualTo(true)
+        assertThat(f.matches("m/n/o")).isEqualTo(false)
+        assertThat(f.matches("m/n/o\$p")).isEqualTo(true)
+        assertThat(f.matches("m/n/o\$p\$r")).isEqualTo(true)
+        assertThat(f.matches("m/n/o\$p\$r\$")).isEqualTo(true)
+    }
+
+    @Test
     fun testBadFilter1() {
         try {
             ClassFilter.buildFromString("""