Merge "Fix test failing with NPE due to null List" into main
diff --git a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
index 8bd3ef4..637c726 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/controllers/QuotaController.java
@@ -36,10 +36,14 @@
 import android.app.ActivityManager;
 import android.app.AlarmManager;
 import android.app.UidObserver;
+import android.app.compat.CompatChanges;
 import android.app.job.JobInfo;
 import android.app.usage.UsageEvents;
 import android.app.usage.UsageStatsManagerInternal;
 import android.app.usage.UsageStatsManagerInternal.UsageEventListener;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.Disabled;
+import android.compat.annotation.Overridable;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
@@ -132,6 +136,27 @@
         return (int) (val ^ (val >>> 32));
     }
 
+    /**
+     * When enabled this change id overrides the default quota policy enforcement to the jobs
+     * running in the foreground process state.
+     */
+    // TODO: b/379681266 - Might need some refactoring for a better app-compat strategy.
+    @VisibleForTesting
+    @ChangeId
+    @Disabled // Disabled by default
+    @Overridable // The change can be overridden in user build
+    static final long OVERRIDE_QUOTA_ENFORCEMENT_TO_FGS_JOBS = 341201311L;
+
+    /**
+     * When enabled this change id overrides the default quota policy enforcement policy
+     * the jobs started when app was in the TOP state.
+     */
+    @VisibleForTesting
+    @ChangeId
+    @Disabled // Disabled by default
+    @Overridable // The change can be overridden in user build.
+    static final long OVERRIDE_QUOTA_ENFORCEMENT_TO_TOP_STARTED_JOBS = 374323858L;
+
     @VisibleForTesting
     static class ExecutionStats {
         /**
@@ -622,7 +647,9 @@
         }
 
         final int uid = jobStatus.getSourceUid();
-        if (!Flags.enforceQuotaPolicyToTopStartedJobs() && mTopAppCache.get(uid)) {
+        if ((!Flags.enforceQuotaPolicyToTopStartedJobs()
+                || CompatChanges.isChangeEnabled(OVERRIDE_QUOTA_ENFORCEMENT_TO_TOP_STARTED_JOBS,
+                        uid)) && mTopAppCache.get(uid)) {
             if (DEBUG) {
                 Slog.d(TAG, jobStatus.toShortString() + " is top started job");
             }
@@ -659,7 +686,9 @@
                 timer.stopTrackingJob(jobStatus);
             }
         }
-        if (!Flags.enforceQuotaPolicyToTopStartedJobs()) {
+        if (!Flags.enforceQuotaPolicyToTopStartedJobs()
+                || CompatChanges.isChangeEnabled(OVERRIDE_QUOTA_ENFORCEMENT_TO_TOP_STARTED_JOBS,
+                        jobStatus.getSourceUid())) {
             mTopStartedJobs.remove(jobStatus);
         }
     }
@@ -772,7 +801,13 @@
 
     /** @return true if the job was started while the app was in the TOP state. */
     private boolean isTopStartedJobLocked(@NonNull final JobStatus jobStatus) {
-        return !Flags.enforceQuotaPolicyToTopStartedJobs() && mTopStartedJobs.contains(jobStatus);
+        if (!Flags.enforceQuotaPolicyToTopStartedJobs()
+                || CompatChanges.isChangeEnabled(OVERRIDE_QUOTA_ENFORCEMENT_TO_TOP_STARTED_JOBS,
+                        jobStatus.getSourceUid())) {
+            return mTopStartedJobs.contains(jobStatus);
+        }
+
+        return false;
     }
 
     /** Returns the maximum amount of time this job could run for. */
@@ -2634,9 +2669,13 @@
     }
 
     @VisibleForTesting
-    int getProcessStateQuotaFreeThreshold() {
-        return Flags.enforceQuotaPolicyToFgsJobs() ? ActivityManager.PROCESS_STATE_BOUND_TOP :
-                ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE;
+    int getProcessStateQuotaFreeThreshold(int uid) {
+        if (Flags.enforceQuotaPolicyToFgsJobs()
+                && !CompatChanges.isChangeEnabled(OVERRIDE_QUOTA_ENFORCEMENT_TO_FGS_JOBS, uid)) {
+            return ActivityManager.PROCESS_STATE_BOUND_TOP;
+        }
+
+        return ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE;
     }
 
     private class QcHandler extends Handler {
@@ -2776,7 +2815,7 @@
                                 isQuotaFree = true;
                             } else {
                                 final boolean reprocess;
-                                if (procState <= getProcessStateQuotaFreeThreshold()) {
+                                if (procState <= getProcessStateQuotaFreeThreshold(uid)) {
                                     reprocess = !mForegroundUids.get(uid);
                                     mForegroundUids.put(uid, true);
                                     isQuotaFree = true;
diff --git a/core/api/current.txt b/core/api/current.txt
index 6cf572b..5a7f2bd 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -4614,6 +4614,7 @@
     method public void reportFullyDrawn();
     method public android.view.DragAndDropPermissions requestDragAndDropPermissions(android.view.DragEvent);
     method public void requestFullscreenMode(int, @Nullable android.os.OutcomeReceiver<java.lang.Void,java.lang.Throwable>);
+    method @FlaggedApi("com.android.window.flags.enable_desktop_windowing_app_to_web_education") public final void requestOpenInBrowserEducation();
     method public final void requestPermissions(@NonNull String[], int);
     method @FlaggedApi("android.permission.flags.device_aware_permission_apis_enabled") public final void requestPermissions(@NonNull String[], int, int);
     method public final void requestShowKeyboardShortcuts();
@@ -4639,7 +4640,6 @@
     method public void setInheritShowWhenLocked(boolean);
     method public void setIntent(android.content.Intent);
     method @FlaggedApi("android.security.content_uri_permission_apis") public void setIntent(@Nullable android.content.Intent, @Nullable android.app.ComponentCaller);
-    method @FlaggedApi("com.android.window.flags.enable_desktop_windowing_app_to_web_education") public final void setLimitSystemEducationDialogs(boolean);
     method public void setLocusContext(@Nullable android.content.LocusId, @Nullable android.os.Bundle);
     method public final void setMediaController(android.media.session.MediaController);
     method public void setPictureInPictureParams(@NonNull android.app.PictureInPictureParams);
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index bc0037d..14a6c8c 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -103,6 +103,7 @@
     field public static final String BRIGHTNESS_SLIDER_USAGE = "android.permission.BRIGHTNESS_SLIDER_USAGE";
     field public static final String BROADCAST_CLOSE_SYSTEM_DIALOGS = "android.permission.BROADCAST_CLOSE_SYSTEM_DIALOGS";
     field @Deprecated public static final String BROADCAST_NETWORK_PRIVILEGED = "android.permission.BROADCAST_NETWORK_PRIVILEGED";
+    field @FlaggedApi("android.media.audio.concurrent_audio_record_bypass_permission") public static final String BYPASS_CONCURRENT_RECORD_AUDIO_RESTRICTION = "android.permission.BYPASS_CONCURRENT_RECORD_AUDIO_RESTRICTION";
     field public static final String BYPASS_ROLE_QUALIFICATION = "android.permission.BYPASS_ROLE_QUALIFICATION";
     field public static final String CALL_AUDIO_INTERCEPTION = "android.permission.CALL_AUDIO_INTERCEPTION";
     field public static final String CAMERA_DISABLE_TRANSMIT_LED = "android.permission.CAMERA_DISABLE_TRANSMIT_LED";
@@ -229,6 +230,7 @@
     field public static final String MANAGE_ROTATION_RESOLVER = "android.permission.MANAGE_ROTATION_RESOLVER";
     field public static final String MANAGE_SAFETY_CENTER = "android.permission.MANAGE_SAFETY_CENTER";
     field public static final String MANAGE_SEARCH_UI = "android.permission.MANAGE_SEARCH_UI";
+    field @FlaggedApi("android.security.secure_lockdown") public static final String MANAGE_SECURE_LOCK_DEVICE = "android.permission.MANAGE_SECURE_LOCK_DEVICE";
     field public static final String MANAGE_SENSOR_PRIVACY = "android.permission.MANAGE_SENSOR_PRIVACY";
     field public static final String MANAGE_SMARTSPACE = "android.permission.MANAGE_SMARTSPACE";
     field public static final String MANAGE_SOUND_TRIGGER = "android.permission.MANAGE_SOUND_TRIGGER";
@@ -3869,6 +3871,7 @@
     field public static final String APP_INTEGRITY_SERVICE = "app_integrity";
     field public static final String APP_PREDICTION_SERVICE = "app_prediction";
     field public static final String AUDIO_DEVICE_VOLUME_SERVICE = "audio_device_volume";
+    field @FlaggedApi("android.security.secure_lockdown") public static final String AUTHENTICATION_POLICY_SERVICE = "authentication_policy";
     field public static final String BACKUP_SERVICE = "backup";
     field public static final String BATTERY_STATS_SERVICE = "batterystats";
     field public static final int BIND_ALLOW_BACKGROUND_ACTIVITY_STARTS = 1048576; // 0x100000
@@ -12858,6 +12861,36 @@
 
 }
 
+package android.security.authenticationpolicy {
+
+  @FlaggedApi("android.security.secure_lockdown") public final class AuthenticationPolicyManager {
+    method @FlaggedApi("android.security.secure_lockdown") @RequiresPermission(android.Manifest.permission.MANAGE_SECURE_LOCK_DEVICE) public int disableSecureLockDevice(@NonNull android.security.authenticationpolicy.DisableSecureLockDeviceParams);
+    method @FlaggedApi("android.security.secure_lockdown") @RequiresPermission(android.Manifest.permission.MANAGE_SECURE_LOCK_DEVICE) public int enableSecureLockDevice(@NonNull android.security.authenticationpolicy.EnableSecureLockDeviceParams);
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_ALREADY_ENABLED = 6; // 0x6
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_INSUFFICIENT_BIOMETRICS = 5; // 0x5
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_INVALID_PARAMS = 3; // 0x3
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_NO_BIOMETRICS_ENROLLED = 4; // 0x4
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_UNKNOWN = 0; // 0x0
+    field @FlaggedApi("android.security.secure_lockdown") public static final int ERROR_UNSUPPORTED = 2; // 0x2
+    field @FlaggedApi("android.security.secure_lockdown") public static final int SUCCESS = 1; // 0x1
+  }
+
+  @FlaggedApi("android.security.secure_lockdown") public final class DisableSecureLockDeviceParams implements android.os.Parcelable {
+    ctor public DisableSecureLockDeviceParams(@NonNull String);
+    method public int describeContents();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.security.authenticationpolicy.DisableSecureLockDeviceParams> CREATOR;
+  }
+
+  @FlaggedApi("android.security.secure_lockdown") public final class EnableSecureLockDeviceParams implements android.os.Parcelable {
+    ctor public EnableSecureLockDeviceParams(@NonNull String);
+    method public int describeContents();
+    method public void writeToParcel(@NonNull android.os.Parcel, int);
+    field @NonNull public static final android.os.Parcelable.Creator<android.security.authenticationpolicy.EnableSecureLockDeviceParams> CREATOR;
+  }
+
+}
+
 package android.security.forensic {
 
   @FlaggedApi("android.security.afl_api") public class ForensicManager {
diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index 44bcc2a..ad1d937 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -1805,6 +1805,7 @@
     method @RequiresPermission(android.Manifest.permission.REMAP_MODIFIER_KEYS) public void remapModifierKey(int, int);
     method @FlaggedApi("com.android.input.flags.device_associations") @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void removeUniqueIdAssociationByDescriptor(@NonNull String);
     method @RequiresPermission("android.permission.ASSOCIATE_INPUT_DEVICE_TO_DISPLAY") public void removeUniqueIdAssociationByPort(@NonNull String);
+    method public void resetLockedModifierState();
     field public static final long BLOCK_UNTRUSTED_TOUCHES = 158002302L; // 0x96aec7eL
   }
 
diff --git a/core/java/android/app/Activity.java b/core/java/android/app/Activity.java
index 419eb7d..38aea64 100644
--- a/core/java/android/app/Activity.java
+++ b/core/java/android/app/Activity.java
@@ -1270,27 +1270,22 @@
     }
 
     /**
-     * To make users aware of system features such as the app header menu and its various
-     * functionalities, educational dialogs are shown to demonstrate how to find and utilize these
-     * features. Using this method, an activity can specify if it wants these educational dialogs to
-     * be shown. When set to {@code true}, these dialogs are not completely blocked; however, the
-     * system will be notified that they should not be shown unless necessary. If this API is not
-     * called, the system's educational dialogs are not limited by default.
+     * Requests to show the “Open in browser” education. “Open in browser” is a feature
+     * within the app header that allows users to switch from an app to the web. The feature
+     * is made available when an application is opened by a user clicking a link or when a
+     * link is provided by an application. Links can be provided by utilizing
+     * {@link AssistContent#EXTRA_AUTHENTICATING_USER_WEB_URI} or
+     * {@link AssistContent#setWebUri}.
      *
-     * <p>This method can be utilized when activities have states where showing an
-     * educational dialog would be disruptive to the user. For example, if a game application is
-     * expecting prompt user input, this method can be used to limit educational dialogs such as the
-     * dialogs that showcase the app header's features which, in this instance, would disrupt the
-     * user's experience if shown.</p>
-     *
-     * <p>Note that educational dialogs may be shown soon after this activity is launched, so
-     * this method must be called early if the intent is to limit the dialogs from the start.</p>
+     * <p>This method should be utilized when an activity wants to nudge the user to switch
+     * to the web application in cases where the web may provide the user with a better
+     * experience. Note that this method does not guarantee that the education will be shown.</p>
      */
     @FlaggedApi(com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_APP_TO_WEB_EDUCATION)
-    public final void setLimitSystemEducationDialogs(boolean limitSystemEducationDialogs) {
+    public final void requestOpenInBrowserEducation() {
         try {
             ActivityTaskManager
-                  .getService().setLimitSystemEducationDialogs(mToken, limitSystemEducationDialogs);
+                  .getService().requestOpenInBrowserEducation(mToken);
         } catch (RemoteException e) {
             // Empty
         }
diff --git a/core/java/android/app/IActivityTaskManager.aidl b/core/java/android/app/IActivityTaskManager.aidl
index ec7b72e..c6f62a2 100644
--- a/core/java/android/app/IActivityTaskManager.aidl
+++ b/core/java/android/app/IActivityTaskManager.aidl
@@ -242,8 +242,8 @@
 
     boolean supportsLocalVoiceInteraction();
 
-    // Sets whether system educational dialogs should be limited
-    void setLimitSystemEducationDialogs(IBinder appToken, boolean limitSystemEducationDialogs);
+    // Requests the "Open in browser" education to be shown
+    void requestOpenInBrowserEducation(IBinder appToken);
 
     // Get device configuration
     ConfigurationInfo getDeviceConfigurationInfo();
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index eb161ea..ff73382 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -814,8 +814,8 @@
         if (Flags.notificationsRedesignTemplates()) {
             return switch (layoutId) {
                 case R.layout.notification_2025_template_collapsed_base,
+                     R.layout.notification_2025_template_heads_up_base,
                      R.layout.notification_2025_template_header,
-                     R.layout.notification_template_material_heads_up_base,
                      R.layout.notification_template_material_big_base,
                      R.layout.notification_template_material_big_picture,
                      R.layout.notification_template_material_big_text,
@@ -7517,7 +7517,11 @@
         }
 
         private int getHeadsUpBaseLayoutResource() {
-            return R.layout.notification_template_material_heads_up_base;
+            if (Flags.notificationsRedesignTemplates()) {
+                return R.layout.notification_2025_template_heads_up_base;
+            } else {
+                return R.layout.notification_template_material_heads_up_base;
+            }
         }
 
         private int getCompactHeadsUpBaseLayoutResource() {
diff --git a/core/java/android/app/SystemServiceRegistry.java b/core/java/android/app/SystemServiceRegistry.java
index 53a7dad..6a23349 100644
--- a/core/java/android/app/SystemServiceRegistry.java
+++ b/core/java/android/app/SystemServiceRegistry.java
@@ -241,6 +241,8 @@
 import android.security.advancedprotection.IAdvancedProtectionService;
 import android.security.attestationverification.AttestationVerificationManager;
 import android.security.attestationverification.IAttestationVerificationManagerService;
+import android.security.authenticationpolicy.AuthenticationPolicyManager;
+import android.security.authenticationpolicy.IAuthenticationPolicyService;
 import android.security.forensic.ForensicManager;
 import android.security.forensic.IForensicService;
 import android.security.keystore.KeyStoreManager;
@@ -1025,6 +1027,25 @@
                     }
                 });
 
+        registerService(Context.AUTHENTICATION_POLICY_SERVICE,
+                AuthenticationPolicyManager.class,
+                new CachedServiceFetcher<AuthenticationPolicyManager>() {
+                    @Override
+                    public AuthenticationPolicyManager createService(ContextImpl ctx)
+                            throws ServiceNotFoundException {
+                        if (!android.security.Flags.secureLockdown()) {
+                            throw new ServiceNotFoundException(
+                                    Context.AUTHENTICATION_POLICY_SERVICE);
+                        }
+
+                        final IBinder binder = ServiceManager.getServiceOrThrow(
+                                Context.AUTHENTICATION_POLICY_SERVICE);
+                        final IAuthenticationPolicyService service =
+                                IAuthenticationPolicyService.Stub.asInterface(binder);
+                        return new AuthenticationPolicyManager(ctx.getOuterContext(), service);
+                    }
+                });
+
         registerService(Context.TV_INTERACTIVE_APP_SERVICE, TvInteractiveAppManager.class,
                 new CachedServiceFetcher<TvInteractiveAppManager>() {
             @Override
diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index aac963a..01cc9d8 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -340,10 +340,10 @@
     public int requestedVisibleTypes;
 
     /**
-     * Whether the top activity has requested to limit educational dialogs shown by the system.
+     * The timestamp of the top activity's last request to show the "Open in Browser" education.
      * @hide
      */
-    public boolean isTopActivityLimitSystemEducationDialogs;
+    public long topActivityRequestOpenInBrowserEducationTimestamp;
 
     /**
      * Encapsulate specific App Compat information.
@@ -493,8 +493,8 @@
                 && Objects.equals(capturedLink, that.capturedLink)
                 && capturedLinkTimestamp == that.capturedLinkTimestamp
                 && requestedVisibleTypes == that.requestedVisibleTypes
-                && isTopActivityLimitSystemEducationDialogs
-                    == that.isTopActivityLimitSystemEducationDialogs
+                && topActivityRequestOpenInBrowserEducationTimestamp
+                    == that.topActivityRequestOpenInBrowserEducationTimestamp
                 && appCompatTaskInfo.equalsForTaskOrganizer(that.appCompatTaskInfo)
                 && Objects.equals(topActivityMainWindowFrame, that.topActivityMainWindowFrame);
     }
@@ -571,7 +571,7 @@
         capturedLink = source.readTypedObject(Uri.CREATOR);
         capturedLinkTimestamp = source.readLong();
         requestedVisibleTypes = source.readInt();
-        isTopActivityLimitSystemEducationDialogs = source.readBoolean();
+        topActivityRequestOpenInBrowserEducationTimestamp = source.readLong();
         appCompatTaskInfo = source.readTypedObject(AppCompatTaskInfo.CREATOR);
         topActivityMainWindowFrame = source.readTypedObject(Rect.CREATOR);
     }
@@ -627,7 +627,7 @@
         dest.writeTypedObject(capturedLink, flags);
         dest.writeLong(capturedLinkTimestamp);
         dest.writeInt(requestedVisibleTypes);
-        dest.writeBoolean(isTopActivityLimitSystemEducationDialogs);
+        dest.writeLong(topActivityRequestOpenInBrowserEducationTimestamp);
         dest.writeTypedObject(appCompatTaskInfo, flags);
         dest.writeTypedObject(topActivityMainWindowFrame, flags);
     }
@@ -672,8 +672,8 @@
                 + " capturedLink=" + capturedLink
                 + " capturedLinkTimestamp=" + capturedLinkTimestamp
                 + " requestedVisibleTypes=" + requestedVisibleTypes
-                + " isTopActivityLimitSystemEducationDialogs="
-                + isTopActivityLimitSystemEducationDialogs
+                + " topActivityRequestOpenInBrowserEducationTimestamp="
+                + topActivityRequestOpenInBrowserEducationTimestamp
                 + " appCompatTaskInfo=" + appCompatTaskInfo
                 + " topActivityMainWindowFrame=" + topActivityMainWindowFrame
                 + "}";
diff --git a/core/java/android/app/jank/FrameOverrunHistogram.java b/core/java/android/app/jank/FrameOverrunHistogram.java
index e28ac12..3ad6531 100644
--- a/core/java/android/app/jank/FrameOverrunHistogram.java
+++ b/core/java/android/app/jank/FrameOverrunHistogram.java
@@ -39,7 +39,7 @@
      * Create a new instance of FrameOverrunHistogram.
      */
     public FrameOverrunHistogram() {
-        mBucketCounts = new int[sBucketEndpoints.length - 1];
+        mBucketCounts = new int[sBucketEndpoints.length];
     }
 
     /**
diff --git a/core/java/android/app/jank/JankDataProcessor.java b/core/java/android/app/jank/JankDataProcessor.java
index 7525d04..7ceaeb3 100644
--- a/core/java/android/app/jank/JankDataProcessor.java
+++ b/core/java/android/app/jank/JankDataProcessor.java
@@ -59,7 +59,6 @@
      * @param appUid the uid of the app.
      */
     public void processJankData(List<JankData> jankData, String activityName, int appUid) {
-        mCurrentBatchCount++;
         // add all the previous and active states to the pending states list.
         mStateTracker.retrieveAllStates(mPendingStates);
 
@@ -79,9 +78,8 @@
             }
         }
         // At this point we have attributed all frames to a state.
-        if (mCurrentBatchCount >= LOG_BATCH_FREQUENCY) {
-            logMetricCounts();
-        }
+        incrementBatchCountAndMaybeLogStats();
+
         // return the StatData object back to the pool to be reused.
         jankDataProcessingComplete();
     }
@@ -91,7 +89,73 @@
      * stats
      */
     public void mergeJankStats(AppJankStats jankStats, String activityName) {
-        // TODO b/377572463 Add Merging Logic
+        // Each state has a key which is a combination of widget category, widget id and widget
+        // state, this key is also used to identify pending stats, a pending stat is essentially a
+        // state with frames associated with it.
+        String stateKey = mStateTracker.getStateKey(jankStats.getWidgetCategory(),
+                jankStats.getWidgetId(), jankStats.getWidgetState());
+
+        if (mPendingJankStats.containsKey(stateKey)) {
+            mergeExistingStat(stateKey, jankStats);
+        } else {
+            mergeNewStat(stateKey, activityName, jankStats);
+        }
+
+        incrementBatchCountAndMaybeLogStats();
+    }
+
+    private void mergeExistingStat(String stateKey, AppJankStats jankStat) {
+        PendingJankStat pendingStat = mPendingJankStats.get(stateKey);
+
+        pendingStat.mJankyFrames += jankStat.getJankyFrameCount();
+        pendingStat.mTotalFrames += jankStat.getTotalFrameCount();
+
+        mergeOverrunHistograms(pendingStat.mFrameOverrunBuckets,
+                jankStat.getFrameOverrunHistogram().getBucketCounters());
+    }
+
+    private void mergeNewStat(String stateKey, String activityName, AppJankStats jankStats) {
+        // Check if we have space for a new stat
+        if (mPendingJankStats.size() > MAX_IN_MEMORY_STATS) {
+            return;
+        }
+
+        PendingJankStat pendingStat = mPendingJankStatsPool.acquire();
+        if (pendingStat == null) {
+            pendingStat = new PendingJankStat();
+
+        }
+        pendingStat.clearStats();
+
+        pendingStat.mActivityName = activityName;
+        pendingStat.mUid = jankStats.getUid();
+        pendingStat.mWidgetId = jankStats.getWidgetId();
+        pendingStat.mWidgetCategory = jankStats.getWidgetCategory();
+        pendingStat.mWidgetState = jankStats.getWidgetState();
+        pendingStat.mTotalFrames = jankStats.getTotalFrameCount();
+        pendingStat.mJankyFrames = jankStats.getJankyFrameCount();
+
+        mergeOverrunHistograms(pendingStat.mFrameOverrunBuckets,
+                jankStats.getFrameOverrunHistogram().getBucketCounters());
+
+        mPendingJankStats.put(stateKey, pendingStat);
+    }
+
+    private void mergeOverrunHistograms(int[] mergeTarget, int[] mergeSource) {
+        // The length of each histogram should be identical, if they are not then its possible the
+        // buckets are not in sync, these records should not be recorded.
+        if (mergeTarget.length != mergeSource.length) return;
+
+        for (int i = 0; i < mergeTarget.length; i++) {
+            mergeTarget[i] += mergeSource[i];
+        }
+    }
+
+    private void incrementBatchCountAndMaybeLogStats() {
+        mCurrentBatchCount++;
+        if (mCurrentBatchCount >= LOG_BATCH_FREQUENCY) {
+            logMetricCounts();
+        }
     }
 
     /**
diff --git a/core/java/android/app/jank/JankTracker.java b/core/java/android/app/jank/JankTracker.java
index 202281f..4695216 100644
--- a/core/java/android/app/jank/JankTracker.java
+++ b/core/java/android/app/jank/JankTracker.java
@@ -89,7 +89,12 @@
      * stats
      */
     public void mergeAppJankStats(AppJankStats appJankStats) {
-        mJankDataProcessor.mergeJankStats(appJankStats, mActivityName);
+        getHandler().post(new Runnable() {
+            @Override
+            public void run() {
+                mJankDataProcessor.mergeJankStats(appJankStats, mActivityName);
+            }
+        });
     }
 
     public void setActivityName(@NonNull String activityName) {
diff --git a/core/java/android/app/jank/StateTracker.java b/core/java/android/app/jank/StateTracker.java
index c86d5a5..21bb5e8 100644
--- a/core/java/android/app/jank/StateTracker.java
+++ b/core/java/android/app/jank/StateTracker.java
@@ -180,7 +180,11 @@
         }
     }
 
-    private String getStateKey(String widgetCategory, String widgetId, String widgetState) {
+    /**
+     * Returns a concatenated string of the inputs. This key can be used to retrieve both pending
+     * stats and the state that was used to create the pending stat.
+     */
+    public String getStateKey(String widgetCategory, String widgetId, String widgetState) {
         return widgetCategory + widgetId + widgetState;
     }
 
diff --git a/core/java/android/content/Context.java b/core/java/android/content/Context.java
index 6086f24..19cd2e6 100644
--- a/core/java/android/content/Context.java
+++ b/core/java/android/content/Context.java
@@ -18,6 +18,7 @@
 
 import static android.app.appfunctions.flags.Flags.FLAG_ENABLE_APP_FUNCTION_MANAGER;
 import static android.content.flags.Flags.FLAG_ENABLE_BIND_PACKAGE_ISOLATED_PROCESS;
+import static android.security.Flags.FLAG_SECURE_LOCKDOWN;
 
 import android.annotation.AttrRes;
 import android.annotation.CallbackExecutor;
@@ -4256,6 +4257,7 @@
             FINGERPRINT_SERVICE,
             //@hide: FACE_SERVICE,
             BIOMETRIC_SERVICE,
+            AUTHENTICATION_POLICY_SERVICE,
             MEDIA_ROUTER_SERVICE,
             TELEPHONY_SERVICE,
             TELEPHONY_SUBSCRIPTION_SERVICE,
@@ -4437,6 +4439,9 @@
      * web domain approval state.
      * <dt> {@link #DISPLAY_HASH_SERVICE} ("display_hash")
      * <dd> A {@link android.view.displayhash.DisplayHashManager} for management of display hashes.
+     * <dt> {@link #AUTHENTICATION_POLICY_SERVICE} ("authentication_policy")
+     * <dd> A {@link android.security.authenticationpolicy.AuthenticationPolicyManager}
+     * for managing authentication related policies on the device.
      * </dl>
      *
      * <p>Note:  System services obtained via this API may be closely associated with
@@ -4521,8 +4526,9 @@
      * @see android.content.pm.verify.domain.DomainVerificationManager
      * @see #DISPLAY_HASH_SERVICE
      * @see android.view.displayhash.DisplayHashManager
+     * @see #AUTHENTICATION_POLICY_SERVICE
+     * @see android.security.authenticationpolicy.AuthenticationPolicyManager
      */
-    // TODO(b/347269120): Re-add @Nullable
     public abstract Object getSystemService(@ServiceName @NonNull String name);
 
     /**
@@ -4543,7 +4549,8 @@
      * {@link android.os.BatteryManager}, {@link android.app.job.JobScheduler},
      * {@link android.app.usage.NetworkStatsManager},
      * {@link android.content.pm.verify.domain.DomainVerificationManager},
-     * {@link android.view.displayhash.DisplayHashManager}.
+     * {@link android.view.displayhash.DisplayHashManager}
+     * {@link android.security.authenticationpolicy.AuthenticationPolicyManager}.
      * </p>
      *
      * <p>
@@ -4568,7 +4575,6 @@
      */
     @SuppressWarnings("unchecked")
     @RavenwoodKeep
-    // TODO(b/347269120): Re-add @Nullable
     public final <T> T getSystemService(@NonNull Class<T> serviceClass) {
         // Because subclasses may override getSystemService(String) we cannot
         // perform a lookup by class alone.  We must first map the class to its
@@ -5183,6 +5189,18 @@
     public static final String AUTH_SERVICE = "auth";
 
     /**
+     * Use with {@link #getSystemService(String)} to retrieve an {@link
+     * android.security.authenticationpolicy.AuthenticationPolicyManager}.
+     * @see #getSystemService
+     * @see android.security.authenticationpolicy.AuthenticationPolicyManager
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final String AUTHENTICATION_POLICY_SERVICE = "authentication_policy";
+
+    /**
      * Use with {@link #getSystemService(String)} to retrieve a
      * {@link android.hardware.fingerprint.FingerprintManager} for handling management
      * of fingerprints.
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index 23d17cb..413eb98 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -960,7 +960,6 @@
     }
 
     @Override
-    // TODO(b/347269120): Re-add @Nullable
     public Object getSystemService(String name) {
         return mBase.getSystemService(name);
     }
diff --git a/core/java/android/content/pm/ActivityInfo.java b/core/java/android/content/pm/ActivityInfo.java
index ce52825..b10f5e4 100644
--- a/core/java/android/content/pm/ActivityInfo.java
+++ b/core/java/android/content/pm/ActivityInfo.java
@@ -1320,23 +1320,23 @@
             264301586L; // buganizer id
 
     /**
-     * Excludes the packages the override is applied to from the camera compatibility treatment
-     * in free-form windowing mode for fixed-orientation apps.
+     * Includes the packages the override is applied to in the camera compatibility treatment in
+     * free-form windowing mode for fixed-orientation apps.
      *
      * <p>In free-form windowing mode, the compatibility treatment emulates running on a portrait
      * device by letterboxing the app window and changing the camera characteristics to what apps
      * commonly expect in a portrait device: 90 and 270 degree sensor rotation for back and front
      * cameras, respectively, and setting display rotation to 0.
      *
-     * <p>Use this flag to disable the compatibility treatment for apps that do not respond well to
-     * the treatment.
+     * <p>Use this flag to enable the compatibility treatment for apps in which camera doesn't work
+     * well in freeform windowing.
      *
      * @hide
      */
     @ChangeId
     @Overridable
     @Disabled
-    public static final long OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT =
+    public static final long OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT =
             314961188L;
 
     /**
diff --git a/core/java/android/hardware/input/IInputManager.aidl b/core/java/android/hardware/input/IInputManager.aidl
index 3284761..ed510e4 100644
--- a/core/java/android/hardware/input/IInputManager.aidl
+++ b/core/java/android/hardware/input/IInputManager.aidl
@@ -281,4 +281,6 @@
     AidlInputGestureData[] getCustomInputGestures(int userId, int tag);
 
     AidlInputGestureData[] getAppLaunchBookmarks();
+
+    void resetLockedModifierState();
 }
diff --git a/core/java/android/hardware/input/InputManager.java b/core/java/android/hardware/input/InputManager.java
index f824192..10224c1 100644
--- a/core/java/android/hardware/input/InputManager.java
+++ b/core/java/android/hardware/input/InputManager.java
@@ -260,7 +260,7 @@
     }
 
     /**
-     * Custom input gesture error: Input gesture already exists
+     * Custom input gesture result success
      *
      * @hide
      */
@@ -1590,6 +1590,21 @@
     }
 
     /**
+     * Resets locked modifier state (i.e.. Caps Lock, Num Lock, Scroll Lock state)
+     *
+     * @hide
+     */
+    @TestApi
+    @SuppressLint("UnflaggedApi") // @TestApi without associated feature.
+    public void resetLockedModifierState() {
+        try {
+            mIm.resetLockedModifierState();
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * A callback used to be notified about battery state changes for an input device. The
      * {@link #onBatteryStateChanged(int, long, BatteryState)} method will be called once after the
      * listener is successfully registered to provide the initial battery state of the device.
diff --git a/core/java/android/os/CombinedMessageQueue/MessageQueue.java b/core/java/android/os/CombinedMessageQueue/MessageQueue.java
index 036ccd8..4c9f08d 100644
--- a/core/java/android/os/CombinedMessageQueue/MessageQueue.java
+++ b/core/java/android/os/CombinedMessageQueue/MessageQueue.java
@@ -123,11 +123,6 @@
         // We can lift this restriction in the future after we've made it possible for test authors
         // to test Looper and MessageQueue without resorting to reflection.
 
-        // Holdback study.
-        if (mUseConcurrent && Flags.messageQueueForceLegacy()) {
-            mUseConcurrent = false;
-        }
-
         mQuitAllowed = quitAllowed;
         mPtr = nativeInit();
     }
diff --git a/core/java/android/os/flags.aconfig b/core/java/android/os/flags.aconfig
index 9b1bf05..2ef8764 100644
--- a/core/java/android/os/flags.aconfig
+++ b/core/java/android/os/flags.aconfig
@@ -4,15 +4,6 @@
 
 # keep-sorted start block=yes newline_separated=yes
 flag {
-     # Holdback study for concurrent MessageQueue.
-     # Do not promote beyond trunkfood.
-     namespace: "system_performance"
-     name: "message_queue_force_legacy"
-     description: "Whether to holdback concurrent MessageQueue (force legacy)."
-     bug: "336880969"
-}
-
-flag {
     name: "adpf_gpu_report_actual_work_duration"
     is_exported: true
     namespace: "game"
diff --git a/core/java/android/security/authenticationpolicy/AuthenticationPolicyManager.java b/core/java/android/security/authenticationpolicy/AuthenticationPolicyManager.java
new file mode 100644
index 0000000..75abd5f
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/AuthenticationPolicyManager.java
@@ -0,0 +1,237 @@
+/*
+ * 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.security.authenticationpolicy;
+
+import static android.Manifest.permission.MANAGE_SECURE_LOCK_DEVICE;
+import static android.security.Flags.FLAG_SECURE_LOCKDOWN;
+
+import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.RequiresPermission;
+import android.annotation.SystemApi;
+import android.annotation.SystemService;
+import android.content.Context;
+import android.os.RemoteException;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+/**
+ * AuthenticationPolicyManager is a centralized interface for managing authentication related
+ * policies on the device. This includes device locking capabilities to protect users in "at risk"
+ * environments.
+ *
+ * AuthenticationPolicyManager is designed to protect Android users by integrating with apps and
+ * key system components, such as the lock screen. It is not related to enterprise control surfaces
+ * and does not offer additional administrative controls.
+ *
+ * <p>
+ * To use this class, call {@link #enableSecureLockDevice} to enable secure lock on the device.
+ * This will require the caller to have the
+ * {@link android.Manifest.permission#MANAGE_SECURE_LOCK_DEVICE} permission.
+ *
+ * <p>
+ * To disable secure lock on the device, call {@link #disableSecureLockDevice}. This will require
+ * the caller to have the {@link android.Manifest.permission#MANAGE_SECURE_LOCK_DEVICE} permission.
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_SECURE_LOCKDOWN)
+@SystemService(Context.AUTHENTICATION_POLICY_SERVICE)
+public final class AuthenticationPolicyManager {
+    private static final String TAG = "AuthenticationPolicyManager";
+
+    @NonNull private final IAuthenticationPolicyService mAuthenticationPolicyService;
+    @NonNull private final Context mContext;
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice} and {@link
+     * #disableSecureLockDevice}.
+     *
+     * Secure lock device request status unknown.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_UNKNOWN = 0;
+
+    /**
+     * Success result code for {@link #enableSecureLockDevice} and {@link #disableSecureLockDevice}.
+     *
+     * Secure lock device request successful.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int SUCCESS = 1;
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice} and {@link #disableSecureLockDevice}.
+     *
+     * Secure lock device is unsupported.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_UNSUPPORTED = 2;
+
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice} and {@link #disableSecureLockDevice}.
+     *
+     * Invalid secure lock device request params provided.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_INVALID_PARAMS = 3;
+
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice} and {@link #disableSecureLockDevice}.
+     *
+     * Secure lock device is unavailable because there are no biometrics enrolled on the device.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_NO_BIOMETRICS_ENROLLED = 4;
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice} and {@link #disableSecureLockDevice}.
+     *
+     * Secure lock device is unavailable because the device has no biometric hardware or the
+     * biometric sensors do not meet
+     * {@link android.hardware.biometrics.BiometricManager.Authenticators#BIOMETRIC_STRONG}
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_INSUFFICIENT_BIOMETRICS = 5;
+
+    /**
+     * Error result code for {@link #enableSecureLockDevice}.
+     *
+     * Secure lock is already enabled.
+     *
+     * @hide
+     */
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public static final int ERROR_ALREADY_ENABLED = 6;
+
+    /**
+     * Communicates the current status of a request to enable secure lock on the device.
+     *
+     * @hide
+     */
+    @IntDef(prefix = {"ENABLE_SECURE_LOCK_DEVICE_STATUS_"}, value = {
+            ERROR_UNKNOWN,
+            SUCCESS,
+            ERROR_UNSUPPORTED,
+            ERROR_INVALID_PARAMS,
+            ERROR_NO_BIOMETRICS_ENROLLED,
+            ERROR_INSUFFICIENT_BIOMETRICS,
+            ERROR_ALREADY_ENABLED
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface EnableSecureLockDeviceRequestStatus {}
+
+    /**
+     * Communicates the current status of a request to disable secure lock on the device.
+     *
+     * @hide
+     */
+    @IntDef(prefix = {"DISABLE_SECURE_LOCK_DEVICE_STATUS_"}, value = {
+            ERROR_UNKNOWN,
+            SUCCESS,
+            ERROR_UNSUPPORTED,
+            ERROR_INVALID_PARAMS,
+    })
+    @Retention(RetentionPolicy.SOURCE)
+    public @interface DisableSecureLockDeviceRequestStatus {}
+
+    /** @hide */
+    public AuthenticationPolicyManager(@NonNull Context context,
+            @NonNull IAuthenticationPolicyService authenticationPolicyService) {
+        mContext = context;
+        mAuthenticationPolicyService = authenticationPolicyService;
+    }
+
+    /**
+     * Called by a privileged component to remotely enable secure lock on the device.
+     *
+     * Secure lock is an enhanced security state that restricts access to sensitive data (app
+     * notifications, widgets, quick settings, assistant, etc) and requires multi-factor
+     * authentication for device entry, such as
+     * {@link android.hardware.biometrics.BiometricManager.Authenticators#DEVICE_CREDENTIAL} and
+     * {@link android.hardware.biometrics.BiometricManager.Authenticators#BIOMETRIC_STRONG}.
+     *
+     * If secure lock is already enabled when this method is called, it will return
+     * {@link ERROR_ALREADY_ENABLED}.
+     *
+     * @param params EnableSecureLockDeviceParams for caller to supply params related to the secure
+     *               lock device request
+     * @return @EnableSecureLockDeviceRequestStatus int indicating the result of the secure lock
+     * device request
+     *
+     * @hide
+     */
+    @EnableSecureLockDeviceRequestStatus
+    @RequiresPermission(MANAGE_SECURE_LOCK_DEVICE)
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public int enableSecureLockDevice(@NonNull EnableSecureLockDeviceParams params) {
+        try {
+            return mAuthenticationPolicyService.enableSecureLockDevice(params);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
+     * Called by a privileged component to disable secure lock on the device.
+     *
+     * If secure lock is already disabled when this method is called, it will return
+     * {@link SUCCESS}.
+     *
+     * @param params @DisableSecureLockDeviceParams for caller to supply params related to the
+     *               secure lock device request
+     * @return @DisableSecureLockDeviceRequestStatus int indicating the result of the secure lock
+     * device request
+     *
+     * @hide
+     */
+    @DisableSecureLockDeviceRequestStatus
+    @RequiresPermission(MANAGE_SECURE_LOCK_DEVICE)
+    @SystemApi
+    @FlaggedApi(FLAG_SECURE_LOCKDOWN)
+    public int disableSecureLockDevice(@NonNull DisableSecureLockDeviceParams params) {
+        try {
+            return mAuthenticationPolicyService.disableSecureLockDevice(params);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+}
diff --git a/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.aidl b/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.aidl
new file mode 100644
index 0000000..81f7726
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.aidl
@@ -0,0 +1,21 @@
+/*
+ * 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.security.authenticationpolicy;
+
+/**
+ * @hide
+ */
+parcelable DisableSecureLockDeviceParams;
\ No newline at end of file
diff --git a/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.java b/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.java
new file mode 100644
index 0000000..64a3f0f
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/DisableSecureLockDeviceParams.java
@@ -0,0 +1,82 @@
+/*
+ * 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.security.authenticationpolicy;
+
+import static android.security.Flags.FLAG_SECURE_LOCKDOWN;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.SystemApi;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import java.util.Objects;
+
+/**
+ * Parameters related to a request to disable secure lock on the device.
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_SECURE_LOCKDOWN)
+public final class DisableSecureLockDeviceParams implements Parcelable {
+
+    /**
+     * Client message associated with the request to disable secure lock on the device. This message
+     * will be shown on the device when secure lock mode is disabled.
+     */
+    private final @NonNull String mMessage;
+
+    /**
+     * Creates DisableSecureLockDeviceParams with the given params.
+     *
+     * @param message Allows clients to pass in a message with information about the request to
+     *                disable secure lock on the device. This message will be shown to the user when
+     *                secure lock mode is disabled. If an empty string is provided, it will default
+     *                to a system-defined string (e.g. "Secure lock mode has been disabled.")
+     */
+    public DisableSecureLockDeviceParams(@NonNull String message) {
+        mMessage = message;
+    }
+
+    private DisableSecureLockDeviceParams(@NonNull Parcel in) {
+        mMessage = Objects.requireNonNull(in.readString8());
+    }
+
+    public static final @NonNull Creator<DisableSecureLockDeviceParams> CREATOR =
+            new Creator<DisableSecureLockDeviceParams>() {
+                @Override
+                public DisableSecureLockDeviceParams createFromParcel(Parcel in) {
+                    return new DisableSecureLockDeviceParams(in);
+                }
+
+                @Override
+                public DisableSecureLockDeviceParams[] newArray(int size) {
+                    return new DisableSecureLockDeviceParams[size];
+                }
+            };
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeString8(mMessage);
+    }
+}
diff --git a/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.aidl b/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.aidl
new file mode 100644
index 0000000..9e496f8
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.aidl
@@ -0,0 +1,21 @@
+/*
+ * 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.security.authenticationpolicy;
+
+/**
+ * @hide
+ */
+parcelable EnableSecureLockDeviceParams;
\ No newline at end of file
diff --git a/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.java b/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.java
new file mode 100644
index 0000000..1d72772
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/EnableSecureLockDeviceParams.java
@@ -0,0 +1,82 @@
+/*
+ * 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.security.authenticationpolicy;
+
+import static android.security.Flags.FLAG_SECURE_LOCKDOWN;
+
+import android.annotation.FlaggedApi;
+import android.annotation.NonNull;
+import android.annotation.SystemApi;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+import java.util.Objects;
+
+
+/**
+ * Parameters related to a request to enable secure lock on the device.
+ *
+ * @hide
+ */
+@SystemApi
+@FlaggedApi(FLAG_SECURE_LOCKDOWN)
+public final class EnableSecureLockDeviceParams implements Parcelable {
+
+    /**
+     * Client message associated with the request to enable secure lock on the device. This message
+     * will be shown on the device when secure lock mode is enabled.
+     */
+    private final @NonNull String mMessage;
+
+    /**
+     * Creates EnableSecureLockDeviceParams with the given params.
+     *
+     * @param message Allows clients to pass in a message with information about the request to
+     *                enable secure lock on the device. This message will be shown to the user when
+     *                secure lock mode is enabled. If an empty string is provided, it will default
+     *                to a system-defined string (e.g. "Device is securely locked remotely.")
+     */
+    public EnableSecureLockDeviceParams(@NonNull String message) {
+        mMessage = message;
+    }
+
+    private EnableSecureLockDeviceParams(@NonNull Parcel in) {
+        mMessage = Objects.requireNonNull(in.readString8());
+    }
+
+    public static final @NonNull Creator<EnableSecureLockDeviceParams> CREATOR =
+            new Creator<EnableSecureLockDeviceParams>() {
+                @Override
+                public EnableSecureLockDeviceParams createFromParcel(Parcel in) {
+                    return new EnableSecureLockDeviceParams(in);
+                }
+
+                @Override
+                public EnableSecureLockDeviceParams[] newArray(int size) {
+                    return new EnableSecureLockDeviceParams[size];
+                }
+            };
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
+        dest.writeString8(mMessage);
+    }
+}
diff --git a/core/java/android/security/authenticationpolicy/IAuthenticationPolicyService.aidl b/core/java/android/security/authenticationpolicy/IAuthenticationPolicyService.aidl
new file mode 100644
index 0000000..5ad4534
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/IAuthenticationPolicyService.aidl
@@ -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 android.security.authenticationpolicy;
+
+import android.security.authenticationpolicy.EnableSecureLockDeviceParams;
+import android.security.authenticationpolicy.DisableSecureLockDeviceParams;
+
+/**
+ * Communication channel from AuthenticationPolicyManager to AuthenticationPolicyService.
+ * @hide
+ */
+interface IAuthenticationPolicyService {
+    @EnforcePermission("MANAGE_SECURE_LOCK_DEVICE")
+    int enableSecureLockDevice(in EnableSecureLockDeviceParams params);
+
+    @EnforcePermission("MANAGE_SECURE_LOCK_DEVICE")
+    int disableSecureLockDevice(in DisableSecureLockDeviceParams params);
+}
\ No newline at end of file
diff --git a/core/java/android/security/authenticationpolicy/OWNERS b/core/java/android/security/authenticationpolicy/OWNERS
new file mode 100644
index 0000000..4310d1a
--- /dev/null
+++ b/core/java/android/security/authenticationpolicy/OWNERS
@@ -0,0 +1 @@
+include /services/core/java/com/android/server/security/authenticationpolicy/OWNERS
\ No newline at end of file
diff --git a/core/java/android/service/notification/ZenModeConfig.java b/core/java/android/service/notification/ZenModeConfig.java
index 24328eb..1388778 100644
--- a/core/java/android/service/notification/ZenModeConfig.java
+++ b/core/java/android/service/notification/ZenModeConfig.java
@@ -2532,7 +2532,7 @@
 
     /** Returns whether the rule id corresponds to an implicit rule. */
     public static boolean isImplicitRuleId(@NonNull String ruleId) {
-        return ruleId.startsWith(IMPLICIT_RULE_ID_PREFIX);
+        return ruleId != null && ruleId.startsWith(IMPLICIT_RULE_ID_PREFIX);
     }
 
     private static int[] tryParseHourAndMinute(String value) {
diff --git a/core/java/android/window/flags/lse_desktop_experience.aconfig b/core/java/android/window/flags/lse_desktop_experience.aconfig
index eebdeadc..c7d5a9f 100644
--- a/core/java/android/window/flags/lse_desktop_experience.aconfig
+++ b/core/java/android/window/flags/lse_desktop_experience.aconfig
@@ -253,6 +253,9 @@
     namespace: "lse_desktop_experience"
     description: "Enables enter desktop windowing transition & motion polish changes"
     bug: "369763947"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
 
 flag {
@@ -260,6 +263,9 @@
     namespace: "lse_desktop_experience"
     description: "Enables exit desktop windowing transition & motion polish changes"
     bug: "353650462"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
 
 flag {
@@ -343,6 +349,9 @@
     namespace: "lse_desktop_experience"
     description: "Enables custom transitions for alt-tab app launches in Desktop Mode."
     bug: "370735595"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
 
 flag {
@@ -350,6 +359,9 @@
     namespace: "lse_desktop_experience"
     description: "Enables custom transitions for app launches in Desktop Mode."
     bug: "375992828"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
 }
 
 flag {
diff --git a/core/java/android/window/flags/window_surfaces.aconfig b/core/java/android/window/flags/window_surfaces.aconfig
index 96b9dc7..bb47707 100644
--- a/core/java/android/window/flags/window_surfaces.aconfig
+++ b/core/java/android/window/flags/window_surfaces.aconfig
@@ -29,14 +29,6 @@
 
 flag {
     namespace: "window_surfaces"
-    name: "secure_window_state"
-    description: "Move SC secure flag to WindowState level"
-    is_fixed_read_only: true
-    bug: "308662081"
-}
-
-flag {
-    namespace: "window_surfaces"
     name: "trusted_presentation_listener_for_window"
     is_exported: true
     description: "Enable trustedPresentationListener on windows public API"
diff --git a/core/res/Android.bp b/core/res/Android.bp
index 0e4e22b..26e63bc 100644
--- a/core/res/Android.bp
+++ b/core/res/Android.bp
@@ -161,6 +161,7 @@
         "android.app.flags-aconfig",
         "android.appwidget.flags-aconfig",
         "android.content.pm.flags-aconfig",
+        "android.media.audio-aconfig",
         "android.provider.flags-aconfig",
         "camera_platform_flags",
         "android.net.platform.flags-aconfig",
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index 7fcbf19..8cc7b0b 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -6492,6 +6492,15 @@
     <permission android:name="android.permission.CAPTURE_VOICE_COMMUNICATION_OUTPUT"
         android:protectionLevel="signature|privileged|role" />
 
+    <!-- @SystemApi Allows an application to bypass concurrency restrictions while
+        recording audio. For example, apps with this permission can continue to record
+        while a voice call is active.</p>
+     @FlaggedApi(android.media.audio.Flags.FLAG_CONCURRENT_AUDIO_RECORD_BYPASS_PERMISSION)
+     @hide -->
+    <permission android:name="android.permission.BYPASS_CONCURRENT_RECORD_AUDIO_RESTRICTION"
+        android:featureFlag="android.media.audio.concurrent_audio_record_bypass_permission"
+        android:protectionLevel="signature|privileged|role" />
+
     <!-- @SystemApi Allows an application to capture audio for hotword detection.
          <p>Not for use by third-party applications.</p>
          @hide -->
@@ -8651,6 +8660,17 @@
     <permission android:name="android.permission.SETUP_FSVERITY"
                 android:protectionLevel="signature|privileged"/>
 
+    <!-- @SystemApi
+        @FlaggedApi(android.security.Flags.FLAG_SECURE_LOCKDOWN)
+        Allows an application to lock down the device into an enhanced security state.
+        <p>Not for use by third-party applications.
+        <p>Protection level: signature|privileged
+        @hide
+    -->
+    <permission android:name="android.permission.MANAGE_SECURE_LOCK_DEVICE"
+        android:protectionLevel="signature|privileged"
+        android:featureFlag="android.security.secure_lockdown" />
+
     <!-- Allows app to enter trade-in-mode.
         <p>Protection level: signature|privileged
         @hide
diff --git a/core/res/res/layout/notification_2025_template_heads_up_base.xml b/core/res/res/layout/notification_2025_template_heads_up_base.xml
new file mode 100644
index 0000000..e4ff835
--- /dev/null
+++ b/core/res/res/layout/notification_2025_template_heads_up_base.xml
@@ -0,0 +1,66 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2014 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
+  -->
+
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/status_bar_latest_event_content"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:orientation="vertical"
+    android:clipChildren="false"
+    android:tag="headsUp"
+    >
+
+    <LinearLayout
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:clipChildren="false"
+        android:orientation="vertical"
+        >
+
+        <include
+            layout="@layout/notification_2025_template_collapsed_base"
+            android:id="@null"
+            />
+
+        <LinearLayout
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            android:layout_marginTop="-20dp"
+            android:clipChildren="false"
+            android:orientation="vertical"
+            >
+
+            <ViewStub
+                android:layout="@layout/notification_material_reply_text"
+                android:id="@+id/notification_material_reply_container"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                />
+
+            <include
+                layout="@layout/notification_template_smart_reply_container"
+                android:layout_width="match_parent"
+                android:layout_height="wrap_content"
+                android:layout_marginStart="@dimen/notification_2025_content_margin_start"
+                android:layout_marginEnd="@dimen/notification_content_margin_end"
+                android:layout_marginTop="@dimen/notification_content_margin"
+                />
+
+            <include layout="@layout/notification_material_action_list" />
+        </LinearLayout>
+    </LinearLayout>
+</FrameLayout>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index a2a19a2..5a6b66c 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -2390,6 +2390,7 @@
   <java-symbol type="layout" name="notification_material_action_list" />
   <java-symbol type="layout" name="notification_material_action_tombstone" />
   <java-symbol type="layout" name="notification_2025_template_collapsed_base" />
+  <java-symbol type="layout" name="notification_2025_template_heads_up_base" />
   <java-symbol type="layout" name="notification_2025_template_header" />
   <java-symbol type="layout" name="notification_template_material_base" />
   <java-symbol type="layout" name="notification_template_material_heads_up_base" />
diff --git a/data/etc/privapp-permissions-platform.xml b/data/etc/privapp-permissions-platform.xml
index fea7cb4..329e5de 100644
--- a/data/etc/privapp-permissions-platform.xml
+++ b/data/etc/privapp-permissions-platform.xml
@@ -533,6 +533,8 @@
         <permission name="com.android.cellbroadcastservice.FULL_ACCESS_CELL_BROADCAST_HISTORY" />
         <!-- Permission required for ATS test - CarDevicePolicyManagerTest -->
         <permission name="android.permission.LOCK_DEVICE" />
+        <!-- Permission required for AuthenticationPolicyManagerTest -->
+        <permission name="android.permission.MANAGE_SECURE_LOCK_DEVICE" />
         <!-- Permissions required for CTS test - CtsSafetyCenterTestCases -->
         <permission name="android.permission.SEND_SAFETY_CENTER_UPDATE" />
         <permission name="android.permission.READ_SAFETY_CENTER_STATUS" />
diff --git a/libs/WindowManager/Shell/res/values/styles.xml b/libs/WindowManager/Shell/res/values/styles.xml
index 55cda78..597a921 100644
--- a/libs/WindowManager/Shell/res/values/styles.xml
+++ b/libs/WindowManager/Shell/res/values/styles.xml
@@ -29,6 +29,7 @@
         <item name="android:navigationBarColor">@android:color/transparent</item>
         <item name="android:windowDrawsSystemBarBackgrounds">true</item>
         <item name="android:windowAnimationStyle">@null</item>
+        <item name="android:windowIsTranslucent">true</item>
     </style>
 
     <style name="Animation.ForcedResizable" parent="@android:style/Animation">
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 0a39076..6fc6eb7 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
@@ -70,6 +70,8 @@
 import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.internal.protolog.ProtoLog
 import com.android.window.flags.Flags
+import com.android.window.flags.Flags.enableMoveToNextDisplayShortcut
+import com.android.wm.shell.Flags.enableFlexibleSplit
 import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
 import com.android.wm.shell.ShellTaskOrganizer
@@ -102,6 +104,7 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.DESKTOP_DENSITY_OVERRIDE
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus.useDesktopOverrideDensity
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED
 import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT
 import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT
 import com.android.wm.shell.splitscreen.SplitScreenController
@@ -1512,11 +1515,15 @@
             WINDOWING_MODE_MULTI_WINDOW -> {
                 val splitPosition = splitScreenController
                     .determineNewInstancePosition(callingTaskInfo)
+                // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+                //  specific cases in the future.
+                val splitIndex = if (enableFlexibleSplit())
+                    splitScreenController.determineNewInstanceIndex(callingTaskInfo) else
+                    SPLIT_INDEX_UNDEFINED
                 splitScreenController.startIntent(
                     launchIntent, context.userId, fillIn, splitPosition,
                     options.toBundle(), null /* hideTaskToken */,
-                    true /* forceLaunchNewTask */
-                )
+                    true /* forceLaunchNewTask */, splitIndex)
             }
             WINDOWING_MODE_FREEFORM -> {
                 val wct = WindowContainerTransaction()
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
index 5d22c1e..ae9d21f62 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/SplitDragPolicy.java
@@ -41,6 +41,9 @@
 import static com.android.wm.shell.draganddrop.SplitDragPolicy.Target.TYPE_SPLIT_TOP;
 import static com.android.wm.shell.shared.draganddrop.DragAndDropConstants.EXTRA_DISALLOW_HIT_REGION;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_1;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
@@ -81,6 +84,7 @@
 import com.android.wm.shell.draganddrop.anim.TwoFiftyFiftyTargetAnimator;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.shared.split.SplitScreenConstants;
+import com.android.wm.shell.shared.split.SplitScreenConstants.SplitIndex;
 import com.android.wm.shell.shared.split.SplitScreenConstants.SplitPosition;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 
@@ -219,8 +223,10 @@
                         displayRegion.splitHorizontally(startHitRegion, endHitRegion);
                     }
 
-                    mTargets.add(new Target(TYPE_SPLIT_LEFT, startHitRegion, startBounds, -1));
-                    mTargets.add(new Target(TYPE_SPLIT_RIGHT, endHitRegion, endBounds, -1));
+                    mTargets.add(new Target(TYPE_SPLIT_LEFT, startHitRegion, startBounds,
+                            SPLIT_INDEX_0));
+                    mTargets.add(new Target(TYPE_SPLIT_RIGHT, endHitRegion, endBounds,
+                            SPLIT_INDEX_1));
                 } else {
                     // TODO(b/349828130), move this into init function and/or the insets updating
                     //  callback
@@ -287,9 +293,10 @@
                         displayRegion.splitVertically(leftHitRegion, rightHitRegion);
                     }
 
-                    mTargets.add(new Target(TYPE_SPLIT_LEFT, leftHitRegion, topOrLeftBounds, -1));
+                    mTargets.add(new Target(TYPE_SPLIT_LEFT, leftHitRegion, topOrLeftBounds,
+                            SPLIT_INDEX_UNDEFINED));
                     mTargets.add(new Target(TYPE_SPLIT_RIGHT, rightHitRegion, bottomOrRightBounds,
-                            -1));
+                            SPLIT_INDEX_UNDEFINED));
                 } else {
                     final Rect topHitRegion = new Rect();
                     final Rect bottomHitRegion = new Rect();
@@ -308,9 +315,10 @@
                         displayRegion.splitHorizontally(topHitRegion, bottomHitRegion);
                     }
 
-                    mTargets.add(new Target(TYPE_SPLIT_TOP, topHitRegion, topOrLeftBounds, -1));
+                    mTargets.add(new Target(TYPE_SPLIT_TOP, topHitRegion, topOrLeftBounds,
+                            SPLIT_INDEX_UNDEFINED));
                     mTargets.add(new Target(TYPE_SPLIT_BOTTOM, bottomHitRegion, bottomOrRightBounds,
-                            -1));
+                            SPLIT_INDEX_UNDEFINED));
                 }
             }
         } else {
@@ -378,9 +386,9 @@
                 ? mFullscreenStarter
                 : mSplitscreenStarter;
         if (mSession.appData != null) {
-            launchApp(mSession, starter, position, hideTaskToken);
+            launchApp(mSession, starter, position, hideTaskToken, target.index);
         } else {
-            launchIntent(mSession, starter, position, hideTaskToken);
+            launchIntent(mSession, starter, position, hideTaskToken, target.index);
         }
 
         if (enableFlexibleSplit()) {
@@ -392,9 +400,10 @@
      * Launches an app provided by SysUI.
      */
     private void launchApp(DragSession session, Starter starter, @SplitPosition int position,
-            @Nullable WindowContainerToken hideTaskToken) {
-        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP, "Launching app data at position=%d",
-                position);
+            @Nullable WindowContainerToken hideTaskToken, @SplitIndex int splitIndex) {
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP,
+                "Launching app data at position=%d index=%d",
+                position, splitIndex);
         final ClipDescription description = session.getClipDescription();
         final boolean isTask = description.hasMimeType(MIMETYPE_APPLICATION_TASK);
         final boolean isShortcut = description.hasMimeType(MIMETYPE_APPLICATION_SHORTCUT);
@@ -429,7 +438,7 @@
                 }
             }
             starter.startIntent(launchIntent, user.getIdentifier(), null /* fillIntent */,
-                    position, opts, hideTaskToken);
+                    position, opts, hideTaskToken, splitIndex);
         }
     }
 
@@ -437,7 +446,7 @@
      * Launches an intent sender provided by an application.
      */
     private void launchIntent(DragSession session, Starter starter, @SplitPosition int position,
-            @Nullable WindowContainerToken hideTaskToken) {
+            @Nullable WindowContainerToken hideTaskToken, @SplitIndex int index) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP, "Launching intent at position=%d",
                 position);
         final ActivityOptions baseActivityOpts = ActivityOptions.makeBasic();
@@ -452,7 +461,7 @@
         final Bundle opts = baseActivityOpts.toBundle();
         starter.startIntent(session.launchableIntent,
                 session.launchableIntent.getCreatorUserHandle().getIdentifier(),
-                null /* fillIntent */, position, opts, hideTaskToken);
+                null /* fillIntent */, position, opts, hideTaskToken, index);
     }
 
     @Override
@@ -541,7 +550,7 @@
                 @Nullable Bundle options, UserHandle user);
         void startIntent(PendingIntent intent, int userId, Intent fillInIntent,
                 @SplitPosition int position, @Nullable Bundle options,
-                @Nullable WindowContainerToken hideTaskToken);
+                @Nullable WindowContainerToken hideTaskToken, @SplitIndex int index);
         void enterSplitScreen(int taskId, boolean leftOrTop);
 
         /**
@@ -592,7 +601,7 @@
         @Override
         public void startIntent(PendingIntent intent, int userId, @Nullable Intent fillInIntent,
                 int position, @Nullable Bundle options,
-                @Nullable WindowContainerToken hideTaskToken) {
+                @Nullable WindowContainerToken hideTaskToken, @SplitIndex int index) {
             if (hideTaskToken != null) {
                 ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DRAG_AND_DROP,
                         "Default starter does not support hide task token");
@@ -641,13 +650,13 @@
         final Rect hitRegion;
         // The approximate visual region for where the task will start
         final Rect drawRegion;
-        int index;
+        @SplitIndex int index;
 
         /**
          * @param index 0-indexed, represents which position of drop target this object represents,
          *              0 to N for left to right, top to bottom
          */
-        public Target(@Type int t, Rect hit, Rect draw, int index) {
+        public Target(@Type int t, Rect hit, Rect draw, @SplitIndex int index) {
             type = t;
             hitRegion = hit;
             drawRegion = draw;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/anim/TwoFiftyFiftyTargetAnimator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/anim/TwoFiftyFiftyTargetAnimator.kt
index 9f532f5..5461952 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/anim/TwoFiftyFiftyTargetAnimator.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/anim/TwoFiftyFiftyTargetAnimator.kt
@@ -22,6 +22,10 @@
 import com.android.wm.shell.R
 import com.android.wm.shell.common.DisplayLayout
 import com.android.wm.shell.draganddrop.SplitDragPolicy.Target
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_1
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_2
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_3
 
 /**
  * Represents Drop Zone targets and animations for when the system is currently in a 2 app 50/50
@@ -98,7 +102,7 @@
                     farStartBounds.right + halfDividerWidth,
                     farStartBounds.bottom
                 ),
-                farStartBounds, 0
+                farStartBounds, SPLIT_INDEX_0
             )
         )
         targets.add(
@@ -110,7 +114,7 @@
                     startBounds.right + halfDividerWidth,
                     startBounds.bottom
                 ),
-                startBounds, 1
+                startBounds, SPLIT_INDEX_1
             )
         )
         targets.add(
@@ -120,7 +124,7 @@
                     endBounds.left - halfDividerWidth,
                     endBounds.top, endBounds.right, endBounds.bottom
                 ),
-                endBounds, 2
+                endBounds, SPLIT_INDEX_2
             )
         )
         targets.add(
@@ -130,7 +134,7 @@
                     farEndBounds.left - halfDividerWidth,
                     farEndBounds.top, farEndBounds.right, farEndBounds.bottom
                 ),
-                farEndBounds, 3
+                farEndBounds, SPLIT_INDEX_3
             )
         )
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index 3e6d36c..39ed9ab 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -54,10 +54,27 @@
      */
     int STAGE_TYPE_SIDE = 1;
 
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    int STAGE_TYPE_A = 2;
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    int STAGE_TYPE_B = 3;
+    /**
+     * Position independent stage identifier for a given Stage
+     */
+    int STAGE_TYPE_C = 4;
+
     @IntDef(prefix = { "STAGE_TYPE_" }, value = {
             STAGE_TYPE_UNDEFINED,
             STAGE_TYPE_MAIN,
-            STAGE_TYPE_SIDE
+            STAGE_TYPE_SIDE,
+            // Used for flexible split
+            STAGE_TYPE_A,
+            STAGE_TYPE_B,
+            STAGE_TYPE_C
     })
     @interface StageType {}
 
@@ -128,6 +145,9 @@
             case STAGE_TYPE_UNDEFINED: return "UNDEFINED";
             case STAGE_TYPE_MAIN: return "MAIN";
             case STAGE_TYPE_SIDE: return "SIDE";
+            case STAGE_TYPE_A: return "STAGE_A";
+            case STAGE_TYPE_B: return "STAGE_B";
+            case STAGE_TYPE_C: return "STAGE_C";
             default: return "UNKNOWN(" + stage + ")";
         }
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 6398d31..4f0f676 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -24,6 +24,7 @@
 import static android.content.Intent.FLAG_ACTIVITY_NO_USER_ACTION;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.common.MultiInstanceHelper.getComponent;
 import static com.android.wm.shell.common.MultiInstanceHelper.getShortcutComponent;
 import static com.android.wm.shell.common.MultiInstanceHelper.samePackage;
@@ -33,6 +34,9 @@
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
 import static com.android.wm.shell.shared.ShellSharedConstants.KEY_EXTRA_SHELL_SPLIT_SCREEN;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.KEY_EXTRA_WIDGET_INTENT;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_1;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
@@ -98,6 +102,7 @@
 import com.android.wm.shell.shared.TransactionPool;
 import com.android.wm.shell.shared.annotations.ExternalThread;
 import com.android.wm.shell.shared.split.SplitScreenConstants.PersistentSnapPosition;
+import com.android.wm.shell.shared.split.SplitScreenConstants.SplitIndex;
 import com.android.wm.shell.shared.split.SplitScreenConstants.SplitPosition;
 import com.android.wm.shell.splitscreen.SplitScreen.StageType;
 import com.android.wm.shell.sysui.KeyguardChangeListener;
@@ -325,7 +330,6 @@
     /**
      * @return an Array of RunningTaskInfo's ordered by leftToRight or topTopBottom
      */
-    @Nullable
     public ActivityManager.RunningTaskInfo[] getAllTaskInfos() {
         // TODO(b/349828130) Add the third stage task info and not rely on positions
         ActivityManager.RunningTaskInfo topLeftTask = getTaskInfo(SPLIT_POSITION_TOP_OR_LEFT);
@@ -335,7 +339,7 @@
             return new ActivityManager.RunningTaskInfo[]{topLeftTask, bottomRightTask};
         }
 
-        return null;
+        return new ActivityManager.RunningTaskInfo[0];
     }
 
     /** Check task is under split or not by taskId. */
@@ -405,7 +409,7 @@
     public void prepareEnterSplitScreen(WindowContainerTransaction wct,
             ActivityManager.RunningTaskInfo taskInfo, int startPosition) {
         mStageCoordinator.prepareEnterSplitScreen(wct, taskInfo, startPosition,
-                false /* resizeAnim */);
+                false /* resizeAnim */, SPLIT_INDEX_UNDEFINED);
     }
 
     /**
@@ -451,6 +455,24 @@
         }
     }
 
+    /**
+     * Determines which split index a new instance of a task should take.
+     * @param callingTask The task requesting a new instance.
+     * @return the split index of the new instance
+     */
+    @SplitIndex
+    public int determineNewInstanceIndex(@NonNull ActivityManager.RunningTaskInfo callingTask) {
+        if (!enableFlexibleSplit()) {
+            throw new IllegalStateException("Use determineNewInstancePosition");
+        }
+        if (callingTask.getWindowingMode() == WINDOWING_MODE_FULLSCREEN
+                || getSplitPosition(callingTask.taskId) == SPLIT_POSITION_TOP_OR_LEFT) {
+            return SPLIT_INDEX_1;
+        } else {
+            return SPLIT_INDEX_0;
+        }
+    }
+
     public void enterSplitScreen(int taskId, boolean leftOrTop) {
         enterSplitScreen(taskId, leftOrTop, new WindowContainerTransaction());
     }
@@ -685,7 +707,10 @@
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "startIntentWithInstanceId: reason=%d",
                 ENTER_REASON_LAUNCHER);
         mStageCoordinator.getLogger().enterRequested(instanceId, ENTER_REASON_LAUNCHER);
-        startIntent(intent, userId, fillInIntent, position, options, null /* hideTaskToken */);
+        // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+        //  specific cases in the future. Only focusing on parity with starting intent/task
+        startIntent(intent, userId, fillInIntent, position, options, null /* hideTaskToken */,
+                SPLIT_INDEX_UNDEFINED);
     }
 
     private void startIntentAndTask(PendingIntent pendingIntent, int userId1,
@@ -775,9 +800,9 @@
     @Override
     public void startIntent(PendingIntent intent, int userId1, @Nullable Intent fillInIntent,
             @SplitPosition int position, @Nullable Bundle options,
-            @Nullable WindowContainerToken hideTaskToken) {
+            @Nullable WindowContainerToken hideTaskToken, @SplitIndex int index) {
         startIntent(intent, userId1, fillInIntent, position, options, hideTaskToken,
-                false /* forceLaunchNewTask */);
+                false /* forceLaunchNewTask */, index);
     }
 
     /**
@@ -790,7 +815,8 @@
      */
     public void startIntent(PendingIntent intent, int userId1, @Nullable Intent fillInIntent,
             @SplitPosition int position, @Nullable Bundle options,
-            @Nullable WindowContainerToken hideTaskToken, boolean forceLaunchNewTask) {
+            @Nullable WindowContainerToken hideTaskToken, boolean forceLaunchNewTask,
+            @SplitIndex int index) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
                 "startIntent(): intent=%s user=%d fillInIntent=%s position=%d", intent, userId1,
                 fillInIntent, position);
@@ -816,7 +842,7 @@
         if (taskInfo != null) {
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
                     "Found suitable background task=%s", taskInfo);
-            mStageCoordinator.startTask(taskInfo.taskId, position, options, hideTaskToken);
+            mStageCoordinator.startTask(taskInfo.taskId, position, options, hideTaskToken, index);
 
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Start task in background");
             return;
@@ -841,7 +867,8 @@
             }
         }
 
-        mStageCoordinator.startIntent(intent, fillInIntent, position, options, hideTaskToken);
+        mStageCoordinator.startIntent(intent, fillInIntent, position, options, hideTaskToken,
+                index);
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
index 84004941..3091be5 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
@@ -21,6 +21,7 @@
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
 
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_TRANSITIONS;
 import static com.android.wm.shell.shared.animation.Interpolators.ALPHA_IN;
@@ -55,6 +56,9 @@
 import com.android.wm.shell.transition.Transitions;
 
 import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.Executor;
 
 /** Manages transition animations for split-screen. */
@@ -268,22 +272,21 @@
             @NonNull SurfaceControl.Transaction startTransaction,
             @NonNull SurfaceControl.Transaction finishTransaction,
             @NonNull Transitions.TransitionFinishCallback finishCallback,
-            @NonNull WindowContainerToken mainRoot, @NonNull WindowContainerToken sideRoot,
-            @NonNull SplitDecorManager mainDecor, @NonNull SplitDecorManager sideDecor) {
+            @NonNull Map<WindowContainerToken, SplitDecorManager> rootDecorMap) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "playResizeAnimation: transition=%d", info.getDebugId());
         initTransition(transition, finishTransaction, finishCallback);
 
+        Set<WindowContainerToken> rootDecorKeys = rootDecorMap.keySet();
         for (int i = info.getChanges().size() - 1; i >= 0; --i) {
             final TransitionInfo.Change change = info.getChanges().get(i);
-            if (mainRoot.equals(change.getContainer()) || sideRoot.equals(change.getContainer())) {
+            if (rootDecorKeys.contains(change.getContainer())) {
                 final SurfaceControl leash = change.getLeash();
                 startTransaction.setPosition(leash, change.getEndAbsBounds().left,
                         change.getEndAbsBounds().top);
                 startTransaction.setWindowCrop(leash, change.getEndAbsBounds().width(),
                         change.getEndAbsBounds().height());
 
-                SplitDecorManager decor = mainRoot.equals(change.getContainer())
-                        ? mainDecor : sideDecor;
+                SplitDecorManager decor = rootDecorMap.get(change.getContainer());
 
                 // This is to ensure onFinished be called after all animations ended.
                 ValueAnimator va = new ValueAnimator();
@@ -433,15 +436,22 @@
             Transitions.TransitionHandler handler,
             @Nullable TransitionConsumedCallback consumedCallback,
             @Nullable TransitionFinishedCallback finishCallback,
-            @NonNull SplitDecorManager mainDecor, @NonNull SplitDecorManager sideDecor) {
+            @Nullable SplitDecorManager mainDecor, @Nullable SplitDecorManager sideDecor,
+            @Nullable List<SplitDecorManager> decorManagers) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS,
                 "  splitTransition deduced Resize split screen.");
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "setResizeTransition: hasPendingResize=%b",
                 mPendingResize != null);
         if (mPendingResize != null) {
             mPendingResize.cancel(null);
-            mainDecor.cancelRunningAnimations();
-            sideDecor.cancelRunningAnimations();
+            if (enableFlexibleSplit()) {
+                for (SplitDecorManager stage : decorManagers) {
+                    stage.cancelRunningAnimations();
+                }
+            } else {
+                mainDecor.cancelRunningAnimations();
+                sideDecor.cancelRunningAnimations();
+            }
             mAnimations.clear();
             onFinish(null /* wct */);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index e692c61..f07ae44 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -33,6 +33,7 @@
 import static android.window.TransitionInfo.FLAG_IS_DISPLAY;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.common.split.SplitLayout.PARALLAX_ALIGN_CENTER;
 import static com.android.wm.shell.common.split.SplitScreenUtils.isPartiallyOffscreen;
 import static com.android.wm.shell.common.split.SplitScreenUtils.reverseSplitPosition;
@@ -43,6 +44,7 @@
 import static com.android.wm.shell.shared.TransitionUtil.isOrderOnly;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.FLAG_IS_DIVIDER_BAR;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_10_90;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_90_10;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_1;
@@ -51,9 +53,12 @@
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.splitPositionToString;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_A;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_B;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_MAIN;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_SIDE;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
+import static com.android.wm.shell.splitscreen.SplitScreen.stageTypeToString;
 import static com.android.wm.shell.splitscreen.SplitScreenController.ENTER_REASON_LAUNCHER;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_APP_DOES_NOT_SUPPORT_MULTIWINDOW;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_APP_FINISHED;
@@ -142,6 +147,7 @@
 import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.shared.split.SplitBounds;
 import com.android.wm.shell.shared.split.SplitScreenConstants.PersistentSnapPosition;
+import com.android.wm.shell.shared.split.SplitScreenConstants.SplitIndex;
 import com.android.wm.shell.shared.split.SplitScreenConstants.SplitPosition;
 import com.android.wm.shell.splitscreen.SplitScreen.StageType;
 import com.android.wm.shell.splitscreen.SplitScreenController.ExitReason;
@@ -153,11 +159,15 @@
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.function.Consumer;
+import java.util.function.Predicate;
 
 /**
  * Coordinates the staging (visibility, sizing, ...) of the split-screen stages.
@@ -178,10 +188,11 @@
     // entered
     private static final int DISABLE_LAUNCH_ADJACENT_AFTER_ENTER_TIMEOUT_MS = 1000;
 
-    private final StageTaskListener mMainStage;
-    private final StageTaskListener mSideStage;
+    private StageTaskListener mMainStage;
+    private StageTaskListener mSideStage;
     @SplitPosition
     private int mSideStagePosition = SPLIT_POSITION_BOTTOM_OR_RIGHT;
+    private StageOrderOperator mStageOrderOperator;
 
     private final int mDisplayId;
     private SplitLayout mSplitLayout;
@@ -335,22 +346,32 @@
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_FULLSCREEN, this /* listener */);
 
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "Creating main/side root task");
-        mMainStage = new StageTaskListener(
-                mContext,
-                mTaskOrganizer,
-                mDisplayId,
-                this /*stageListenerCallbacks*/,
-                mSyncQueue,
-                iconProvider,
-                mWindowDecorViewModel);
-        mSideStage = new StageTaskListener(
-                mContext,
-                mTaskOrganizer,
-                mDisplayId,
-                this /*stageListenerCallbacks*/,
-                mSyncQueue,
-                iconProvider,
-                mWindowDecorViewModel);
+        if (enableFlexibleSplit()) {
+            mStageOrderOperator = new StageOrderOperator(mContext,
+                    mTaskOrganizer,
+                    mDisplayId,
+                    this /*stageListenerCallbacks*/,
+                    mSyncQueue,
+                    iconProvider,
+                    mWindowDecorViewModel);
+        } else {
+            mMainStage = new StageTaskListener(
+                    mContext,
+                    mTaskOrganizer,
+                    mDisplayId,
+                    this /*stageListenerCallbacks*/,
+                    mSyncQueue,
+                    iconProvider,
+                    mWindowDecorViewModel, STAGE_TYPE_MAIN);
+            mSideStage = new StageTaskListener(
+                    mContext,
+                    mTaskOrganizer,
+                    mDisplayId,
+                    this /*stageListenerCallbacks*/,
+                    mSyncQueue,
+                    iconProvider,
+                    mWindowDecorViewModel, STAGE_TYPE_SIDE);
+        }
         mDisplayController = displayController;
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
@@ -422,24 +443,63 @@
     }
 
     public boolean isSplitScreenVisible() {
-        return mSideStage.mVisible && mMainStage.mVisible;
+        if (enableFlexibleSplit()) {
+            return runForActiveStagesAllMatch((stage) -> stage.mVisible);
+        } else {
+            return mSideStage.mVisible && mMainStage.mVisible;
+        }
     }
 
-    private void activateSplit(WindowContainerTransaction wct, boolean includingTopTask) {
-        mMainStage.activate(wct, includingTopTask);
+    /**
+     * @param includingTopTask reparents the current top task into the stage defined by index
+     *                         (or mainStage in legacy split)
+     * @param index the index to move the current visible task into, if undefined will arbitrarily
+     *              choose a stage to launch into
+     */
+    private void activateSplit(WindowContainerTransaction wct, boolean includingTopTask,
+            int index) {
+        if (enableFlexibleSplit()) {
+            mStageOrderOperator.onEnteringSplit(SNAP_TO_2_50_50);
+            if (index == SPLIT_INDEX_UNDEFINED || !includingTopTask) {
+                // If we aren't includingTopTask, then the call to activate on the stage is
+                // effectively a no-op. Previously the stage kept track of the "isActive" state,
+                // but now that gets set in the "onEnteringSplit" call above.
+                //
+                // index == UNDEFINED case might change, but as of now no use case where we activate
+                // without an index specified.
+                return;
+            }
+            @SplitIndex int oppositeIndex = index == SPLIT_INDEX_0 ? SPLIT_INDEX_1 : SPLIT_INDEX_0;
+            StageTaskListener activatingStage = mStageOrderOperator.getStageForIndex(oppositeIndex);
+            activatingStage.activate(wct, includingTopTask);
+        } else {
+            mMainStage.activate(wct, includingTopTask);
+        }
     }
 
     public boolean isSplitActive() {
-        return mMainStage.isActive();
+        if (enableFlexibleSplit()) {
+            return mStageOrderOperator.isActive();
+        } else {
+            return mMainStage.isActive();
+        }
     }
 
     /**
      * Deactivates main stage by removing the stage from the top level split root (usually when a
      * task underneath gets removed from the stage root).
-     * @param reparentToTop whether we want to put the stage root back on top
+     * @param stageToTop stage which we want to put on top
      */
-    private void deactivateSplit(WindowContainerTransaction wct, boolean reparentToTop) {
-        mMainStage.deactivate(wct, reparentToTop);
+    private void deactivateSplit(WindowContainerTransaction wct, @StageType int stageToTop) {
+        if (enableFlexibleSplit()) {
+            StageTaskListener stageToDeactivate = mStageOrderOperator.getAllStages().stream()
+                    .filter(stage -> stage.getId() == stageToTop)
+                    .findFirst().orElseThrow();
+            stageToDeactivate.deactivate(wct, true /*toTop*/);
+            mStageOrderOperator.onExitingSplit();
+        } else {
+            mMainStage.deactivate(wct, stageToTop == STAGE_TYPE_MAIN);
+        }
     }
 
     /** @return whether this transition-request has the launch-adjacent flag. */
@@ -463,11 +523,12 @@
         // If one of the splitting tasks support auto-pip, wm-core might reparent the task to TDA
         // and file a TRANSIT_PIP transition when finishing transitions.
         // @see com.android.server.wm.RootWindowContainer#moveActivityToPinnedRootTask
-        if (mMainStage.getChildCount() == 0 || mSideStage.getChildCount() == 0) {
-            return true;
+        if (enableFlexibleSplit()) {
+            return mStageOrderOperator.getActiveStages().stream()
+                    .anyMatch(stage -> stage.getChildCount() == 0);
+        } else {
+            return mMainStage.getChildCount() == 0 || mSideStage.getChildCount() == 0;
         }
-
-        return false;
     }
 
     /** Checks if `transition` is a pending enter-split transition. */
@@ -477,10 +538,19 @@
 
     @StageType
     int getStageOfTask(int taskId) {
-        if (mMainStage.containsTask(taskId)) {
-            return STAGE_TYPE_MAIN;
-        } else if (mSideStage.containsTask(taskId)) {
-            return STAGE_TYPE_SIDE;
+        if (enableFlexibleSplit()) {
+            StageTaskListener stageTaskListener = mStageOrderOperator.getActiveStages().stream()
+                    .filter(stage -> stage.containsTask(taskId))
+                    .findFirst().orElse(null);
+            if (stageTaskListener != null) {
+                return stageTaskListener.getId();
+            }
+        } else {
+            if (mMainStage.containsTask(taskId)) {
+                return STAGE_TYPE_MAIN;
+            } else if (mSideStage.containsTask(taskId)) {
+                return STAGE_TYPE_SIDE;
+            }
         }
 
         return STAGE_TYPE_UNDEFINED;
@@ -490,14 +560,22 @@
         if (mRootTaskInfo != null && mRootTaskInfo.taskId == taskId) {
             return true;
         }
-        return mMainStage.isRootTaskId(taskId) || mSideStage.isRootTaskId(taskId);
+        if (enableFlexibleSplit()) {
+            return mStageOrderOperator.getActiveStages().stream()
+                    .anyMatch((stage) -> stage.isRootTaskId(taskId));
+        } else {
+            return mMainStage.isRootTaskId(taskId) || mSideStage.isRootTaskId(taskId);
+        }
     }
 
     boolean moveToStage(ActivityManager.RunningTaskInfo task, @SplitPosition int stagePosition,
             WindowContainerTransaction wct) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "moveToStage: task=%d position=%d", task.taskId,
                 stagePosition);
-        prepareEnterSplitScreen(wct, task, stagePosition, false /* resizeAnim */);
+        // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+        //  specific cases in the future. Only focusing on parity with starting intent/task
+        prepareEnterSplitScreen(wct, task, stagePosition, false /* resizeAnim */,
+                SPLIT_INDEX_UNDEFINED);
         mSplitTransitions.startEnterTransition(TRANSIT_TO_FRONT, wct,
                 null, this,
                 isSplitScreenVisible()
@@ -595,11 +673,14 @@
      *                      same window container transaction as the starting of the intent.
      */
     void startTask(int taskId, @SplitPosition int position, @Nullable Bundle options,
-            @Nullable WindowContainerToken hideTaskToken) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "startTask: task=%d position=%d", taskId, position);
+            @Nullable WindowContainerToken hideTaskToken, @SplitIndex int index) {
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "startTask: task=%d position=%d index=%d",
+                taskId, position, index);
         mSplitRequest = new SplitRequest(taskId, position);
         final WindowContainerTransaction wct = new WindowContainerTransaction();
-        options = resolveStartStage(STAGE_TYPE_UNDEFINED, position, options, null /* wct */);
+        options = enableFlexibleSplit()
+                ? resolveStartStageForIndex(options, null /*wct*/, index)
+                : resolveStartStage(STAGE_TYPE_UNDEFINED, position, options, null /* wct */);
         if (hideTaskToken != null) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "Reordering hide-task to bottom");
             wct.reorder(hideTaskToken, false /* onTop */);
@@ -623,7 +704,7 @@
         // If split screen is not activated, we're expecting to open a pair of apps to split.
         final int extraTransitType = isSplitActive()
                 ? TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE : TRANSIT_SPLIT_SCREEN_PAIR_OPEN;
-        prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering);
+        prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering, index);
 
         mSplitTransitions.startEnterTransition(TRANSIT_TO_FRONT, wct, null, this,
                 extraTransitType, !mIsDropEntering);
@@ -635,13 +716,16 @@
      *                      same window container transaction as the starting of the intent.
      */
     void startIntent(PendingIntent intent, Intent fillInIntent, @SplitPosition int position,
-            @Nullable Bundle options, @Nullable WindowContainerToken hideTaskToken) {
+            @Nullable Bundle options, @Nullable WindowContainerToken hideTaskToken,
+            @SplitIndex int index) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "startIntent: intent=%s position=%d", intent.getIntent(),
                 position);
         mSplitRequest = new SplitRequest(intent.getIntent(), position);
 
         final WindowContainerTransaction wct = new WindowContainerTransaction();
-        options = resolveStartStage(STAGE_TYPE_UNDEFINED, position, options, null /* wct */);
+        options = enableFlexibleSplit()
+                ? resolveStartStageForIndex(options, null /*wct*/, index)
+                : resolveStartStage(STAGE_TYPE_UNDEFINED, position, options, null /* wct */);
         if (hideTaskToken != null) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "Reordering hide-task to bottom");
             wct.reorder(hideTaskToken, false /* onTop */);
@@ -666,13 +750,17 @@
         // If split screen is not activated, we're expecting to open a pair of apps to split.
         final int extraTransitType = isSplitActive()
                 ? TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE : TRANSIT_SPLIT_SCREEN_PAIR_OPEN;
-        prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering);
+        prepareEnterSplitScreen(wct, null /* taskInfo */, position, !mIsDropEntering, index);
 
         mSplitTransitions.startEnterTransition(TRANSIT_TO_FRONT, wct, null, this,
                 extraTransitType, !mIsDropEntering);
     }
 
-    /** Starts 2 tasks in one transition. */
+    /**
+     * Starts 2 tasks in one transition.
+     * @param taskId1 starts in the mSideStage
+     * @param taskId2 starts in the mainStage #startWithTask()
+     */
     void startTasks(int taskId1, @Nullable Bundle options1, int taskId2, @Nullable Bundle options2,
             @SplitPosition int splitPosition, @PersistentSnapPosition int snapPosition,
             @Nullable RemoteTransition remoteTransition, InstanceId instanceId) {
@@ -687,11 +775,19 @@
 
         setSideStagePosition(splitPosition, wct);
         options1 = options1 != null ? options1 : new Bundle();
-        addActivityOptions(options1, mSideStage);
+        StageTaskListener stageForTask1;
+        if (enableFlexibleSplit()) {
+            stageForTask1 = mStageOrderOperator.getStageForLegacyPosition(splitPosition,
+                    true /*checkAllStagesIfNotActive*/);
+        } else {
+            stageForTask1 = mSideStage;
+        }
+        addActivityOptions(options1, stageForTask1);
         prepareTasksForSplitScreen(new int[] {taskId1, taskId2}, wct);
         wct.startTask(taskId1, options1);
 
-        startWithTask(wct, taskId2, options2, snapPosition, remoteTransition, instanceId);
+        startWithTask(wct, taskId2, options2, snapPosition, remoteTransition, instanceId,
+                splitPosition);
     }
 
     /** Start an intent and a task to a split pair in one transition. */
@@ -721,7 +817,8 @@
         wct.sendPendingIntent(pendingIntent, fillInIntent, options1);
         prepareTasksForSplitScreen(new int[] {taskId}, wct);
 
-        startWithTask(wct, taskId, options2, snapPosition, remoteTransition, instanceId);
+        startWithTask(wct, taskId, options2, snapPosition, remoteTransition, instanceId,
+                splitPosition);
     }
 
     /**
@@ -765,7 +862,8 @@
         wct.startShortcut(mContext.getPackageName(), shortcutInfo, options1);
         prepareTasksForSplitScreen(new int[] {taskId}, wct);
 
-        startWithTask(wct, taskId, options2, snapPosition, remoteTransition, instanceId);
+        startWithTask(wct, taskId, options2, snapPosition, remoteTransition, instanceId,
+                splitPosition);
     }
 
     /**
@@ -795,11 +893,14 @@
      */
     private void startWithTask(WindowContainerTransaction wct, int mainTaskId,
             @Nullable Bundle mainOptions, @PersistentSnapPosition int snapPosition,
-            @Nullable RemoteTransition remoteTransition, InstanceId instanceId) {
+            @Nullable RemoteTransition remoteTransition, InstanceId instanceId,
+            @SplitPosition int splitPosition) {
         if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
-            activateSplit(wct, false /* reparentToTop */);
+            // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+            //  specific cases in the future. Only focusing on parity with starting intent/task
+            activateSplit(wct, false /* reparentToTop */, SPLIT_INDEX_UNDEFINED);
         }
         mSplitLayout.setDivideRatio(snapPosition);
         updateWindowBounds(mSplitLayout, wct);
@@ -807,10 +908,19 @@
         wct.setReparentLeafTaskIfRelaunch(mRootTaskInfo.token,
                 false /* reparentLeafTaskIfRelaunch */);
         setRootForceTranslucent(false, wct);
-
+        // All callers of this method set the correct activity options on mSideStage,
+        // so we choose the opposite stage for this method
+        StageTaskListener stage;
+        if (enableFlexibleSplit()) {
+            stage = mStageOrderOperator
+                    .getStageForLegacyPosition(reverseSplitPosition(splitPosition),
+                            false /*checkAllStagesIfNotActive*/);
+        } else {
+            stage = mMainStage;
+        }
         // Make sure the launch options will put tasks in the corresponding split roots
         mainOptions = mainOptions != null ? mainOptions : new Bundle();
-        addActivityOptions(mainOptions, mMainStage);
+        addActivityOptions(mainOptions, stage);
 
         // Add task launch requests
         wct.startTask(mainTaskId, mainOptions);
@@ -866,7 +976,9 @@
         if (!isSplitActive()) {
             // Build a request WCT that will launch both apps such that task 0 is on the main stage
             // while task 1 is on the side stage.
-            activateSplit(wct, false /* reparentToTop */);
+            // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+            //  specific cases in the future. Only focusing on parity with starting intent/task
+            activateSplit(wct, false /* reparentToTop */, SPLIT_INDEX_UNDEFINED);
         }
 
         setSideStagePosition(splitPosition, wct);
@@ -974,6 +1086,31 @@
         mSideStage.evictInvisibleChildren(wct);
     }
 
+    /**
+     * @param index for the new stage that will be opening. Ex. if app is dragged to
+     *              index=1, then this will tell the stage at index=1 to launch the task
+     *              in the wct in that stage. This doesn't verify that the non-specified
+     *              indices' stages have their tasks correctly set/re-parented.
+     */
+    Bundle resolveStartStageForIndex(@Nullable Bundle options,
+            @Nullable WindowContainerTransaction wct,
+            @SplitIndex int index) {
+        StageTaskListener oppositeStage;
+        if (index == SPLIT_INDEX_UNDEFINED) {
+            // Arbitrarily choose a stage
+            oppositeStage = mStageOrderOperator.getStageForIndex(SPLIT_INDEX_1);
+        } else {
+            oppositeStage = mStageOrderOperator.getStageForIndex(index);
+        }
+        if (options == null) {
+            options = new Bundle();
+        }
+        updateStageWindowBoundsForIndex(wct, index);
+        addActivityOptions(options, oppositeStage);
+
+        return options;
+    }
+
     Bundle resolveStartStage(@StageType int stage, @SplitPosition int position,
             @Nullable Bundle options, @Nullable WindowContainerTransaction wct) {
         switch (stage) {
@@ -1041,20 +1178,35 @@
             return INVALID_TASK_ID;
         }
 
-        return mSideStagePosition == splitPosition
-                ? mSideStage.getTopVisibleChildTaskId()
-                : mMainStage.getTopVisibleChildTaskId();
+        if (enableFlexibleSplit()) {
+            StageTaskListener stage = mStageOrderOperator.getStageForLegacyPosition(splitPosition,
+                    true /*checkAllStagesIfNotActive*/);
+            return stage != null ? stage.getTopVisibleChildTaskId() : INVALID_TASK_ID;
+        } else {
+            return mSideStagePosition == splitPosition
+                    ? mSideStage.getTopVisibleChildTaskId()
+                    : mMainStage.getTopVisibleChildTaskId();
+        }
     }
 
     void switchSplitPosition(String reason) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "switchSplitPosition");
         final SurfaceControl.Transaction t = mTransactionPool.acquire();
         mTempRect1.setEmpty();
-        final StageTaskListener topLeftStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mSideStage : mMainStage;
-        final StageTaskListener bottomRightStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mMainStage : mSideStage;
-
+        final StageTaskListener topLeftStage;
+        final StageTaskListener bottomRightStage;
+        if (enableFlexibleSplit()) {
+            topLeftStage = mStageOrderOperator.getStageForLegacyPosition(SPLIT_POSITION_TOP_OR_LEFT,
+                            false /*checkAllStagesIfNotActive*/);
+            bottomRightStage = mStageOrderOperator
+                    .getStageForLegacyPosition(SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                    false /*checkAllStagesIfNotActive*/);
+        } else {
+            topLeftStage =
+                    mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mSideStage : mMainStage;
+            bottomRightStage =
+                    mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mMainStage : mSideStage;
+        }
         // Don't allow windows or divider to be focused during animation (mRootTaskInfo is the
         // parent of all 3 leaves). We don't want the user to be able to tap and focus a window
         // while it is moving across the screen, because granting focus also recalculates the
@@ -1091,9 +1243,13 @@
                 });
 
         ProtoLog.v(WM_SHELL_SPLIT_SCREEN, "Switch split position: %s", reason);
-        mLogger.logSwap(getMainStagePosition(), mMainStage.getTopChildTaskUid(),
-                getSideStagePosition(), mSideStage.getTopChildTaskUid(),
-                mSplitLayout.isLeftRightSplit());
+        if (enableFlexibleSplit()) {
+            // TODO(b/374825718) update logging for 2+ apps
+        } else {
+            mLogger.logSwap(getMainStagePosition(), mMainStage.getTopChildTaskUid(),
+                    getSideStagePosition(), mSideStage.getTopChildTaskUid(),
+                    mSplitLayout.isLeftRightSplit());
+        }
     }
 
     void setSideStagePosition(@SplitPosition int sideStagePosition,
@@ -1101,8 +1257,26 @@
         if (mSideStagePosition == sideStagePosition) return;
         mSideStagePosition = sideStagePosition;
         sendOnStagePositionChanged();
+        StageTaskListener stage = enableFlexibleSplit()
+                ? mStageOrderOperator.getStageForLegacyPosition(mSideStagePosition,
+                true /*checkAllStagesIfNotActive*/)
+                : mSideStage;
 
-        if (mSideStage.mVisible) {
+        if (stage.mVisible) {
+            if (wct == null) {
+                // onLayoutChanged builds/applies a wct with the contents of updateWindowBounds.
+                onLayoutSizeChanged(mSplitLayout);
+            } else {
+                updateWindowBounds(mSplitLayout, wct);
+                sendOnBoundsChanged();
+            }
+        }
+    }
+
+    private void updateStageWindowBoundsForIndex(@Nullable WindowContainerTransaction wct,
+            @SplitIndex int index) {
+        StageTaskListener stage = mStageOrderOperator.getStageForIndex(index);
+        if (stage.mVisible) {
             if (wct == null) {
                 // onLayoutChanged builds/applies a wct with the contents of updateWindowBounds.
                 onLayoutSizeChanged(mSplitLayout);
@@ -1152,10 +1326,17 @@
     void recordLastActiveStage() {
         if (!isSplitActive() || !isSplitScreenVisible()) {
             mLastActiveStage = STAGE_TYPE_UNDEFINED;
-        } else if (mMainStage.isFocused()) {
-            mLastActiveStage = STAGE_TYPE_MAIN;
-        } else if (mSideStage.isFocused()) {
-            mLastActiveStage = STAGE_TYPE_SIDE;
+        } else if (enableFlexibleSplit()) {
+            mStageOrderOperator.getActiveStages().stream()
+                    .filter(StageTaskListener::isFocused)
+                    .findFirst()
+                    .ifPresent(stage -> mLastActiveStage = stage.getId());
+        } else {
+            if (mMainStage.isFocused()) {
+                mLastActiveStage = STAGE_TYPE_MAIN;
+            } else if (mSideStage.isFocused()) {
+                mLastActiveStage = STAGE_TYPE_SIDE;
+            }
         }
     }
 
@@ -1212,7 +1393,7 @@
         mSplitLayout.getInvisibleBounds(mTempRect1);
         if (childrenToTop == null || childrenToTop.getTopVisibleChildTaskId() == INVALID_TASK_ID) {
             mSideStage.removeAllTasks(wct, false /* toTop */);
-            deactivateSplit(wct, false /* reparentToTop */);
+            deactivateSplit(wct, STAGE_TYPE_UNDEFINED);
             wct.reorder(mRootTaskInfo.token, false /* onTop */);
             setRootForceTranslucent(true, wct);
             wct.setBounds(mSideStage.mRootTaskInfo.token, mTempRect1);
@@ -1241,7 +1422,7 @@
                 childrenToTop.fadeOutDecor(() -> {
                     WindowContainerTransaction finishedWCT = new WindowContainerTransaction();
                     mIsExiting = false;
-                    deactivateSplit(finishedWCT, childrenToTop == mMainStage /* reparentToTop */);
+                    deactivateSplit(finishedWCT, childrenToTop.getId());
                     mSideStage.removeAllTasks(finishedWCT, childrenToTop == mSideStage /* toTop */);
                     finishedWCT.reorder(mRootTaskInfo.token, false /* toTop */);
                     setRootForceTranslucent(true, finishedWCT);
@@ -1369,8 +1550,13 @@
         mRecentTasks.ifPresent(recentTasks -> {
             // Notify recents if we are exiting in a way that breaks the pair, and disable further
             // updates to splits in the recents until we enter split again
-            mMainStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
-            mSideStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
+            if (enableFlexibleSplit()) {
+                runForActiveStages((stage) ->
+                        stage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId)));
+            } else {
+                mMainStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
+                mSideStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
+            }
         });
         logExit(exitReason);
     }
@@ -1383,15 +1569,22 @@
     void prepareExitSplitScreen(@StageType int stageToTop,
             @NonNull WindowContainerTransaction wct) {
         if (!isSplitActive()) return;
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareExitSplitScreen: stageToTop=%d", stageToTop);
-        mSideStage.removeAllTasks(wct, stageToTop == STAGE_TYPE_SIDE);
-        deactivateSplit(wct, stageToTop == STAGE_TYPE_MAIN);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareExitSplitScreen: stageToTop=%s",
+                stageTypeToString(stageToTop));
+        if (enableFlexibleSplit()) {
+            mStageOrderOperator.getActiveStages().stream()
+                    .filter(stage -> stage.getId() != stageToTop)
+                    .forEach(stage -> stage.removeAllTasks(wct, false /*toTop*/));
+        } else {
+            mSideStage.removeAllTasks(wct, stageToTop == STAGE_TYPE_SIDE);
+        }
+        deactivateSplit(wct, stageToTop);
     }
 
     private void prepareEnterSplitScreen(WindowContainerTransaction wct) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareEnterSplitScreen");
         prepareEnterSplitScreen(wct, null /* taskInfo */, SPLIT_POSITION_UNDEFINED,
-                !mIsDropEntering);
+                !mIsDropEntering, SPLIT_INDEX_UNDEFINED);
     }
 
     /**
@@ -1400,7 +1593,7 @@
      */
     void prepareEnterSplitScreen(WindowContainerTransaction wct,
             @Nullable ActivityManager.RunningTaskInfo taskInfo, @SplitPosition int startPosition,
-            boolean resizeAnim) {
+            boolean resizeAnim, @SplitIndex int index) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareEnterSplitScreen: position=%d resize=%b",
                 startPosition, resizeAnim);
         onSplitScreenEnter();
@@ -1412,7 +1605,7 @@
         if (isSplitActive()) {
             prepareBringSplit(wct, taskInfo, startPosition, resizeAnim);
         } else {
-            prepareActiveSplit(wct, taskInfo, startPosition, resizeAnim);
+            prepareActiveSplit(wct, taskInfo, startPosition, resizeAnim, index);
         }
     }
 
@@ -1433,14 +1626,22 @@
             if (!mSkipEvictingMainStageChildren) {
                 mMainStage.evictAllChildren(wct);
             }
-            mMainStage.reparentTopTask(wct);
+            // TODO(b/349828130) revisit bring split from BG to FG scenarios
+            if (enableFlexibleSplit()) {
+                runForActiveStages(stage -> stage.reparentTopTask(wct));
+            } else {
+                mMainStage.reparentTopTask(wct);
+            }
             prepareSplitLayout(wct, resizeAnim);
         }
     }
 
+    /**
+     * @param index The index that has already been assigned a stage
+     */
     private void prepareActiveSplit(WindowContainerTransaction wct,
             @Nullable ActivityManager.RunningTaskInfo taskInfo, @SplitPosition int startPosition,
-            boolean resizeAnim) {
+            boolean resizeAnim, @SplitIndex int index) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "prepareActiveSplit: task=%d isSplitVisible=%b",
                 taskInfo != null ? taskInfo.taskId : -1, isSplitScreenVisible());
         // We handle split visibility itself on shell transition, but sometimes we didn't
@@ -1450,7 +1651,7 @@
             setSideStagePosition(startPosition, wct);
             mSideStage.addTask(taskInfo, wct);
         }
-        activateSplit(wct, true /* reparentToTop */);
+        activateSplit(wct, true /* reparentToTop */, index);
         prepareSplitLayout(wct, resizeAnim);
     }
 
@@ -1477,8 +1678,13 @@
     void finishEnterSplitScreen(SurfaceControl.Transaction finishT) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "finishEnterSplitScreen");
         mSplitLayout.update(null, true /* resetImePosition */);
-        mMainStage.getSplitDecorManager().inflate(mContext, mMainStage.mRootLeash);
-        mSideStage.getSplitDecorManager().inflate(mContext, mSideStage.mRootLeash);
+        if (enableFlexibleSplit()) {
+            runForActiveStages((stage) ->
+                    stage.getSplitDecorManager().inflate(mContext, stage.mRootLeash));
+        } else {
+            mMainStage.getSplitDecorManager().inflate(mContext, mMainStage.mRootLeash);
+            mSideStage.getSplitDecorManager().inflate(mContext, mSideStage.mRootLeash);
+        }
         setDividerVisibility(true, finishT);
         // Ensure divider surface are re-parented back into the hierarchy at the end of the
         // transition. See Transition#buildFinishTransaction for more detail.
@@ -1492,6 +1698,10 @@
         mSplitRequest = null;
         updateRecentTasksSplitPair();
 
+        if (enableFlexibleSplit()) {
+            // TODO(b/374825718) log 2+ apps
+            return;
+        }
         mLogger.logEnter(mSplitLayout.getDividerPositionAsFraction(),
                 getMainStagePosition(), mMainStage.getTopChildTaskUid(),
                 getSideStagePosition(), mSideStage.getTopChildTaskUid(),
@@ -1503,6 +1713,15 @@
         outBottomOrRightBounds.set(mSplitLayout.getBounds2());
     }
 
+    private void runForActiveStages(Consumer<StageTaskListener> consumer) {
+        mStageOrderOperator.getActiveStages().forEach(consumer);
+    }
+
+    private boolean runForActiveStagesAllMatch(Predicate<StageTaskListener> predicate) {
+        List<StageTaskListener> activeStages = mStageOrderOperator.getActiveStages();
+        return !activeStages.isEmpty() && activeStages.stream().allMatch(predicate);
+    }
+
     @SplitPosition
     int getSplitPosition(int taskId) {
         if (mSideStage.getTopVisibleChildTaskId() == taskId) {
@@ -1516,6 +1735,9 @@
     private void addActivityOptions(Bundle opts, @Nullable StageTaskListener launchTarget) {
         ActivityOptions options = ActivityOptions.fromBundle(opts);
         if (launchTarget != null) {
+            ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
+                    "addActivityOptions setting launch root for stage=%s",
+                    stageTypeToString(launchTarget.getId()));
             options.setLaunchRootTask(launchTarget.mRootTaskInfo.token);
         }
         // Put BAL flags to avoid activity start aborted. Otherwise, flows like shortcut to split
@@ -1559,8 +1781,15 @@
             listener.onSplitBoundsChanged(mSplitLayout.getRootBounds(), getMainStageBounds(),
                     getSideStageBounds());
         }
-        mSideStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_SIDE);
-        mMainStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_MAIN);
+        if (enableFlexibleSplit()) {
+            // TODO(b/349828130) replace w/ stageID
+            mStageOrderOperator.getAllStages().forEach(
+                    stage -> stage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_UNDEFINED)
+            );
+        } else {
+            mSideStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_SIDE);
+            mMainStage.onSplitScreenListenerRegistered(listener, STAGE_TYPE_MAIN);
+        }
     }
 
     private void sendOnStagePositionChanged() {
@@ -1584,17 +1813,25 @@
             boolean present, boolean visible) {
         int stage;
         if (present) {
-            stage = stageListener == mSideStage ? STAGE_TYPE_SIDE : STAGE_TYPE_MAIN;
+            if (enableFlexibleSplit()) {
+                stage = stageListener.getId();
+            } else {
+                stage = stageListener == mSideStage ? STAGE_TYPE_SIDE : STAGE_TYPE_MAIN;
+            }
         } else {
             // No longer on any stage
             stage = STAGE_TYPE_UNDEFINED;
         }
-        if (stage == STAGE_TYPE_MAIN) {
-            mLogger.logMainStageAppChange(getMainStagePosition(), mMainStage.getTopChildTaskUid(),
-                    mSplitLayout.isLeftRightSplit());
-        } else if (stage == STAGE_TYPE_SIDE) {
-            mLogger.logSideStageAppChange(getSideStagePosition(), mSideStage.getTopChildTaskUid(),
-                    mSplitLayout.isLeftRightSplit());
+        if (!enableFlexibleSplit()) {
+            if (stage == STAGE_TYPE_MAIN) {
+                mLogger.logMainStageAppChange(getMainStagePosition(),
+                        mMainStage.getTopChildTaskUid(),
+                        mSplitLayout.isLeftRightSplit());
+            } else if (stage == STAGE_TYPE_SIDE) {
+                mLogger.logSideStageAppChange(getSideStagePosition(),
+                        mSideStage.getTopChildTaskUid(),
+                        mSplitLayout.isLeftRightSplit());
+            }
         }
         if (present) {
             updateRecentTasksSplitPair();
@@ -1622,17 +1859,35 @@
         mRecentTasks.ifPresent(recentTasks -> {
             Rect topLeftBounds = mSplitLayout.getBounds1();
             Rect bottomRightBounds = mSplitLayout.getBounds2();
-            int mainStageTopTaskId = mMainStage.getTopVisibleChildTaskId();
-            int sideStageTopTaskId = mSideStage.getTopVisibleChildTaskId();
+            int sideStageTopTaskId;
+            int mainStageTopTaskId;
+            if (enableFlexibleSplit()) {
+                List<StageTaskListener> activeStages = mStageOrderOperator.getActiveStages();
+                if (activeStages.size() != 2) {
+                    sideStageTopTaskId = mainStageTopTaskId = INVALID_TASK_ID;
+                } else {
+                    // doesn't matter which one we assign to? What matters is the order of 0 and 1?
+                    mainStageTopTaskId = activeStages.get(0).getTopVisibleChildTaskId();
+                    sideStageTopTaskId = activeStages.get(1).getTopVisibleChildTaskId();
+                }
+            } else {
+                mainStageTopTaskId = mMainStage.getTopVisibleChildTaskId();
+                sideStageTopTaskId= mSideStage.getTopVisibleChildTaskId();
+            }
             boolean sideStageTopLeft = mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT;
             int leftTopTaskId;
             int rightBottomTaskId;
-            if (sideStageTopLeft) {
-                leftTopTaskId = sideStageTopTaskId;
-                rightBottomTaskId = mainStageTopTaskId;
-            } else {
+            if (enableFlexibleSplit()) {
                 leftTopTaskId = mainStageTopTaskId;
                 rightBottomTaskId = sideStageTopTaskId;
+            } else {
+                if (sideStageTopLeft) {
+                    leftTopTaskId = sideStageTopTaskId;
+                    rightBottomTaskId = mainStageTopTaskId;
+                } else {
+                    leftTopTaskId = mainStageTopTaskId;
+                    rightBottomTaskId = sideStageTopTaskId;
+                }
             }
 
             if (Flags.enableFlexibleTwoAppSplit()) {
@@ -1741,29 +1996,59 @@
     @VisibleForTesting
     @Override
     public void onRootTaskAppeared() {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onRootTaskAppeared: rootTask=%s mainRoot=%b sideRoot=%b",
-                mRootTaskInfo, mMainStage.mHasRootTask, mSideStage.mHasRootTask);
+        if (enableFlexibleSplit()) {
+            ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onRootTaskAppeared: rootTask=%s",
+                mRootTaskInfo);
+            mStageOrderOperator.getAllStages().forEach(stage -> {
+                ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
+                        "    onRootStageAppeared stageId=%s hasRoot=%b",
+                        stageTypeToString(stage.getId()), stage.mHasRootTask);
+            });
+        } else {
+            ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
+                    "onRootTaskAppeared: rootTask=%s mainRoot=%b sideRoot=%b",
+                    mRootTaskInfo, mMainStage.mHasRootTask, mSideStage.mHasRootTask);
+        }
+        boolean notAllStagesHaveRootTask;
+        if (enableFlexibleSplit()) {
+            notAllStagesHaveRootTask = mStageOrderOperator.getAllStages().stream()
+                    .anyMatch((stage) -> !stage.mHasRootTask);
+        } else {
+            notAllStagesHaveRootTask = !mMainStage.mHasRootTask
+                    || !mSideStage.mHasRootTask;
+        }
         // Wait unit all root tasks appeared.
-        if (mRootTaskInfo == null
-                || !mMainStage.mHasRootTask
-                || !mSideStage.mHasRootTask) {
+        if (mRootTaskInfo == null || notAllStagesHaveRootTask) {
             return;
         }
 
         final WindowContainerTransaction wct = new WindowContainerTransaction();
-        wct.reparent(mMainStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
-        wct.reparent(mSideStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
+        if (enableFlexibleSplit()) {
+            mStageOrderOperator.getAllStages().forEach(stage ->
+                    wct.reparent(stage.mRootTaskInfo.token, mRootTaskInfo.token, true));
+        } else {
+            wct.reparent(mMainStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
+            wct.reparent(mSideStage.mRootTaskInfo.token, mRootTaskInfo.token, true);
+        }
 
-        // Make the stages adjacent to each other so they occlude what's behind them.
-        wct.setAdjacentRoots(mMainStage.mRootTaskInfo.token, mSideStage.mRootTaskInfo.token);
         setRootForceTranslucent(true, wct);
-        mSplitLayout.getInvisibleBounds(mTempRect1);
-        wct.setBounds(mSideStage.mRootTaskInfo.token, mTempRect1);
+        if (!enableFlexibleSplit()) {
+            //TODO(b/373709676) Need to figure out how adjacentRoots work for flex split
+
+            // Make the stages adjacent to each other so they occlude what's behind them.
+            wct.setAdjacentRoots(mMainStage.mRootTaskInfo.token, mSideStage.mRootTaskInfo.token);
+            mSplitLayout.getInvisibleBounds(mTempRect1);
+            wct.setBounds(mSideStage.mRootTaskInfo.token, mTempRect1);
+        }
         mSyncQueue.queue(wct);
-        mSyncQueue.runInSync(t -> {
-            t.setPosition(mSideStage.mRootLeash, mTempRect1.left, mTempRect1.top);
-        });
-        mLaunchAdjacentController.setLaunchAdjacentRoot(mSideStage.mRootTaskInfo.token);
+        if (!enableFlexibleSplit()) {
+            mSyncQueue.runInSync(t -> {
+                t.setPosition(mSideStage.mRootLeash, mTempRect1.left, mTempRect1.top);
+            });
+            mLaunchAdjacentController.setLaunchAdjacentRoot(mSideStage.mRootTaskInfo.token);
+        } else {
+            // TODO(b/373709676) Need to figure out how adjacentRoots work for flex split
+        }
     }
 
     @Override
@@ -1958,15 +2243,21 @@
     }
 
     @Override
-    public void onSnappedToDismiss(boolean bottomOrRight, @ExitReason int exitReason) {
+    public void onSnappedToDismiss(boolean closedBottomRightStage, @ExitReason int exitReason) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onSnappedToDismiss: bottomOrRight=%b reason=%s",
-                bottomOrRight, exitReasonToString(exitReason));
-        final boolean mainStageToTop =
-                bottomOrRight ? mSideStagePosition == SPLIT_POSITION_BOTTOM_OR_RIGHT
+                closedBottomRightStage, exitReasonToString(exitReason));
+        boolean mainStageToTop =
+                closedBottomRightStage ? mSideStagePosition == SPLIT_POSITION_BOTTOM_OR_RIGHT
                         : mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT;
-        final StageTaskListener toTopStage = mainStageToTop ? mMainStage : mSideStage;
-
-        final int dismissTop = mainStageToTop ? STAGE_TYPE_MAIN : STAGE_TYPE_SIDE;
+        StageTaskListener toTopStage = mainStageToTop ? mMainStage : mSideStage;
+        int dismissTop = mainStageToTop ? STAGE_TYPE_MAIN : STAGE_TYPE_SIDE;
+        if (enableFlexibleSplit()) {
+            toTopStage = mStageOrderOperator.getStageForLegacyPosition(closedBottomRightStage
+                            ? SPLIT_POSITION_TOP_OR_LEFT
+                            : SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                    false /*checkAllStagesIfNotActive*/);
+            dismissTop = toTopStage.getId();
+        }
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         toTopStage.resetBounds(wct);
         prepareExitSplitScreen(dismissTop, wct);
@@ -1998,8 +2289,21 @@
         updateSurfaceBounds(layout, t, shouldUseParallaxEffect);
         getMainStageBounds(mTempRect1);
         getSideStageBounds(mTempRect2);
-        mMainStage.onResizing(mTempRect1, mTempRect2, t, offsetX, offsetY, mShowDecorImmediately);
-        mSideStage.onResizing(mTempRect2, mTempRect1, t, offsetX, offsetY, mShowDecorImmediately);
+        if (enableFlexibleSplit()) {
+            StageTaskListener ltStage =
+                    mStageOrderOperator.getStageForLegacyPosition(SPLIT_POSITION_TOP_OR_LEFT,
+                    false /*checkAllStagesIfNotActive*/);
+            StageTaskListener brStage =
+                    mStageOrderOperator.getStageForLegacyPosition(SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                            false /*checkAllStagesIfNotActive*/);
+            ltStage.onResizing(mTempRect1, mTempRect2, t, offsetX, offsetY, mShowDecorImmediately);
+            brStage.onResizing(mTempRect2, mTempRect1, t, offsetX, offsetY, mShowDecorImmediately);
+        } else {
+            mMainStage.onResizing(mTempRect1, mTempRect2, t, offsetX, offsetY,
+                    mShowDecorImmediately);
+            mSideStage.onResizing(mTempRect2, mTempRect1, t, offsetX, offsetY,
+                    mShowDecorImmediately);
+        }
         t.apply();
         mTransactionPool.release(t);
     }
@@ -2015,19 +2319,33 @@
         if (!sizeChanged) {
             // We still need to resize on decor for ensure all current status clear.
             final SurfaceControl.Transaction t = mTransactionPool.acquire();
-            mMainStage.onResized(t);
-            mSideStage.onResized(t);
+            if (enableFlexibleSplit()) {
+                runForActiveStages(stage -> stage.onResized(t));
+            } else {
+                mMainStage.onResized(t);
+                mSideStage.onResized(t);
+            }
             mTransactionPool.release(t);
             return;
         }
-
+        List<SplitDecorManager> decorManagers = new ArrayList<>();
+        SplitDecorManager mainDecor = null;
+        SplitDecorManager sideDecor = null;
+        if (enableFlexibleSplit()) {
+            decorManagers = mStageOrderOperator.getActiveStages().stream()
+                    .map(StageTaskListener::getSplitDecorManager)
+                    .toList();
+        } else {
+            mainDecor = mMainStage.getSplitDecorManager();
+            sideDecor = mSideStage.getSplitDecorManager();
+        }
         sendOnBoundsChanged();
         mSplitLayout.setDividerInteractive(false, false, "onSplitResizeStart");
         mSplitTransitions.startResizeTransition(wct, this, (aborted) -> {
             mSplitLayout.setDividerInteractive(true, false, "onSplitResizeConsumed");
         }, (finishWct, t) -> {
             mSplitLayout.setDividerInteractive(true, false, "onSplitResizeFinish");
-        }, mMainStage.getSplitDecorManager(), mSideStage.getSplitDecorManager());
+        }, mainDecor, sideDecor, decorManagers);
 
         if (Flags.enableFlexibleTwoAppSplit()) {
             switch (layout.calculateCurrentSnapPosition()) {
@@ -2054,10 +2372,23 @@
      * @return true if stage bounds actually .
      */
     private boolean updateWindowBounds(SplitLayout layout, WindowContainerTransaction wct) {
-        final StageTaskListener topLeftStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mSideStage : mMainStage;
-        final StageTaskListener bottomRightStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mMainStage : mSideStage;
+        final StageTaskListener topLeftStage;
+        final StageTaskListener bottomRightStage;
+        if (enableFlexibleSplit()) {
+            topLeftStage = mStageOrderOperator
+                    .getStageForLegacyPosition(SPLIT_POSITION_TOP_OR_LEFT,
+                            true /*checkAllStagesIfNotActive*/);
+            bottomRightStage = mStageOrderOperator
+                    .getStageForLegacyPosition(SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                            true /*checkAllStagesIfNotActive*/);
+        } else {
+            topLeftStage = mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
+                    ? mSideStage
+                    : mMainStage;
+            bottomRightStage = mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
+                    ? mMainStage
+                    : mSideStage;
+        }
         boolean updated = layout.applyTaskChanges(wct, topLeftStage.mRootTaskInfo,
                 bottomRightStage.mRootTaskInfo);
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "updateWindowBounds: topLeftStage=%s bottomRightStage=%s",
@@ -2067,10 +2398,23 @@
 
     void updateSurfaceBounds(@Nullable SplitLayout layout, @NonNull SurfaceControl.Transaction t,
             boolean applyResizingOffset) {
-        final StageTaskListener topLeftStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mSideStage : mMainStage;
-        final StageTaskListener bottomRightStage =
-                mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT ? mMainStage : mSideStage;
+        final StageTaskListener topLeftStage;
+        final StageTaskListener bottomRightStage;
+        if (enableFlexibleSplit()) {
+            topLeftStage = mStageOrderOperator
+                    .getStageForLegacyPosition(SPLIT_POSITION_TOP_OR_LEFT,
+                            true /*checkAllStagesIfNotActive*/);
+            bottomRightStage = mStageOrderOperator
+                    .getStageForLegacyPosition(SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                            true /*checkAllStagesIfNotActive*/);
+        } else {
+            topLeftStage = mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
+                    ? mSideStage
+                    : mMainStage;
+            bottomRightStage = mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
+                    ? mMainStage
+                    : mSideStage;
+        }
         (layout != null ? layout : mSplitLayout).applySurfaceChanges(t, topLeftStage.mRootLeash,
                 bottomRightStage.mRootLeash, topLeftStage.mDimLayer, bottomRightStage.mDimLayer,
                 applyResizingOffset);
@@ -2085,10 +2429,22 @@
             return SPLIT_POSITION_UNDEFINED;
         }
 
-        if (mMainStage.containsToken(token)) {
-            return getMainStagePosition();
-        } else if (mSideStage.containsToken(token)) {
-            return getSideStagePosition();
+        if (enableFlexibleSplit()) {
+            // We could migrate to/return the new INDEX enums here since most callers just care that
+            // this value isn't SPLIT_POSITION_UNDEFINED, but
+            // ImePositionProcessor#getImeTargetPosition actually uses the leftTop/bottomRight value
+            StageTaskListener stageForToken = mStageOrderOperator.getAllStages().stream()
+                    .filter(stage -> stage.containsToken(token))
+                    .findFirst().orElse(null);
+            return stageForToken == null
+                    ? SPLIT_POSITION_UNDEFINED
+                    : mStageOrderOperator.getLegacyPositionForStage(stageForToken);
+        } else {
+            if (mMainStage.containsToken(token)) {
+                return getMainStagePosition();
+            } else if (mSideStage.containsToken(token)) {
+                return getSideStagePosition();
+            }
         }
 
         return SPLIT_POSITION_UNDEFINED;
@@ -2193,19 +2549,41 @@
                 ? mSplitLayout.getBounds2() : mSplitLayout.getBounds1();
     }
 
+    /**
+     * TODO(b/349828130) Currently the way this is being used is only to to get the bottomRight
+     *  stage. Eventually we'll need to rename and for now we'll repurpose the method to return
+     *  the bottomRight bounds under the flex split flag
+     */
     private void getSideStageBounds(Rect rect) {
-        if (mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT) {
+        if (enableFlexibleSplit()) {
+            // Split Layout doesn't actually keep track of the bounds based on the stage,
+            // it only knows that bounds1 is leftTop position and bounds2 is bottomRight position
+            // We'll then assume this method is to get bounds of bottomRight stage
+            mSplitLayout.getBounds2(rect);
+        }  else if (mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT) {
             mSplitLayout.getBounds1(rect);
         } else {
             mSplitLayout.getBounds2(rect);
         }
     }
 
+    /**
+     * TODO(b/349828130) Currently the way this is being used is only to to get the leftTop
+     *  stage. Eventually we'll need to rename and for now we'll repurpose the method to return
+     *  the leftTop bounds under the flex split flag
+     */
     private void getMainStageBounds(Rect rect) {
-        if (mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT) {
-            mSplitLayout.getBounds2(rect);
-        } else {
+        if (enableFlexibleSplit()) {
+            // Split Layout doesn't actually keep track of the bounds based on the stage,
+            // it only knows that bounds1 is leftTop position and bounds2 is bottomRight position
+            // We'll then assume this method is to get bounds of topLeft stage
             mSplitLayout.getBounds1(rect);
+        } else {
+            if (mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT) {
+                mSplitLayout.getBounds2(rect);
+            } else {
+                mSplitLayout.getBounds1(rect);
+            }
         }
     }
 
@@ -2214,14 +2592,23 @@
      * this task (yet) so this can also be used to identify which stage to put a task into.
      */
     private StageTaskListener getStageOfTask(ActivityManager.RunningTaskInfo taskInfo) {
-        // TODO(b/184679596): Find a way to either include task-org information in the transition,
-        //                    or synchronize task-org callbacks so we can use stage.containsTask
-        if (mMainStage.mRootTaskInfo != null
-                && taskInfo.parentTaskId == mMainStage.mRootTaskInfo.taskId) {
-            return mMainStage;
-        } else if (mSideStage.mRootTaskInfo != null
-                && taskInfo.parentTaskId == mSideStage.mRootTaskInfo.taskId) {
-            return mSideStage;
+        if (enableFlexibleSplit()) {
+            return mStageOrderOperator.getActiveStages().stream()
+                    .filter((stage) -> stage.mRootTaskInfo != null &&
+                            taskInfo.parentTaskId == stage.mRootTaskInfo.taskId
+                    )
+                    .findFirst()
+                    .orElse(null);
+        } else {
+            // TODO(b/184679596): Find a way to either include task-org information in the
+            //  transition, or synchronize task-org callbacks so we can use stage.containsTask
+            if (mMainStage.mRootTaskInfo != null
+                    && taskInfo.parentTaskId == mMainStage.mRootTaskInfo.taskId) {
+                return mMainStage;
+            } else if (mSideStage.mRootTaskInfo != null
+                    && taskInfo.parentTaskId == mSideStage.mRootTaskInfo.taskId) {
+                return mSideStage;
+            }
         }
         return null;
     }
@@ -2229,7 +2616,11 @@
     @StageType
     private int getStageType(StageTaskListener stage) {
         if (stage == null) return STAGE_TYPE_UNDEFINED;
-        return stage == mMainStage ? STAGE_TYPE_MAIN : STAGE_TYPE_SIDE;
+        if (enableFlexibleSplit()) {
+            return stage.getId();
+        } else {
+            return stage == mMainStage ? STAGE_TYPE_MAIN : STAGE_TYPE_SIDE;
+        }
     }
 
     @Override
@@ -2276,11 +2667,17 @@
         if (isSplitActive()) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "handleRequest: transition=%d split active",
                     request.getDebugId());
+            StageTaskListener primaryStage = enableFlexibleSplit()
+                    ? mStageOrderOperator.getActiveStages().get(0)
+                    : mMainStage;
+            StageTaskListener secondaryStage = enableFlexibleSplit()
+                    ? mStageOrderOperator.getActiveStages().get(1)
+                    : mSideStage;
             // Try to handle everything while in split-screen, so return a WCT even if it's empty.
             ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "  split is active so using split"
                             + "Transition to handle request. triggerTask=%d type=%s mainChildren=%d"
                             + " sideChildren=%d", triggerTask.taskId, transitTypeToString(type),
-                    mMainStage.getChildCount(), mSideStage.getChildCount());
+                    primaryStage.getChildCount(), secondaryStage.getChildCount());
             out = new WindowContainerTransaction();
             if (stage != null) {
                 if (isClosingType(type) && stage.getChildCount() == 1) {
@@ -2320,11 +2717,21 @@
                     // the remote handler.
                     return null;
                 }
-
-                if ((mMainStage.containsTask(triggerTask.taskId)
-                            && mMainStage.getChildCount() == 1)
-                        || (mSideStage.containsTask(triggerTask.taskId)
-                            && mSideStage.getChildCount() == 1)) {
+                boolean anyStageContainsSingleFullscreenTask;
+                if (enableFlexibleSplit()) {
+                    anyStageContainsSingleFullscreenTask =
+                            mStageOrderOperator.getActiveStages().stream()
+                                    .anyMatch(stageListener ->
+                                            stageListener.containsTask(triggerTask.taskId)
+                                                    && stageListener.getChildCount() == 1);
+                } else {
+                    anyStageContainsSingleFullscreenTask =
+                            (mMainStage.containsTask(triggerTask.taskId)
+                                    && mMainStage.getChildCount() == 1)
+                                    || (mSideStage.containsTask(triggerTask.taskId)
+                                    && mSideStage.getChildCount() == 1);
+                }
+                if (anyStageContainsSingleFullscreenTask) {
                     // A splitting task is opening to fullscreen causes one side of the split empty,
                     // so appends operations to exit split.
                     prepareExitSplitScreen(STAGE_TYPE_UNDEFINED, out);
@@ -2344,11 +2751,19 @@
                 // One of the cases above handled it
                 return out;
             } else if (isSplitScreenVisible()) {
+                boolean allStagesHaveChildren;
+                if (enableFlexibleSplit()) {
+                    allStagesHaveChildren = runForActiveStagesAllMatch(stageTaskListener ->
+                            stageTaskListener.getChildCount() != 0);
+                } else {
+                    allStagesHaveChildren = mMainStage.getChildCount() != 0
+                            && mSideStage.getChildCount() != 0;
+                }
                 // If split is visible, only defer handling this transition if it's launching
                 // adjacent while there is already a split pair -- this may trigger PIP and
                 // that should be handled by the mixed handler.
                 final boolean deferTransition = requestHasLaunchAdjacentFlag(request)
-                    && mMainStage.getChildCount() != 0 && mSideStage.getChildCount() != 0;
+                    && allStagesHaveChildren;
                 return !deferTransition ? out : null;
             }
             // Don't intercept the transition if we are not handling it as a part of one of the
@@ -2588,8 +3003,15 @@
             }
 
             final ArraySet<StageTaskListener> dismissStages = record.getShouldDismissedStage();
-            if (mMainStage.getChildCount() == 0 || mSideStage.getChildCount() == 0
-                    || dismissStages.size() == 1) {
+            boolean anyStageHasNoChildren;
+            if (enableFlexibleSplit()) {
+                anyStageHasNoChildren = mStageOrderOperator.getActiveStages().stream()
+                        .anyMatch(stage -> stage.getChildCount() == 0);
+            } else {
+                anyStageHasNoChildren = mMainStage.getChildCount() == 0
+                        || mSideStage.getChildCount() == 0;
+            }
+            if (anyStageHasNoChildren || dismissStages.size() == 1) {
                 // If the size of dismissStages == 1, one of the task is closed without prepare
                 // pending transition, which could happen if all activities were finished after
                 // finish top activity in a task, so the trigger task is null when handleRequest.
@@ -2735,24 +3157,48 @@
             shouldAnimate = startPendingDismissAnimation(
                     dismiss, info, startTransaction, finishTransaction);
             if (shouldAnimate && dismiss.mReason == EXIT_REASON_DRAG_DIVIDER) {
-                final StageTaskListener toTopStage =
-                        dismiss.mDismissTop == STAGE_TYPE_MAIN ? mMainStage : mSideStage;
+                StageTaskListener toTopStage;
+                if (enableFlexibleSplit()) {
+                    toTopStage = mStageOrderOperator.getAllStages().stream()
+                            .filter(stage -> stage.getId() == dismiss.mDismissTop)
+                            .findFirst().orElseThrow();
+                } else {
+                    toTopStage = dismiss.mDismissTop == STAGE_TYPE_MAIN ? mMainStage : mSideStage;
+                }
                 mSplitTransitions.playDragDismissAnimation(transition, info, startTransaction,
                         finishTransaction, finishCallback, toTopStage.mRootTaskInfo.token,
                         toTopStage.getSplitDecorManager(), mRootTaskInfo.token);
                 return true;
             }
         } else if (mSplitTransitions.isPendingResize(transition)) {
+            Map<WindowContainerToken, SplitDecorManager> tokenDecorMap = new HashMap<>();
+            if (enableFlexibleSplit()) {
+                runForActiveStages(stageTaskListener ->
+                        tokenDecorMap.put(stageTaskListener.mRootTaskInfo.getToken(),
+                                stageTaskListener.getSplitDecorManager()));
+            } else {
+                tokenDecorMap.put(mMainStage.mRootTaskInfo.getToken(),
+                        mMainStage.getSplitDecorManager());
+                tokenDecorMap.put(mSideStage.mRootTaskInfo.getToken(),
+                        mSideStage.getSplitDecorManager());
+            }
             mSplitTransitions.playResizeAnimation(transition, info, startTransaction,
-                    finishTransaction, finishCallback, mMainStage.mRootTaskInfo.token,
-                    mSideStage.mRootTaskInfo.token, mMainStage.getSplitDecorManager(),
-                    mSideStage.getSplitDecorManager());
+                    finishTransaction, finishCallback, tokenDecorMap);
             return true;
         }
         if (!shouldAnimate) return false;
 
+        WindowContainerToken mainToken;
+        WindowContainerToken sideToken;
+        if (enableFlexibleSplit()) {
+            mainToken = mStageOrderOperator.getActiveStages().get(0).mRootTaskInfo.token;
+            sideToken = mStageOrderOperator.getActiveStages().get(1).mRootTaskInfo.token;
+        } else {
+            mainToken = mMainStage.mRootTaskInfo.token;
+            sideToken = mSideStage.mRootTaskInfo.token;
+        }
         mSplitTransitions.playAnimation(transition, info, startTransaction, finishTransaction,
-                finishCallback, mMainStage.mRootTaskInfo.token, mSideStage.mRootTaskInfo.token,
+                finishCallback, mainToken, sideToken,
                 mRootTaskInfo.token);
         return true;
     }
@@ -2777,6 +3223,8 @@
         // First, verify that we actually have opened apps in both splits.
         TransitionInfo.Change mainChild = null;
         TransitionInfo.Change sideChild = null;
+        StageTaskListener firstAppStage = null;
+        StageTaskListener secondAppStage = null;
         final WindowContainerTransaction evictWct = new WindowContainerTransaction();
         for (int iC = 0; iC < info.getChanges().size(); ++iC) {
             final TransitionInfo.Change change = info.getChanges().get(iC);
@@ -2785,14 +3233,19 @@
             if (mPausingTasks.contains(taskInfo.taskId)) {
                 continue;
             }
-            final @StageType int stageType = getStageType(getStageOfTask(taskInfo));
-            if (mainChild == null && stageType == STAGE_TYPE_MAIN
+            StageTaskListener stage = getStageOfTask(taskInfo);
+            final @StageType int stageType = getStageType(stage);
+            if (mainChild == null
+                    && stageType == (enableFlexibleSplit() ? STAGE_TYPE_A : STAGE_TYPE_MAIN)
                     && (isOpeningType(change.getMode()) || change.getMode() == TRANSIT_CHANGE)) {
                 // Includes TRANSIT_CHANGE to cover reparenting top-most task to split.
                 mainChild = change;
-            } else if (sideChild == null && stageType == STAGE_TYPE_SIDE
+                firstAppStage = getStageOfTask(taskInfo);
+            } else if (sideChild == null
+                    && stageType == (enableFlexibleSplit() ? STAGE_TYPE_B : STAGE_TYPE_SIDE)
                     && (isOpeningType(change.getMode()) || change.getMode() == TRANSIT_CHANGE)) {
                 sideChild = change;
+                secondAppStage = stage;
             } else if (stageType != STAGE_TYPE_UNDEFINED && change.getMode() == TRANSIT_TO_BACK) {
                 // Collect all to back task's and evict them when transition finished.
                 evictWct.reparent(taskInfo.token, null /* parent */, false /* onTop */);
@@ -2848,9 +3301,9 @@
         // TODO(b/184679596): Find a way to either include task-org information in
         //                    the transition, or synchronize task-org callbacks.
         final boolean mainNotContainOpenTask =
-                mainChild != null && !mMainStage.containsTask(mainChild.getTaskInfo().taskId);
+                mainChild != null && !firstAppStage.containsTask(mainChild.getTaskInfo().taskId);
         final boolean sideNotContainOpenTask =
-                sideChild != null && !mSideStage.containsTask(sideChild.getTaskInfo().taskId);
+                sideChild != null && !secondAppStage.containsTask(sideChild.getTaskInfo().taskId);
         if (mainNotContainOpenTask) {
             Log.w(TAG, "Expected onTaskAppeared on " + mMainStage
                     + " to have been called with " + mainChild.getTaskInfo().taskId
@@ -2863,6 +3316,8 @@
         }
         final TransitionInfo.Change finalMainChild = mainChild;
         final TransitionInfo.Change finalSideChild = sideChild;
+        final StageTaskListener finalFirstAppStage = firstAppStage;
+        final StageTaskListener finalSecondAppStage = secondAppStage;
         enterTransition.setFinishedCallback((callbackWct, callbackT) -> {
             if (!enterTransition.mResizeAnim) {
                 // If resizing, we'll call notify at the end of the resizing animation (below)
@@ -2870,16 +3325,18 @@
             }
             if (finalMainChild != null) {
                 if (!mainNotContainOpenTask) {
-                    mMainStage.evictOtherChildren(callbackWct, finalMainChild.getTaskInfo().taskId);
+                    finalFirstAppStage.evictOtherChildren(callbackWct,
+                            finalMainChild.getTaskInfo().taskId);
                 } else {
-                    mMainStage.evictInvisibleChildren(callbackWct);
+                    finalFirstAppStage.evictInvisibleChildren(callbackWct);
                 }
             }
             if (finalSideChild != null) {
                 if (!sideNotContainOpenTask) {
-                    mSideStage.evictOtherChildren(callbackWct, finalSideChild.getTaskInfo().taskId);
+                    finalSecondAppStage.evictOtherChildren(callbackWct,
+                            finalSideChild.getTaskInfo().taskId);
                 } else {
-                    mSideStage.evictInvisibleChildren(callbackWct);
+                    finalSecondAppStage.evictInvisibleChildren(callbackWct);
                 }
             }
             if (!evictWct.isEmpty()) {
@@ -2961,8 +3418,10 @@
     public void onPipExpandToSplit(WindowContainerTransaction wct,
             ActivityManager.RunningTaskInfo taskInfo) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onPipExpandToSplit: task=%s", taskInfo);
+        // TODO(b/349828130) currently pass in index_undefined until we can revisit these
+        //  flex split + pip interactions in the future
         prepareEnterSplitScreen(wct, taskInfo, getActivateSplitPosition(taskInfo),
-                false /*resizeAnim*/);
+                false /*resizeAnim*/, SPLIT_INDEX_UNDEFINED);
 
         if (!isSplitScreenVisible() || mSplitRequest == null) {
             return;
@@ -3067,13 +3526,28 @@
         // Wait until after animation to update divider
 
         // Reset crops so they don't interfere with subsequent launches
-        t.setCrop(mMainStage.mRootLeash, null);
-        t.setCrop(mSideStage.mRootLeash, null);
+        if (enableFlexibleSplit()) {
+            runForActiveStages(stage -> t.setCrop(stage.mRootLeash, null /*crop*/));
+        } else {
+            t.setCrop(mMainStage.mRootLeash, null);
+            t.setCrop(mSideStage.mRootLeash, null);
+        }
         // Hide the non-top stage and set the top one to the fullscreen position.
         if (toStage != STAGE_TYPE_UNDEFINED) {
-            t.hide(toStage == STAGE_TYPE_MAIN ? mSideStage.mRootLeash : mMainStage.mRootLeash);
-            t.setPosition(toStage == STAGE_TYPE_MAIN
-                    ? mMainStage.mRootLeash : mSideStage.mRootLeash, 0, 0);
+            if (enableFlexibleSplit()) {
+                StageTaskListener stageToKeep = mStageOrderOperator.getAllStages().stream()
+                        .filter(stage -> stage.getId() == toStage)
+                        .findFirst().orElseThrow();
+                List<StageTaskListener> stagesToHide = mStageOrderOperator.getAllStages().stream()
+                        .filter(stage -> stage.getId() != toStage)
+                        .toList();
+                stagesToHide.forEach(stage -> t.hide(stage.mRootLeash));
+                t.setPosition(stageToKeep.mRootLeash, 0, 0);
+            } else {
+                t.hide(toStage == STAGE_TYPE_MAIN ? mSideStage.mRootLeash : mMainStage.mRootLeash);
+                t.setPosition(toStage == STAGE_TYPE_MAIN
+                        ? mMainStage.mRootLeash : mSideStage.mRootLeash, 0, 0);
+            }
         } else {
             for (int i = dismissingTasks.keySet().size() - 1; i >= 0; --i) {
                 finishT.hide(dismissingTasks.valueAt(i));
@@ -3088,8 +3562,12 @@
 
         // Hide divider and dim layer on transition finished.
         setDividerVisibility(false, t);
-        finishT.hide(mMainStage.mDimLayer);
-        finishT.hide(mSideStage.mDimLayer);
+        if (enableFlexibleSplit()) {
+            runForActiveStages(stage -> finishT.hide(stage.mRootLeash));
+        } else {
+            finishT.hide(mMainStage.mDimLayer);
+            finishT.hide(mSideStage.mDimLayer);
+        }
     }
 
     private boolean startPendingDismissAnimation(
@@ -3110,8 +3588,12 @@
             return false;
         }
         dismissTransition.setFinishedCallback((callbackWct, callbackT) -> {
-            mMainStage.getSplitDecorManager().release(callbackT);
-            mSideStage.getSplitDecorManager().release(callbackT);
+            if (enableFlexibleSplit()) {
+                runForActiveStages(stage -> stage.getSplitDecorManager().release(callbackT));
+            } else {
+                mMainStage.getSplitDecorManager().release(callbackT);
+                mSideStage.getSplitDecorManager().release(callbackT);
+            }
             callbackWct.setReparentLeafTaskIfRelaunch(mRootTaskInfo.token, false);
         });
         return true;
@@ -3128,8 +3610,15 @@
                 if (TransitionUtil.isClosingType(change.getMode())
                         && change.getTaskInfo() != null) {
                     final int taskId = change.getTaskInfo().taskId;
-                    if (mMainStage.getTopVisibleChildTaskId() == taskId
-                            || mSideStage.getTopVisibleChildTaskId() == taskId) {
+                    boolean anyStagesHaveTask;
+                    if (enableFlexibleSplit()) {
+                        anyStagesHaveTask = mStageOrderOperator.getActiveStages().stream()
+                                .anyMatch(stage -> stage.getTopVisibleChildTaskId() == taskId);
+                    } else {
+                        anyStagesHaveTask = mMainStage.getTopVisibleChildTaskId() == taskId
+                                || mSideStage.getTopVisibleChildTaskId() == taskId;
+                    }
+                    if (anyStagesHaveTask) {
                         mPausingTasks.add(taskId);
                     }
                 }
@@ -3161,9 +3650,16 @@
             final WindowContainerTransaction.HierarchyOp op =
                     finishWct.getHierarchyOps().get(i);
             final IBinder container = op.getContainer();
+            boolean anyStageContainsContainer;
+            if (enableFlexibleSplit()) {
+                anyStageContainsContainer = mStageOrderOperator.getActiveStages().stream()
+                        .anyMatch(stage -> stage.containsContainer(container));
+            } else {
+                anyStageContainsContainer = mMainStage.containsContainer(container)
+                        || mSideStage.containsContainer(container);
+            }
             if (op.getType() == HIERARCHY_OP_TYPE_REORDER && op.getToTop()
-                    && (mMainStage.containsContainer(container)
-                    || mSideStage.containsContainer(container))) {
+                    && anyStageContainsContainer) {
                 updateSurfaceBounds(mSplitLayout, finishT,
                         false /* applyResizingOffset */);
                 finishT.reparent(mSplitLayout.getDividerLeash(), mRootTaskLeash);
@@ -3190,10 +3686,18 @@
         // user entering recents.
         for (int i = mPausingTasks.size() - 1; i >= 0; --i) {
             final int taskId = mPausingTasks.get(i);
-            if (mMainStage.containsTask(taskId)) {
-                mMainStage.evictChild(finishWct, taskId, "recentsPairToPair");
-            } else if (mSideStage.containsTask(taskId)) {
-                mSideStage.evictChild(finishWct, taskId, "recentsPairToPair");
+            if (enableFlexibleSplit()) {
+                mStageOrderOperator.getActiveStages().stream()
+                        .filter(stage -> stage.containsTask(taskId))
+                        .findFirst()
+                        .ifPresent(stageToEvict ->
+                            stageToEvict.evictChild(finishWct, taskId, "recentsPairToPair"));
+            } else {
+                if (mMainStage.containsTask(taskId)) {
+                    mMainStage.evictChild(finishWct, taskId, "recentsPairToPair");
+                } else if (mSideStage.containsTask(taskId)) {
+                    mSideStage.evictChild(finishWct, taskId, "recentsPairToPair");
+                }
             }
         }
         // If pending enter hasn't consumed, the mix handler will invoke start pending
@@ -3256,8 +3760,15 @@
      */
     private void setSplitsVisible(boolean visible) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "setSplitsVisible: visible=%b", visible);
-        mMainStage.mVisible = mSideStage.mVisible = visible;
-        mMainStage.mHasChildren = mSideStage.mHasChildren = visible;
+        if (enableFlexibleSplit()) {
+            runForActiveStages(stage -> {
+                stage.mVisible = visible;
+                stage.mHasChildren = visible;
+            });
+        } else {
+            mMainStage.mVisible = mSideStage.mVisible = visible;
+            mMainStage.mHasChildren = mSideStage.mHasChildren = visible;
+        }
     }
 
     /**
@@ -3300,6 +3811,10 @@
      * executed.
      */
     private void logExitToStage(@ExitReason int exitReason, boolean toMainStage) {
+        if (enableFlexibleSplit()) {
+            // TODO(b/374825718) update logging for 2+ apps
+            return;
+        }
         mLogger.logExit(exitReason,
                 toMainStage ? getMainStagePosition() : SPLIT_POSITION_UNDEFINED,
                 toMainStage ? mMainStage.getTopChildTaskUid() : 0 /* mainStageUid */,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageOrderOperator.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageOrderOperator.kt
new file mode 100644
index 0000000..b7b3c9b
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageOrderOperator.kt
@@ -0,0 +1,185 @@
+/*
+ * 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.splitscreen
+
+import android.content.Context
+import com.android.internal.protolog.ProtoLog
+import com.android.launcher3.icons.IconProvider
+import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.common.SyncTransactionQueue
+import com.android.wm.shell.protolog.ShellProtoLogGroup
+import com.android.wm.shell.shared.split.SplitScreenConstants
+import com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_NONE
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_1
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_2
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED
+import com.android.wm.shell.shared.split.SplitScreenConstants.SnapPosition
+import com.android.wm.shell.shared.split.SplitScreenConstants.SplitIndex
+import com.android.wm.shell.shared.split.SplitScreenConstants.SplitPosition
+import com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_A
+import com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_B
+import com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_C
+import com.android.wm.shell.splitscreen.SplitScreen.stageTypeToString
+import com.android.wm.shell.windowdecor.WindowDecorViewModel
+import java.util.Optional
+
+/**
+ * Responsible for creating [StageTaskListener]s and maintaining their ordering on screen.
+ * Must be notified whenever stages positions change via swapping or starting/ending tasks
+ */
+class StageOrderOperator (
+        context: Context,
+        taskOrganizer: ShellTaskOrganizer,
+        displayId: Int,
+        stageCallbacks: StageTaskListener.StageListenerCallbacks,
+        syncQueue: SyncTransactionQueue,
+        iconProvider: IconProvider,
+        windowDecorViewModel: Optional<WindowDecorViewModel>
+    ) {
+
+    private val MAX_STAGES = 3
+    /**
+     * This somewhat acts as a replacement to stageTypes in the intermediary, so we want to start
+     * it after the @StageType constant values just to be safe and avoid potentially subtle bugs.
+     */
+    private var stageIds = listOf(STAGE_TYPE_A, STAGE_TYPE_B, STAGE_TYPE_C)
+
+    /**
+     * Active Stages, this list represent the current, ordered list of stages that are
+     * currently visible to the user. This map should be empty if the user is currently
+     * not in split screen. Note that this is different than if split screen is visible, which
+     * is determined by [StageListenerImpl.mVisible].
+     * Split stages can be active and in the background
+     */
+    val activeStages = mutableListOf<StageTaskListener>()
+    val allStages = mutableListOf<StageTaskListener>()
+    var isActive: Boolean = false
+    var isVisible: Boolean = false
+    @SnapPosition private var currentLayout: Int = SNAP_TO_NONE
+
+    init {
+        for(i in 0 until MAX_STAGES) {
+            allStages.add(StageTaskListener(context,
+                taskOrganizer,
+                displayId,
+                stageCallbacks,
+                syncQueue,
+                iconProvider,
+                windowDecorViewModel,
+                stageIds[i])
+            )
+        }
+    }
+
+    /**
+     * Updates internal state to keep record of "active" stages. Note that this does NOT call
+     * [StageTaskListener.activate] on the stages.
+     */
+    fun onEnteringSplit(@SnapPosition goingToLayout: Int) {
+        if (goingToLayout == currentLayout) {
+            // Add protolog here. Return for now, but maybe we want to handle swap case, TBD
+            return
+        }
+        val freeStages: List<StageTaskListener> =
+            allStages.filterNot { activeStages.contains(it) }
+        when(goingToLayout) {
+            SplitScreenConstants.SNAP_TO_2_50_50 -> {
+                if (activeStages.size < 2) {
+                    // take from allStages and add into activeStages
+                    for (i in 0 until (2 - activeStages.size)) {
+                        val stage = freeStages[i]
+                        activeStages.add(stage)
+                    }
+                }
+            }
+        }
+        ProtoLog.d(
+            ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+            "Activated stages: %d ids=%s",
+            activeStages.size,
+            activeStages.joinToString(",") { stageTypeToString(it.id) }
+        )
+        isActive = true
+    }
+
+    fun onExitingSplit() {
+        activeStages.clear()
+        isActive = false
+    }
+
+    /**
+     * Given a legacy [SplitPosition] returns one of the stages from the actives stages.
+     * If there are no active stages and [checkAllStagesIfNotActive] is not true, then will return
+     * null
+     */
+    fun getStageForLegacyPosition(@SplitPosition position: Int,
+                                  checkAllStagesIfNotActive : Boolean = false) :
+            StageTaskListener? {
+        if (activeStages.size != 2 && !checkAllStagesIfNotActive) {
+            return null
+        }
+        val listToCheck = if (activeStages.isEmpty() and checkAllStagesIfNotActive)
+            allStages else
+            activeStages
+        if (position == SPLIT_POSITION_TOP_OR_LEFT) {
+            return listToCheck[0]
+        } else if (position == SPLIT_POSITION_BOTTOM_OR_RIGHT) {
+            return listToCheck[1]
+        } else {
+            throw IllegalArgumentException("No stage for invalid position")
+        }
+    }
+
+    /**
+     * Returns a legacy split position for the given stage. If no stages are active then this will
+     * return [SPLIT_POSITION_UNDEFINED]
+     */
+    @SplitPosition
+    fun getLegacyPositionForStage(stage: StageTaskListener) : Int {
+        if (allStages[0] == stage) {
+            return SPLIT_POSITION_TOP_OR_LEFT
+        } else if (allStages[1] == stage) {
+            return SPLIT_POSITION_BOTTOM_OR_RIGHT
+        } else {
+            return SPLIT_POSITION_UNDEFINED
+        }
+    }
+
+    /**
+     * Returns the stageId from a given splitIndex. This will default to checking from all stages if
+     * [isActive] is false, otherwise will only check active stages.
+     */
+    fun getStageForIndex(@SplitIndex splitIndex: Int) : StageTaskListener {
+        // Probably should do a check for index to be w/in the bounds of the current split layout
+        // that we're currently in
+        val listToCheck = if (isActive) activeStages else allStages
+        if (splitIndex == SPLIT_INDEX_0) {
+            return listToCheck[0]
+        } else if (splitIndex == SPLIT_INDEX_1) {
+            return listToCheck[1]
+        } else if (splitIndex == SPLIT_INDEX_2) {
+            return listToCheck[2]
+        } else {
+            // Though I guess what if we're adding to the end? Maybe that indexing needs to be
+            // resolved elsewhere
+            throw IllegalStateException("No stage for the given splitIndex")
+        }
+    }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 08cdfdb..4a37169 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -22,10 +22,12 @@
 import static android.content.res.Configuration.SMALLEST_SCREEN_WIDTH_DP_UNDEFINED;
 import static android.view.RemoteAnimationTarget.MODE_OPENING;
 
+import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.CONTROLLED_ACTIVITY_TYPES;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.CONTROLLED_WINDOWING_MODES;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.CONTROLLED_WINDOWING_MODES_WHEN_ACTIVE;
+import static com.android.wm.shell.splitscreen.SplitScreen.stageTypeToString;
 
 import android.annotation.CallSuper;
 import android.annotation.Nullable;
@@ -72,6 +74,8 @@
     // No current way to enforce this but if enableFlexibleSplit() is enabled, then only 1 of the
     // stages should have this be set/being used
     private boolean mIsActive;
+    /** Unique identifier for this state, > 0 */
+    @StageType private final int mId;
     /** Callback interface for listening to changes in a split-screen stage. */
     public interface StageListenerCallbacks {
         void onRootTaskAppeared();
@@ -110,13 +114,14 @@
     StageTaskListener(Context context, ShellTaskOrganizer taskOrganizer, int displayId,
             StageListenerCallbacks callbacks, SyncTransactionQueue syncQueue,
             IconProvider iconProvider,
-            Optional<WindowDecorViewModel> windowDecorViewModel) {
+            Optional<WindowDecorViewModel> windowDecorViewModel, int id) {
         mContext = context;
         mCallbacks = callbacks;
         mSyncQueue = syncQueue;
         mIconProvider = iconProvider;
         mWindowDecorViewModel = windowDecorViewModel;
         taskOrganizer.createRootTask(displayId, WINDOWING_MODE_MULTI_WINDOW, this);
+        mId = id;
     }
 
     int getChildCount() {
@@ -161,6 +166,11 @@
         return contains(t -> t.isFocused);
     }
 
+    @StageType
+    int getId() {
+        return mId;
+    }
+
     private boolean contains(Predicate<ActivityManager.RunningTaskInfo> predicate) {
         if (mRootTaskInfo != null && predicate.test(mRootTaskInfo)) {
             return true;
@@ -197,10 +207,10 @@
     @CallSuper
     public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskAppeared: taskId=%d taskParent=%d rootTask=%d "
-                        + "taskActivity=%s",
+                        + "stageId=%s taskActivity=%s",
                 taskInfo.taskId, taskInfo.parentTaskId,
                 mRootTaskInfo != null ? mRootTaskInfo.taskId : -1,
-                taskInfo.baseActivity);
+                stageTypeToString(mId), taskInfo.baseActivity);
         if (mRootTaskInfo == null) {
             mRootLeash = leash;
             mRootTaskInfo = taskInfo;
@@ -230,8 +240,9 @@
     @Override
     @CallSuper
     public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskInfoChanged: taskId=%d taskAct=%s",
-                taskInfo.taskId, taskInfo.baseActivity);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskInfoChanged: taskId=%d taskAct=%s "
+                        + "stageId=%s",
+                taskInfo.taskId, taskInfo.baseActivity, stageTypeToString(mId));
         mWindowDecorViewModel.ifPresent(viewModel -> viewModel.onTaskInfoChanged(taskInfo));
         if (mRootTaskInfo.taskId == taskInfo.taskId) {
             mRootTaskInfo = taskInfo;
@@ -261,7 +272,8 @@
     @Override
     @CallSuper
     public void onTaskVanished(ActivityManager.RunningTaskInfo taskInfo) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskVanished: task=%d", taskInfo.taskId);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskVanished: task=%d stageId=%s",
+                taskInfo.taskId, stageTypeToString(mId));
         final int taskId = taskInfo.taskId;
         mWindowDecorViewModel.ifPresent(vm -> vm.onTaskVanished(taskInfo));
         if (mRootTaskInfo.taskId == taskId) {
@@ -466,14 +478,17 @@
     }
 
     void activate(WindowContainerTransaction wct, boolean includingTopTask) {
-        if (mIsActive) return;
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "activate: includingTopTask=%b",
-                includingTopTask);
+        if (mIsActive && !enableFlexibleSplit()) return;
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "activate: includingTopTask=%b stage=%s",
+                includingTopTask, stageTypeToString(mId));
 
         if (includingTopTask) {
             reparentTopTask(wct);
         }
 
+        if (enableFlexibleSplit()) {
+            return;
+        }
         mIsActive = true;
     }
 
@@ -481,11 +496,14 @@
         deactivate(wct, false /* toTop */);
     }
 
-    void deactivate(WindowContainerTransaction wct, boolean toTop) {
-        if (!mIsActive) return;
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "deactivate: toTop=%b rootTaskInfo=%s",
-                toTop, mRootTaskInfo);
-        mIsActive = false;
+    void deactivate(WindowContainerTransaction wct, boolean reparentTasksToTop) {
+        if (!mIsActive && !enableFlexibleSplit()) return;
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "deactivate: reparentTasksToTop=%b "
+                        + "rootTaskInfo=%s stage=%s",
+                reparentTasksToTop, mRootTaskInfo, stageTypeToString(mId));
+        if (!enableFlexibleSplit()) {
+            mIsActive = false;
+        }
 
         if (mRootTaskInfo == null) return;
         final WindowContainerToken rootToken = mRootTaskInfo.token;
@@ -494,14 +512,15 @@
                 null /* newParent */,
                 null /* windowingModes */,
                 null /* activityTypes */,
-                toTop);
+                reparentTasksToTop);
     }
 
     // --------
-    // Previously only used in SideStage
+    // Previously only used in SideStage. With flexible split this is called for all stages
     boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove all side stage tasks: childCount=%d toTop=%b",
-                mChildrenTaskInfo.size(), toTop);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "remove all side stage tasks: childCount=%d toTop=%b "
+                        + " stageI=%s",
+                mChildrenTaskInfo.size(), toTop, stageTypeToString(mId));
         if (mChildrenTaskInfo.size() == 0) return false;
         wct.reparentTasks(
                 mRootTaskInfo.token,
@@ -522,6 +541,15 @@
     }
 
     @Override
+    public String toString() {
+        return "mId: " + stageTypeToString(mId)
+                + " mVisible: " + mVisible
+                + " mActive: " + mIsActive
+                + " mHasRootTask: " + mHasRootTask
+                + " childSize: " + mChildrenTaskInfo.size();
+    }
+
+    @Override
     @CallSuper
     public void dump(@NonNull PrintWriter pw, String prefix) {
         final String innerPrefix = prefix + "  ";
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 0f546cd..8d04749 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
@@ -49,7 +49,8 @@
 
     @Test
     open fun enterDesktopWithDrag() {
-        testApp.enterDesktopModeWithDrag(wmHelper, device)
+        // By default this method uses drag to desktop
+        testApp.enterDesktopMode(wmHelper, device)
     }
 
     @After
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 936835c..5df3957 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
@@ -116,6 +116,7 @@
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource.UNKNOWN
 import com.android.wm.shell.shared.split.SplitScreenConstants
+import com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED
 import com.android.wm.shell.splitscreen.SplitScreenController
 import com.android.wm.shell.sysui.ShellCommandHandler
 import com.android.wm.shell.sysui.ShellController
@@ -3158,7 +3159,7 @@
     runOpenNewWindow(task)
     verify(splitScreenController)
       .startIntent(any(), anyInt(), any(), any(),
-        optionsCaptor.capture(), anyOrNull(), eq(true)
+        optionsCaptor.capture(), anyOrNull(), eq(true), eq(SPLIT_INDEX_UNDEFINED)
       )
     assertThat(ActivityOptions.fromBundle(optionsCaptor.value).launchWindowingMode)
       .isEqualTo(WINDOWING_MODE_MULTI_WINDOW)
@@ -3174,7 +3175,7 @@
     verify(splitScreenController)
       .startIntent(
         any(), anyInt(), any(), any(),
-        optionsCaptor.capture(), anyOrNull(), eq(true)
+        optionsCaptor.capture(), anyOrNull(), eq(true), eq(SPLIT_INDEX_UNDEFINED)
       )
     assertThat(ActivityOptions.fromBundle(optionsCaptor.value).launchWindowingMode)
       .isEqualTo(WINDOWING_MODE_MULTI_WINDOW)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
index 2cfce69..0cf15ba 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/SplitDragPolicyTest.java
@@ -24,6 +24,7 @@
 import static android.content.ClipDescription.MIMETYPE_APPLICATION_TASK;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED;
 import static com.android.wm.shell.draganddrop.DragTestUtils.createAppClipData;
 import static com.android.wm.shell.draganddrop.DragTestUtils.createIntentClipData;
 import static com.android.wm.shell.draganddrop.DragTestUtils.createTaskInfo;
@@ -226,7 +227,7 @@
 
         mPolicy.onDropped(filterTargetByType(targets, TYPE_FULLSCREEN), null /* hideTaskToken */);
         verify(mFullscreenStarter).startIntent(any(), anyInt(), any(),
-                eq(SPLIT_POSITION_UNDEFINED), any(), any());
+                eq(SPLIT_POSITION_UNDEFINED), any(), any(), eq(SPLIT_INDEX_UNDEFINED));
     }
 
     private void dragOverFullscreenApp_expectSplitScreenTargets(ClipData data) {
@@ -241,12 +242,12 @@
 
         mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_LEFT), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
-                eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any());
+                eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any(), eq(SPLIT_INDEX_UNDEFINED));
         reset(mSplitScreenStarter);
 
         mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_RIGHT), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
-                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any());
+                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any(), eq(SPLIT_INDEX_UNDEFINED));
     }
 
     private void dragOverFullscreenAppPhone_expectVerticalSplitScreenTargets(ClipData data) {
@@ -261,13 +262,13 @@
 
         mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_TOP), null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
-                eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any());
+                eq(SPLIT_POSITION_TOP_OR_LEFT), any(), any(), eq(SPLIT_INDEX_UNDEFINED));
         reset(mSplitScreenStarter);
 
         mPolicy.onDropped(filterTargetByType(targets, TYPE_SPLIT_BOTTOM),
                 null /* hideTaskToken */);
         verify(mSplitScreenStarter).startIntent(any(), anyInt(), any(),
-                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any());
+                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), any(), any(), eq(SPLIT_INDEX_UNDEFINED));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
index 966651f..72a7a3f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
@@ -23,6 +23,7 @@
 import static android.content.Intent.FLAG_ACTIVITY_MULTIPLE_TASK;
 import static android.content.Intent.FLAG_ACTIVITY_NO_USER_ACTION;
 
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_0;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 
@@ -200,10 +201,11 @@
                 PendingIntent.getActivity(mContext, 0, startIntent, FLAG_IMMUTABLE);
 
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
-                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */);
+                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
+                SPLIT_INDEX_0);
 
         verify(mStageCoordinator).startIntent(eq(pendingIntent), mIntentCaptor.capture(),
-                eq(SPLIT_POSITION_TOP_OR_LEFT), isNull(), isNull());
+                eq(SPLIT_POSITION_TOP_OR_LEFT), isNull(), isNull(), eq(SPLIT_INDEX_0));
         assertEquals(FLAG_ACTIVITY_NO_USER_ACTION,
                 mIntentCaptor.getValue().getFlags() & FLAG_ACTIVITY_NO_USER_ACTION);
     }
@@ -220,10 +222,11 @@
         doReturn(topRunningTask).when(mRecentTasks).getTopRunningTask(any());
 
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
-                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */);
+                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
+                SPLIT_INDEX_0);
 
         verify(mStageCoordinator).startIntent(eq(pendingIntent), mIntentCaptor.capture(),
-                eq(SPLIT_POSITION_TOP_OR_LEFT), isNull(), isNull());
+                eq(SPLIT_POSITION_TOP_OR_LEFT), isNull(), isNull(), eq(SPLIT_INDEX_0));
         assertEquals(FLAG_ACTIVITY_MULTIPLE_TASK,
                 mIntentCaptor.getValue().getFlags() & FLAG_ACTIVITY_MULTIPLE_TASK);
     }
@@ -243,10 +246,11 @@
         doReturn(sameTaskInfo).when(mRecentTasks).findTaskInBackground(any(), anyInt(), any());
 
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
-                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */);
+                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
+                SPLIT_INDEX_0);
 
         verify(mStageCoordinator).startTask(anyInt(), eq(SPLIT_POSITION_TOP_OR_LEFT),
-                isNull(), isNull());
+                isNull(), isNull(), eq(SPLIT_INDEX_0));
         verify(mMultiInstanceHelper, never()).supportsMultiInstanceSplit(any());
         verify(mStageCoordinator, never()).switchSplitPosition(any());
     }
@@ -269,10 +273,11 @@
                 .findTaskInBackground(any(), anyInt(), any());
 
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
-                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */);
+                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
+                SPLIT_INDEX_0);
         verify(mMultiInstanceHelper, never()).supportsMultiInstanceSplit(any());
         verify(mStageCoordinator).startTask(anyInt(), eq(SPLIT_POSITION_TOP_OR_LEFT),
-                isNull(), isNull());
+                isNull(), isNull(), eq(SPLIT_INDEX_0));
     }
 
     @Test
@@ -289,7 +294,8 @@
                 SPLIT_POSITION_BOTTOM_OR_RIGHT);
 
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
-                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */);
+                SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
+                SPLIT_INDEX_0);
 
         verify(mStageCoordinator).switchSplitPosition(anyString());
     }
@@ -301,7 +307,7 @@
                 PendingIntent.getActivity(mContext, 0, startIntent, FLAG_IMMUTABLE);
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
                 SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
-                true /* forceLaunchNewTask */);
+                true /* forceLaunchNewTask */, SPLIT_INDEX_0);
         verify(mRecentTasks, never()).findTaskInBackground(any(), anyInt(), any());
     }
 
@@ -312,7 +318,7 @@
                 PendingIntent.getActivity(mContext, 0, startIntent, FLAG_IMMUTABLE);
         mSplitScreenController.startIntent(pendingIntent, mContext.getUserId(), null,
                 SPLIT_POSITION_TOP_OR_LEFT, null /* options */, null /* hideTaskToken */,
-                false /* forceLaunchNewTask */);
+                false /* forceLaunchNewTask */, SPLIT_INDEX_0);
         verify(mRecentTasks).findTaskInBackground(any(), anyInt(), any());
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index ce3944a..e32cf38 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -27,6 +27,7 @@
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_CHILDREN_TASKS_REPARENT;
 import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_MAIN;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_SIDE;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_APP_DOES_NOT_SUPPORT_MULTIWINDOW;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_DRAG_DIVIDER;
@@ -133,11 +134,11 @@
         mSplitLayout = SplitTestUtils.createMockSplitLayout();
         mMainStage = spy(new StageTaskListener(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue,
-                mIconProvider, Optional.of(mWindowDecorViewModel)));
+                mIconProvider, Optional.of(mWindowDecorViewModel), STAGE_TYPE_MAIN));
         mMainStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mSideStage = spy(new StageTaskListener(mContext, mTaskOrganizer, DEFAULT_DISPLAY, mock(
                 StageTaskListener.StageListenerCallbacks.class), mSyncQueue,
-                mIconProvider, Optional.of(mWindowDecorViewModel)));
+                mIconProvider, Optional.of(mWindowDecorViewModel), STAGE_TYPE_SIDE));
         mSideStage.onTaskAppeared(new TestRunningTaskInfoBuilder().build(), createMockSurface());
         mStageCoordinator = new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
                 mSyncQueue, mTaskOrganizer, mMainStage, mSideStage, mDisplayController,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index a252a9d..2d102fb 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -19,6 +19,7 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_INDEX_UNDEFINED;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
@@ -165,8 +166,9 @@
         final WindowContainerTransaction wct = spy(new WindowContainerTransaction());
 
         mStageCoordinator.moveToStage(task, SPLIT_POSITION_BOTTOM_OR_RIGHT, wct);
+        // TODO(b/349828130) Address this once we remove index_undefined called
         verify(mStageCoordinator).prepareEnterSplitScreen(eq(wct), eq(task),
-                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false));
+                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false), eq(SPLIT_INDEX_UNDEFINED));
         verify(mMainStage).reparentTopTask(eq(wct));
         assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getSideStagePosition());
         assertEquals(SPLIT_POSITION_TOP_OR_LEFT, mStageCoordinator.getMainStagePosition());
@@ -183,8 +185,9 @@
         final WindowContainerTransaction wct = new WindowContainerTransaction();
 
         mStageCoordinator.moveToStage(task, SPLIT_POSITION_BOTTOM_OR_RIGHT, wct);
+        // TODO(b/349828130) Address this once we remove index_undefined called
         verify(mStageCoordinator).prepareEnterSplitScreen(eq(wct), eq(task),
-                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false));
+                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false), eq(SPLIT_INDEX_UNDEFINED));
         assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getMainStagePosition());
         assertEquals(SPLIT_POSITION_TOP_OR_LEFT, mStageCoordinator.getSideStagePosition());
     }
@@ -195,8 +198,9 @@
         final WindowContainerTransaction wct = new WindowContainerTransaction();
 
         mStageCoordinator.moveToStage(task, SPLIT_POSITION_BOTTOM_OR_RIGHT, wct);
+        // TODO(b/349828130) Address this once we remove index_undefined called
         verify(mStageCoordinator).prepareEnterSplitScreen(eq(wct), eq(task),
-                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false));
+                eq(SPLIT_POSITION_BOTTOM_OR_RIGHT), eq(false), eq(SPLIT_INDEX_UNDEFINED));
         assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getSideStagePosition());
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
index 7144a1e..fe91440 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageTaskListenerTests.java
@@ -19,6 +19,8 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertFalse;
@@ -93,7 +95,8 @@
                 mCallbacks,
                 mSyncQueue,
                 mIconProvider,
-                Optional.of(mWindowDecorViewModel));
+                Optional.of(mWindowDecorViewModel),
+                STAGE_TYPE_UNDEFINED);
         mRootTask = new TestRunningTaskInfoBuilder().build();
         mRootTask.parentTaskId = INVALID_TASK_ID;
         mSurfaceControl = new SurfaceControl.Builder().setName("test").build();
diff --git a/packages/Shell/AndroidManifest.xml b/packages/Shell/AndroidManifest.xml
index 526320d..1070ebd 100644
--- a/packages/Shell/AndroidManifest.xml
+++ b/packages/Shell/AndroidManifest.xml
@@ -746,6 +746,9 @@
     <!-- Permission required for ATS test - CarDevicePolicyManagerTest -->
     <uses-permission android:name="android.permission.LOCK_DEVICE" />
 
+    <!-- Permission required for AuthenticationPolicyManagerTest -->
+    <uses-permission android:name="android.permission.MANAGE_SECURE_LOCK_DEVICE" />
+
     <!-- Permissions required for CTS test - CtsSafetyCenterTestCases -->
     <uses-permission android:name="android.permission.SEND_SAFETY_CENTER_UPDATE" />
     <uses-permission android:name="android.permission.READ_SAFETY_CENTER_STATUS" />
diff --git a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/LockSection.kt b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/LockSection.kt
index 9390664..597cbf2 100644
--- a/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/LockSection.kt
+++ b/packages/SystemUI/compose/features/src/com/android/systemui/keyguard/ui/composable/section/LockSection.kt
@@ -31,6 +31,7 @@
 import com.android.compose.animation.scene.ElementKey
 import com.android.compose.animation.scene.SceneScope
 import com.android.systemui.biometrics.AuthController
+import com.android.systemui.customization.R as customR
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.flags.FeatureFlagsClassic
 import com.android.systemui.flags.Flags
@@ -147,7 +148,7 @@
             } else {
                 val scaleFactor = authController.scaleFactor
                 val bottomPaddingPx =
-                    context.resources.getDimensionPixelSize(R.dimen.lock_icon_margin_bottom)
+                    context.resources.getDimensionPixelSize(customR.dimen.lock_icon_margin_bottom)
                 val heightPx = windowViewBounds.bottom.toFloat()
 
                 Pair(
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 1603267..0db545f 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
@@ -519,8 +519,19 @@
                 // we intercept an ongoing swipe transition (i.e. startDragImmediately() returned
                 // true).
                 if (overSlop == 0f) {
-                    val delta = (drag.position - consumablePointer.position).toFloat()
-                    check(delta != 0f) { "delta is equal to 0" }
+                    // If the user drags in the opposite direction, the delta becomes zero because
+                    // we return to the original point. Therefore, we should use the previous event
+                    // to calculate the direction.
+                    val delta = (drag.position - drag.previousPosition).toFloat()
+                    check(delta != 0f) {
+                        buildString {
+                            append("delta is equal to 0 ")
+                            append("touchSlop ${currentValueOf(LocalViewConfiguration).touchSlop} ")
+                            append("consumablePointer.position ${consumablePointer.position} ")
+                            append("drag.position ${drag.position} ")
+                            append("drag.previousPosition ${drag.previousPosition}")
+                        }
+                    }
                     overSlop = delta.sign
                 }
                 drag
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
index 5ec74f8..85d8b60 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/MultiPointerDraggableTest.kt
@@ -593,7 +593,7 @@
             }
         }
 
-        fun continueDraggingDown() {
+        fun dragDown() {
             rule.onRoot().performTouchInput { moveBy(Offset(0f, touchSlop)) }
         }
 
@@ -603,11 +603,78 @@
         assertThat(started).isFalse()
 
         swipeConsume = true
-        continueDraggingDown()
+        // Drag in same direction
+        dragDown()
         assertThat(capturedChange).isNotNull()
         capturedChange = null
 
-        continueDraggingDown()
+        dragDown()
+        assertThat(capturedChange).isNull()
+
+        assertThat(started).isTrue()
+    }
+
+    @Test
+    fun multiPointerSwipeDetectorInteractionZeroOffsetFromStartPosition() {
+        val size = 200f
+        val middle = Offset(size / 2f, size / 2f)
+
+        var started = false
+
+        var capturedChange: PointerInputChange? = null
+        var swipeConsume = false
+
+        var touchSlop = 0f
+        rule.setContent {
+            touchSlop = LocalViewConfiguration.current.touchSlop
+            Box(
+                Modifier.size(with(LocalDensity.current) { Size(size, size).toDpSize() })
+                    .nestedScrollDispatcher()
+                    .multiPointerDraggable(
+                        orientation = Orientation.Vertical,
+                        startDragImmediately = { false },
+                        swipeDetector =
+                            object : SwipeDetector {
+                                override fun detectSwipe(change: PointerInputChange): Boolean {
+                                    capturedChange = change
+                                    return swipeConsume
+                                }
+                            },
+                        onDragStarted = { _, _ ->
+                            started = true
+                            SimpleDragController(
+                                onDrag = { /* do nothing */ },
+                                onStop = { /* do nothing */ },
+                            )
+                        },
+                        dispatcher = defaultDispatcher,
+                    )
+            ) {}
+        }
+
+        fun startDraggingDown() {
+            rule.onRoot().performTouchInput {
+                down(middle)
+                moveBy(Offset(0f, touchSlop))
+            }
+        }
+
+        fun dragUp() {
+            rule.onRoot().performTouchInput { moveBy(Offset(0f, -touchSlop)) }
+        }
+
+        startDraggingDown()
+        assertThat(capturedChange).isNotNull()
+        capturedChange = null
+        assertThat(started).isFalse()
+
+        swipeConsume = true
+        // Drag in the opposite direction
+        dragUp()
+        assertThat(capturedChange).isNotNull()
+        capturedChange = null
+
+        dragUp()
         assertThat(capturedChange).isNull()
 
         assertThat(started).isTrue()
diff --git a/packages/SystemUI/customization/res/values-land/dimens.xml b/packages/SystemUI/customization/res/values-land/dimens.xml
new file mode 100644
index 0000000..50f220c8
--- /dev/null
+++ b/packages/SystemUI/customization/res/values-land/dimens.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<resources>
+    <dimen name="lock_icon_margin_bottom">24dp</dimen>
+</resources>
\ No newline at end of file
diff --git a/packages/SystemUI/customization/res/values-sw600dp-land/dimens.xml b/packages/SystemUI/customization/res/values-sw600dp-land/dimens.xml
index 18073ad..8760281 100644
--- a/packages/SystemUI/customization/res/values-sw600dp-land/dimens.xml
+++ b/packages/SystemUI/customization/res/values-sw600dp-land/dimens.xml
@@ -17,4 +17,5 @@
 <resources>
     <dimen name="keyguard_smartspace_top_offset">0dp</dimen>
     <dimen name="status_view_margin_horizontal">8dp</dimen>
+    <dimen name="lock_icon_margin_bottom">60dp</dimen>
 </resources>
\ No newline at end of file
diff --git a/packages/SystemUI/customization/res/values/dimens.xml b/packages/SystemUI/customization/res/values/dimens.xml
index 041ae62..21b4c71 100644
--- a/packages/SystemUI/customization/res/values/dimens.xml
+++ b/packages/SystemUI/customization/res/values/dimens.xml
@@ -40,4 +40,5 @@
     <dimen name="date_weather_view_height">24dp</dimen>
     <dimen name="enhanced_smartspace_height">104dp</dimen>
     <dimen name="status_view_margin_horizontal">0dp</dimen>
+    <dimen name="lock_icon_margin_bottom">74dp</dimen>
 </resources>
\ No newline at end of file
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractorTest.kt
index f0d79bb..47cba07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractorTest.kt
@@ -568,6 +568,41 @@
             assertThat(isUnlocked).isFalse()
         }
 
+    @Test
+    fun lockNow() =
+        testScope.runTest {
+            setLockAfterScreenTimeout(5000)
+            val isUnlocked by collectLastValue(underTest.deviceUnlockStatus.map { it.isUnlocked })
+            unlockDevice()
+            assertThat(isUnlocked).isTrue()
+
+            underTest.lockNow()
+            runCurrent()
+
+            assertThat(isUnlocked).isFalse()
+        }
+
+    @Test
+    fun deviceUnlockStatus_isResetToFalse_whenDeviceGoesToSleep_fromSleepButton() =
+        testScope.runTest {
+            setLockAfterScreenTimeout(5000)
+            kosmos.fakeAuthenticationRepository.powerButtonInstantlyLocks = false
+            val deviceUnlockStatus by collectLastValue(underTest.deviceUnlockStatus)
+
+            kosmos.fakeDeviceEntryFingerprintAuthRepository.setAuthenticationStatus(
+                SuccessFingerprintAuthenticationStatus(0, true)
+            )
+            runCurrent()
+            assertThat(deviceUnlockStatus?.isUnlocked).isTrue()
+
+            kosmos.powerInteractor.setAsleepForTest(
+                sleepReason = PowerManager.GO_TO_SLEEP_REASON_SLEEP_BUTTON
+            )
+            runCurrent()
+
+            assertThat(deviceUnlockStatus?.isUnlocked).isFalse()
+        }
+
     private fun TestScope.unlockDevice() {
         val deviceUnlockStatus by collectLastValue(underTest.deviceUnlockStatus)
 
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
index 92764ae..74a0bafda 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/ui/binder/WindowManagerLockscreenVisibilityManagerTest.kt
@@ -17,6 +17,10 @@
 package com.android.systemui.keyguard.ui.binder
 
 import android.app.IActivityTaskManager
+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 androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -25,8 +29,10 @@
 import com.android.systemui.statusbar.policy.KeyguardStateController
 import com.android.systemui.util.concurrency.FakeExecutor
 import com.android.systemui.util.time.FakeSystemClock
+import com.android.window.flags.Flags
 import com.android.wm.shell.keyguard.KeyguardTransitions
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.ArgumentMatchers.eq
@@ -41,6 +47,9 @@
 @RunWith(AndroidJUnit4::class)
 @kotlinx.coroutines.ExperimentalCoroutinesApi
 class WindowManagerLockscreenVisibilityManagerTest : SysuiTestCase() {
+
+    @get:Rule val checkFlagsRule: CheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule()
+
     private lateinit var underTest: WindowManagerLockscreenVisibilityManager
     private lateinit var executor: FakeExecutor
 
@@ -68,32 +77,62 @@
     }
 
     @Test
-    fun testLockscreenVisible_andAodVisible() {
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testLockscreenVisible_andAodVisible_without_keyguard_shell_transitions() {
         underTest.setLockscreenShown(true)
-        underTest.setAodVisible(true)
-
         verify(activityTaskManagerService).setLockScreenShown(true, false)
+        underTest.setAodVisible(true)
         verify(activityTaskManagerService).setLockScreenShown(true, true)
+
         verifyNoMoreInteractions(activityTaskManagerService)
     }
 
     @Test
-    fun testGoingAway_whenLockscreenVisible_thenSurfaceMadeVisible() {
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testLockscreenVisible_andAodVisible_with_keyguard_shell_transitions() {
         underTest.setLockscreenShown(true)
+        verify(keyguardTransitions).startKeyguardTransition(true, false)
         underTest.setAodVisible(true)
+        verify(keyguardTransitions).startKeyguardTransition(true, true)
 
+        verifyNoMoreInteractions(keyguardTransitions)
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testGoingAway_whenLockscreenVisible_thenSurfaceMadeVisible_without_keyguard_shell_transitions() {
+        underTest.setLockscreenShown(true)
         verify(activityTaskManagerService).setLockScreenShown(true, false)
+        underTest.setAodVisible(true)
         verify(activityTaskManagerService).setLockScreenShown(true, true)
+
         verifyNoMoreInteractions(activityTaskManagerService)
 
         underTest.setSurfaceBehindVisibility(true)
-
         verify(activityTaskManagerService).keyguardGoingAway(anyInt())
+
         verifyNoMoreInteractions(activityTaskManagerService)
     }
 
     @Test
-    fun testSurfaceVisible_whenLockscreenNotShowing_doesNotTriggerGoingAway() {
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testGoingAway_whenLockscreenVisible_thenSurfaceMadeVisible_with_keyguard_shell_transitions() {
+        underTest.setLockscreenShown(true)
+        verify(keyguardTransitions).startKeyguardTransition(true, false)
+        underTest.setAodVisible(true)
+        verify(keyguardTransitions).startKeyguardTransition(true, true)
+
+        verifyNoMoreInteractions(keyguardTransitions)
+
+        underTest.setSurfaceBehindVisibility(true)
+        verify(keyguardTransitions).startKeyguardTransition(false, false)
+
+        verifyNoMoreInteractions(keyguardTransitions)
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testSurfaceVisible_whenLockscreenNotShowing_doesNotTriggerGoingAway_without_keyguard_shell_transitions() {
         underTest.setLockscreenShown(false)
         underTest.setAodVisible(false)
 
@@ -106,7 +145,22 @@
     }
 
     @Test
-    fun testAodVisible_noLockscreenShownCallYet_doesNotShowLockscreenUntilLater() {
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testSurfaceVisible_whenLockscreenNotShowing_doesNotTriggerGoingAway_with_keyguard_shell_transitions() {
+        underTest.setLockscreenShown(false)
+        underTest.setAodVisible(false)
+
+        verify(keyguardTransitions).startKeyguardTransition(false, false)
+        verifyNoMoreInteractions(keyguardTransitions)
+
+        underTest.setSurfaceBehindVisibility(true)
+
+        verifyNoMoreInteractions(keyguardTransitions)
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testAodVisible_noLockscreenShownCallYet_doesNotShowLockscreenUntilLater_without_keyguard_shell_transitions() {
         underTest.setAodVisible(false)
         verifyNoMoreInteractions(activityTaskManagerService)
 
@@ -116,7 +170,19 @@
     }
 
     @Test
-    fun setSurfaceBehindVisibility_goesAwayFirst_andIgnoresSecondCall() {
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun testAodVisible_noLockscreenShownCallYet_doesNotShowLockscreenUntilLater_with_keyguard_shell_transitions() {
+        underTest.setAodVisible(false)
+        verifyNoMoreInteractions(keyguardTransitions)
+
+        underTest.setLockscreenShown(true)
+        verify(keyguardTransitions).startKeyguardTransition(true, false)
+        verifyNoMoreInteractions(activityTaskManagerService)
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun setSurfaceBehindVisibility_goesAwayFirst_andIgnoresSecondCall_without_keyguard_shell_transitions() {
         underTest.setLockscreenShown(true)
         underTest.setSurfaceBehindVisibility(true)
         verify(activityTaskManagerService).keyguardGoingAway(0)
@@ -126,8 +192,27 @@
     }
 
     @Test
-    fun setSurfaceBehindVisibility_falseSetsLockscreenVisibility() {
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun setSurfaceBehindVisibility_goesAwayFirst_andIgnoresSecondCall_with_keyguard_shell_transitions() {
+        underTest.setLockscreenShown(true)
+        underTest.setSurfaceBehindVisibility(true)
+        verify(keyguardTransitions).startKeyguardTransition(false, false)
+
+        underTest.setSurfaceBehindVisibility(true)
+        verifyNoMoreInteractions(keyguardTransitions)
+    }
+
+    @Test
+    @RequiresFlagsDisabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun setSurfaceBehindVisibility_falseSetsLockscreenVisibility_without_keyguard_shell_transitions() {
         underTest.setSurfaceBehindVisibility(false)
         verify(activityTaskManagerService).setLockScreenShown(eq(true), any())
     }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_ENSURE_KEYGUARD_DOES_TRANSITION_STARTING)
+    fun setSurfaceBehindVisibility_falseSetsLockscreenVisibility_with_keyguard_shell_transitions() {
+        underTest.setSurfaceBehindVisibility(false)
+        verify(keyguardTransitions).startKeyguardTransition(eq(true), any())
+    }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/FakeNoteTaskBubbleController.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/FakeNoteTaskBubbleController.kt
index 450aadd..ebc00c3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/FakeNoteTaskBubbleController.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/FakeNoteTaskBubbleController.kt
@@ -20,6 +20,7 @@
 import android.content.Intent
 import android.graphics.drawable.Icon
 import android.os.UserHandle
+import com.android.wm.shell.bubbles.Bubble
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
 import kotlinx.coroutines.CoroutineDispatcher
@@ -33,14 +34,29 @@
 class FakeNoteTaskBubbleController(
     unUsed1: Context,
     unsUsed2: CoroutineDispatcher,
-    private val optionalBubbles: Optional<Bubbles>
+    private val optionalBubbles: Optional<Bubbles>,
 ) : NoteTaskBubblesController(unUsed1, unsUsed2) {
     override suspend fun areBubblesAvailable() = optionalBubbles.isPresent
 
-    override suspend fun showOrHideAppBubble(intent: Intent, userHandle: UserHandle, icon: Icon) {
+    override suspend fun showOrHideAppBubble(
+        intent: Intent,
+        userHandle: UserHandle,
+        icon: Icon,
+        bubbleExpandBehavior: NoteTaskBubbleExpandBehavior,
+    ) {
         optionalBubbles.ifPresentOrElse(
-            { bubbles -> bubbles.showOrHideAppBubble(intent, userHandle, icon) },
-            { throw IllegalAccessException() }
+            { bubbles ->
+                if (
+                    bubbleExpandBehavior == NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED &&
+                        bubbles.isBubbleExpanded(
+                            Bubble.getAppBubbleKeyForApp(intent.`package`, userHandle)
+                        )
+                ) {
+                    return@ifPresentOrElse
+                }
+                bubbles.showOrHideAppBubble(intent, userHandle, icon)
+            },
+            { throw IllegalAccessException() },
         )
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskBubblesServiceTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskBubblesServiceTest.kt
index 9ef6b9c..e55d6ad 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskBubblesServiceTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskBubblesServiceTest.kt
@@ -21,9 +21,9 @@
 import android.os.UserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
-import com.android.systemui.res.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.notetask.NoteTaskBubblesController.NoteTaskBubblesService
+import com.android.systemui.res.R
 import com.android.wm.shell.bubbles.Bubbles
 import com.google.common.truth.Truth.assertThat
 import java.util.Optional
@@ -33,6 +33,9 @@
 import org.mockito.Mock
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.never
+import org.mockito.kotlin.whenever
 
 /** atest SystemUITests:NoteTaskBubblesServiceTest */
 @SmallTest
@@ -61,12 +64,40 @@
     }
 
     @Test
-    fun showOrHideAppBubble() {
+    fun showOrHideAppBubble_defaultExpandBehavior_shouldCallBubblesApi() {
         val intent = Intent()
         val user = UserHandle.SYSTEM
         val icon = Icon.createWithResource(context, R.drawable.ic_note_task_shortcut_widget)
+        val bubbleExpandBehavior = NoteTaskBubbleExpandBehavior.DEFAULT
+        whenever(bubbles.isBubbleExpanded(any())).thenReturn(false)
 
-        createServiceBinder().showOrHideAppBubble(intent, user, icon)
+        createServiceBinder().showOrHideAppBubble(intent, user, icon, bubbleExpandBehavior)
+
+        verify(bubbles).showOrHideAppBubble(intent, user, icon)
+    }
+
+    @Test
+    fun showOrHideAppBubble_keepIfExpanded_bubbleShown_shouldNotCallBubblesApi() {
+        val intent = Intent().apply { setPackage("test") }
+        val user = UserHandle.SYSTEM
+        val icon = Icon.createWithResource(context, R.drawable.ic_note_task_shortcut_widget)
+        val bubbleExpandBehavior = NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED
+        whenever(bubbles.isBubbleExpanded(any())).thenReturn(true)
+
+        createServiceBinder().showOrHideAppBubble(intent, user, icon, bubbleExpandBehavior)
+
+        verify(bubbles, never()).showOrHideAppBubble(intent, user, icon)
+    }
+
+    @Test
+    fun showOrHideAppBubble_keepIfExpanded_bubbleNotShown_shouldCallBubblesApi() {
+        val intent = Intent().apply { setPackage("test") }
+        val user = UserHandle.SYSTEM
+        val icon = Icon.createWithResource(context, R.drawable.ic_note_task_shortcut_widget)
+        val bubbleExpandBehavior = NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED
+        whenever(bubbles.isBubbleExpanded(any())).thenReturn(false)
+
+        createServiceBinder().showOrHideAppBubble(intent, user, icon, bubbleExpandBehavior)
 
         verify(bubbles).showOrHideAppBubble(intent, user, icon)
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
index 8f4078b..d3578fb 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/notetask/NoteTaskInfoTest.kt
@@ -19,6 +19,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.notetask.NoteTaskEntryPoint.QS_NOTES_TILE
 import com.android.systemui.notetask.NoteTaskEntryPoint.WIDGET_PICKER_SHORTCUT_IN_MULTI_WINDOW_MODE
 import com.google.common.truth.Truth.assertThat
 import org.junit.Test
@@ -44,10 +45,19 @@
     }
 
     @Test
-    fun launchMode_keyguardUnlocked_launchModeAppBubble() {
+    fun launchMode_keyguardUnlocked_launchModeAppBubble_withDefaultExpandBehavior() {
         val underTest = DEFAULT_INFO.copy(isKeyguardLocked = false)
 
-        assertThat(underTest.launchMode).isEqualTo(NoteTaskLaunchMode.AppBubble)
+        assertThat(underTest.launchMode)
+            .isEqualTo(NoteTaskLaunchMode.AppBubble(NoteTaskBubbleExpandBehavior.DEFAULT))
+    }
+
+    @Test
+    fun launchMode_keyguardUnlocked_qsTileEntryPoint_launchModeAppBubble_withKeepIfExpandedExpandBehavior() {
+        val underTest = DEFAULT_INFO.copy(isKeyguardLocked = false, entryPoint = QS_NOTES_TILE)
+
+        assertThat(underTest.launchMode)
+            .isEqualTo(NoteTaskLaunchMode.AppBubble(NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED))
     }
 
     private companion object {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
index fb7252b..60a1855 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/StatusBarIconViewTest.java
@@ -16,8 +16,6 @@
 
 package com.android.systemui.statusbar;
 
-import static android.app.Notification.CATEGORY_CALL;
-
 import static com.google.common.truth.Truth.assertThat;
 
 import static junit.framework.Assert.assertEquals;
@@ -199,19 +197,6 @@
     }
 
     @Test
-    public void testContentDescForNotification_noNotifContent() {
-        Notification n = new Notification.Builder(mContext, "test")
-                .setSmallIcon(0)
-                .setContentTitle("hello")
-                .setCategory(CATEGORY_CALL)
-                .build();
-        assertThat(NotificationContentDescription.contentDescForNotification(mContext, n)
-                .toString()).startsWith("com.android.systemui.tests notification");
-        assertThat(NotificationContentDescription.contentDescForNotification(mContext, n)
-                .toString()).doesNotContain("hello");
-    }
-
-    @Test
     @EnableFlags({Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS})
     public void setIcon_withPreloaded_usesPreloaded() {
         Icon mockIcon = mock(Icon.class);
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractorTest.kt
index 7c5a487..3f995c6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractorTest.kt
@@ -155,6 +155,30 @@
             }
         }
 
+    @Test
+    fun activeStreamChanges_showBoth() {
+        with(kosmos) {
+            testScope.runTest {
+                runCurrent()
+                fakeVolumeDialogController.updateState {
+                    activeStream = AudioManager.STREAM_SYSTEM
+                    states.put(AudioManager.STREAM_MUSIC, buildStreamState())
+                    states.put(AudioManager.STREAM_SYSTEM, buildStreamState())
+                }
+                val slidersModel by collectLastValue(underTest.sliders)
+                runCurrent()
+
+                fakeVolumeDialogController.updateState { activeStream = AudioManager.STREAM_MUSIC }
+                runCurrent()
+
+                assertThat(slidersModel!!.slider)
+                    .isEqualTo(VolumeDialogSliderType.Stream(AudioManager.STREAM_MUSIC))
+                assertThat(slidersModel!!.floatingSliders)
+                    .containsExactly(VolumeDialogSliderType.Stream(AudioManager.STREAM_SYSTEM))
+            }
+        }
+    }
+
     private fun buildStreamState(
         build: VolumeDialogController.StreamState.() -> Unit = {}
     ): VolumeDialogController.StreamState {
diff --git a/packages/SystemUI/res/values-land/dimens.xml b/packages/SystemUI/res/values-land/dimens.xml
index 235015b..70ae5c1 100644
--- a/packages/SystemUI/res/values-land/dimens.xml
+++ b/packages/SystemUI/res/values-land/dimens.xml
@@ -91,7 +91,6 @@
     <dimen name="notification_blocker_channel_list_height">128dp</dimen>
 
     <dimen name="keyguard_indication_margin_bottom">8dp</dimen>
-    <dimen name="lock_icon_margin_bottom">24dp</dimen>
 
     <!-- Keyboard shortcuts helper -->
     <dimen name="ksh_container_horizontal_margin">48dp</dimen>
diff --git a/packages/SystemUI/res/values-sw600dp-land/dimens.xml b/packages/SystemUI/res/values-sw600dp-land/dimens.xml
index 75bee9f..055c3a6 100644
--- a/packages/SystemUI/res/values-sw600dp-land/dimens.xml
+++ b/packages/SystemUI/res/values-sw600dp-land/dimens.xml
@@ -20,7 +20,6 @@
     <!-- keyguard-->
     <dimen name="keyguard_indication_margin_bottom">25dp</dimen>
     <dimen name="ambient_indication_margin_bottom">115dp</dimen>
-    <dimen name="lock_icon_margin_bottom">60dp</dimen>
     <!-- margin from keyguard status bar to clock. For split shade it should be
          keyguard_split_shade_top_margin - status_bar_header_height_keyguard = 8dp -->
     <dimen name="keyguard_clock_top_margin">8dp</dimen>
diff --git a/packages/SystemUI/res/values/dimens.xml b/packages/SystemUI/res/values/dimens.xml
index 580f6d3..813bb9c 100644
--- a/packages/SystemUI/res/values/dimens.xml
+++ b/packages/SystemUI/res/values/dimens.xml
@@ -948,7 +948,6 @@
     <dimen name="keyguard_translate_distance_on_swipe_up">-200dp</dimen>
 
     <dimen name="keyguard_indication_margin_bottom">32dp</dimen>
-    <dimen name="lock_icon_margin_bottom">74dp</dimen>
     <dimen name="ambient_indication_margin_bottom">71dp</dimen>
 
 
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 1a56541..400d097 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,6 +16,7 @@
 
 package com.android.systemui.deviceentry.domain.interactor
 
+import com.android.app.tracing.coroutines.launchTraced as launch
 import com.android.internal.policy.IKeyguardDismissCallback
 import com.android.systemui.authentication.domain.interactor.AuthenticationInteractor
 import com.android.systemui.authentication.shared.model.AuthenticationMethodModel
@@ -43,7 +44,6 @@
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.stateIn
-import com.android.app.tracing.coroutines.launchTraced as launch
 
 /**
  * Hosts application business logic related to device entry.
@@ -174,6 +174,14 @@
     }
 
     /**
+     * Whether lockscreen bypass is enabled. When enabled, the lockscreen will be automatically
+     * dismissed once the authentication challenge is completed. For example, completing a biometric
+     * authentication challenge via face unlock or fingerprint sensor can automatically bypass the
+     * lockscreen.
+     */
+    val isBypassEnabled: StateFlow<Boolean> = repository.isBypassEnabled
+
+    /**
      * Attempt to enter the device and dismiss the lockscreen. If authentication is required to
      * unlock the device it will transition to bouncer.
      *
@@ -238,11 +246,8 @@
         isLockscreenEnabled()
     }
 
-    /**
-     * Whether lockscreen bypass is enabled. When enabled, the lockscreen will be automatically
-     * dismissed once the authentication challenge is completed. For example, completing a biometric
-     * authentication challenge via face unlock or fingerprint sensor can automatically bypass the
-     * lockscreen.
-     */
-    val isBypassEnabled: StateFlow<Boolean> = repository.isBypassEnabled
+    /** Locks the device instantly. */
+    fun lockNow() {
+        deviceUnlockedInteractor.lockNow()
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractor.kt b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractor.kt
index 35eed5e..7d684ca 100644
--- a/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/deviceentry/domain/interactor/DeviceUnlockedInteractor.kt
@@ -43,6 +43,7 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.awaitCancellation
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
@@ -57,6 +58,7 @@
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.launch
 
 @OptIn(ExperimentalCoroutinesApi::class)
@@ -178,6 +180,8 @@
     val deviceUnlockStatus: StateFlow<DeviceUnlockStatus> =
         repository.deviceUnlockStatus.asStateFlow()
 
+    private val lockNowRequests = Channel<Unit>()
+
     override suspend fun onActivated(): Nothing {
         authenticationInteractor.authenticationMethod.collectLatest { authMethod ->
             if (!authMethod.isSecure) {
@@ -196,6 +200,11 @@
         awaitCancellation()
     }
 
+    /** Locks the device instantly. */
+    fun lockNow() {
+        lockNowRequests.trySend(Unit)
+    }
+
     private suspend fun handleLockAndUnlockEvents() {
         try {
             Log.d(TAG, "started watching for lock and unlock events")
@@ -225,10 +234,12 @@
                     .map { (isAsleep, lastSleepReason) ->
                         if (isAsleep) {
                             if (
-                                lastSleepReason == WakeSleepReason.POWER_BUTTON &&
+                                (lastSleepReason == WakeSleepReason.POWER_BUTTON) &&
                                     authenticationInteractor.getPowerButtonInstantlyLocks()
                             ) {
                                 LockImmediately("locked instantly from power button")
+                            } else if (lastSleepReason == WakeSleepReason.SLEEP_BUTTON) {
+                                LockImmediately("locked instantly from sleep button")
                             } else {
                                 LockWithDelay("entering sleep")
                             }
@@ -256,6 +267,7 @@
                         emptyFlow()
                     }
                 },
+                lockNowRequests.receiveAsFlow().map { LockImmediately("lockNow") },
             )
             .collectLatest(::onLockEvent)
     }
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
index 32c2bc7..7097c1d 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardService.java
@@ -662,7 +662,9 @@
             trace("doKeyguardTimeout");
             checkPermission();
 
-            if (KeyguardWmStateRefactor.isEnabled()) {
+            if (SceneContainerFlag.isEnabled()) {
+                mDeviceEntryInteractorLazy.get().lockNow();
+            } else if (KeyguardWmStateRefactor.isEnabled()) {
                 mKeyguardLockWhileAwakeInteractor.onKeyguardServiceDoKeyguardTimeout(options);
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultDeviceEntrySection.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultDeviceEntrySection.kt
index b51bb7b..aa7eb29 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultDeviceEntrySection.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/sections/DefaultDeviceEntrySection.kt
@@ -28,6 +28,7 @@
 import androidx.constraintlayout.widget.ConstraintLayout
 import androidx.constraintlayout.widget.ConstraintSet
 import com.android.systemui.biometrics.AuthController
+import com.android.systemui.customization.R as customR
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
@@ -115,7 +116,7 @@
 
         val scaleFactor: Float = authController.scaleFactor
         val mBottomPaddingPx =
-            context.resources.getDimensionPixelSize(R.dimen.lock_icon_margin_bottom)
+            context.resources.getDimensionPixelSize(customR.dimen.lock_icon_margin_bottom)
         val bounds = windowManager.currentWindowMetrics.bounds
         var widthPixels = bounds.right.toFloat()
         if (featureFlags.isEnabled(Flags.LOCKSCREEN_ENABLE_LANDSCAPE)) {
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/INoteTaskBubblesService.aidl b/packages/SystemUI/src/com/android/systemui/notetask/INoteTaskBubblesService.aidl
index 3e947d9..7803f22 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/INoteTaskBubblesService.aidl
+++ b/packages/SystemUI/src/com/android/systemui/notetask/INoteTaskBubblesService.aidl
@@ -16,6 +16,7 @@
 
 package com.android.systemui.notetask;
 
+import com.android.systemui.notetask.NoteTaskBubbleExpandBehavior;
 import android.content.Intent;
 import android.graphics.drawable.Icon;
 import android.os.UserHandle;
@@ -25,5 +26,6 @@
 
     boolean areBubblesAvailable();
 
-    void showOrHideAppBubble(in Intent intent, in UserHandle userHandle, in Icon icon);
+    void showOrHideAppBubble(in Intent intent, in UserHandle userHandle, in Icon icon,
+     in NoteTaskBubbleExpandBehavior bubbleExpandBehavior);
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.aidl b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.aidl
new file mode 100644
index 0000000..86a06a5
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.aidl
@@ -0,0 +1,3 @@
+package com.android.systemui.notetask;
+
+parcelable NoteTaskBubbleExpandBehavior;
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.kt
new file mode 100644
index 0000000..63b38a1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubbleExpandBehavior.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.notetask
+
+import android.os.Parcel
+import android.os.Parcelable
+
+enum class NoteTaskBubbleExpandBehavior : Parcelable {
+    /**
+     * The default bubble expand behavior for note task bubble: The bubble will collapse if there is
+     * already an expanded bubble, The bubble will expand if there is a collapsed bubble.
+     */
+    DEFAULT,
+    /**
+     * The special bubble expand behavior for note task bubble: The bubble will stay expanded, not
+     * collapse, if there is already an expanded bubble, The bubble will expand if there is a
+     * collapsed bubble.
+     */
+    KEEP_IF_EXPANDED;
+
+    override fun describeContents(): Int {
+        return 0
+    }
+
+    override fun writeToParcel(dest: Parcel, flags: Int) {
+        dest.writeString(name)
+    }
+
+    companion object CREATOR : Parcelable.Creator<NoteTaskBubbleExpandBehavior> {
+        override fun createFromParcel(parcel: Parcel?): NoteTaskBubbleExpandBehavior {
+            return parcel?.readString()?.let { valueOf(it) } ?: DEFAULT
+        }
+
+        override fun newArray(size: Int) = arrayOfNulls<NoteTaskBubbleExpandBehavior>(size)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubblesController.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubblesController.kt
index ec205f8..169285f 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubblesController.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskBubblesController.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.log.DebugLogger.debugLog
+import com.android.wm.shell.bubbles.Bubble
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
 import javax.inject.Inject
@@ -48,7 +49,7 @@
 @Inject
 constructor(
     @Application private val context: Context,
-    @Background private val bgDispatcher: CoroutineDispatcher
+    @Background private val bgDispatcher: CoroutineDispatcher,
 ) {
 
     private val serviceConnector: ServiceConnector<INoteTaskBubblesService> =
@@ -57,7 +58,7 @@
             Intent(context, NoteTaskBubblesService::class.java),
             Context.BIND_AUTO_CREATE or Context.BIND_WAIVE_PRIORITY or Context.BIND_NOT_VISIBLE,
             UserHandle.USER_SYSTEM,
-            INoteTaskBubblesService.Stub::asInterface
+            INoteTaskBubblesService.Stub::asInterface,
         )
 
     /** Returns whether notes app bubble is supported. */
@@ -79,11 +80,12 @@
     open suspend fun showOrHideAppBubble(
         intent: Intent,
         userHandle: UserHandle,
-        icon: Icon
+        icon: Icon,
+        bubbleExpandBehavior: NoteTaskBubbleExpandBehavior,
     ) {
         withContext(bgDispatcher) {
             serviceConnector
-                .post { it.showOrHideAppBubble(intent, userHandle, icon) }
+                .post { it.showOrHideAppBubble(intent, userHandle, icon, bubbleExpandBehavior) }
                 .whenComplete { _, error ->
                     if (error != null) {
                         debugLog(error = error) {
@@ -120,16 +122,28 @@
                 override fun showOrHideAppBubble(
                     intent: Intent,
                     userHandle: UserHandle,
-                    icon: Icon
+                    icon: Icon,
+                    bubbleExpandBehavior: NoteTaskBubbleExpandBehavior,
                 ) {
                     mOptionalBubbles.ifPresentOrElse(
-                        { bubbles -> bubbles.showOrHideAppBubble(intent, userHandle, icon) },
+                        { bubbles ->
+                            if (
+                                bubbleExpandBehavior ==
+                                    NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED &&
+                                    bubbles.isBubbleExpanded(
+                                        Bubble.getAppBubbleKeyForApp(intent.`package`, userHandle)
+                                    )
+                            ) {
+                                return@ifPresentOrElse
+                            }
+                            bubbles.showOrHideAppBubble(intent, userHandle, icon)
+                        },
                         {
                             debugLog {
                                 "Failed to show or hide bubble for intent $intent," +
                                     "user $user, and icon $icon as bubble is empty."
                             }
-                        }
+                        },
                     )
                 }
             }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
index 1fa5baa..a615963 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
@@ -84,7 +84,7 @@
     private val userTracker: UserTracker,
     private val secureSettings: SecureSettings,
     @Application private val applicationScope: CoroutineScope,
-    @Background private val bgCoroutineContext: CoroutineContext
+    @Background private val bgCoroutineContext: CoroutineContext,
 ) {
 
     @VisibleForTesting val infoReference = AtomicReference<NoteTaskInfo?>()
@@ -98,7 +98,7 @@
         if (key != Bubble.getAppBubbleKeyForApp(info.packageName, info.user)) return
 
         // Safe guard mechanism, this callback should only be called for app bubbles.
-        if (info.launchMode != NoteTaskLaunchMode.AppBubble) return
+        if (info.launchMode !is NoteTaskLaunchMode.AppBubble) return
 
         if (isExpanding) {
             debugLog { "onBubbleExpandChanged - expanding: $info" }
@@ -117,10 +117,8 @@
             } else {
                 getUserForHandlingNotesTaking(entryPoint)
             }
-        activityContext.startActivityAsUser(
-            createNotesRoleHolderSettingsIntent(),
-            user
-        )
+
+        activityContext.startActivityAsUser(createNotesRoleHolderSettingsIntent(), user)
     }
 
     /**
@@ -140,8 +138,7 @@
                 entryPoint == QUICK_AFFORDANCE -> {
                 userTracker.userProfiles
                     .firstOrNull { userManager.isManagedProfile(it.id) }
-                    ?.userHandle
-                    ?: userTracker.userHandle
+                    ?.userHandle ?: userTracker.userHandle
             }
             // On work profile devices, SysUI always run in the main user.
             else -> userTracker.userHandle
@@ -158,19 +155,14 @@
      *
      * That will let users open other apps in full screen, and take contextual notes.
      */
-    fun showNoteTask(
-        entryPoint: NoteTaskEntryPoint,
-    ) {
+    fun showNoteTask(entryPoint: NoteTaskEntryPoint) {
         if (!isEnabled) return
 
         showNoteTaskAsUser(entryPoint, getUserForHandlingNotesTaking(entryPoint))
     }
 
     /** A variant of [showNoteTask] which launches note task in the given [user]. */
-    fun showNoteTaskAsUser(
-        entryPoint: NoteTaskEntryPoint,
-        user: UserHandle,
-    ) {
+    fun showNoteTaskAsUser(entryPoint: NoteTaskEntryPoint, user: UserHandle) {
         if (!isEnabled) return
 
         applicationScope.launch("$TAG#showNoteTaskAsUser") {
@@ -178,10 +170,7 @@
         }
     }
 
-    private suspend fun awaitShowNoteTaskAsUser(
-        entryPoint: NoteTaskEntryPoint,
-        user: UserHandle,
-    ) {
+    private suspend fun awaitShowNoteTaskAsUser(entryPoint: NoteTaskEntryPoint, user: UserHandle) {
         if (!isEnabled) return
 
         if (!noteTaskBubblesController.areBubblesAvailable()) {
@@ -222,7 +211,13 @@
                     val intent = createNoteTaskIntent(info)
                     val icon =
                         Icon.createWithResource(context, R.drawable.ic_note_task_shortcut_widget)
-                    noteTaskBubblesController.showOrHideAppBubble(intent, user, icon)
+                    noteTaskBubblesController.showOrHideAppBubble(
+                        intent,
+                        user,
+                        icon,
+                        info.launchMode.bubbleExpandBehavior,
+                    )
+
                     // App bubble logging happens on `onBubbleExpandChanged`.
                     debugLog { "onShowNoteTask - opened as app bubble: $info" }
                 }
@@ -399,8 +394,8 @@
         const val EXTRA_SHORTCUT_BADGE_OVERRIDE_PACKAGE = "extra_shortcut_badge_override_package"
 
         /** Returns notes role holder settings intent. */
-        fun createNotesRoleHolderSettingsIntent() = Intent(Intent.ACTION_MANAGE_DEFAULT_APP).
-            putExtra(Intent.EXTRA_ROLE_NAME, ROLE_NOTES)
+        fun createNotesRoleHolderSettingsIntent() =
+            Intent(Intent.ACTION_MANAGE_DEFAULT_APP).putExtra(Intent.EXTRA_ROLE_NAME, ROLE_NOTES)
     }
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
index 269eb87..8319e07 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInfo.kt
@@ -31,6 +31,10 @@
         if (isKeyguardLocked || entryPoint == WIDGET_PICKER_SHORTCUT_IN_MULTI_WINDOW_MODE) {
             NoteTaskLaunchMode.Activity
         } else {
-            NoteTaskLaunchMode.AppBubble
+            if (entryPoint == NoteTaskEntryPoint.QS_NOTES_TILE) {
+                NoteTaskLaunchMode.AppBubble(NoteTaskBubbleExpandBehavior.KEEP_IF_EXPANDED)
+            } else {
+                NoteTaskLaunchMode.AppBubble(NoteTaskBubbleExpandBehavior.DEFAULT)
+            }
         }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskLaunchMode.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskLaunchMode.kt
index 836e103f..6c85f20 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskLaunchMode.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskLaunchMode.kt
@@ -26,7 +26,8 @@
 sealed class NoteTaskLaunchMode {
 
     /** @see Bubbles.showOrHideAppBubble */
-    object AppBubble : NoteTaskLaunchMode()
+    data class AppBubble(val bubbleExpandBehavior: NoteTaskBubbleExpandBehavior) :
+        NoteTaskLaunchMode()
 
     /** @see Context.startActivity */
     object Activity : NoteTaskLaunchMode()
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
index a1c5c9c..5d54656 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
@@ -41,6 +41,7 @@
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfig
 import com.android.systemui.qs.tiles.viewmodel.QSTileUIConfig
 import com.android.systemui.qs.tiles.viewmodel.QSTileViewModel
+import com.android.systemui.qs.tiles.viewmodel.StubQSTileViewModel
 import com.android.systemui.res.R
 import dagger.Binds
 import dagger.Module
@@ -106,12 +107,14 @@
             stateInteractor: NotesTileDataInteractor,
             userActionInteractor: NotesTileUserActionInteractor,
         ): QSTileViewModel =
-            factory.create(
-                TileSpec.create(NOTES_TILE_SPEC),
-                userActionInteractor,
-                stateInteractor,
-                mapper,
-            )
+            if (com.android.systemui.Flags.qsNewTilesFuture())
+                factory.create(
+                    TileSpec.create(NOTES_TILE_SPEC),
+                    userActionInteractor,
+                    stateInteractor,
+                    mapper,
+                )
+            else StubQSTileViewModel
 
         @Provides
         @IntoMap
@@ -120,10 +123,10 @@
             QSTileConfig(
                 tileSpec = TileSpec.create(NOTES_TILE_SPEC),
                 uiConfig =
-                QSTileUIConfig.Resource(
-                    iconRes = R.drawable.ic_qs_notes,
-                    labelRes = R.string.quick_settings_notes_label,
-                ),
+                    QSTileUIConfig.Resource(
+                        iconRes = R.drawable.ic_qs_notes,
+                        labelRes = R.string.quick_settings_notes_label,
+                    ),
                 instanceId = uiEventLogger.getNewInstanceId(),
                 category = TileCategory.UTILITIES,
             )
diff --git a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
index 776a8f4..c57b53b 100644
--- a/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
+++ b/packages/SystemUI/src/com/android/systemui/power/shared/model/WakeSleepReason.kt
@@ -26,6 +26,9 @@
     /** The physical power button was pressed to wake up or sleep the device. */
     POWER_BUTTON(isTouch = false, PowerManager.WAKE_REASON_POWER_BUTTON),
 
+    /** The sleep button was pressed to sleep the device. */
+    SLEEP_BUTTON(isTouch = false, PowerManager.GO_TO_SLEEP_REASON_SLEEP_BUTTON),
+
     /** The user has tapped or double tapped to wake the screen. */
     TAP(isTouch = true, PowerManager.WAKE_REASON_TAP),
 
@@ -78,6 +81,7 @@
         fun fromPowerManagerSleepReason(reason: Int): WakeSleepReason {
             return when (reason) {
                 PowerManager.GO_TO_SLEEP_REASON_POWER_BUTTON -> POWER_BUTTON
+                PowerManager.GO_TO_SLEEP_REASON_SLEEP_BUTTON -> SLEEP_BUTTON
                 PowerManager.GO_TO_SLEEP_REASON_TIMEOUT -> TIMEOUT
                 PowerManager.GO_TO_SLEEP_REASON_DEVICE_FOLD -> FOLD
                 else -> OTHER
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
index 2dcb706..d0c033b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationRowContentBinderImpl.kt
@@ -834,6 +834,15 @@
             public?.let {
                 it.layoutInflaterFactory = provider.provide(row, FLAG_CONTENT_VIEW_PUBLIC)
             }
+            if (android.app.Flags.notificationsRedesignAppIcons()) {
+                normalGroupHeader?.let {
+                    it.layoutInflaterFactory = provider.provide(row, FLAG_GROUP_SUMMARY_HEADER)
+                }
+                minimizedGroupHeader?.let {
+                    it.layoutInflaterFactory =
+                        provider.provide(row, FLAG_LOW_PRIORITY_GROUP_SUMMARY_HEADER)
+                }
+            }
             return this
         }
 
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 7ef1e41..5837a49 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
@@ -31,6 +31,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.function.Consumer;
 
 import javax.inject.Inject;
 
@@ -63,17 +64,21 @@
 
     @Override
     public void addCallback(@NonNull Callback callback) {
-        mCallbacks.add(callback);
-        if (mCallbacks.size() == 1) {
-            setListening(true);
+        synchronized (mCallbacks) {
+            mCallbacks.add(callback);
+            if (mCallbacks.size() == 1) {
+                setListening(true);
+            }
+            callback.onManagedProfileChanged();
         }
-        callback.onManagedProfileChanged();
     }
 
     @Override
     public void removeCallback(@NonNull Callback callback) {
-        if (mCallbacks.remove(callback) && mCallbacks.size() == 0) {
-            setListening(false);
+        synchronized (mCallbacks) {
+            if (mCallbacks.remove(callback) && mCallbacks.size() == 0) {
+                setListening(false);
+            }
         }
     }
 
@@ -109,10 +114,7 @@
     }
 
     private void notifyManagedProfileRemoved() {
-        ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
-        for (Callback callback : copy) {
-            callback.onManagedProfileRemoved();
-        }
+        notifyCallbacks(Callback::onManagedProfileRemoved);
     }
 
     public boolean hasActiveProfile() {
@@ -136,6 +138,16 @@
         }
     }
 
+    private void notifyCallbacks(Consumer<Callback> method) {
+        ArrayList<Callback> copy;
+        synchronized (mCallbacks) {
+            copy = new ArrayList<>(mCallbacks);
+        }
+        for (Callback callback : copy) {
+            method.accept(callback);
+        }
+    }
+
     private void setListening(boolean listening) {
         if (mListening == listening) {
             return;
@@ -154,19 +166,13 @@
         @Override
         public void onUserChanged(int newUser, @NonNull Context userContext) {
             reloadManagedProfiles();
-            ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
-            for (Callback callback : copy) {
-                callback.onManagedProfileChanged();
-            }
+            notifyCallbacks(Callback::onManagedProfileChanged);
         }
 
         @Override
         public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
             reloadManagedProfiles();
-            ArrayList<Callback> copy = new ArrayList<>(mCallbacks);
-            for (Callback callback : copy) {
-                callback.onManagedProfileChanged();
-            }
+            notifyCallbacks(Callback::onManagedProfileChanged);
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/TutorialSelectionScreen.kt b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/TutorialSelectionScreen.kt
index d371acf..66a900b 100644
--- a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/TutorialSelectionScreen.kt
+++ b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/composable/TutorialSelectionScreen.kt
@@ -18,6 +18,7 @@
 
 import android.content.res.Configuration
 import androidx.compose.foundation.background
+import androidx.compose.foundation.focusable
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
@@ -36,8 +37,12 @@
 import androidx.compose.material3.MaterialTheme
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
+import androidx.compose.runtime.LaunchedEffect
+import androidx.compose.runtime.remember
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.focus.FocusRequester
+import androidx.compose.ui.focus.focusRequester
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.vector.ImageVector
 import androidx.compose.ui.input.pointer.pointerInteropFilter
@@ -49,6 +54,7 @@
 import com.android.systemui.res.R
 import com.android.systemui.touchpad.tutorial.ui.gesture.isFourFingerTouchpadSwipe
 import com.android.systemui.touchpad.tutorial.ui.gesture.isThreeFingerTouchpadSwipe
+import com.android.systemui.touchpad.tutorial.ui.viewmodel.Screen
 
 @Composable
 fun TutorialSelectionScreen(
@@ -56,6 +62,7 @@
     onHomeTutorialClicked: () -> Unit,
     onRecentAppsTutorialClicked: () -> Unit,
     onDoneButtonClicked: () -> Unit,
+    lastSelectedScreen: Screen,
 ) {
     Column(
         verticalArrangement = Arrangement.Center,
@@ -80,6 +87,7 @@
                     onHomeTutorialClicked = onHomeTutorialClicked,
                     onRecentAppsTutorialClicked = onRecentAppsTutorialClicked,
                     modifier = Modifier.weight(1f).padding(60.dp),
+                    lastSelectedScreen,
                 )
             }
             else -> {
@@ -88,6 +96,7 @@
                     onHomeTutorialClicked = onHomeTutorialClicked,
                     onRecentAppsTutorialClicked = onRecentAppsTutorialClicked,
                     modifier = Modifier.weight(1f).padding(60.dp),
+                    lastSelectedScreen,
                 )
             }
         }
@@ -105,6 +114,7 @@
     onHomeTutorialClicked: () -> Unit,
     onRecentAppsTutorialClicked: () -> Unit,
     modifier: Modifier = Modifier,
+    lastSelectedScreen: Screen,
 ) {
     Row(
         horizontalArrangement = Arrangement.spacedBy(20.dp),
@@ -116,6 +126,7 @@
             onHomeTutorialClicked,
             onRecentAppsTutorialClicked,
             modifier = Modifier.weight(1f).fillMaxSize(),
+            lastSelectedScreen,
         )
     }
 }
@@ -126,6 +137,7 @@
     onHomeTutorialClicked: () -> Unit,
     onRecentAppsTutorialClicked: () -> Unit,
     modifier: Modifier = Modifier,
+    lastSelectedScreen: Screen,
 ) {
     Column(
         verticalArrangement = Arrangement.spacedBy(20.dp),
@@ -137,6 +149,7 @@
             onHomeTutorialClicked,
             onRecentAppsTutorialClicked,
             modifier = Modifier.weight(1f).fillMaxSize(),
+            lastSelectedScreen,
         )
     }
 }
@@ -147,14 +160,26 @@
     onHomeTutorialClicked: () -> Unit,
     onRecentAppsTutorialClicked: () -> Unit,
     modifier: Modifier = Modifier,
+    lastSelectedScreen: Screen,
 ) {
+    val homeFocusRequester = remember { FocusRequester() }
+    val backFocusRequester = remember { FocusRequester() }
+    val recentAppsFocusRequester = remember { FocusRequester() }
+    LaunchedEffect(Unit) {
+        when (lastSelectedScreen) {
+            Screen.HOME_GESTURE -> homeFocusRequester.requestFocus()
+            Screen.BACK_GESTURE -> backFocusRequester.requestFocus()
+            Screen.RECENT_APPS_GESTURE -> recentAppsFocusRequester.requestFocus()
+            else -> {} // No-Op.
+        }
+    }
     TutorialButton(
         text = stringResource(R.string.touchpad_tutorial_home_gesture_button),
         icon = ImageVector.vectorResource(id = R.drawable.touchpad_tutorial_home_icon),
         iconColor = MaterialTheme.colorScheme.onPrimary,
         onClick = onHomeTutorialClicked,
         backgroundColor = MaterialTheme.colorScheme.primary,
-        modifier = modifier,
+        modifier = modifier.focusRequester(homeFocusRequester).focusable(),
     )
     TutorialButton(
         text = stringResource(R.string.touchpad_tutorial_back_gesture_button),
@@ -162,7 +187,7 @@
         iconColor = MaterialTheme.colorScheme.onTertiary,
         onClick = onBackTutorialClicked,
         backgroundColor = MaterialTheme.colorScheme.tertiary,
-        modifier = modifier,
+        modifier = modifier.focusRequester(backFocusRequester).focusable(),
     )
     TutorialButton(
         text = stringResource(R.string.touchpad_tutorial_recent_apps_gesture_button),
@@ -170,7 +195,7 @@
         iconColor = MaterialTheme.colorScheme.onSecondary,
         onClick = onRecentAppsTutorialClicked,
         backgroundColor = MaterialTheme.colorScheme.secondary,
-        modifier = modifier,
+        modifier = modifier.focusRequester(recentAppsFocusRequester).focusable(),
     )
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/view/TouchpadTutorialActivity.kt b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/view/TouchpadTutorialActivity.kt
index e1f7bd5..6662fc5 100644
--- a/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/view/TouchpadTutorialActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/touchpad/tutorial/ui/view/TouchpadTutorialActivity.kt
@@ -24,12 +24,16 @@
 import androidx.activity.viewModels
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
 import androidx.lifecycle.Lifecycle.State.STARTED
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import com.android.compose.theme.PlatformTheme
 import com.android.systemui.inputdevice.tutorial.InputDeviceTutorialLogger
 import com.android.systemui.inputdevice.tutorial.InputDeviceTutorialLogger.TutorialContext
 import com.android.systemui.inputdevice.tutorial.KeyboardTouchpadTutorialMetricsLogger
+import com.android.systemui.res.R
 import com.android.systemui.touchpad.tutorial.ui.composable.BackGestureTutorialScreen
 import com.android.systemui.touchpad.tutorial.ui.composable.HomeGestureTutorialScreen
 import com.android.systemui.touchpad.tutorial.ui.composable.RecentAppsGestureTutorialScreen
@@ -54,6 +58,7 @@
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         enableEdgeToEdge()
+        setTitle(getString(R.string.launch_touchpad_tutorial_notification_content))
         setContent {
             PlatformTheme { TouchpadTutorialScreen(vm, closeTutorial = ::finishTutorial) }
         }
@@ -82,13 +87,24 @@
 @Composable
 fun TouchpadTutorialScreen(vm: TouchpadTutorialViewModel, closeTutorial: () -> Unit) {
     val activeScreen by vm.screen.collectAsStateWithLifecycle(STARTED)
+    var lastSelectedScreen by remember { mutableStateOf(TUTORIAL_SELECTION) }
     when (activeScreen) {
         TUTORIAL_SELECTION ->
             TutorialSelectionScreen(
-                onBackTutorialClicked = { vm.goTo(BACK_GESTURE) },
-                onHomeTutorialClicked = { vm.goTo(HOME_GESTURE) },
-                onRecentAppsTutorialClicked = { vm.goTo(RECENT_APPS_GESTURE) },
+                onBackTutorialClicked = {
+                    lastSelectedScreen = BACK_GESTURE
+                    vm.goTo(BACK_GESTURE)
+                },
+                onHomeTutorialClicked = {
+                    lastSelectedScreen = HOME_GESTURE
+                    vm.goTo(HOME_GESTURE)
+                },
+                onRecentAppsTutorialClicked = {
+                    lastSelectedScreen = RECENT_APPS_GESTURE
+                    vm.goTo(RECENT_APPS_GESTURE)
+                },
                 onDoneButtonClicked = closeTutorial,
+                lastSelectedScreen,
             )
         BACK_GESTURE ->
             BackGestureTutorialScreen(
diff --git a/packages/SystemUI/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractor.kt
index 66a598b..c904ac5 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/dialog/sliders/domain/interactor/VolumeDialogSlidersInteractor.kt
@@ -35,6 +35,7 @@
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.runningReduce
 import kotlinx.coroutines.flow.stateIn
 
 private const val DEFAULT_STREAM = AudioManager.STREAM_MUSIC
@@ -54,13 +55,17 @@
         volumeDialogStateInteractor.volumeDialogState
             .filter { it.streamModels.isNotEmpty() }
             .map { stateModel ->
-                stateModel.streamModels.values
-                    .filter { streamModel -> shouldShowSliders(stateModel, streamModel) }
-                    .sortedWith(streamsSorter)
+                val sliderTypes =
+                    stateModel.streamModels.values
+                        .filter { streamModel -> shouldShowSliders(stateModel, streamModel) }
+                        .sortedWith(streamsSorter)
+                        .map { model -> model.toType() }
+                LinkedHashSet(sliderTypes)
             }
-            .map { models ->
-                val sliderTypes: List<VolumeDialogSliderType> =
-                    models.map { model -> model.toType() }
+            .runningReduce { sliderTypes, newSliderTypes ->
+                newSliderTypes.apply { addAll(sliderTypes) }
+            }
+            .map { sliderTypes ->
                 VolumeDialogSlidersModel(
                     slider = sliderTypes.first(),
                     floatingSliders = sliderTypes.drop(1),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
index 6e9b24f..4ca84c58 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
@@ -37,7 +37,6 @@
 import java.io.File
 import java.io.FileInputStream
 import java.io.FileOutputStream
-import org.junit.After
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -63,22 +62,18 @@
             Room.inMemoryDatabaseBuilder(context, CommunalDatabase::class.java)
                 .allowMainThreadQueries()
                 .build()
+        onTeardown { database.close() }
         CommunalDatabase.setInstance(database)
 
         dao = database.communalWidgetDao()
         backupUtils = CommunalBackupUtils(context)
 
         backupDataFile = File(context.cacheDir, "backup_data_file")
+        onTeardown { backupDataFile.delete() }
 
         underTest = CommunalBackupHelper(UserHandle.SYSTEM, backupUtils)
     }
 
-    @After
-    fun teardown() {
-        backupDataFile.delete()
-        database.close()
-    }
-
     @Test
     @EnableFlags(Flags.FLAG_COMMUNAL_HUB)
     fun backupAndRestoreCommunalHub() {
diff --git a/services/autofill/java/com/android/server/autofill/Helper.java b/services/autofill/java/com/android/server/autofill/Helper.java
index e59bb42..4f632c9 100644
--- a/services/autofill/java/com/android/server/autofill/Helper.java
+++ b/services/autofill/java/com/android/server/autofill/Helper.java
@@ -353,7 +353,10 @@
     private static void addAutofillableIds(@NonNull ViewNode node,
             @NonNull ArrayList<AutofillId> ids, boolean autofillableOnly) {
         if (!autofillableOnly || node.getAutofillType() != View.AUTOFILL_TYPE_NONE) {
-            ids.add(node.getAutofillId());
+            AutofillId id = node.getAutofillId();
+            if (id != null) {
+                ids.add(id);
+            }
         }
         final int size = node.getChildCount();
         for (int i = 0; i < size; i++) {
diff --git a/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
new file mode 100644
index 0000000..5e4bab15
--- /dev/null
+++ b/services/backup/java/com/android/server/backup/BackupAgentConnectionManager.java
@@ -0,0 +1,318 @@
+/*
+ * Copyright (C) 2017 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.backup;
+
+import static com.android.server.backup.BackupManagerService.MORE_DEBUG;
+import static com.android.server.backup.BackupManagerService.TAG;
+
+import android.annotation.Nullable;
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
+import android.app.ApplicationThreadConstants;
+import android.app.IActivityManager;
+import android.app.IBackupAgent;
+import android.app.backup.BackupAnnotations;
+import android.app.compat.CompatChanges;
+import android.compat.annotation.ChangeId;
+import android.compat.annotation.EnabledSince;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.content.pm.PackageManager.NameNotFoundException;
+import android.os.Binder;
+import android.os.Build;
+import android.os.IBinder;
+import android.os.Process;
+import android.os.RemoteException;
+import android.os.UserHandle;
+import android.util.ArraySet;
+import android.util.Slog;
+
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.LocalServices;
+import com.android.server.backup.internal.LifecycleOperationStorage;
+
+import java.util.Set;
+
+/**
+ * Handles the lifecycle of {@link IBackupAgent}s that the {@link UserBackupManagerService}
+ * communicates with.
+ *
+ * <p>There can only be one agent that's connected to at a time.
+ *
+ * <p>There should be only one instance of this class per {@link UserBackupManagerService}.
+ */
+public class BackupAgentConnectionManager {
+
+    /**
+     * Enables the OS making a decision on whether backup restricted mode should be used for apps
+     * that haven't explicitly opted in or out. See
+     * {@link android.content.pm.PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE} for details.
+     */
+    @ChangeId
+    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.BAKLAVA)
+    public static final long OS_DECIDES_BACKUP_RESTRICTED_MODE = 376661510;
+
+    // The thread performing the sequence of queued backups binds to each app's agent
+    // in succession.  Bind notifications are asynchronously delivered through the
+    // Activity Manager; use this lock object to signal when a requested binding has
+    // completed.
+    private final Object mAgentConnectLock = new Object();
+    private IBackupAgent mConnectedAgent;
+    private volatile boolean mConnecting;
+    private final ArraySet<String> mRestoreNoRestrictedModePackages = new ArraySet<>();
+    private final ArraySet<String> mBackupNoRestrictedModePackages = new ArraySet<>();
+
+    private final IActivityManager mActivityManager;
+    private final ActivityManagerInternal mActivityManagerInternal;
+    private final LifecycleOperationStorage mOperationStorage;
+    private final PackageManager mPackageManager;
+    private final UserBackupManagerService mUserBackupManagerService;
+    private final int mUserId;
+    private final String mUserIdMsg;
+
+    BackupAgentConnectionManager(LifecycleOperationStorage operationStorage,
+            PackageManager packageManager, UserBackupManagerService userBackupManagerService,
+            int userId) {
+        mActivityManager = ActivityManager.getService();
+        mActivityManagerInternal = LocalServices.getService(ActivityManagerInternal.class);
+        mOperationStorage = operationStorage;
+        mPackageManager = packageManager;
+        mUserBackupManagerService = userBackupManagerService;
+        mUserId = userId;
+        mUserIdMsg = "[UserID:" + userId + "] ";
+    }
+
+    /**
+     * Fires off a backup agent, blocking until it attaches (and ActivityManager will call
+     * {@link #agentConnected(String, IBinder)}) or until this operation times out.
+     *
+     * @param mode a {@code BACKUP_MODE} from {@link android.app.ApplicationThreadConstants}.
+     */
+    @Nullable
+    public IBackupAgent bindToAgentSynchronous(ApplicationInfo app, int mode,
+            @BackupAnnotations.BackupDestination int backupDestination) {
+        IBackupAgent agent = null;
+        synchronized (mAgentConnectLock) {
+            mConnecting = true;
+            mConnectedAgent = null;
+            boolean useRestrictedMode = shouldUseRestrictedBackupModeForPackage(mode,
+                    app.packageName);
+            try {
+                if (mActivityManager.bindBackupAgent(app.packageName, mode, mUserId,
+                        backupDestination, useRestrictedMode)) {
+                    Slog.d(TAG, mUserIdMsg + "awaiting agent for " + app);
+
+                    // success; wait for the agent to arrive
+                    // only wait 10 seconds for the bind to happen
+                    long timeoutMark = System.currentTimeMillis() + 10 * 1000;
+                    while (mConnecting && mConnectedAgent == null && (System.currentTimeMillis()
+                            < timeoutMark)) {
+                        try {
+                            mAgentConnectLock.wait(5000);
+                        } catch (InterruptedException e) {
+                            // just bail
+                            Slog.w(TAG, mUserIdMsg + "Interrupted: " + e);
+                            mConnecting = false;
+                            mConnectedAgent = null;
+                        }
+                    }
+
+                    // if we timed out with no connect, abort and move on
+                    if (mConnecting) {
+                        Slog.w(TAG, mUserIdMsg + "Timeout waiting for agent " + app);
+                        mConnectedAgent = null;
+                    }
+                    Slog.i(TAG, mUserIdMsg + "got agent " + mConnectedAgent);
+                    agent = mConnectedAgent;
+                }
+            } catch (RemoteException e) {
+                // can't happen - ActivityManager is local
+            }
+        }
+        if (agent == null) {
+            mActivityManagerInternal.clearPendingBackup(mUserId);
+        }
+        return agent;
+    }
+
+    /**
+     * Tell the ActivityManager that we are done with the {@link IBackupAgent} of this {@code app}.
+     * It will tell the app to destroy the agent.
+     */
+    public void unbindAgent(ApplicationInfo app) {
+        try {
+            mActivityManager.unbindBackupAgent(app);
+        } catch (RemoteException e) {
+            // Can't happen - activity manager is local
+        }
+    }
+
+    /**
+     * Callback: a requested backup agent has been instantiated. This should only be called from
+     * the
+     * {@link ActivityManager} when it's telling us that an agent is ready after a call to
+     * {@link #bindToAgentSynchronous(ApplicationInfo, int, int)}.
+     */
+    public void agentConnected(String packageName, IBinder agentBinder) {
+        synchronized (mAgentConnectLock) {
+            if (getCallingUid() == android.os.Process.SYSTEM_UID) {
+                Slog.d(TAG,
+                        mUserIdMsg + "agentConnected pkg=" + packageName + " agent=" + agentBinder);
+                mConnectedAgent = IBackupAgent.Stub.asInterface(agentBinder);
+                mConnecting = false;
+            } else {
+                Slog.w(TAG, mUserIdMsg + "Non-system process uid=" + getCallingUid()
+                        + " claiming agent connected");
+            }
+            mAgentConnectLock.notifyAll();
+        }
+    }
+
+    /**
+     * Callback: a backup agent has failed to come up, or has unexpectedly quit. If the agent failed
+     * to come up in the first place, the agentBinder argument will be {@code null}. This should
+     * only be called from the {@link ActivityManager}.
+     */
+    public void agentDisconnected(String packageName) {
+        synchronized (mAgentConnectLock) {
+            if (getCallingUid() == Process.SYSTEM_UID) {
+                mConnectedAgent = null;
+                mConnecting = false;
+            } else {
+                Slog.w(TAG, mUserIdMsg + "Non-system process uid=" + getCallingUid()
+                        + " claiming agent disconnected");
+            }
+            Slog.w(TAG, mUserIdMsg + "agentDisconnected: the backup agent for " + packageName
+                    + " died: cancel current operations");
+
+            // Offload operation cancellation off the main thread as the cancellation callbacks
+            // might call out to BackupTransport. Other operations started on the same package
+            // before the cancellation callback has executed will also be cancelled by the callback.
+            Runnable cancellationRunnable = () -> {
+                // handleCancel() causes the PerformFullTransportBackupTask to go on to
+                // tearDownAgentAndKill: that will unbindBackupAgent in the Activity Manager, so
+                // that the package being backed up doesn't get stuck in restricted mode until the
+                // backup time-out elapses.
+                for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
+                    if (MORE_DEBUG) {
+                        Slog.d(TAG,
+                                mUserIdMsg + "agentDisconnected: will handleCancel(all) for token:"
+                                        + Integer.toHexString(token));
+                    }
+                    mUserBackupManagerService.handleCancel(token, true /* cancelAll */);
+                }
+            };
+            getThreadForCancellation(cancellationRunnable).start();
+
+            mAgentConnectLock.notifyAll();
+        }
+    }
+
+    /**
+     * Marks the given set of packages as packages that should not be put into restricted mode if
+     * they are started for the given {@link BackupAnnotations.OperationType}.
+     */
+    public void setNoRestrictedModePackages(Set<String> packageNames,
+            @BackupAnnotations.OperationType int opType) {
+        if (opType == BackupAnnotations.OperationType.BACKUP) {
+            mBackupNoRestrictedModePackages.clear();
+            mBackupNoRestrictedModePackages.addAll(packageNames);
+        } else if (opType == BackupAnnotations.OperationType.RESTORE) {
+            mRestoreNoRestrictedModePackages.clear();
+            mRestoreNoRestrictedModePackages.addAll(packageNames);
+        } else {
+            throw new IllegalArgumentException("opType must be BACKUP or RESTORE");
+        }
+    }
+
+    /**
+     * Clears the list of packages that should not be put into restricted mode for either backup or
+     * restore.
+     */
+    public void clearNoRestrictedModePackages() {
+        mBackupNoRestrictedModePackages.clear();
+        mRestoreNoRestrictedModePackages.clear();
+    }
+
+    /**
+     * If the app has specified {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE}, then
+     * its value is returned. If it hasn't and it targets an SDK below
+     * {@link Build.VERSION_CODES#BAKLAVA} then returns true. If it targets a newer SDK, then
+     * returns the decision made by the {@link android.app.backup.BackupTransport}.
+     *
+     * <p>When this method is called, we should have already asked the transport and cached its
+     * response in {@link #mBackupNoRestrictedModePackages} or
+     * {@link #mRestoreNoRestrictedModePackages} so this method will immediately return without
+     * any IPC to the transport.
+     */
+    private boolean shouldUseRestrictedBackupModeForPackage(
+            @BackupAnnotations.OperationType int mode, String packageName) {
+        if (!Flags.enableRestrictedModeChanges()) {
+            return true;
+        }
+
+        // Key/Value apps are never put in restricted mode.
+        if (mode == ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL
+                || mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE) {
+            return false;
+        }
+
+        try {
+            PackageManager.Property property = mPackageManager.getPropertyAsUser(
+                    PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE,
+                    packageName, /* className= */ null, mUserId);
+            if (property.isBoolean()) {
+                // If the package has explicitly specified, we won't ask the transport.
+                return property.getBoolean();
+            } else {
+                Slog.w(TAG,
+                        PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE + "must be a boolean.");
+            }
+        } catch (NameNotFoundException e) {
+            // This is expected when the package has not defined the property in its manifest.
+        }
+
+        // The package has not specified the property. The behavior depends on the package's
+        // targetSdk.
+        // <36 gets the old behavior of always using restricted mode.
+        if (!CompatChanges.isChangeEnabled(OS_DECIDES_BACKUP_RESTRICTED_MODE, packageName,
+                UserHandle.of(mUserId))) {
+            return true;
+        }
+
+        // Apps targeting >=36 get the behavior decided by the transport.
+        // By this point, we should have asked the transport and cached its decision.
+        if ((mode == ApplicationThreadConstants.BACKUP_MODE_FULL
+                && mBackupNoRestrictedModePackages.contains(packageName)) || (
+                mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL
+                        && mRestoreNoRestrictedModePackages.contains(packageName))) {
+            Slog.d(TAG, "Transport requested no restricted mode for: " + packageName);
+            return false;
+        }
+        return true;
+    }
+
+    @VisibleForTesting
+    Thread getThreadForCancellation(Runnable operation) {
+        return new Thread(operation, /* operationName */ "agent-disconnected");
+    }
+
+    @VisibleForTesting
+    int getCallingUid() {
+        return Binder.getCallingUid();
+    }
+}
diff --git a/services/backup/java/com/android/server/backup/BackupManagerService.java b/services/backup/java/com/android/server/backup/BackupManagerService.java
index 5f0071d..3f6ede9 100644
--- a/services/backup/java/com/android/server/backup/BackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/BackupManagerService.java
@@ -658,7 +658,8 @@
                 getServiceForUserIfCallerHasPermission(userId, "agentConnected()");
 
         if (userBackupManagerService != null) {
-            userBackupManagerService.agentConnected(packageName, agentBinder);
+            userBackupManagerService.getBackupAgentConnectionManager().agentConnected(packageName,
+                    agentBinder);
         }
     }
 
@@ -683,7 +684,8 @@
                 getServiceForUserIfCallerHasPermission(userId, "agentDisconnected()");
 
         if (userBackupManagerService != null) {
-            userBackupManagerService.agentDisconnected(packageName);
+            userBackupManagerService.getBackupAgentConnectionManager().agentDisconnected(
+                    packageName);
         }
     }
 
diff --git a/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java b/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
index f5d6836..136bacd 100644
--- a/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
+++ b/services/backup/java/com/android/server/backup/KeyValueAdbBackupEngine.java
@@ -146,7 +146,8 @@
 
     private IBackupAgent bindToAgent(ApplicationInfo targetApp) {
         try {
-            return mBackupManagerService.bindToAgentSynchronous(targetApp,
+            return mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
+                    targetApp,
                     ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL,
                     BackupAnnotations.BackupDestination.CLOUD);
         } catch (SecurityException e) {
diff --git a/services/backup/java/com/android/server/backup/UserBackupManagerService.java b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
index 5de2fb3..e085f6e 100644
--- a/services/backup/java/com/android/server/backup/UserBackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
@@ -43,9 +43,7 @@
 import android.app.ActivityManagerInternal;
 import android.app.AlarmManager;
 import android.app.AppGlobals;
-import android.app.ApplicationThreadConstants;
 import android.app.IActivityManager;
-import android.app.IBackupAgent;
 import android.app.PendingIntent;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupAnnotations;
@@ -60,9 +58,6 @@
 import android.app.backup.IFullBackupRestoreObserver;
 import android.app.backup.IRestoreSession;
 import android.app.backup.ISelectBackupTransportCallback;
-import android.app.compat.CompatChanges;
-import android.compat.annotation.ChangeId;
-import android.compat.annotation.EnabledSince;
 import android.content.ActivityNotFoundException;
 import android.content.BroadcastReceiver;
 import android.content.ComponentName;
@@ -83,7 +78,6 @@
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.HandlerThread;
-import android.os.IBinder;
 import android.os.Message;
 import android.os.ParcelFileDescriptor;
 import android.os.PowerManager;
@@ -302,21 +296,10 @@
     private static final String BACKUP_FINISHED_ACTION = "android.intent.action.BACKUP_FINISHED";
     private static final String BACKUP_FINISHED_PACKAGE_EXTRA = "packageName";
 
-    /**
-     * Enables the OS making a decision on whether backup restricted mode should be used for apps
-     * that haven't explicitly opted in or out. See
-     * {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE} for details.
-     */
-    @ChangeId
-    @EnabledSince(targetSdkVersion = Build.VERSION_CODES.BAKLAVA)
-    public static final long OS_DECIDES_BACKUP_RESTRICTED_MODE = 376661510;
-
     // Time delay for initialization operations that can be delayed so as not to consume too much
     // CPU on bring-up and increase time-to-UI.
     private static final long INITIALIZATION_DELAY_MILLIS = 3000;
 
-    // Timeout interval for deciding that a bind has taken too long.
-    private static final long BIND_TIMEOUT_INTERVAL = 10 * 1000;
     // Timeout interval for deciding that a clear-data has taken too long.
     private static final long CLEAR_DATA_TIMEOUT_INTERVAL = 30 * 1000;
 
@@ -365,22 +348,9 @@
     // Backups that we haven't started yet.  Keys are package names.
     private final HashMap<String, BackupRequest> mPendingBackups = new HashMap<>();
 
-    private final ArraySet<String> mRestoreNoRestrictedModePackages = new ArraySet<>();
-    private final ArraySet<String> mBackupNoRestrictedModePackages = new ArraySet<>();
-
     // locking around the pending-backup management
     private final Object mQueueLock = new Object();
-
     private final UserBackupPreferences mBackupPreferences;
-
-    // The thread performing the sequence of queued backups binds to each app's agent
-    // in succession.  Bind notifications are asynchronously delivered through the
-    // Activity Manager; use this lock object to signal when a requested binding has
-    // completed.
-    private final Object mAgentConnectLock = new Object();
-    private IBackupAgent mConnectedAgent;
-    private volatile boolean mConnecting;
-
     private volatile boolean mBackupRunning;
     private volatile long mLastBackupPass;
 
@@ -410,6 +380,7 @@
 
     private ActiveRestoreSession mActiveRestoreSession;
 
+    private final BackupAgentConnectionManager mBackupAgentConnectionManager;
     private final LifecycleOperationStorage mOperationStorage;
 
     private final Random mTokenGenerator = new Random();
@@ -547,6 +518,8 @@
         mRegisterTransportsRequestedTime = 0;
         mPackageManager = packageManager;
         mOperationStorage = operationStorage;
+        mBackupAgentConnectionManager = new BackupAgentConnectionManager(mOperationStorage,
+                mPackageManager, this, mUserId);
         mTransportManager = transportManager;
         mFullBackupQueue = new ArrayList<>();
         mBackupHandler = backupHandler;
@@ -599,6 +572,8 @@
         mAgentTimeoutParameters.start();
 
         mOperationStorage = new LifecycleOperationStorage(mUserId);
+        mBackupAgentConnectionManager = new BackupAgentConnectionManager(mOperationStorage,
+                mPackageManager, this, mUserId);
 
         Objects.requireNonNull(userBackupThread, "userBackupThread cannot be null");
         mBackupHandler = new BackupHandler(this, mOperationStorage, userBackupThread);
@@ -1660,67 +1635,6 @@
         }
     }
 
-    /** Fires off a backup agent, blocking until it attaches or times out. */
-    @Nullable
-    public IBackupAgent bindToAgentSynchronous(ApplicationInfo app, int mode,
-            @BackupDestination int backupDestination) {
-        IBackupAgent agent = null;
-        synchronized (mAgentConnectLock) {
-            mConnecting = true;
-            mConnectedAgent = null;
-            boolean useRestrictedMode = shouldUseRestrictedBackupModeForPackage(mode,
-                    app.packageName);
-            try {
-                if (mActivityManager.bindBackupAgent(app.packageName, mode, mUserId,
-                        backupDestination, useRestrictedMode)) {
-                    Slog.d(TAG, addUserIdToLogMessage(mUserId, "awaiting agent for " + app));
-
-                    // success; wait for the agent to arrive
-                    // only wait 10 seconds for the bind to happen
-                    long timeoutMark = System.currentTimeMillis() + BIND_TIMEOUT_INTERVAL;
-                    while (mConnecting && mConnectedAgent == null
-                            && (System.currentTimeMillis() < timeoutMark)) {
-                        try {
-                            mAgentConnectLock.wait(5000);
-                        } catch (InterruptedException e) {
-                            // just bail
-                            Slog.w(TAG, addUserIdToLogMessage(mUserId, "Interrupted: " + e));
-                            mConnecting = false;
-                            mConnectedAgent = null;
-                        }
-                    }
-
-                    // if we timed out with no connect, abort and move on
-                    if (mConnecting) {
-                        Slog.w(
-                                TAG,
-                                addUserIdToLogMessage(mUserId, "Timeout waiting for agent " + app));
-                        mConnectedAgent = null;
-                    }
-                    if (DEBUG) {
-                        Slog.i(TAG, addUserIdToLogMessage(mUserId, "got agent " + mConnectedAgent));
-                    }
-                    agent = mConnectedAgent;
-                }
-            } catch (RemoteException e) {
-                // can't happen - ActivityManager is local
-            }
-        }
-        if (agent == null) {
-            mActivityManagerInternal.clearPendingBackup(mUserId);
-        }
-        return agent;
-    }
-
-    /** Unbind from a backup agent. */
-    public void unbindAgent(ApplicationInfo app) {
-        try {
-            mActivityManager.unbindBackupAgent(app);
-        } catch (RemoteException e) {
-            // Can't happen - activity manager is local
-        }
-    }
-
     /**
      * Clear an application's data after a failed restore, blocking until the operation completes or
      * times out.
@@ -2493,10 +2407,6 @@
         AppWidgetBackupBridge.restoreWidgetState(packageName, widgetData, mUserId);
     }
 
-    // *****************************
-    // NEW UNIFIED RESTORE IMPLEMENTATION
-    // *****************************
-
     /** Schedule a backup pass for {@code packageName}. */
     public void dataChangedImpl(String packageName) {
         HashSet<String> targets = dataChangedTargets(packageName);
@@ -3122,91 +3032,6 @@
         }
     }
 
-    /**
-     * Marks the given set of packages as packages that should not be put into restricted mode if
-     * they are started for the given {@link BackupAnnotations.OperationType}.
-     */
-    public void setNoRestrictedModePackages(Set<String> packageNames,
-            @BackupAnnotations.OperationType int opType) {
-        if (opType == BackupAnnotations.OperationType.BACKUP) {
-            mBackupNoRestrictedModePackages.clear();
-            mBackupNoRestrictedModePackages.addAll(packageNames);
-        } else if (opType == BackupAnnotations.OperationType.RESTORE) {
-            mRestoreNoRestrictedModePackages.clear();
-            mRestoreNoRestrictedModePackages.addAll(packageNames);
-        } else {
-            throw new IllegalArgumentException("opType must be BACKUP or RESTORE");
-        }
-    }
-
-    /**
-     * Clears the list of packages that should not be put into restricted mode for either backup or
-     * restore.
-     */
-    public void clearNoRestrictedModePackages() {
-        mBackupNoRestrictedModePackages.clear();
-        mRestoreNoRestrictedModePackages.clear();
-    }
-
-    /**
-     * If the app has specified {@link PackageManager#PROPERTY_USE_RESTRICTED_BACKUP_MODE}, then
-     * its value is returned. If it hasn't and it targets an SDK below
-     * {@link Build.VERSION_CODES#BAKLAVA} then returns true. If it targets a newer SDK, then
-     * returns the decision made by the {@link android.app.backup.BackupTransport}.
-     *
-     * <p>When this method is called, we should have already asked the transport and cached its
-     * response in {@link #mBackupNoRestrictedModePackages} or
-     * {@link #mRestoreNoRestrictedModePackages} so this method will immediately return without
-     * any IPC to the transport.
-     */
-    private boolean shouldUseRestrictedBackupModeForPackage(
-            @BackupAnnotations.OperationType int mode, String packageName) {
-        if (!Flags.enableRestrictedModeChanges()) {
-            return true;
-        }
-
-        // Key/Value apps are never put in restricted mode.
-        if (mode == ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL
-                || mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE) {
-            return false;
-        }
-
-        try {
-            PackageManager.Property property = mPackageManager.getPropertyAsUser(
-                    PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE,
-                    packageName, /* className= */ null,
-                    mUserId);
-            if (property.isBoolean()) {
-                // If the package has explicitly specified, we won't ask the transport.
-                return property.getBoolean();
-            } else {
-                Slog.w(TAG, PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE
-                        + "must be a boolean.");
-            }
-        } catch (NameNotFoundException e) {
-            // This is expected when the package has not defined the property in its manifest.
-        }
-
-        // The package has not specified the property. The behavior depends on the package's
-        // targetSdk.
-        // <36 gets the old behavior of always using restricted mode.
-        if (!CompatChanges.isChangeEnabled(OS_DECIDES_BACKUP_RESTRICTED_MODE, packageName,
-                UserHandle.of(mUserId))) {
-            return true;
-        }
-
-        // Apps targeting >=36 get the behavior decided by the transport.
-        // By this point, we should have asked the transport and cached its decision.
-        if ((mode == ApplicationThreadConstants.BACKUP_MODE_FULL
-                && mBackupNoRestrictedModePackages.contains(packageName))
-                || (mode == ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL
-                && mRestoreNoRestrictedModePackages.contains(packageName))) {
-            Slog.d(TAG, "Transport requested no restricted mode for: " + packageName);
-            return false;
-        }
-        return true;
-    }
-
     private boolean startConfirmationUi(int token, String action) {
         try {
             Intent confIntent = new Intent(action);
@@ -3898,83 +3723,6 @@
     }
 
     /**
-     * Callback: a requested backup agent has been instantiated. This should only be called from the
-     * {@link ActivityManager}.
-     */
-    public void agentConnected(String packageName, IBinder agentBinder) {
-        synchronized (mAgentConnectLock) {
-            if (Binder.getCallingUid() == Process.SYSTEM_UID) {
-                Slog.d(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "agentConnected pkg=" + packageName + " agent=" + agentBinder));
-                mConnectedAgent = IBackupAgent.Stub.asInterface(agentBinder);
-                mConnecting = false;
-            } else {
-                Slog.w(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "Non-system process uid="
-                                        + Binder.getCallingUid()
-                                        + " claiming agent connected"));
-            }
-            mAgentConnectLock.notifyAll();
-        }
-    }
-
-    /**
-     * Callback: a backup agent has failed to come up, or has unexpectedly quit. If the agent failed
-     * to come up in the first place, the agentBinder argument will be {@code null}. This should
-     * only be called from the {@link ActivityManager}.
-     */
-    public void agentDisconnected(String packageName) {
-        synchronized (mAgentConnectLock) {
-            if (Binder.getCallingUid() == Process.SYSTEM_UID) {
-                mConnectedAgent = null;
-                mConnecting = false;
-            } else {
-                Slog.w(
-                        TAG,
-                        addUserIdToLogMessage(
-                                mUserId,
-                                "Non-system process uid="
-                                        + Binder.getCallingUid()
-                                        + " claiming agent disconnected"));
-            }
-            Slog.w(TAG, "agentDisconnected: the backup agent for " + packageName
-                    + " died: cancel current operations");
-
-            // Offload operation cancellation off the main thread as the cancellation callbacks
-            // might call out to BackupTransport. Other operations started on the same package
-            // before the cancellation callback has executed will also be cancelled by the callback.
-            Runnable cancellationRunnable = () -> {
-                // handleCancel() causes the PerformFullTransportBackupTask to go on to
-                // tearDownAgentAndKill: that will unbindBackupAgent in the Activity Manager, so
-                // that the package being backed up doesn't get stuck in restricted mode until the
-                // backup time-out elapses.
-                for (int token : mOperationStorage.operationTokensForPackage(packageName)) {
-                    if (MORE_DEBUG) {
-                        Slog.d(TAG, "agentDisconnected: will handleCancel(all) for token:"
-                                + Integer.toHexString(token));
-                    }
-                    handleCancel(token, true /* cancelAll */);
-                }
-            };
-            getThreadForAsyncOperation(/* operationName */ "agent-disconnected",
-                    cancellationRunnable).start();
-
-            mAgentConnectLock.notifyAll();
-        }
-    }
-
-    @VisibleForTesting
-    Thread getThreadForAsyncOperation(String operationName, Runnable operation) {
-        return new Thread(operation, operationName);
-    }
-
-    /**
      * An application being installed will need a restore pass, then the {@link PackageManager} will
      * need to be told when the restore is finished.
      */
@@ -4521,4 +4269,8 @@
     public IBackupManager getBackupManagerBinder() {
         return mBackupManagerBinder;
     }
+
+    public BackupAgentConnectionManager getBackupAgentConnectionManager() {
+        return mBackupAgentConnectionManager;
+    }
 }
diff --git a/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java b/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
index 1271206..b98cb10 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/FullBackupEngine.java
@@ -314,7 +314,7 @@
                 Slog.d(TAG, "Binding to full backup agent : " + mPkg.packageName);
             }
             mAgent =
-                    backupManagerService.bindToAgentSynchronous(
+                    backupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                             mPkg.applicationInfo, ApplicationThreadConstants.BACKUP_MODE_FULL,
                             mBackupEligibilityRules.getBackupDestination());
         }
diff --git a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
index be9cdc8..65730c9 100644
--- a/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
+++ b/services/backup/java/com/android/server/backup/fullbackup/PerformFullTransportBackupTask.java
@@ -702,7 +702,8 @@
             }
 
             // Clear this to avoid using the memory until reboot.
-            mUserBackupManagerService.clearNoRestrictedModePackages();
+            mUserBackupManagerService
+                    .getBackupAgentConnectionManager().clearNoRestrictedModePackages();
 
             Slog.i(TAG, "Full data backup pass finished.");
             mUserBackupManagerService.getWakelock().release();
@@ -741,7 +742,8 @@
             }
             packageNames = transport.getPackagesThatShouldNotUseRestrictedMode(packageNames,
                     BACKUP);
-            mUserBackupManagerService.setNoRestrictedModePackages(packageNames, BACKUP);
+            mUserBackupManagerService.getBackupAgentConnectionManager().setNoRestrictedModePackages(
+                    packageNames, BACKUP);
         } catch (RemoteException e) {
             Slog.i(TAG, "Failed to retrieve no restricted mode packages from transport");
         }
diff --git a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
index 3a6e1ca..82232a6 100644
--- a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
+++ b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
@@ -741,7 +741,7 @@
         final IBackupAgent agent;
         try {
             agent =
-                    mBackupManagerService.bindToAgentSynchronous(
+                    mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                             packageInfo.applicationInfo, BACKUP_MODE_INCREMENTAL,
                             mBackupEligibilityRules.getBackupDestination());
             if (agent == null) {
@@ -1302,7 +1302,8 @@
 
         // For PM metadata (for which applicationInfo is null) there is no agent-bound state.
         if (mCurrentPackage.applicationInfo != null) {
-            mBackupManagerService.unbindAgent(mCurrentPackage.applicationInfo);
+            mBackupManagerService.getBackupAgentConnectionManager().unbindAgent(
+                    mCurrentPackage.applicationInfo);
         }
         mAgent = null;
     }
diff --git a/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java b/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
index 2d99c96..b59e860 100644
--- a/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
+++ b/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
@@ -410,11 +410,7 @@
 
                             // All set; now set up the IPC and launch the agent
                             setUpPipes();
-                            mAgent = mBackupManagerService.bindToAgentSynchronous(mTargetApp,
-                                    FullBackup.KEY_VALUE_DATA_TOKEN.equals(info.domain)
-                                            ? ApplicationThreadConstants.BACKUP_MODE_RESTORE
-                                            : ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL,
-                                    mBackupEligibilityRules.getBackupDestination());
+                            mAgent = bindToAgent(info);
                             mAgentPackage = pkg;
                         } catch (IOException | NameNotFoundException e) {
                             // fall through to error handling
@@ -805,15 +801,12 @@
         return packages.contains(packageName);
     }
 
-    void sendOnRestorePackage(String name) {
-        if (mObserver != null) {
-            try {
-                // TODO: use a more user-friendly name string
-                mObserver.onRestorePackage(name);
-            } catch (RemoteException e) {
-                Slog.w(TAG, "full restore observer went away: restorePackage");
-                mObserver = null;
-            }
-        }
+    private IBackupAgent bindToAgent(FileMetadata info) {
+        return mBackupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
+                mTargetApp,
+                FullBackup.KEY_VALUE_DATA_TOKEN.equals(info.domain)
+                        ? ApplicationThreadConstants.BACKUP_MODE_RESTORE
+                        : ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL,
+                mBackupEligibilityRules.getBackupDestination());
     }
 }
diff --git a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
index 5ee51a5..e5c7e5c 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -814,7 +814,7 @@
 
         // Good to go!  Set up and bind the agent...
         mAgent =
-                backupManagerService.bindToAgentSynchronous(
+                backupManagerService.getBackupAgentConnectionManager().bindToAgentSynchronous(
                         mCurrentPackage.applicationInfo,
                         ApplicationThreadConstants.BACKUP_MODE_RESTORE,
                         mBackupEligibilityRules.getBackupDestination());
@@ -1364,7 +1364,7 @@
         backupManagerService.getBackupHandler().removeMessages(MSG_RESTORE_SESSION_TIMEOUT);
 
         // Clear this to avoid using the memory until reboot.
-        backupManagerService.clearNoRestrictedModePackages();
+        backupManagerService.getBackupAgentConnectionManager().clearNoRestrictedModePackages();
 
         // If we have a PM token, we must under all circumstances be sure to
         // handshake when we've finished.
@@ -1838,7 +1838,8 @@
             }
             packageNames = transport.getPackagesThatShouldNotUseRestrictedMode(packageNames,
                     RESTORE);
-            backupManagerService.setNoRestrictedModePackages(packageNames, RESTORE);
+            backupManagerService.getBackupAgentConnectionManager().setNoRestrictedModePackages(
+                    packageNames, RESTORE);
         } catch (RemoteException e) {
             Slog.i(TAG, "Failed to retrieve restricted mode packages from transport");
         }
diff --git a/services/core/java/com/android/server/am/AppStartInfoTracker.java b/services/core/java/com/android/server/am/AppStartInfoTracker.java
index 3913d2f..961022b 100644
--- a/services/core/java/com/android/server/am/AppStartInfoTracker.java
+++ b/services/core/java/com/android/server/am/AppStartInfoTracker.java
@@ -634,12 +634,20 @@
         }
 
         final ApplicationStartInfo info = new ApplicationStartInfo(raw);
+        int uid = raw.getRealUid();
 
-        AppStartInfoContainer container = mData.get(raw.getPackageName(), raw.getRealUid());
+        // Isolated process starts won't be reasonably accessible if stored by their uid, don't
+        // store them.
+        if (com.android.server.am.Flags.appStartInfoIsolatedProcess()
+                && UserHandle.isIsolated(uid)) {
+            return null;
+        }
+
+        AppStartInfoContainer container = mData.get(raw.getPackageName(), uid);
         if (container == null) {
             container = new AppStartInfoContainer(mAppStartInfoHistoryListSize);
-            container.mUid = info.getRealUid();
-            mData.put(raw.getPackageName(), raw.getRealUid(), container);
+            container.mUid = uid;
+            mData.put(raw.getPackageName(), uid, container);
         }
         container.addStartInfoLocked(info);
 
@@ -1010,6 +1018,17 @@
                             new AppStartInfoContainer(mAppStartInfoHistoryListSize);
                     int uid = container.readFromProto(proto, AppsStartInfoProto.Package.USERS,
                             pkgName);
+
+                    // If the isolated process flag is enabled and the uid is that of an isolated
+                    // process, then break early so that the container will not be added to mData.
+                    // This is expected only as a one time mitigation, records added after this flag
+                    // is enabled should always return false for isIsolated and thereby always
+                    // continue on.
+                    if (com.android.server.am.Flags.appStartInfoIsolatedProcess()
+                            && UserHandle.isIsolated(uid)) {
+                        break;
+                    }
+
                     synchronized (mLock) {
                         mData.put(pkgName, uid, container);
                     }
diff --git a/services/core/java/com/android/server/am/flags.aconfig b/services/core/java/com/android/server/am/flags.aconfig
index c59c40f..6d247d2 100644
--- a/services/core/java/com/android/server/am/flags.aconfig
+++ b/services/core/java/com/android/server/am/flags.aconfig
@@ -270,3 +270,13 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "app_start_info_isolated_process"
+    namespace: "system_performance"
+    description: "Adjust handling of isolated process records to be discarded."
+    bug: "374032823"
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index edad247..e5b077d 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -3031,6 +3031,11 @@
         return mKeyGestureController.getAppLaunchBookmarks();
     }
 
+    @Override
+    public void resetLockedModifierState() {
+        mNative.resetLockedModifierState();
+    }
+
     private void handleCurrentUserChanged(@UserIdInt int userId) {
         mCurrentUserId = userId;
         mKeyGestureController.setCurrentUserId(userId);
diff --git a/services/core/java/com/android/server/input/NativeInputManagerService.java b/services/core/java/com/android/server/input/NativeInputManagerService.java
index 8903c27..728e440 100644
--- a/services/core/java/com/android/server/input/NativeInputManagerService.java
+++ b/services/core/java/com/android/server/input/NativeInputManagerService.java
@@ -98,6 +98,8 @@
 
     void toggleCapsLock(int deviceId);
 
+    void resetLockedModifierState();
+
     void displayRemoved(int displayId);
 
     void setInputDispatchMode(boolean enabled, boolean frozen);
@@ -370,6 +372,9 @@
         public native void toggleCapsLock(int deviceId);
 
         @Override
+        public native void resetLockedModifierState();
+
+        @Override
         public native void displayRemoved(int displayId);
 
         @Override
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index d8483f7..b7af9a4 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -1310,7 +1310,7 @@
         // Do not reset the default (current) IME when it is a 3rd-party IME
         String selectedMethodId = bindingController.getSelectedMethodId();
         final InputMethodSettings settings = InputMethodSettingsRepository.get(userId);
-        if (selectedMethodId != null
+        if (selectedMethodId != null && settings.getMethodMap().get(selectedMethodId) != null
                 && !settings.getMethodMap().get(selectedMethodId).isSystem()) {
             return;
         }
diff --git a/services/core/java/com/android/server/location/contexthub/ContextHubClientManager.java b/services/core/java/com/android/server/location/contexthub/ContextHubClientManager.java
index 0fdd0ae..5248a05 100644
--- a/services/core/java/com/android/server/location/contexthub/ContextHubClientManager.java
+++ b/services/core/java/com/android/server/location/contexthub/ContextHubClientManager.java
@@ -237,15 +237,16 @@
 
         if (message.isBroadcastMessage()) {
             if (message.isReliable()) {
-                Log.e(TAG, "Received reliable broadcast message from " + message.getNanoAppId());
+                Log.e(TAG, "Received reliable broadcast message from 0x"
+                        + Long.toHexString(message.getNanoAppId()));
                 return ErrorCode.PERMANENT_ERROR;
             }
 
             // Broadcast messages shouldn't be sent with any permissions tagged per CHRE API
             // requirements.
             if (!messagePermissions.isEmpty()) {
-                Log.e(TAG, "Received broadcast message with permissions from "
-                        + message.getNanoAppId());
+                Log.e(TAG, "Received broadcast message with permissions from 0x"
+                        + Long.toHexString(message.getNanoAppId()));
                 return ErrorCode.PERMANENT_ERROR;
             }
 
diff --git a/services/core/java/com/android/server/location/contexthub/ContextHubService.java b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
index 946e896..7c9d9c5 100644
--- a/services/core/java/com/android/server/location/contexthub/ContextHubService.java
+++ b/services/core/java/com/android/server/location/contexthub/ContextHubService.java
@@ -326,8 +326,15 @@
         }
 
         if (Flags.offloadApi()) {
-            mHubInfoRegistry = new HubInfoRegistry(mContextHubWrapper);
-            Log.i(TAG, "Enabling generic offload API");
+            HubInfoRegistry registry;
+            try {
+                registry = new HubInfoRegistry(mContextHubWrapper);
+                Log.i(TAG, "Enabling generic offload API");
+            } catch (UnsupportedOperationException e) {
+                registry = null;
+                Log.w(TAG, "Generic offload API not supported, disabling");
+            }
+            mHubInfoRegistry = registry;
         } else {
             mHubInfoRegistry = null;
             Log.i(TAG, "Disabling generic offload API");
diff --git a/services/core/java/com/android/server/security/authenticationpolicy/AuthenticationPolicyService.java b/services/core/java/com/android/server/security/authenticationpolicy/AuthenticationPolicyService.java
index b8a4a9c..6798a61 100644
--- a/services/core/java/com/android/server/security/authenticationpolicy/AuthenticationPolicyService.java
+++ b/services/core/java/com/android/server/security/authenticationpolicy/AuthenticationPolicyService.java
@@ -16,8 +16,11 @@
 
 package com.android.server.security.authenticationpolicy;
 
+import static android.Manifest.permission.MANAGE_SECURE_LOCK_DEVICE;
+
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.SOME_AUTH_REQUIRED_AFTER_ADAPTIVE_AUTH_REQUEST;
 
+import android.annotation.EnforcePermission;
 import android.app.KeyguardManager;
 import android.content.Context;
 import android.content.pm.PackageManager;
@@ -32,9 +35,14 @@
 import android.hardware.biometrics.events.AuthenticationSucceededInfo;
 import android.os.Build;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 import android.os.SystemClock;
+import android.security.authenticationpolicy.AuthenticationPolicyManager;
+import android.security.authenticationpolicy.DisableSecureLockDeviceParams;
+import android.security.authenticationpolicy.EnableSecureLockDeviceParams;
+import android.security.authenticationpolicy.IAuthenticationPolicyService;
 import android.util.Log;
 import android.util.Slog;
 import android.util.SparseIntArray;
@@ -74,6 +82,7 @@
     private final KeyguardManager mKeyguardManager;
     private final WindowManagerInternal mWindowManager;
     private final UserManagerInternal mUserManager;
+    private SecureLockDeviceServiceInternal mSecureLockDeviceService;
     @VisibleForTesting
     final SparseIntArray mFailedAttemptsForUser = new SparseIntArray();
     private final SparseLongArray mLastLockedTimestamp = new SparseLongArray();
@@ -94,10 +103,16 @@
         mWindowManager = Objects.requireNonNull(
                 LocalServices.getService(WindowManagerInternal.class));
         mUserManager = Objects.requireNonNull(LocalServices.getService(UserManagerInternal.class));
+        if (android.security.Flags.secureLockdown()) {
+            mSecureLockDeviceService = Objects.requireNonNull(
+                    LocalServices.getService(SecureLockDeviceServiceInternal.class));
+        }
     }
 
     @Override
-    public void onStart() {}
+    public void onStart() {
+        publishBinderService(Context.AUTHENTICATION_POLICY_SERVICE, mService);
+    }
 
     @Override
     public void onBootPhase(int phase) {
@@ -294,4 +309,36 @@
         // next successful primary or biometric auth happens
         mLastLockedTimestamp.put(userId, SystemClock.elapsedRealtime());
     }
+
+    private final IBinder mService = new IAuthenticationPolicyService.Stub() {
+        /**
+         * @see AuthenticationPolicyManager#enableSecureLockDevice(EnableSecureLockDeviceParams)
+         * @param params EnableSecureLockDeviceParams for caller to supply params related
+         *               to the secure lock device request
+         * @return @EnableSecureLockDeviceRequestStatus int indicating the result of the Secure
+         * Lock Device request
+         */
+        @Override
+        @EnforcePermission(MANAGE_SECURE_LOCK_DEVICE)
+        @AuthenticationPolicyManager.EnableSecureLockDeviceRequestStatus
+        public int enableSecureLockDevice(EnableSecureLockDeviceParams params) {
+            enableSecureLockDevice_enforcePermission();
+            return mSecureLockDeviceService.enableSecureLockDevice(params);
+        }
+
+        /**
+         * @see AuthenticationPolicyManager#disableSecureLockDevice(DisableSecureLockDeviceParams)
+         * @param params @DisableSecureLockDeviceParams for caller to supply params related
+         *               to the secure lock device request
+         * @return @DisableSecureLockDeviceRequestStatus int indicating the result of the Secure
+         * Lock Device request
+         */
+        @Override
+        @EnforcePermission(MANAGE_SECURE_LOCK_DEVICE)
+        @AuthenticationPolicyManager.DisableSecureLockDeviceRequestStatus
+        public int disableSecureLockDevice(DisableSecureLockDeviceParams params) {
+            disableSecureLockDevice_enforcePermission();
+            return mSecureLockDeviceService.disableSecureLockDevice(params);
+        }
+    };
 }
diff --git a/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceService.java b/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceService.java
new file mode 100644
index 0000000..7b89723
--- /dev/null
+++ b/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceService.java
@@ -0,0 +1,120 @@
+/*
+ * 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.security.authenticationpolicy;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.security.authenticationpolicy.AuthenticationPolicyManager;
+import android.security.authenticationpolicy.AuthenticationPolicyManager.DisableSecureLockDeviceRequestStatus;
+import android.security.authenticationpolicy.AuthenticationPolicyManager.EnableSecureLockDeviceRequestStatus;
+import android.security.authenticationpolicy.DisableSecureLockDeviceParams;
+import android.security.authenticationpolicy.EnableSecureLockDeviceParams;
+import android.util.Slog;
+
+import com.android.server.LocalServices;
+import com.android.server.SystemService;
+
+/**
+ * System service for remotely calling secure lock on the device.
+ *
+ * Callers will access this class via
+ * {@link com.android.server.security.authenticationpolicy.AuthenticationPolicyService}.
+ *
+ * @see AuthenticationPolicyService
+ * @see AuthenticationPolicyManager#enableSecureLockDevice
+ * @see AuthenticationPolicyManager#disableSecureLockDevice
+ * @hide
+ */
+public class SecureLockDeviceService extends SecureLockDeviceServiceInternal {
+    private static final String TAG = "SecureLockDeviceService";
+    private final Context mContext;
+
+    public SecureLockDeviceService(@NonNull Context context) {
+        mContext = context;
+    }
+
+    private void start() {
+        // Expose private service for system components to use.
+        LocalServices.addService(SecureLockDeviceServiceInternal.class, this);
+    }
+
+    /**
+     * @see AuthenticationPolicyManager#enableSecureLockDevice
+     * @param params EnableSecureLockDeviceParams for caller to supply params related
+     *               to the secure lock device request
+     * @return @EnableSecureLockDeviceRequestStatus int indicating the result of the
+     * secure lock device request
+     *
+     * @hide
+     */
+    @Override
+    @EnableSecureLockDeviceRequestStatus
+    public int enableSecureLockDevice(EnableSecureLockDeviceParams params) {
+        // (1) Call into system_server to lock device, configure allowed auth types
+        // for secure lock
+        // TODO: lock device, configure allowed authentication types for device entry
+        // (2) Call into framework to configure secure lock 2FA lockscreen
+        // update, UI & string updates
+        // TODO: implement 2FA lockscreen when SceneContainerFlag.isEnabled()
+        // TODO: implement 2FA lockscreen when !SceneContainerFlag.isEnabled()
+        // (3) Call into framework to configure keyguard security updates
+        // TODO: implement security updates
+        return AuthenticationPolicyManager.ERROR_UNSUPPORTED;
+    }
+
+    /**
+     * @see AuthenticationPolicyManager#disableSecureLockDevice
+     * @param params @DisableSecureLockDeviceParams for caller to supply params related
+     *               to the secure lock device request
+     * @return @DisableSecureLockDeviceRequestStatus int indicating the result of the
+     * secure lock device request
+     *
+     * @hide
+     */
+    @Override
+    @DisableSecureLockDeviceRequestStatus
+    public int disableSecureLockDevice(DisableSecureLockDeviceParams params) {
+        // (1) Call into system_server to reset allowed auth types
+        // TODO: reset allowed authentication types for device entry;
+        // (2) Call into framework to disable secure lock 2FA lockscreen, reset UI
+        // & string updates
+        // TODO: implement reverting to normal lockscreen when SceneContainerFlag.isEnabled()
+        // TODO: implement reverting to normal lockscreen when !SceneContainerFlag.isEnabled()
+        // (3) Call into framework to revert keyguard security updates
+        // TODO: implement reverting security updates
+        return AuthenticationPolicyManager.ERROR_UNSUPPORTED;
+    }
+
+    /**
+     * System service lifecycle.
+     */
+    public static final class Lifecycle extends SystemService {
+        private final SecureLockDeviceService mService;
+
+        public Lifecycle(@NonNull Context context) {
+            super(context);
+            mService = new SecureLockDeviceService(context);
+        }
+
+        @Override
+        public void onStart() {
+            Slog.i(TAG, "Starting SecureLockDeviceService");
+            mService.start();
+            Slog.i(TAG, "Started SecureLockDeviceService");
+        }
+    }
+}
diff --git a/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceServiceInternal.java b/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceServiceInternal.java
new file mode 100644
index 0000000..b903709
--- /dev/null
+++ b/services/core/java/com/android/server/security/authenticationpolicy/SecureLockDeviceServiceInternal.java
@@ -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.server.security.authenticationpolicy;
+
+import android.security.authenticationpolicy.AuthenticationPolicyManager;
+import android.security.authenticationpolicy.AuthenticationPolicyManager.DisableSecureLockDeviceRequestStatus;
+import android.security.authenticationpolicy.AuthenticationPolicyManager.EnableSecureLockDeviceRequestStatus;
+import android.security.authenticationpolicy.DisableSecureLockDeviceParams;
+import android.security.authenticationpolicy.EnableSecureLockDeviceParams;
+
+/**
+ * Local system service interface for {@link SecureLockDeviceService}.
+ *
+ * <p>No permission / argument checks will be performed inside.
+ * Callers must check the calling app permission and the calling package name.
+ *
+ * @hide
+ */
+public abstract class SecureLockDeviceServiceInternal {
+    private static final String TAG = "SecureLockDeviceServiceInternal";
+
+    /**
+     * @see AuthenticationPolicyManager#enableSecureLockDevice(EnableSecureLockDeviceParams)
+     * @param params EnableSecureLockDeviceParams for caller to supply params related
+     *               to the secure lock request
+     * @return @EnableSecureLockDeviceRequestStatus int indicating the result of the
+     * secure lock request
+     */
+    @EnableSecureLockDeviceRequestStatus
+    public abstract int enableSecureLockDevice(EnableSecureLockDeviceParams params);
+
+    /**
+     * @see AuthenticationPolicyManager#disableSecureLockDevice(DisableSecureLockDeviceParams)
+     * @param params @DisableSecureLockDeviceParams for caller to supply params related
+     *               to the secure lock device request
+     * @return @DisableSecureLockDeviceRequestStatus int indicating the result of the
+     * secure lock device request
+     */
+    @DisableSecureLockDeviceRequestStatus
+    public abstract int disableSecureLockDevice(DisableSecureLockDeviceParams params);
+}
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index bf93f73..f70dec1 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -630,8 +630,8 @@
     // The locusId associated with this activity, if set.
     private LocusId mLocusId;
 
-    // Whether the activity is requesting to limit the system's educational dialogs
-    public boolean mShouldLimitSystemEducationDialogs;
+    // The timestamp of the last request to show the "Open in browser" education
+    public long mRequestOpenInBrowserEducationTimestamp;
 
     // Whether the activity was launched from a bubble.
     private boolean mLaunchedFromBubble;
@@ -7342,9 +7342,8 @@
         return mLocusId;
     }
 
-    void setLimitSystemEducationDialogs(boolean limitSystemEducationDialogs) {
-        if (mShouldLimitSystemEducationDialogs == limitSystemEducationDialogs) return;
-        mShouldLimitSystemEducationDialogs = limitSystemEducationDialogs;
+    void requestOpenInBrowserEducation() {
+        mRequestOpenInBrowserEducationTimestamp = System.currentTimeMillis();
         final Task task = getTask();
         if (task != null) {
             final boolean force = isVisibleRequested() && this == task.getTopNonFinishingActivity();
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 198e14a..8ff0818 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -3915,12 +3915,11 @@
     }
 
     @Override
-    public void setLimitSystemEducationDialogs(
-            IBinder appToken, boolean limitSystemEducationDialogs) {
+    public void requestOpenInBrowserEducation(IBinder appToken) {
         synchronized (mGlobalLock) {
             final ActivityRecord r = ActivityRecord.isInRootTaskLocked(appToken);
             if (r != null) {
-                r.setLimitSystemEducationDialogs(limitSystemEducationDialogs);
+                r.requestOpenInBrowserEducation();
             }
         }
     }
diff --git a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
index f2362b9..9754595 100644
--- a/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
+++ b/services/core/java/com/android/server/wm/AppCompatCameraOverrides.java
@@ -17,8 +17,8 @@
 package com.android.server.wm;
 
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_FORCE_ROTATION;
-import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT;
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH;
+import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT;
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_ENABLE_REFRESH_VIA_PAUSE;
 import static android.content.pm.ActivityInfo.OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA;
 import static android.content.pm.ActivityInfo.OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA;
@@ -164,12 +164,12 @@
      * <p>The treatment is enabled when the following conditions are met:
      * <ul>
      * <li>Feature flag gating the camera compatibility free-form treatment is enabled.
-     * <li>Activity isn't opted out by the device manufacturer with override.
+     * <li>Activity is opted in by the device manufacturer with override.
      * </ul>
      */
     boolean shouldApplyFreeformTreatmentForCameraCompat() {
-        return Flags.enableCameraCompatForDesktopWindowing() && !isChangeEnabled(mActivityRecord,
-                OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT);
+        return Flags.enableCameraCompatForDesktopWindowing() && isChangeEnabled(mActivityRecord,
+                OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT);
     }
 
     boolean isOverrideOrientationOnlyForCameraEnabled() {
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index dbc3b76c2..e8ae280 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -3423,8 +3423,8 @@
                 ? top.getLastParentBeforePip().mTaskId : INVALID_TASK_ID;
         info.shouldDockBigOverlays = top != null && top.shouldDockBigOverlays;
         info.mTopActivityLocusId = top != null ? top.getLocusId() : null;
-        info.isTopActivityLimitSystemEducationDialogs = top != null
-              ? top.mShouldLimitSystemEducationDialogs : false;
+        info.topActivityRequestOpenInBrowserEducationTimestamp = top != null
+              ? top.mRequestOpenInBrowserEducationTimestamp : 0;
         final Task parentTask = getParent() != null ? getParent().asTask() : null;
         info.parentTaskId = parentTask != null && parentTask.mCreatedByOrganizer
                 ? parentTask.mTaskId
@@ -3523,6 +3523,7 @@
 
         info.capturedLink = null;
         info.capturedLinkTimestamp = 0;
+        info.topActivityRequestOpenInBrowserEducationTimestamp = 0;
     }
 
     @Nullable PictureInPictureParams getPictureInPictureParams() {
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 81af78e..6009848 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -182,7 +182,6 @@
 import static com.android.server.wm.WindowStateProto.VIEW_VISIBILITY;
 import static com.android.server.wm.WindowStateProto.WINDOW_CONTAINER;
 import static com.android.server.wm.WindowStateProto.WINDOW_FRAMES;
-import static com.android.window.flags.Flags.secureWindowState;
 import static com.android.window.flags.Flags.surfaceTrustedOverlay;
 
 import android.annotation.CallSuper;
@@ -1187,9 +1186,7 @@
         if (surfaceTrustedOverlay() && isWindowTrustedOverlay()) {
             getPendingTransaction().setTrustedOverlay(mSurfaceControl, true);
         }
-        if (secureWindowState()) {
-            getPendingTransaction().setSecure(mSurfaceControl, isSecureLocked());
-        }
+        getPendingTransaction().setSecure(mSurfaceControl, isSecureLocked());
         // All apps should be considered as occluding when computing TrustedPresentation Thresholds.
         final boolean canOccludePresentation = !mSession.mCanAddInternalSystemWindow;
         getPendingTransaction().setCanOccludePresentation(mSurfaceControl, canOccludePresentation);
@@ -6174,18 +6171,10 @@
 
     void setSecureLocked(boolean isSecure) {
         ProtoLog.i(WM_SHOW_TRANSACTIONS, "SURFACE isSecure=%b: %s", isSecure, getName());
-        if (secureWindowState()) {
-            if (mSurfaceControl == null) {
-                return;
-            }
-            getPendingTransaction().setSecure(mSurfaceControl, isSecure);
-        } else {
-            if (mWinAnimator.mSurfaceControl == null) {
-                return;
-            }
-            getPendingTransaction().setSecure(mWinAnimator.mSurfaceControl,
-                    isSecure);
+        if (mSurfaceControl == null) {
+            return;
         }
+        getPendingTransaction().setSecure(mSurfaceControl, isSecure);
         if (mDisplayContent != null) {
             mDisplayContent.refreshImeSecureFlag(getSyncTransaction());
         }
diff --git a/services/core/java/com/android/server/wm/WindowStateAnimator.java b/services/core/java/com/android/server/wm/WindowStateAnimator.java
index a934eea..0154d95 100644
--- a/services/core/java/com/android/server/wm/WindowStateAnimator.java
+++ b/services/core/java/com/android/server/wm/WindowStateAnimator.java
@@ -49,7 +49,6 @@
 import static com.android.server.wm.WindowStateAnimatorProto.SURFACE;
 import static com.android.server.wm.WindowStateAnimatorProto.SYSTEM_DECOR_RECT;
 import static com.android.server.wm.WindowSurfaceControllerProto.SHOWN;
-import static com.android.window.flags.Flags.secureWindowState;
 import static com.android.window.flags.Flags.setScPropertiesInClient;
 
 import android.content.Context;
@@ -310,12 +309,6 @@
         int flags = SurfaceControl.HIDDEN;
         final WindowManager.LayoutParams attrs = w.mAttrs;
 
-        if (!secureWindowState()) {
-            if (w.isSecureLocked()) {
-                flags |= SurfaceControl.SECURE;
-            }
-        }
-
         if ((mWin.mAttrs.privateFlags & PRIVATE_FLAG_IS_ROUNDED_CORNERS_OVERLAY) != 0) {
             flags |= SurfaceControl.SKIP_SCREENSHOT;
         }
diff --git a/services/core/jni/com_android_server_input_InputManagerService.cpp b/services/core/jni/com_android_server_input_InputManagerService.cpp
index e383375..dece612 100644
--- a/services/core/jni/com_android_server_input_InputManagerService.cpp
+++ b/services/core/jni/com_android_server_input_InputManagerService.cpp
@@ -2330,6 +2330,12 @@
     im->getInputManager()->getReader().toggleCapsLockState(deviceId);
 }
 
+static void resetLockedModifierState(JNIEnv* env, jobject nativeImplObj) {
+    NativeInputManager* im = getNativeInputManager(env, nativeImplObj);
+
+    im->getInputManager()->getReader().resetLockedModifierState();
+}
+
 static void nativeDisplayRemoved(JNIEnv* env, jobject nativeImplObj, jint displayId) {
     NativeInputManager* im = getNativeInputManager(env, nativeImplObj);
 
@@ -3134,6 +3140,7 @@
         {"verifyInputEvent", "(Landroid/view/InputEvent;)Landroid/view/VerifiedInputEvent;",
          (void*)nativeVerifyInputEvent},
         {"toggleCapsLock", "(I)V", (void*)nativeToggleCapsLock},
+        {"resetLockedModifierState", "()V", (void*)resetLockedModifierState},
         {"displayRemoved", "(I)V", (void*)nativeDisplayRemoved},
         {"setFocusedApplication", "(ILandroid/view/InputApplicationHandle;)V",
          (void*)nativeSetFocusedApplication},
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index 9e9dadb..dc10dd9 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -251,6 +251,7 @@
 import com.android.server.security.KeyChainSystemService;
 import com.android.server.security.advancedprotection.AdvancedProtectionService;
 import com.android.server.security.authenticationpolicy.AuthenticationPolicyService;
+import com.android.server.security.authenticationpolicy.SecureLockDeviceService;
 import com.android.server.security.forensic.ForensicService;
 import com.android.server.security.rkp.RemoteProvisioningService;
 import com.android.server.selinux.SelinuxAuditLogsService;
@@ -2659,6 +2660,12 @@
             mSystemServiceManager.startService(AuthService.class);
             t.traceEnd();
 
+            if (android.security.Flags.secureLockdown()) {
+                t.traceBegin("StartSecureLockDeviceService.Lifecycle");
+                mSystemServiceManager.startService(SecureLockDeviceService.Lifecycle.class);
+                t.traceEnd();
+            }
+
             if (android.adaptiveauth.Flags.enableAdaptiveAuth()) {
                 t.traceBegin("StartAuthenticationPolicyService");
                 mSystemServiceManager.startService(AuthenticationPolicyService.class);
diff --git a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
index a547d0f..4e9fff2 100644
--- a/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/BackupManagerServiceRoboTest.java
@@ -31,6 +31,7 @@
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 import static org.robolectric.Shadows.shadowOf;
 import static org.testng.Assert.expectThrows;
 
@@ -96,6 +97,7 @@
     @UserIdInt private int mUserTwoId;
     @Mock private UserBackupManagerService mUserSystemService;
     @Mock private UserBackupManagerService mUserOneService;
+    @Mock private BackupAgentConnectionManager mUserOneBackupAgentConnectionManager;
     @Mock private UserBackupManagerService mUserTwoService;
 
     /** Setup */
@@ -116,6 +118,9 @@
         mShadowContext.grantPermissions(BACKUP);
         mShadowContext.grantPermissions(INTERACT_ACROSS_USERS_FULL);
 
+        when(mUserOneService.getBackupAgentConnectionManager()).thenReturn(
+                mUserOneBackupAgentConnectionManager);
+
         ShadowBinder.setCallingUid(Process.SYSTEM_UID);
     }
 
@@ -226,7 +231,7 @@
 
         backupManagerService.agentConnected(mUserOneId, TEST_PACKAGE, agentBinder);
 
-        verify(mUserOneService).agentConnected(TEST_PACKAGE, agentBinder);
+        verify(mUserOneBackupAgentConnectionManager).agentConnected(TEST_PACKAGE, agentBinder);
     }
 
     /** Test that the backup service does not route methods for non-registered users. */
@@ -239,7 +244,8 @@
 
         backupManagerService.agentConnected(mUserTwoId, TEST_PACKAGE, agentBinder);
 
-        verify(mUserOneService, never()).agentConnected(TEST_PACKAGE, agentBinder);
+        verify(mUserOneBackupAgentConnectionManager, never()).agentConnected(TEST_PACKAGE,
+                agentBinder);
     }
 
     /** Test that the backup service routes methods correctly to the user that requests it. */
diff --git a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
index 7349c14..aeb1ba9 100644
--- a/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
+++ b/services/robotests/backup/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
@@ -107,6 +107,7 @@
 import com.android.internal.infra.AndroidFuture;
 import com.android.server.EventLogTags;
 import com.android.server.LocalServices;
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.BackupRestoreTask;
 import com.android.server.backup.DataChangedJournal;
 import com.android.server.backup.KeyValueBackupJob;
@@ -167,7 +168,6 @@
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Stream;
 
-// TODO: Test agents timing out
 @RunWith(RobolectricTestRunner.class)
 @Config(
         shadows = {
@@ -195,6 +195,7 @@
     @Mock private IBackupManagerMonitor mMonitor;
     @Mock private OnTaskFinishedListener mListener;
     @Mock private PackageManagerInternal mPackageManagerInternal;
+    @Mock private BackupAgentConnectionManager mBackupAgentConnectionManager;
 
     private UserBackupManagerService mBackupManagerService;
     private TransportData mTransport;
@@ -257,6 +258,8 @@
         when(mBackupManagerService.getBaseStateDir()).thenReturn(mBaseStateDir);
         when(mBackupManagerService.getDataDir()).thenReturn(mDataDir);
         when(mBackupManagerService.getBackupManagerBinder()).thenReturn(mBackupManager);
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
 
         mBackupHandler = mBackupManagerService.getBackupHandler();
         mShadowBackupLooper = shadowOf(mBackupHandler.getLooper());
@@ -749,7 +752,8 @@
 
     /**
      * Agent unavailable means {@link
-     * UserBackupManagerService#bindToAgentSynchronous(ApplicationInfo, int)} returns {@code null}.
+     * BackupAgentConnectionManager#bindToAgentSynchronous(ApplicationInfo, int, int)} returns
+     * {@code null}.
      *
      * @see #setUpAgent(PackageData)
      */
@@ -805,7 +809,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgent(PACKAGE_1);
         doThrow(SecurityException.class)
-                .when(mBackupManagerService)
+                .when(mBackupAgentConnectionManager)
                 .bindToAgentSynchronous(argThat(applicationInfo(PACKAGE_1)), anyInt(), anyInt());
         KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
@@ -823,7 +827,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgent(PACKAGE_1);
         doThrow(SecurityException.class)
-                .when(mBackupManagerService)
+                .when(mBackupAgentConnectionManager)
                 .bindToAgentSynchronous(argThat(applicationInfo(PACKAGE_1)), anyInt(), anyInt());
         KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true, PACKAGE_1);
 
@@ -861,7 +865,7 @@
         runTask(task);
 
         verify(mBackupManagerService).setWorkSource(null);
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -1097,7 +1101,7 @@
         runTask(task);
 
         verify(agentMock.agentBinder).fail(any());
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -1418,7 +1422,8 @@
                 .isEqualTo("newState".getBytes());
         assertCleansUpFiles(mTransport, PM_PACKAGE);
         // We don't unbind PM
-        verify(mBackupManagerService, never()).unbindAgent(argThat(applicationInfo(PM_PACKAGE)));
+        verify(mBackupAgentConnectionManager, never()).unbindAgent(
+                argThat(applicationInfo(PM_PACKAGE)));
     }
 
     @Test
@@ -1439,7 +1444,8 @@
 
         runTask(task);
 
-        verify(mBackupManagerService, never()).unbindAgent(argThat(applicationInfo(PM_PACKAGE)));
+        verify(mBackupAgentConnectionManager, never()).unbindAgent(
+                argThat(applicationInfo(PM_PACKAGE)));
     }
 
     @Test
@@ -1642,9 +1648,10 @@
 
         runTask(task);
 
-        InOrder inOrder = inOrder(agentMock.agent, mBackupManagerService);
+        InOrder inOrder = inOrder(agentMock.agent, mBackupAgentConnectionManager);
         inOrder.verify(agentMock.agent).onQuotaExceeded(anyLong(), eq(1234L));
-        inOrder.verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(PACKAGE_1)));
+        inOrder.verify(mBackupAgentConnectionManager).unbindAgent(
+                argThat(applicationInfo(PACKAGE_1)));
     }
 
     @Test
@@ -2634,12 +2641,12 @@
             doNothing().when(backupAgentBinder).fail(any());
             if (packageData.available) {
                 doReturn(backupAgentBinder)
-                        .when(mBackupManagerService)
+                        .when(mBackupAgentConnectionManager)
                         .bindToAgentSynchronous(argThat(applicationInfo(packageData)), anyInt(),
                                 anyInt());
             } else {
                 doReturn(null)
-                        .when(mBackupManagerService)
+                        .when(mBackupAgentConnectionManager)
                         .bindToAgentSynchronous(argThat(applicationInfo(packageData)), anyInt(),
                                 anyInt());
             }
@@ -2976,7 +2983,7 @@
 
     private void assertCleansUpFilesAndAgent(TransportData transport, PackageData packageData) {
         assertCleansUpFiles(transport, packageData);
-        verify(mBackupManagerService).unbindAgent(argThat(applicationInfo(packageData)));
+        verify(mBackupAgentConnectionManager).unbindAgent(argThat(applicationInfo(packageData)));
     }
 
     private void assertCleansUpFiles(TransportData transport, PackageData packageData) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
new file mode 100644
index 0000000..19e43b6
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/BackupAgentConnectionManagerTest.java
@@ -0,0 +1,381 @@
+/*
+ * 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.backup;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.ActivityManager;
+import android.app.ActivityManagerInternal;
+import android.app.ApplicationThreadConstants;
+import android.app.IActivityManager;
+import android.app.IBackupAgent;
+import android.app.backup.BackupAnnotations.BackupDestination;
+import android.app.backup.BackupAnnotations.OperationType;
+import android.compat.testing.PlatformCompatChangeRule;
+import android.content.pm.ApplicationInfo;
+import android.content.pm.PackageManager;
+import android.os.Process;
+import android.os.UserHandle;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.annotations.Presubmit;
+import android.platform.test.flag.junit.SetFlagsRule;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.server.LocalServices;
+import com.android.server.backup.internal.LifecycleOperationStorage;
+
+import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
+import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
+
+import com.google.common.collect.ImmutableSet;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TestRule;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.mockito.MockitoSession;
+import org.mockito.quality.Strictness;
+
+import java.util.Set;
+
+@Presubmit
+@RunWith(AndroidJUnit4.class)
+public class BackupAgentConnectionManagerTest {
+    private static final String TEST_PACKAGE = "com.test.package";
+
+    @Rule
+    public TestRule compatChangeRule = new PlatformCompatChangeRule();
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    @Mock
+    IActivityManager mActivityManager;
+    @Mock
+    ActivityManagerInternal mActivityManagerInternal;
+    @Mock
+    LifecycleOperationStorage mOperationStorage;
+    @Mock
+    UserBackupManagerService mUserBackupManagerService;
+    @Mock
+    IBackupAgent.Stub mBackupAgentStub;
+    @Mock
+    PackageManager mPackageManager;
+
+    private BackupAgentConnectionManager mConnectionManager;
+    private MockitoSession mSession;
+    private ApplicationInfo mTestApplicationInfo;
+    private IBackupAgent mBackupAgentResult;
+    private Thread mTestThread;
+
+    @Before
+    public void setUp() throws Exception {
+        mSession = mockitoSession().initMocks(this).mockStatic(ActivityManager.class).mockStatic(
+                LocalServices.class).strictness(Strictness.LENIENT).startMocking();
+        MockitoAnnotations.initMocks(this);
+
+        doReturn(mActivityManager).when(ActivityManager::getService);
+        doReturn(mActivityManagerInternal).when(
+                () -> LocalServices.getService(ActivityManagerInternal.class));
+        // Real package manager throws if a property is not defined.
+        when(mPackageManager.getPropertyAsUser(any(), any(), any(), anyInt())).thenThrow(
+                new PackageManager.NameNotFoundException());
+
+        mConnectionManager = spy(
+                new BackupAgentConnectionManager(mOperationStorage, mPackageManager,
+                        mUserBackupManagerService, UserHandle.USER_SYSTEM));
+
+        mTestApplicationInfo = new ApplicationInfo();
+        mTestApplicationInfo.packageName = TEST_PACKAGE;
+
+        mBackupAgentResult = null;
+        mTestThread = null;
+    }
+
+    @After
+    public void tearDown() {
+        if (mSession != null) {
+            mSession.finishMocking();
+        }
+    }
+
+    @Test
+    public void bindToAgentSynchronous_amReturnsFailure_returnsNullAndClearsPendingBackups()
+            throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(false);
+
+        IBackupAgent result = mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        assertThat(result).isNull();
+        verify(mActivityManagerInternal).clearPendingBackup(UserHandle.USER_SYSTEM);
+    }
+
+    @Test
+    public void bindToAgentSynchronous_agentDisconnectedCalled_returnsNullAndClearsPendingBackups()
+            throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(true);
+        // This is so that IBackupAgent.Stub.asInterface() works.
+        when(mBackupAgentStub.queryLocalInterface(any())).thenReturn(mBackupAgentStub);
+        when(mConnectionManager.getCallingUid()).thenReturn(Process.SYSTEM_UID);
+
+        // This is going to block until it receives the callback so we need to run it on a
+        // separate thread.
+        Thread testThread = new Thread(() -> setBackupAgentResult(
+                mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                        ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD)),
+                "backup-agent-connection-manager-test");
+        testThread.start();
+        // Give the testThread a head start, otherwise agentConnected() might run before
+        // bindToAgentSynchronous() is called.
+        Thread.sleep(500);
+        mConnectionManager.agentDisconnected(TEST_PACKAGE);
+        testThread.join();
+
+        assertThat(mBackupAgentResult).isNull();
+        verify(mActivityManagerInternal).clearPendingBackup(UserHandle.USER_SYSTEM);
+    }
+
+    @Test
+    public void bindToAgentSynchronous_agentConnectedCalled_returnsBackupAgent() throws Exception {
+        when(mActivityManager.bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(),
+                anyInt(), anyBoolean())).thenReturn(true);
+        // This is so that IBackupAgent.Stub.asInterface() works.
+        when(mBackupAgentStub.queryLocalInterface(any())).thenReturn(mBackupAgentStub);
+        when(mConnectionManager.getCallingUid()).thenReturn(Process.SYSTEM_UID);
+
+        // This is going to block until it receives the callback so we need to run it on a
+        // separate thread.
+        Thread testThread = new Thread(() -> setBackupAgentResult(
+                mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                        ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD)),
+                "backup-agent-connection-manager-test");
+        testThread.start();
+        // Give the testThread a head start, otherwise agentConnected() might run before
+        // bindToAgentSynchronous() is called.
+        Thread.sleep(500);
+        mConnectionManager.agentConnected(TEST_PACKAGE, mBackupAgentStub);
+        testThread.join();
+
+        assertThat(mBackupAgentResult).isEqualTo(mBackupAgentStub);
+        verify(mActivityManagerInternal, never()).clearPendingBackup(anyInt());
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_restrictedModeChangesFlagOff_shouldUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_keyValueBackup_shouldNotUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_keyValueRestore_shouldNotUseRestrictedMode()
+            throws Exception {
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_RESTORE, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+        // Make sure we never hit the code that checks the property.
+        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_packageOptedIn_shouldUseRestrictedMode() throws Exception {
+        reset(mPackageManager);
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenReturn(new PackageManager.Property(
+                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ true,
+                TEST_PACKAGE, /* className= */ null));
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    public void bindToAgentSynchronous_packageOptedOut_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenReturn(new PackageManager.Property(
+                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ false,
+                TEST_PACKAGE, /* className= */ null));
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @DisableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_targetSdkBelowB_shouldUseRestrictedMode() throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_targetSdkB_notInList_shouldUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.clearNoRestrictedModePackages();
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(true));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_forRestore_targetSdkB_inList_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.RESTORE);
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
+    @EnableCompatChanges({BackupAgentConnectionManager.OS_DECIDES_BACKUP_RESTRICTED_MODE})
+    public void bindToAgentSynchronous_forBackup_targetSdkB_inList_shouldNotUseRestrictedMode()
+            throws Exception {
+        reset(mPackageManager);
+        // Mock that the app has not explicitly set the property.
+        when(mPackageManager.getPropertyAsUser(
+                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE), eq(TEST_PACKAGE), any(),
+                anyInt())).thenThrow(new PackageManager.NameNotFoundException());
+        mConnectionManager.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.BACKUP);
+
+        mConnectionManager.bindToAgentSynchronous(mTestApplicationInfo,
+                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
+
+        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
+                /* useRestrictedMode= */ eq(false));
+    }
+
+    @Test
+    public void agentDisconnected_cancelsCurrentOperations() throws Exception {
+        when(mOperationStorage.operationTokensForPackage(eq(TEST_PACKAGE))).thenReturn(
+                ImmutableSet.of(123, 456, 789));
+        when(mConnectionManager.getThreadForCancellation(any())).thenAnswer(invocation -> {
+            Thread testThread = new Thread((Runnable) invocation.getArgument(0),
+                    "agent-disconnected-test");
+            setTestThread(testThread);
+            return testThread;
+        });
+
+        mConnectionManager.agentDisconnected(TEST_PACKAGE);
+
+        mTestThread.join();
+        verify(mUserBackupManagerService).handleCancel(eq(123), eq(true));
+        verify(mUserBackupManagerService).handleCancel(eq(456), eq(true));
+        verify(mUserBackupManagerService).handleCancel(eq(789), eq(true));
+    }
+
+    @Test
+    public void unbindAgent_callsAmUnbindBackupAgent() throws Exception {
+        mConnectionManager.unbindAgent(mTestApplicationInfo);
+
+        verify(mActivityManager).unbindBackupAgent(eq(mTestApplicationInfo));
+    }
+
+    // Needed because variables can't be assigned directly inside lambdas in Java.
+    private void setBackupAgentResult(IBackupAgent result) {
+        mBackupAgentResult = result;
+    }
+
+    // Needed because variables can't be assigned directly inside lambdas in Java.
+    private void setTestThread(Thread thread) {
+        mTestThread = thread;
+    }
+}
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
index 07f2188..7a9e96f 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/UserBackupManagerServiceTest.java
@@ -18,20 +18,18 @@
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
+
 import static com.google.common.truth.Truth.assertThat;
-import static org.junit.Assert.fail;
+
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.when;
 
 import android.annotation.UserIdInt;
 import android.app.ActivityManagerInternal;
-import android.app.ApplicationThreadConstants;
 import android.app.IActivityManager;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupAnnotations.BackupDestination;
@@ -47,8 +45,6 @@
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.os.Handler;
-import android.platform.test.annotations.DisableFlags;
-import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
@@ -57,7 +53,6 @@
 import android.util.KeyValueListParser;
 
 import androidx.test.core.app.ApplicationProvider;
-import androidx.test.filters.FlakyTest;
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.server.backup.internal.BackupHandler;
@@ -69,8 +64,6 @@
 import com.android.server.backup.utils.BackupEligibilityRules;
 import com.android.server.backup.utils.BackupManagerMonitorEventSender;
 
-import com.google.common.collect.ImmutableSet;
-
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
@@ -84,11 +77,6 @@
 
 import java.util.Arrays;
 import java.util.List;
-import java.util.Set;
-import java.util.function.IntConsumer;
-
-import libcore.junit.util.compat.CoreCompatChangeRule.DisableCompatChanges;
-import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
 
 @Presubmit
 @RunWith(AndroidJUnit4.class)
@@ -96,7 +84,6 @@
     private static final String TEST_PACKAGE = "package1";
     private static final String[] TEST_PACKAGES = new String[] { TEST_PACKAGE };
     private static final String TEST_TRANSPORT = "transport";
-    private static final int WORKER_THREAD_TIMEOUT_MILLISECONDS = 100;
     @UserIdInt private static final int USER_ID = 0;
 
     @Rule
@@ -278,33 +265,6 @@
     }
 
     @Test
-    @FlakyTest
-    public void testAgentDisconnected_cancelsCurrentOperations() throws Exception {
-        when(mOperationStorage.operationTokensForPackage(eq("com.android.foo"))).thenReturn(
-                ImmutableSet.of(123, 456, 789)
-        );
-
-        mService.agentDisconnected("com.android.foo");
-
-        mService.waitForAsyncOperation();
-        verify(mOperationStorage).cancelOperation(eq(123), eq(true), any(IntConsumer.class));
-        verify(mOperationStorage).cancelOperation(eq(456), eq(true), any());
-        verify(mOperationStorage).cancelOperation(eq(789), eq(true), any());
-    }
-
-    @Test
-    public void testAgentDisconnected_unknownPackageName_cancelsNothing() throws Exception {
-        when(mOperationStorage.operationTokensForPackage(eq("com.android.foo"))).thenReturn(
-                ImmutableSet.of()
-        );
-
-        mService.agentDisconnected("com.android.foo");
-
-        verify(mOperationStorage, never())
-                .cancelOperation(anyInt(), anyBoolean(), any(IntConsumer.class));
-    }
-
-    @Test
     public void testReportDelayedRestoreResult_sendsLogsToMonitor() throws Exception {
         PackageInfo packageInfo = getPackageInfo(TEST_PACKAGE);
         when(mPackageManager.getPackageInfoAsUser(anyString(),
@@ -324,158 +284,6 @@
                 eq(packageInfo), eq(results), eq(OperationType.RESTORE));
     }
 
-    @Test
-    @DisableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_restrictedModeChangesFlagOff_shouldUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_keyValueBackup_shouldNotUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_INCREMENTAL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_keyValueRestore_shouldNotUseRestrictedMode()
-            throws Exception {
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_RESTORE, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-        // Make sure we never hit the code that checks the property.
-        verify(mPackageManager, never()).getPropertyAsUser(any(), any(), any(), anyInt());
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_packageOptedIn_shouldUseRestrictedMode()
-            throws Exception {
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenReturn(new PackageManager.Property(
-                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ true,
-                TEST_PACKAGE, /* className= */ null));
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    public void bindToAgentSynchronous_packageOptedOut_shouldNotUseRestrictedMode()
-            throws Exception {
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenReturn(new PackageManager.Property(
-                PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE, /* value= */ false,
-                TEST_PACKAGE, /* className= */ null));
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @DisableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_targetSdkBelowB_shouldUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_targetSdkB_notInList_shouldUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.clearNoRestrictedModePackages();
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(true));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_forRestore_targetSdkB_inList_shouldNotUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.RESTORE);
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_RESTORE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_RESTRICTED_MODE_CHANGES)
-    @EnableCompatChanges({UserBackupManagerService.OS_DECIDES_BACKUP_RESTRICTED_MODE})
-    public void bindToAgentSynchronous_forBackup_targetSdkB_inList_shouldNotUseRestrictedMode()
-            throws Exception {
-        // Mock that the app has not explicitly set the property.
-        when(mPackageManager.getPropertyAsUser(
-                eq(PackageManager.PROPERTY_USE_RESTRICTED_BACKUP_MODE),
-                eq(TEST_PACKAGE), any(), anyInt())).thenThrow(
-                    new PackageManager.NameNotFoundException()
-        );
-        mService.setNoRestrictedModePackages(Set.of(TEST_PACKAGE), OperationType.BACKUP);
-
-        mService.bindToAgentSynchronous(mTestPackageApplicationInfo,
-                ApplicationThreadConstants.BACKUP_MODE_FULL, BackupDestination.CLOUD);
-
-        verify(mActivityManager).bindBackupAgent(eq(TEST_PACKAGE), anyInt(), anyInt(), anyInt(),
-                /* useRestrictedMode= */ eq(false));
-    }
-
     private static PackageInfo getPackageInfo(String packageName) {
         PackageInfo packageInfo = new PackageInfo();
         packageInfo.applicationInfo = new ApplicationInfo();
@@ -487,8 +295,6 @@
         boolean isEnabledStatePersisted = false;
         boolean shouldUseNewBackupEligibilityRules = false;
 
-        private volatile Thread mWorkerThread = null;
-
         TestBackupService() {
             super(mContext, mPackageManager, mOperationStorage, mTransportManager, mBackupHandler,
                     createConstants(mContext), mActivityManager, mActivityManagerInternal);
@@ -523,26 +329,8 @@
         }
 
         @Override
-        Thread getThreadForAsyncOperation(String operationName, Runnable operation) {
-            mWorkerThread = super.getThreadForAsyncOperation(operationName, operation);
-            return mWorkerThread;
-        }
-
-        @Override
         BackupManagerMonitorEventSender getBMMEventSender(IBackupManagerMonitor monitor) {
             return mBackupManagerMonitorEventSender;
         }
-
-        private void waitForAsyncOperation() {
-            if (mWorkerThread == null) {
-                return;
-            }
-
-            try {
-                mWorkerThread.join(/* millis */ WORKER_THREAD_TIMEOUT_MILLISECONDS);
-            } catch (InterruptedException e) {
-                fail("Failed waiting for worker thread to complete: " + e.getMessage());
-            }
-        }
     }
 }
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
index 3310573..e618433 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/fullbackup/PerformFullTransportBackupTaskTest.java
@@ -33,6 +33,7 @@
 
 import androidx.test.runner.AndroidJUnit4;
 
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.BackupAgentTimeoutParameters;
 import com.android.server.backup.OperationStorage;
 import com.android.server.backup.TransportManager;
@@ -66,6 +67,8 @@
     @Mock
     UserBackupManagerService mBackupManagerService;
     @Mock
+    BackupAgentConnectionManager mBackupAgentConnectionManager;
+    @Mock
     BackupTransportClient mBackupTransportClient;
     @Mock
     CountDownLatch mLatch;
@@ -95,6 +98,8 @@
         when(mBackupManagerService.isSetupComplete()).thenReturn(true);
         when(mBackupManagerService.getAgentTimeoutParameters()).thenReturn(
                 mBackupAgentTimeoutParameters);
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
         when(mBackupManagerService.getTransportManager()).thenReturn(mTransportManager);
         when(mTransportManager.getCurrentTransportClient(any())).thenReturn(mTransportConnection);
         when(mTransportConnection.connectOrThrow(any())).thenReturn(mBackupTransportClient);
@@ -142,11 +147,11 @@
 
         mTask.run();
 
-        InOrder inOrder = inOrder(mBackupManagerService);
-        inOrder.verify(mBackupManagerService).setNoRestrictedModePackages(
+        InOrder inOrder = inOrder(mBackupAgentConnectionManager);
+        inOrder.verify(mBackupAgentConnectionManager).setNoRestrictedModePackages(
                 eq(Set.of("package1")),
                 eq(BackupAnnotations.OperationType.BACKUP));
-        inOrder.verify(mBackupManagerService).clearNoRestrictedModePackages();
+        inOrder.verify(mBackupAgentConnectionManager).clearNoRestrictedModePackages();
     }
 
     private void createTask(String[] packageNames) {
diff --git a/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java b/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
index 055adf6..351aac3 100644
--- a/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/backup/restore/PerformUnifiedRestoreTaskTest.java
@@ -44,6 +44,7 @@
 import androidx.test.runner.AndroidJUnit4;
 
 import com.android.modules.utils.testing.TestableDeviceConfig;
+import com.android.server.backup.BackupAgentConnectionManager;
 import com.android.server.backup.Flags;
 import com.android.server.backup.UserBackupManagerService;
 import com.android.server.backup.internal.BackupHandler;
@@ -95,6 +96,8 @@
     private TransportConnection mTransportConnection;
     @Mock
     private BackupTransportClient mBackupTransportClient;
+    @Mock
+    private BackupAgentConnectionManager mBackupAgentConnectionManager;
 
     private Set<String> mExcludedkeys = new HashSet<>();
     private Map<String, String> mBackupData = new HashMap<>();
@@ -122,6 +125,9 @@
 
         mContext = InstrumentationRegistry.getInstrumentation().getTargetContext();
 
+        when(mBackupManagerService.getBackupAgentConnectionManager()).thenReturn(
+                mBackupAgentConnectionManager);
+
         mBackupDataSource = new ArrayDeque<>(mBackupData.keySet());
         when(mBackupDataInput.readNextHeader())
                 .then((Answer<Boolean>) invocation -> !mBackupDataSource.isEmpty());
@@ -166,7 +172,7 @@
         mRestoreTask.setNoRestrictedModePackages(mBackupTransportClient,
                 new PackageInfo[]{packageInfo1, packageInfo2});
 
-        verify(mBackupManagerService).setNoRestrictedModePackages(
+        verify(mBackupAgentConnectionManager).setNoRestrictedModePackages(
                 eq(Set.of("package1")),
                 eq(BackupAnnotations.OperationType.RESTORE));
     }
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
index d66bb00..c6870ad 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/controllers/QuotaControllerTest.java
@@ -65,6 +65,7 @@
 import android.app.usage.UsageEvents;
 import android.app.usage.UsageStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
+import android.compat.testing.PlatformCompatChangeRule;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
@@ -103,10 +104,13 @@
 import com.android.server.job.controllers.QuotaController.TimingSession;
 import com.android.server.usage.AppStandbyInternal;
 
+import libcore.junit.util.compat.CoreCompatChangeRule.EnableCompatChanges;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.TestRule;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatchers;
@@ -135,6 +139,9 @@
     private static final int SOURCE_USER_ID = 0;
 
     @Rule public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    @Rule
+    public TestRule compatChangeRule = new PlatformCompatChangeRule();
     private QuotaController mQuotaController;
     private QuotaController.QcConstants mQcConstants;
     private JobSchedulerService.Constants mConstants = new JobSchedulerService.Constants();
@@ -303,7 +310,7 @@
 
     private int getProcessStateQuotaFreeThreshold() {
         synchronized (mQuotaController.mLock) {
-            return mQuotaController.getProcessStateQuotaFreeThreshold();
+            return mQuotaController.getProcessStateQuotaFreeThreshold(mSourceUid);
         }
     }
 
@@ -5197,6 +5204,101 @@
         assertFalse(jobBg2.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
     }
 
+    @Test
+    @EnableCompatChanges({QuotaController.OVERRIDE_QUOTA_ENFORCEMENT_TO_TOP_STARTED_JOBS,
+            QuotaController.OVERRIDE_QUOTA_ENFORCEMENT_TO_FGS_JOBS})
+    @RequiresFlagsEnabled({Flags.FLAG_ENFORCE_QUOTA_POLICY_TO_TOP_STARTED_JOBS,
+            Flags.FLAG_ENFORCE_QUOTA_POLICY_TO_FGS_JOBS})
+    public void testTracking_OutOfQuota_ForegroundAndBackground_CompactChangeOverrides() {
+        setDischarging();
+
+        JobStatus jobBg = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 1);
+        JobStatus jobTop = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 2);
+        trackJobs(jobBg, jobTop);
+        setStandbyBucket(WORKING_INDEX, jobTop, jobBg); // 2 hour window
+        // Now the package only has 20 seconds to run.
+        final long remainingTimeMs = 20 * SECOND_IN_MILLIS;
+        mQuotaController.saveTimingSession(SOURCE_USER_ID, SOURCE_PACKAGE,
+                createTimingSession(
+                        JobSchedulerService.sElapsedRealtimeClock.millis() - HOUR_IN_MILLIS,
+                        10 * MINUTE_IN_MILLIS - remainingTimeMs, 1), false);
+
+        InOrder inOrder = inOrder(mJobSchedulerService);
+
+        // UID starts out inactive.
+        setProcessState(ActivityManager.PROCESS_STATE_SERVICE);
+        // Start the job.
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.prepareForExecutionLocked(jobBg);
+        }
+        advanceElapsedClock(remainingTimeMs / 2);
+        // New job starts after UID is in the foreground. Since the app is now in the foreground, it
+        // should continue to have remainingTimeMs / 2 time remaining.
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.prepareForExecutionLocked(jobTop);
+        }
+        advanceElapsedClock(remainingTimeMs);
+
+        // Wait for some extra time to allow for job processing.
+        inOrder.verify(mJobSchedulerService,
+                        timeout(remainingTimeMs + 2 * SECOND_IN_MILLIS).times(0))
+                .onControllerStateChanged(argThat(jobs -> jobs.size() > 0));
+        synchronized (mQuotaController.mLock) {
+            assertEquals(remainingTimeMs / 2,
+                    mQuotaController.getRemainingExecutionTimeLocked(jobBg));
+            assertEquals(remainingTimeMs / 2,
+                    mQuotaController.getRemainingExecutionTimeLocked(jobTop));
+        }
+        // Go to a background state.
+        setProcessState(ActivityManager.PROCESS_STATE_TOP_SLEEPING);
+        advanceElapsedClock(remainingTimeMs / 2 + 1);
+        // Only Bg job will be changed from in-quota to out-of-quota.
+        inOrder.verify(mJobSchedulerService,
+                        timeout(remainingTimeMs / 2 + 2 * SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged(argThat(jobs -> jobs.size() == 1));
+        // Top job should still be allowed to run.
+        assertFalse(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        // New jobs to run.
+        JobStatus jobBg2 = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 3);
+        JobStatus jobTop2 = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 4);
+        JobStatus jobFg = createJobStatus("testTracking_OutOfQuota_ForegroundAndBackground", 5);
+        setStandbyBucket(WORKING_INDEX, jobBg2, jobTop2, jobFg);
+
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_TOP);
+        inOrder.verify(mJobSchedulerService, timeout(SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged(argThat(jobs -> jobs.size() == 1));
+        trackJobs(jobFg, jobTop);
+        synchronized (mQuotaController.mLock) {
+            mQuotaController.prepareForExecutionLocked(jobTop);
+        }
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        // App still in foreground so everything should be in quota.
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_FOREGROUND_SERVICE);
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertTrue(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+
+        advanceElapsedClock(20 * SECOND_IN_MILLIS);
+        setProcessState(ActivityManager.PROCESS_STATE_SERVICE);
+        inOrder.verify(mJobSchedulerService, timeout(SECOND_IN_MILLIS).times(1))
+                .onControllerStateChanged(argThat(jobs -> jobs.size() == 2));
+        // App is now in background and out of quota. Fg should now change to out of quota since it
+        // wasn't started. Top should remain in quota since it started when the app was in TOP.
+        assertTrue(jobTop.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertFalse(jobFg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        assertFalse(jobBg.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+        trackJobs(jobBg2);
+        assertFalse(jobBg2.isConstraintSatisfied(JobStatus.CONSTRAINT_WITHIN_QUOTA));
+    }
+
     /**
      * Tests that TOP jobs are stopped when an app runs out of quota.
      */
diff --git a/services/tests/servicestests/src/com/android/server/security/authenticationpolicy/AuthenticationPolicyServiceTest.java b/services/tests/servicestests/src/com/android/server/security/authenticationpolicy/AuthenticationPolicyServiceTest.java
index 2238a1b..ee8eb9b 100644
--- a/services/tests/servicestests/src/com/android/server/security/authenticationpolicy/AuthenticationPolicyServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/security/authenticationpolicy/AuthenticationPolicyServiceTest.java
@@ -19,12 +19,14 @@
 import static android.adaptiveauth.Flags.FLAG_ENABLE_ADAPTIVE_AUTH;
 import static android.adaptiveauth.Flags.FLAG_REPORT_BIOMETRIC_AUTH_ATTEMPTS;
 import static android.security.Flags.FLAG_REPORT_PRIMARY_AUTH_ATTEMPTS;
+import static android.security.authenticationpolicy.AuthenticationPolicyManager.ERROR_UNSUPPORTED;
 
 import static com.android.internal.widget.LockPatternUtils.StrongAuthTracker.SOME_AUTH_REQUIRED_AFTER_ADAPTIVE_AUTH_REQUEST;
 import static com.android.server.security.authenticationpolicy.AuthenticationPolicyService.MAX_ALLOWED_FAILED_AUTH_ATTEMPTS;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assume.assumeTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
@@ -95,6 +97,8 @@
     private WindowManagerInternal mWindowManager;
     @Mock
     private UserManagerInternal mUserManager;
+    @Mock
+    private SecureLockDeviceServiceInternal mSecureLockDeviceService;
 
     @Captor
     ArgumentCaptor<LockSettingsStateListener> mLockSettingsStateListenerCaptor;
@@ -123,6 +127,11 @@
         LocalServices.addService(WindowManagerInternal.class, mWindowManager);
         LocalServices.removeServiceForTest(UserManagerInternal.class);
         LocalServices.addService(UserManagerInternal.class, mUserManager);
+        if (android.security.Flags.secureLockdown()) {
+            LocalServices.removeServiceForTest(SecureLockDeviceServiceInternal.class);
+            LocalServices.addService(SecureLockDeviceServiceInternal.class,
+                    mSecureLockDeviceService);
+        }
 
         mAuthenticationPolicyService = new AuthenticationPolicyService(
                 mContext, mLockPatternUtils);
@@ -136,6 +145,12 @@
         // Set PRIMARY_USER_ID as the parent of MANAGED_PROFILE_USER_ID
         when(mUserManager.getProfileParentId(eq(MANAGED_PROFILE_USER_ID)))
                 .thenReturn(PRIMARY_USER_ID);
+        if (android.security.Flags.secureLockdown()) {
+            when(mSecureLockDeviceService.enableSecureLockDevice(any()))
+                    .thenReturn(ERROR_UNSUPPORTED);
+            when(mSecureLockDeviceService.disableSecureLockDevice(any()))
+                    .thenReturn(ERROR_UNSUPPORTED);
+        }
     }
 
     @After
@@ -143,6 +158,9 @@
         LocalServices.removeServiceForTest(LockSettingsInternal.class);
         LocalServices.removeServiceForTest(WindowManagerInternal.class);
         LocalServices.removeServiceForTest(UserManagerInternal.class);
+        if (android.security.Flags.secureLockdown()) {
+            LocalServices.removeServiceForTest(SecureLockDeviceServiceInternal.class);
+        }
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
index b91a5b7..d5ed048 100644
--- a/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/AppCompatCameraOverridesTest.java
@@ -17,8 +17,8 @@
 package com.android.server.wm;
 
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_FORCE_ROTATION;
-import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT;
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_REFRESH;
+import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT;
 import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_ENABLE_REFRESH_VIA_PAUSE;
 import static android.content.pm.ActivityInfo.OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA;
 import static android.content.pm.ActivityInfo.OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA;
@@ -228,9 +228,8 @@
     }
 
     @Test
-    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT})
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
-    public void testShouldApplyCameraCompatFreeformTreatment_overrideEnabled_returnsFalse() {
+    public void testShouldApplyCameraCompatFreeformTreatment_notEnabledByOverride_returnsFalse() {
         runTestScenario((robot) -> {
             robot.activity().createActivityWithComponentInNewTask();
 
@@ -239,19 +238,9 @@
     }
 
     @Test
-    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT})
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
-    public void testShouldApplyCameraCompatFreeformTreatment_disabledByOverride_returnsFalse() {
-        runTestScenario((robot) -> {
-            robot.activity().createActivityWithComponentInNewTask();
-
-            robot.checkShouldApplyFreeformTreatmentForCameraCompat(false);
-        });
-    }
-
-    @Test
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
-    public void testShouldApplyCameraCompatFreeformTreatment_notDisabledByOverride_returnsTrue() {
+    public void testShouldApplyCameraCompatFreeformTreatment_overrideAndFlagEnabled_returnsTrue() {
         runTestScenario((robot) -> {
             robot.activity().createActivityWithComponentInNewTask();
 
@@ -261,8 +250,9 @@
 
     @Test
     @EnableCompatChanges({OVERRIDE_ORIENTATION_ONLY_FOR_CAMERA,
-            OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA})
-    public void testShouldRecomputeConfigurationForCameraCompat() {
+            OVERRIDE_MIN_ASPECT_RATIO_ONLY_FOR_CAMERA,
+            OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
+    public void testShouldRecomputeConfigurationForFreeformTreatment() {
         runTestScenario((robot) -> {
             robot.conf().enableCameraCompatSplitScreenAspectRatio(true);
             robot.applyOnActivity((a) -> {
diff --git a/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
index 2635bb3..c427583 100644
--- a/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/CameraCompatFreeformPolicyTests.java
@@ -25,7 +25,7 @@
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.app.servertransaction.ActivityLifecycleItem.ON_PAUSE;
 import static android.app.servertransaction.ActivityLifecycleItem.ON_STOP;
-import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT;
+import static android.content.pm.ActivityInfo.OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_FULL_USER;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_PORTRAIT;
@@ -142,8 +142,9 @@
         cameraStateMonitor.startListeningToCameraState();
     }
 
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testFullscreen_doesNotActivateCameraCompatMode() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT, WINDOWING_MODE_FULLSCREEN);
         doReturn(false).when(mActivity).inFreeformWindowingMode();
@@ -153,23 +154,26 @@
         assertNotInCameraCompatMode();
     }
 
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testOrientationUnspecified_doesNotActivateCameraCompatMode() {
         configureActivity(SCREEN_ORIENTATION_UNSPECIFIED);
 
         assertNotInCameraCompatMode();
     }
 
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testNoCameraConnection_doesNotActivateCameraCompatMode() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         assertNotInCameraCompatMode();
     }
 
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraConnected_deviceInPortrait_portraitCameraCompatMode() throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         setDisplayRotation(Surface.ROTATION_0);
@@ -180,6 +184,8 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraConnected_deviceInLandscape_portraitCameraCompatMode() throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         setDisplayRotation(Surface.ROTATION_270);
@@ -190,6 +196,8 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraConnected_deviceInPortrait_landscapeCameraCompatMode() throws Exception {
         configureActivity(SCREEN_ORIENTATION_LANDSCAPE);
         setDisplayRotation(Surface.ROTATION_0);
@@ -200,6 +208,8 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraConnected_deviceInLandscape_landscapeCameraCompatMode() throws Exception {
         configureActivity(SCREEN_ORIENTATION_LANDSCAPE);
         setDisplayRotation(Surface.ROTATION_270);
@@ -209,8 +219,9 @@
         assertActivityRefreshRequested(/* refreshRequested */ false);
     }
 
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraReconnected_cameraCompatModeAndRefresh() throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         setDisplayRotation(Surface.ROTATION_270);
@@ -229,6 +240,8 @@
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testCameraOpenedForDifferentPackage_notInCameraCompatMode() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
 
@@ -239,31 +252,32 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
-    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT})
-    public void testShouldApplyCameraCompatFreeformTreatment_overrideEnabled_returnsFalse() {
+    public void testShouldApplyCameraCompatFreeformTreatment_overrideNotEnabled_returnsFalse() {
+        configureActivity(SCREEN_ORIENTATION_PORTRAIT);
+
+        mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
+
+        assertFalse(mCameraCompatFreeformPolicy.isTreatmentEnabledForActivity(mActivity,
+                /* checkOrientation */ true));
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges(OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT)
+    public void testShouldApplyCameraCompatFreeformTreatment_enabledByOverride_returnsTrue() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
 
         mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
 
         assertTrue(mActivity.info
-                .isChangeEnabled(OVERRIDE_CAMERA_COMPAT_DISABLE_FREEFORM_WINDOWING_TREATMENT));
-        assertFalse(mCameraCompatFreeformPolicy.isTreatmentEnabledForActivity(
-                mActivity, /*  checkOrientation= */ true));
+                .isChangeEnabled(OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT));
+        assertTrue(mCameraCompatFreeformPolicy.isTreatmentEnabledForActivity(mActivity,
+                /* checkOrientation */ true));
     }
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
-    public void testShouldApplyCameraCompatFreeformTreatment_notDisabledByOverride_returnsTrue() {
-        configureActivity(SCREEN_ORIENTATION_PORTRAIT);
-
-        mCameraAvailabilityCallback.onCameraOpened(CAMERA_ID_1, TEST_PACKAGE_1);
-
-        assertTrue(mCameraCompatFreeformPolicy.isTreatmentEnabledForActivity(
-                mActivity, /*  checkOrientation= */ true));
-    }
-
-    @Test
-    @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testShouldRefreshActivity_appBoundsChanged_returnsTrue() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         Configuration oldConfiguration = createConfiguration(/* letterbox= */ false);
@@ -276,6 +290,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testShouldRefreshActivity_displayRotationChanged_returnsTrue() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         Configuration oldConfiguration = createConfiguration(/* letterbox= */ true);
@@ -291,6 +306,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testShouldRefreshActivity_appBoundsNorDisplayChanged_returnsFalse() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         Configuration oldConfiguration = createConfiguration(/* letterbox= */ true);
@@ -306,6 +322,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testOnActivityConfigurationChanging_refreshDisabledViaFlag_noRefresh()
             throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
@@ -321,6 +338,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testOnActivityConfigurationChanging_cycleThroughStopDisabled() throws Exception {
         when(mAppCompatConfiguration.isCameraCompatRefreshCycleThroughStopEnabled())
                 .thenReturn(false);
@@ -335,6 +353,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testOnActivityConfigurationChanging_cycleThroughStopDisabledForApp()
             throws Exception {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
@@ -349,6 +368,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testGetCameraCompatAspectRatio_activityNotInCameraCompat_returnsDefaultAspRatio() {
         configureActivity(SCREEN_ORIENTATION_FULL_USER);
 
@@ -362,6 +382,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testGetCameraCompatAspectRatio_activityInCameraCompat_returnsConfigAspectRatio() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         final float configAspectRatio = 1.5f;
@@ -377,6 +398,7 @@
 
     @Test
     @EnableFlags(FLAG_ENABLE_CAMERA_COMPAT_FOR_DESKTOP_WINDOWING)
+    @EnableCompatChanges({OVERRIDE_CAMERA_COMPAT_ENABLE_FREEFORM_WINDOWING_TREATMENT})
     public void testGetCameraCompatAspectRatio_inCameraCompatPerAppOverride_returnDefAspectRatio() {
         configureActivity(SCREEN_ORIENTATION_PORTRAIT);
         final float configAspectRatio = 1.5f;
diff --git a/telephony/java/android/telephony/CarrierConfigManager.java b/telephony/java/android/telephony/CarrierConfigManager.java
index 6490cbe..7cfdec6 100644
--- a/telephony/java/android/telephony/CarrierConfigManager.java
+++ b/telephony/java/android/telephony/CarrierConfigManager.java
@@ -9734,6 +9734,35 @@
             "carrier_supported_satellite_services_per_provider_bundle";
 
     /**
+     * A PersistableBundle that contains a list of key-value pairs, where the values are integer
+     * arrays.
+     * <p>
+     * Keys are the IDs of regional satellite configs as strings and values are
+     * integer arrays of earfcns in the corresponding regions.
+     *
+     * An example config for two regions "1" and "2":
+     * <pre>{@code
+     * <carrier_config>
+     *   <pbundle_as_map name="regional_satellite_earfcn_bundle">
+     *     <int-array name = "1" num = "2">
+     *       <item value = "100"/>
+     *       <item value = "200"/>
+     *     </int-array>
+     *     <int-array name = "2" num = "1">
+     *       <item value = "200"/>
+     *     </int-array>
+     *   </pbundle_as_map>
+     * </carrier_config>
+     * }</pre>
+     * <p>
+     * This config is empty by default.
+     * @hide
+     */
+    @FlaggedApi(Flags.FLAG_CARRIER_ROAMING_NB_IOT_NTN)
+    public static final String KEY_REGIONAL_SATELLITE_EARFCN_BUNDLE =
+            "regional_satellite_earfcn_bundle";
+
+    /**
      * This config enables modem to scan satellite PLMNs specified as per
      * {@link #KEY_CARRIER_SUPPORTED_SATELLITE_SERVICES_PER_PROVIDER_BUNDLE} and attach to same
      * in case cellular networks are not enabled. This will need specific agreement between
@@ -11264,6 +11293,9 @@
         sDefaults.putPersistableBundle(
                 KEY_CARRIER_SUPPORTED_SATELLITE_SERVICES_PER_PROVIDER_BUNDLE,
                 PersistableBundle.EMPTY);
+        sDefaults.putPersistableBundle(
+                KEY_REGIONAL_SATELLITE_EARFCN_BUNDLE,
+                PersistableBundle.EMPTY);
         sDefaults.putBoolean(KEY_SATELLITE_ATTACH_SUPPORTED_BOOL, false);
         sDefaults.putInt(KEY_SATELLITE_CONNECTION_HYSTERESIS_SEC_INT, 180);
         sDefaults.putIntArray(KEY_NTN_LTE_RSRP_THRESHOLDS_INT_ARRAY,
diff --git a/tests/AppJankTest/src/android/app/jank/tests/JankDataProcessorTest.java b/tests/AppJankTest/src/android/app/jank/tests/JankDataProcessorTest.java
index 2cd625e..4d495ad 100644
--- a/tests/AppJankTest/src/android/app/jank/tests/JankDataProcessorTest.java
+++ b/tests/AppJankTest/src/android/app/jank/tests/JankDataProcessorTest.java
@@ -18,7 +18,9 @@
 
 import static org.junit.Assert.assertEquals;
 
+import android.app.jank.AppJankStats;
 import android.app.jank.Flags;
+import android.app.jank.FrameOverrunHistogram;
 import android.app.jank.JankDataProcessor;
 import android.app.jank.StateTracker;
 import android.platform.test.annotations.RequiresFlagsEnabled;
@@ -39,6 +41,7 @@
 import org.junit.runner.RunWith;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 
 @RunWith(AndroidJUnit4.class)
@@ -154,6 +157,73 @@
         assertEquals(totalFrames, histogramFrames);
     }
 
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_DETAILED_APP_JANK_METRICS_API)
+    public void mergeAppJankStats_confirmStatAddedToPendingStats() {
+        HashMap<String, JankDataProcessor.PendingJankStat> pendingStats =
+                mJankDataProcessor.getPendingJankStats();
+
+        assertEquals(pendingStats.size(), 0);
+
+        AppJankStats jankStats = getAppJankStats();
+        mJankDataProcessor.mergeJankStats(jankStats, sActivityName);
+
+        pendingStats = mJankDataProcessor.getPendingJankStats();
+
+        assertEquals(pendingStats.size(), 1);
+    }
+
+    /**
+     * This test confirms matching states are combined into one pending stat.  When JankStats are
+     * merged from outside the platform they will contain widget category, widget id and widget
+     * state. If an incoming JankStats matches a pending stat on all those fields the incoming
+     * JankStat will be merged into the existing stat.
+     */
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_DETAILED_APP_JANK_METRICS_API)
+    public void mergeAppJankStats_confirmStatsWithMatchingStatesAreCombinedIntoOnePendingStat() {
+        AppJankStats jankStats = getAppJankStats();
+        mJankDataProcessor.mergeJankStats(jankStats, sActivityName);
+
+        HashMap<String, JankDataProcessor.PendingJankStat> pendingStats =
+                mJankDataProcessor.getPendingJankStats();
+        assertEquals(pendingStats.size(), 1);
+
+        AppJankStats secondJankStat = getAppJankStats();
+        mJankDataProcessor.mergeJankStats(secondJankStat, sActivityName);
+
+        pendingStats = mJankDataProcessor.getPendingJankStats();
+
+        assertEquals(pendingStats.size(), 1);
+    }
+
+    @Test
+    @RequiresFlagsEnabled(Flags.FLAG_DETAILED_APP_JANK_METRICS_API)
+    public void mergeAppJankStats_whenStatsWithMatchingStatesMerge_confirmFrameCountsAdded() {
+        AppJankStats jankStats = getAppJankStats();
+        mJankDataProcessor.mergeJankStats(jankStats, sActivityName);
+        mJankDataProcessor.mergeJankStats(jankStats, sActivityName);
+
+        HashMap<String, JankDataProcessor.PendingJankStat> pendingStats =
+                mJankDataProcessor.getPendingJankStats();
+
+        String statKey = pendingStats.keySet().iterator().next();
+        JankDataProcessor.PendingJankStat pendingStat = pendingStats.get(statKey);
+
+        assertEquals(pendingStats.size(), 1);
+        // The same jankStats objects are merged twice, this should result in the frame counts being
+        // doubled.
+        assertEquals(jankStats.getJankyFrameCount() * 2, pendingStat.getJankyFrames());
+        assertEquals(jankStats.getTotalFrameCount() * 2, pendingStat.getTotalFrames());
+
+        int[] originalHistogramBuckets = jankStats.getFrameOverrunHistogram().getBucketCounters();
+        int[] frameOverrunBuckets = pendingStat.getFrameOverrunBuckets();
+
+        for (int i = 0; i < frameOverrunBuckets.length; i++) {
+            assertEquals(originalHistogramBuckets[i] * 2, frameOverrunBuckets[i]);
+        }
+    }
+
     // TODO b/375005277 add tests that cover logging and releasing resources back to pool.
 
     private long getTotalFramesCounted() {
@@ -276,4 +346,26 @@
         return mockData;
     }
 
+    private AppJankStats getAppJankStats() {
+        AppJankStats jankStats = new AppJankStats(
+                /*App Uid*/APP_ID,
+                /*Widget Id*/"test widget id",
+                /*Widget Category*/AppJankStats.SCROLL,
+                /*Widget State*/AppJankStats.SCROLLING,
+                /*Total Frames*/100,
+                /*Janky Frames*/25,
+                getOverrunHistogram()
+        );
+        return jankStats;
+    }
+
+    private FrameOverrunHistogram getOverrunHistogram() {
+        FrameOverrunHistogram overrunHistogram = new FrameOverrunHistogram();
+        overrunHistogram.addFrameOverrunMillis(-2);
+        overrunHistogram.addFrameOverrunMillis(1);
+        overrunHistogram.addFrameOverrunMillis(5);
+        overrunHistogram.addFrameOverrunMillis(25);
+        return overrunHistogram;
+    }
+
 }
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 d2c9eb3..6432827 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
@@ -92,7 +92,7 @@
     }
 
     /** Move an app to Desktop by dragging the app handle at the top. */
-    fun enterDesktopModeWithDrag(
+    private fun enterDesktopModeWithDrag(
         wmHelper: WindowManagerStateHelper,
         device: UiDevice,
         motionEventHelper: MotionEventHelper = MotionEventHelper(getInstrumentation(), TOUCH)