Merge "Add bundle validations on request/response and state params." into main
diff --git a/cmds/bmgr/src/com/android/commands/bmgr/Bmgr.java b/cmds/bmgr/src/com/android/commands/bmgr/Bmgr.java
index 0d3dc49..6310d32 100644
--- a/cmds/bmgr/src/com/android/commands/bmgr/Bmgr.java
+++ b/cmds/bmgr/src/com/android/commands/bmgr/Bmgr.java
@@ -1280,6 +1280,24 @@
                 return "START_PACKAGE_RESTORE";
             case BackupManagerMonitor.LOG_EVENT_ID_AGENT_FAILURE:
                 return "AGENT_FAILURE";
+            case BackupManagerMonitor.LOG_EVENT_ID_RESTORE_AT_INSTALL_INVOKED:
+                return "RESTORE_AT_INSTALL_INVOKED";
+            case BackupManagerMonitor.LOG_EVENT_ID_SKIP_RESTORE_AT_INSTALL:
+                return "SKIP_RESTORE_AT_INSTALL";
+            case BackupManagerMonitor.LOG_EVENT_ID_PACKAGE_ACCEPTED_FOR_RESTORE:
+                return "PACKAGE_ACCEPTED_FOR_RESTORE";
+            case BackupManagerMonitor.LOG_EVENT_ID_RESTORE_DATA_DOES_NOT_BELONG_TO_PACKAGE:
+                return "RESTORE_DATA_DOES_NOT_BELONG_TO_PACKAGE";
+            case BackupManagerMonitor.LOG_EVENT_ID_UNABLE_TO_CREATE_AGENT_FOR_RESTORE:
+                return "UNABLE_TO_CREATE_AGENT_FOR_RESTORE";
+            case BackupManagerMonitor.LOG_EVENT_ID_AGENT_CRASHED_BEFORE_RESTORE_DATA_IS_SENT:
+                return "AGENT_CRASHED_BEFORE_RESTORE_DATA_IS_SEN";
+            case BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_SEND_DATA_TO_AGENT_DURING_RESTORE:
+                return "FAILED_TO_SEND_DATA_TO_AGENT_DURING_RESTORE";
+            case BackupManagerMonitor.LOG_EVENT_ID_AGENT_FAILURE_DURING_RESTORE:
+                return "AGENT_FAILURE_DURING_RESTORE";
+            case BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_READ_DATA_FROM_TRANSPORT:
+                return "FAILED_TO_READ_DATA_FROM_TRANSPORT";
             default:
                 return "UNKNOWN_ID";
         }
diff --git a/core/java/android/app/NotificationChannel.java b/core/java/android/app/NotificationChannel.java
index 7c803eb..193c524 100644
--- a/core/java/android/app/NotificationChannel.java
+++ b/core/java/android/app/NotificationChannel.java
@@ -434,6 +434,40 @@
     /**
      * @hide
      */
+    public NotificationChannel copy() {
+        NotificationChannel copy = new NotificationChannel(mId, mName, mImportance);
+        copy.setDescription(mDesc);
+        copy.setBypassDnd(mBypassDnd);
+        copy.setLockscreenVisibility(mLockscreenVisibility);
+        copy.setSound(mSound, mAudioAttributes);
+        copy.setLightColor(mLightColor);
+        copy.enableLights(mLights);
+        copy.setVibrationPattern(mVibrationPattern);
+        if (Flags.notificationChannelVibrationEffectApi()) {
+            copy.setVibrationEffect(mVibrationEffect);
+        }
+        copy.lockFields(mUserLockedFields);
+        copy.setUserVisibleTaskShown(mUserVisibleTaskShown);
+        copy.enableVibration(mVibrationEnabled);
+        copy.setShowBadge(mShowBadge);
+        copy.setDeleted(mDeleted);
+        copy.setGroup(mGroup);
+        copy.setBlockable(mBlockableSystem);
+        copy.setAllowBubbles(mAllowBubbles);
+        copy.setOriginalImportance(mOriginalImportance);
+        copy.setConversationId(mParentId, mConversationId);
+        copy.setDemoted(mDemoted);
+        copy.setImportantConversation(mImportantConvo);
+        copy.setDeletedTimeMs(mDeletedTime);
+        copy.setImportanceLockedByCriticalDeviceFunction(mImportanceLockedDefaultApp);
+        copy.setLastNotificationUpdateTimeMs(mLastNotificationUpdateTimeMs);
+
+        return copy;
+    }
+
+    /**
+     * @hide
+     */
     @TestApi
     public void lockFields(int field) {
         mUserLockedFields |= field;
diff --git a/core/java/android/app/backup/BackupManagerMonitor.java b/core/java/android/app/backup/BackupManagerMonitor.java
index c66478f..e741bc2 100644
--- a/core/java/android/app/backup/BackupManagerMonitor.java
+++ b/core/java/android/app/backup/BackupManagerMonitor.java
@@ -269,6 +269,33 @@
   /** V to U restore attempt, allowlist and denlist are set
    @hide */
   public static final int LOG_EVENT_ID_V_TO_U_RESTORE_SET_LIST = 72;
+  /** As part of package install, {@link PackageManager} invoked restore.
+   @hide */
+  public static final int LOG_EVENT_ID_RESTORE_AT_INSTALL_INVOKED = 73;
+  /** Skipping restore at package install
+   @hide */
+  public static final int LOG_EVENT_ID_SKIP_RESTORE_AT_INSTALL = 74;
+  /** Package is eligible and is accepted for restore
+   @hide */
+  public static final int LOG_EVENT_ID_PACKAGE_ACCEPTED_FOR_RESTORE = 75;
+  /** Restore data doesn't belong to the package for which restore is started
+   @hide */
+  public static final int LOG_EVENT_ID_RESTORE_DATA_DOES_NOT_BELONG_TO_PACKAGE = 76;
+  /** Unable to create BackupAgent for package for restore
+   @hide */
+  public static final int LOG_EVENT_ID_UNABLE_TO_CREATE_AGENT_FOR_RESTORE = 77;
+  /** BackupAgent crashed after creation but before accepting any data
+   @hide */
+  public static final int LOG_EVENT_ID_AGENT_CRASHED_BEFORE_RESTORE_DATA_IS_SENT = 78;
+  /** Failure in streaming restore data to BackupAgent
+   @hide */
+  public static final int LOG_EVENT_ID_FAILED_TO_SEND_DATA_TO_AGENT_DURING_RESTORE = 79;
+  /** BackupAgent related failure during restore
+   @hide */
+  public static final int LOG_EVENT_ID_AGENT_FAILURE_DURING_RESTORE = 80;
+  /** Failure in reading data from TransportPackage during restore
+   @hide */
+  public static final int LOG_EVENT_ID_FAILED_TO_READ_DATA_FROM_TRANSPORT = 81;
 
   /**
    * This method will be called each time something important happens on BackupManager.
diff --git a/core/java/android/app/notification.aconfig b/core/java/android/app/notification.aconfig
index ce06772b..0082732 100644
--- a/core/java/android/app/notification.aconfig
+++ b/core/java/android/app/notification.aconfig
@@ -89,4 +89,25 @@
   namespace: "systemui"
   description: "Changes notification sort order to be by time within a section"
   bug: "330193582"
+}
+
+flag {
+  name: "restrict_audio_attributes_call"
+  namespace: "systemui"
+  description: "Only CallStyle notifs can use USAGE_NOTIFICATION_RINGTONE"
+  bug: "331793339"
+}
+
+flag {
+  name: "restrict_audio_attributes_alarm"
+  namespace: "systemui"
+  description: "Only alarm category notifs can use USAGE_ALARM"
+  bug: "331793339"
+}
+
+flag {
+  name: "restrict_audio_attributes_media"
+  namespace: "systemui"
+  description: "No notifs can use USAGE_UNKNOWN or USAGE_MEDIA"
+  bug: "331793339"
 }
\ No newline at end of file
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index a82c9a8..bd8e9c6 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -43,6 +43,7 @@
 import static android.view.flags.Flags.toolkitFrameRateBySizeReadOnly;
 import static android.view.flags.Flags.toolkitFrameRateDefaultNormalReadOnly;
 import static android.view.flags.Flags.toolkitFrameRateSmallUsesPercentReadOnly;
+import static android.view.flags.Flags.toolkitFrameRateViewEnablingReadOnly;
 import static android.view.flags.Flags.toolkitMetricsForFrameRateDecision;
 import static android.view.flags.Flags.toolkitSetFrameRateReadOnly;
 import static android.view.flags.Flags.viewVelocityApi;
@@ -2439,6 +2440,8 @@
 
     private static final boolean sToolkitFrameRateSmallUsesPercentReadOnlyFlagValue =
             toolkitFrameRateSmallUsesPercentReadOnly();
+    private static final boolean sToolkitFrameRateViewEnablingReadOnlyFlagValue =
+            toolkitFrameRateViewEnablingReadOnly();
 
     // Used to set frame rate compatibility.
     @Surface.FrameRateCompatibility int mFrameRateCompatibility =
@@ -20794,7 +20797,8 @@
         }
 
         // For VRR to vote the preferred frame rate
-        if (sToolkitSetFrameRateReadOnlyFlagValue) {
+        if (sToolkitSetFrameRateReadOnlyFlagValue
+                && sToolkitFrameRateViewEnablingReadOnlyFlagValue) {
             votePreferredFrameRate();
         }
 
@@ -20901,7 +20905,8 @@
     protected void damageInParent() {
         if (mParent != null && mAttachInfo != null) {
             // For VRR to vote the preferred frame rate
-            if (sToolkitSetFrameRateReadOnlyFlagValue) {
+            if (sToolkitSetFrameRateReadOnlyFlagValue
+                    && sToolkitFrameRateViewEnablingReadOnlyFlagValue) {
                 votePreferredFrameRate();
             }
             mParent.onDescendantInvalidated(this, this);
@@ -23592,7 +23597,8 @@
         }
 
         mPrivateFlags4 = (mPrivateFlags4 & ~PFLAG4_HAS_MOVED) | PFLAG4_HAS_DRAWN;
-        if (sToolkitSetFrameRateReadOnlyFlagValue) {
+        if (sToolkitSetFrameRateReadOnlyFlagValue
+                && sToolkitFrameRateViewEnablingReadOnlyFlagValue) {
             updateInfrequentCount();
         }
 
@@ -25508,7 +25514,7 @@
     }
 
     private void sizeChange(int newWidth, int newHeight, int oldWidth, int oldHeight) {
-        if (mAttachInfo != null) {
+        if (mAttachInfo != null && sToolkitFrameRateViewEnablingReadOnlyFlagValue) {
             boolean isSmall;
             if (sToolkitFrameRateSmallUsesPercentReadOnlyFlagValue) {
                 int size = newWidth * newHeight;
diff --git a/core/java/android/view/flags/refresh_rate_flags.aconfig b/core/java/android/view/flags/refresh_rate_flags.aconfig
index b296e44..d0fe3e0 100644
--- a/core/java/android/view/flags/refresh_rate_flags.aconfig
+++ b/core/java/android/view/flags/refresh_rate_flags.aconfig
@@ -103,4 +103,12 @@
     description: "Feature flag to enable the functionality of the dVRR feature"
     bug: "239979904"
     is_fixed_read_only: true
+}
+
+flag {
+    name: "toolkit_frame_rate_view_enabling_read_only"
+    namespace: "toolkit"
+    description: "Feature flag to enable the functionality on views for the dVRR feature"
+    bug: "239979904"
+    is_fixed_read_only: true
 }
\ No newline at end of file
diff --git a/core/tests/coretests/src/android/app/NotificationChannelTest.java b/core/tests/coretests/src/android/app/NotificationChannelTest.java
index 18209b5..504f98f 100644
--- a/core/tests/coretests/src/android/app/NotificationChannelTest.java
+++ b/core/tests/coretests/src/android/app/NotificationChannelTest.java
@@ -31,6 +31,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import android.annotation.FlaggedApi;
 import android.content.AttributionSource;
 import android.content.ContentProvider;
 import android.content.ContentResolver;
@@ -46,6 +47,7 @@
 import android.os.RemoteCallback;
 import android.os.RemoteException;
 import android.os.VibrationEffect;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.annotations.Presubmit;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.MediaStore.Audio.AudioColumns;
@@ -577,6 +579,40 @@
         assertNull(channel.getVibrationEffect());
     }
 
+    @Test
+    @EnableFlags({Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_MEDIA,
+            Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_CALL, Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_ALARM})
+    public void testCopy() {
+        NotificationChannel original = new NotificationChannel("id", "name", 2);
+        original.setDescription("desc");
+        original.setBypassDnd(true);
+        original.setLockscreenVisibility(7);
+        original.setSound(Uri.EMPTY, new AudioAttributes.Builder().build());
+        original.setLightColor(5);
+        original.enableLights(false);
+        original.setVibrationPattern(new long[] {1, 9, 3});
+        if (Flags.notificationChannelVibrationEffectApi()) {
+            original.setVibrationEffect(VibrationEffect.createOneShot(100, 5));
+        }
+        original.lockFields(9999);
+        original.setUserVisibleTaskShown(true);
+        original.enableVibration(false);
+        original.setShowBadge(true);
+        original.setDeleted(false);
+        original.setGroup("group");
+        original.setBlockable(false);
+        original.setAllowBubbles(true);
+        original.setOriginalImportance(6);
+        original.setConversationId("parent", "convo");
+        original.setDemoted(false);
+        original.setImportantConversation(true);
+        original.setDeletedTimeMs(100);
+        original.setImportanceLockedByCriticalDeviceFunction(false);
+
+        NotificationChannel parcelCopy = writeToAndReadFromParcel(original);
+        assertThat(original.copy()).isEqualTo(parcelCopy);
+    }
+
     /** Backs up a given channel to an XML, and returns the channel read from the XML. */
     private NotificationChannel backUpAndRestore(NotificationChannel channel) throws Exception {
         TypedXmlSerializer serializer = Xml.newFastSerializer();
diff --git a/core/tests/coretests/src/android/view/ViewFrameRateTest.java b/core/tests/coretests/src/android/view/ViewFrameRateTest.java
index 3e172c1..dcfbf64 100644
--- a/core/tests/coretests/src/android/view/ViewFrameRateTest.java
+++ b/core/tests/coretests/src/android/view/ViewFrameRateTest.java
@@ -21,6 +21,7 @@
 import static android.view.Surface.FRAME_RATE_CATEGORY_NORMAL;
 import static android.view.Surface.FRAME_RATE_COMPATIBILITY_GTE;
 import static android.view.flags.Flags.FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY;
+import static android.view.flags.Flags.FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY;
 import static android.view.flags.Flags.FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY;
 import static android.view.flags.Flags.FLAG_VIEW_VELOCITY_API;
 import static android.view.flags.Flags.toolkitFrameRateBySizeReadOnly;
@@ -35,6 +36,8 @@
 import android.app.Activity;
 import android.os.SystemClock;
 import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.util.DisplayMetrics;
 
 import androidx.test.annotation.UiThreadTest;
@@ -60,6 +63,9 @@
     public ActivityTestRule<ViewCaptureTestActivity> mActivityRule = new ActivityTestRule<>(
             ViewCaptureTestActivity.class);
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private Activity mActivity;
     private View mMovingView;
     private ViewRootImpl mViewRoot;
@@ -80,7 +86,8 @@
 
     @UiThreadTest
     @Test
-    @RequiresFlagsEnabled(FLAG_VIEW_VELOCITY_API)
+    @RequiresFlagsEnabled({FLAG_VIEW_VELOCITY_API,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void frameRateChangesWhenContentMoves() {
         mMovingView.offsetLeftAndRight(100);
         float frameRate = mViewRoot.getPreferredFrameRate();
@@ -121,7 +128,8 @@
 
     @Test
     @RequiresFlagsEnabled({FLAG_VIEW_VELOCITY_API,
-            FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void lowVelocity60() throws Throwable {
         mActivityRule.runOnUiThread(() -> {
             ViewGroup.LayoutParams layoutParams = mMovingView.getLayoutParams();
@@ -140,7 +148,8 @@
 
     @Test
     @RequiresFlagsEnabled({FLAG_VIEW_VELOCITY_API,
-            FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_VELOCITY_MAPPING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void highVelocity140() throws Throwable {
         mActivityRule.runOnUiThread(() -> {
             ViewGroup.LayoutParams layoutParams = mMovingView.getLayoutParams();
@@ -172,7 +181,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void noVelocityUsesCategorySmall() throws Throwable {
         final CountDownLatch drawLatch1 = new CountDownLatch(1);
         mActivityRule.runOnUiThread(() -> {
@@ -206,7 +216,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void noVelocityUsesCategoryNarrowWidth() throws Throwable {
         final CountDownLatch drawLatch1 = new CountDownLatch(1);
         mActivityRule.runOnUiThread(() -> {
@@ -239,7 +250,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void noVelocityUsesCategoryNarrowHeight() throws Throwable {
         final CountDownLatch drawLatch1 = new CountDownLatch(1);
         mActivityRule.runOnUiThread(() -> {
@@ -272,7 +284,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void noVelocityUsesCategoryLargeWidth() throws Throwable {
         final CountDownLatch drawLatch1 = new CountDownLatch(1);
         mActivityRule.runOnUiThread(() -> {
@@ -305,7 +318,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void noVelocityUsesCategoryLargeHeight() throws Throwable {
         final CountDownLatch drawLatch1 = new CountDownLatch(1);
         mActivityRule.runOnUiThread(() -> {
@@ -338,7 +352,8 @@
     }
 
     @Test
-    @RequiresFlagsEnabled(FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY)
+    @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void defaultNormal() throws Throwable {
         mActivityRule.runOnUiThread(() -> {
             View parent = (View) mMovingView.getParent();
diff --git a/core/tests/coretests/src/android/view/ViewRootImplTest.java b/core/tests/coretests/src/android/view/ViewRootImplTest.java
index 6f107a9..a86e568 100644
--- a/core/tests/coretests/src/android/view/ViewRootImplTest.java
+++ b/core/tests/coretests/src/android/view/ViewRootImplTest.java
@@ -20,6 +20,7 @@
 import static android.view.flags.Flags.FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY;
 import static android.view.flags.Flags.FLAG_TOOLKIT_FRAME_RATE_BY_SIZE_READ_ONLY;
 import static android.view.flags.Flags.FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY;
+import static android.view.flags.Flags.FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY;
 import static android.view.flags.Flags.FLAG_VIEW_VELOCITY_API;
 import static android.view.Surface.FRAME_RATE_CATEGORY_HIGH;
 import static android.view.Surface.FRAME_RATE_CATEGORY_HIGH_HINT;
@@ -578,7 +579,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRateCategory_visibility_defaultHigh() {
         View view = new View(sContext);
         attachViewToWindow(view);
@@ -612,7 +614,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRateCategory_smallSize_defaultHigh() {
         View view = new View(sContext);
         WindowManager.LayoutParams wmlp = new WindowManager.LayoutParams(TYPE_APPLICATION_OVERLAY);
@@ -640,7 +643,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRateCategory_normalSize_defaultHigh() {
         View view = new View(sContext);
         WindowManager.LayoutParams wmlp = new WindowManager.LayoutParams(TYPE_APPLICATION_OVERLAY);
@@ -777,7 +781,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRate_category() {
         View view = new View(sContext);
         attachViewToWindow(view);
@@ -817,7 +822,9 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_VIEW_VELOCITY_API, FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_VIEW_VELOCITY_API,
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRateCategory_velocityToHigh() {
         View view = new View(sContext);
         WindowManager.LayoutParams wmlp = new WindowManager.LayoutParams(TYPE_APPLICATION_OVERLAY);
@@ -957,7 +964,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_voteFrameRateOnly() {
         View view = new View(sContext);
         float frameRate = 20;
@@ -1000,7 +1008,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_infrequentLayer_defaultHigh() throws InterruptedException {
         final long delay = 200L;
 
@@ -1102,7 +1111,8 @@
      */
     @Test
     @RequiresFlagsEnabled({FLAG_TOOLKIT_SET_FRAME_RATE_READ_ONLY,
-            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY})
+            FLAG_TOOLKIT_FRAME_RATE_FUNCTION_ENABLING_READ_ONLY,
+            FLAG_TOOLKIT_FRAME_RATE_VIEW_ENABLING_READ_ONLY})
     public void votePreferredFrameRate_applyTextureViewHeuristic() throws InterruptedException {
         final long delay = 30L;
 
diff --git a/nfc/Android.bp b/nfc/Android.bp
index 0b3f291..c186804 100644
--- a/nfc/Android.bp
+++ b/nfc/Android.bp
@@ -39,6 +39,7 @@
     libs: [
         "unsupportedappusage", // for android.compat.annotation.UnsupportedAppUsage
         "framework-permission-s",
+        "framework-permission",
     ],
     static_libs: [
         "android.nfc.flags-aconfig-java",
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index c730d49..74d61ca 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -197,12 +197,29 @@
 }
 
 filegroup {
+    name: "kosmos-src",
+    srcs: ["tests/utils/kosmos/src/**/*.kt"],
+    path: "tests/utils/kosmos",
+}
+
+java_library {
+    name: "kosmos",
+    host_supported: true,
+    srcs: [":kosmos-src"],
+    static_libs: [
+        "kotlin-reflect",
+        "kotlin-stdlib",
+    ],
+}
+
+filegroup {
     name: "SystemUI-tests-utils",
     srcs: [
         "tests/utils/src/**/*.java",
         "tests/utils/src/**/*.kt",
+        ":kosmos-src",
     ],
-    path: "tests/utils/src",
+    path: "tests/utils",
 }
 
 filegroup {
@@ -334,6 +351,7 @@
         "kotlin-test",
         "SystemUICustomizationTestUtils",
         "androidx.compose.runtime_runtime",
+        "kosmos",
     ],
     libs: [
         "android.test.runner",
@@ -416,6 +434,7 @@
         "inline-mockito-robolectric-prebuilt",
         "platform-parametric-runner-lib",
         "SystemUICustomizationTestUtils",
+        "kosmos",
     ],
     libs: [
         "android.test.runner",
@@ -449,6 +468,7 @@
         "androidx.test.uiautomator_uiautomator",
         "androidx.core_core-animation-testing",
         "androidx.test.ext.junit",
+        "kosmos",
     ],
     libs: [
         "android.test.runner",
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/ToAodFoldTransitionInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/ToAodFoldTransitionInteractor.kt
index 3b25128..6729246 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/ToAodFoldTransitionInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/ToAodFoldTransitionInteractor.kt
@@ -75,7 +75,7 @@
 
     fun initialize(parentAnimator: ShadeFoldAnimator) {
         this.parentAnimator =
-            parentAnimator as NotificationPanelViewController.ShadeFoldAnimatorImpl?
+            parentAnimator as? NotificationPanelViewController.ShadeFoldAnimatorImpl?
     }
 
     /** Forces the keyguard into AOD or Doze */
diff --git a/packages/SystemUI/tests/utils/kosmos/src/com/android/systemui/kosmos/Kosmos.kt b/packages/SystemUI/tests/utils/kosmos/src/com/android/systemui/kosmos/Kosmos.kt
new file mode 100644
index 0000000..685bb0a
--- /dev/null
+++ b/packages/SystemUI/tests/utils/kosmos/src/com/android/systemui/kosmos/Kosmos.kt
@@ -0,0 +1,97 @@
+/*
+ * 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.kosmos
+
+import kotlin.reflect.KProperty
+
+// (Historical note: The name Kosmos is meant to invoke "Kotlin", the "Object Mother" pattern
+//   (https://martinfowler.com/bliki/ObjectMother.html), and of course the Greek word "kosmos" for
+//   the "order of the world" (https://en.wiktionary.org/wiki/%CE%BA%CF%8C%CF%83%CE%BC%CE%BF%CF%82)
+
+/**
+ * Each Kosmos is its own self-contained set of fixtures, which may reference each other. Fixtures
+ * can be defined through extension properties in any file:
+ * ```
+ * // fixture that must be set:
+ * var Kosmos.context by Fixture<Context>()
+ *
+ * // fixture with overrideable default.
+ * var Kosmos.landscapeMode by Fixture { false }
+ *
+ * // fixture forbidding override (note `val`, and referencing context fixture from above)
+ * val Kosmos.lifecycleScope by Fixture { context.lifecycleScope }
+ * ```
+ *
+ * To use the fixtures, create an instance of Kosmos and retrieve the values you need:
+ * ```
+ * val k = Kosmos()
+ * k.context = mContext
+ * val underTest = YourInteractor(
+ *     context = k.context,
+ *     landscapeMode = k.landscapeMode,
+ * )
+ * ```
+ */
+interface Kosmos {
+    /**
+     * Lookup a fixture in the Kosmos by [name], using [creator] to instantiate and store one if
+     * there is none present.
+     */
+    fun <T> get(name: String, creator: (Kosmos.() -> T)?): T
+
+    /** Sets the [value] of a fixture with the given [name]. */
+    fun set(name: String, value: Any?)
+
+    /**
+     * A value in the kosmos that has a single value once it's read. It can be overridden before
+     * first use only; all objects that are dependent on this fixture will get the same value.
+     *
+     * Example classic uses would be a clock, filesystem, or singleton controller.
+     *
+     * If no [creator] parameter is provided, the fixture must be set before use.
+     */
+    class Fixture<T>(private val creator: (Kosmos.() -> T)? = null) {
+        operator fun getValue(thisRef: Kosmos, property: KProperty<*>): T =
+            thisRef.get(property.name, creator)
+
+        operator fun setValue(thisRef: Kosmos, property: KProperty<*>, value: T) {
+            thisRef.set(property.name, value)
+        }
+    }
+}
+
+/** Constructs a fresh Kosmos. */
+fun Kosmos(): Kosmos = KosmosRegistry()
+
+private class KosmosRegistry : Kosmos {
+    val map: MutableMap<String, Any?> = mutableMapOf()
+    val gotten: MutableSet<String> = mutableSetOf()
+
+    override fun <T> get(name: String, creator: (Kosmos.() -> T)?): T {
+        gotten.add(name)
+        if (name !in map) {
+            checkNotNull(creator) { "Fixture $name has no default, and is read before set." }
+            map[name] = creator()
+        }
+        @Suppress("UNCHECKED_CAST") return map[name] as T
+    }
+
+    override fun set(name: String, value: Any?) {
+        check(name !in gotten) { "Tried to set fixture '$name' after it's already been read." }
+        map[name] = value
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/Kosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/Kosmos.kt
deleted file mode 100644
index c74cdd4..0000000
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/kosmos/Kosmos.kt
+++ /dev/null
@@ -1,69 +0,0 @@
-package com.android.systemui.kosmos
-
-import kotlin.reflect.KProperty
-
-// (Historical note: The name Kosmos is meant to invoke "Kotlin", the "Object Mother" pattern
-//   (https://martinfowler.com/bliki/ObjectMother.html), and of course the Greek word "kosmos" for
-//   the "order of the world" (https://en.wiktionary.org/wiki/%CE%BA%CF%8C%CF%83%CE%BC%CE%BF%CF%82)
-/**
- * Each Kosmos is its own self-contained set of fixtures, which may reference each other. Fixtures
- * can be defined through extension properties in any file:
- * ```
- * // fixture that must be set:
- * var Kosmos.context by Fixture<Context>()
- *
- * // fixture with overrideable default.
- * var Kosmos.landscapeMode by Fixture { false }
- *
- * // fixture forbidding override (note `val`, and referencing context fixture from above)
- * val Kosmos.lifecycleScope by Fixture { context.lifecycleScope }
- * ```
- *
- * To use the fixtures, create an instance of Kosmos and retrieve the values you need:
- * ```
- * val k = Kosmos()
- * k.context = mContext
- * val underTest = YourInteractor(
- *     context = k.context,
- *     landscapeMode = k.landscapeMode,
- * )
- * ```
- */
-class Kosmos {
-    private val map: MutableMap<String, Any?> = mutableMapOf()
-    private val gotten: MutableSet<String> = mutableSetOf()
-
-    /**
-     * A value in the kosmos that has a single value once it's read. It can be overridden before
-     * first use only; all objects that are dependent on this fixture will get the same value.
-     *
-     * Example classic uses would be a clock, filesystem, or singleton controller.
-     *
-     * If no [creator] parameter is provided, the fixture must be set before use.
-     */
-    class Fixture<T>(private val creator: (Kosmos.() -> T)? = null) {
-        operator fun getValue(thisRef: Kosmos, property: KProperty<*>): T {
-            thisRef.gotten.add(property.name)
-            @Suppress("UNCHECKED_CAST")
-            if (!thisRef.map.contains(property.name)) {
-                if (creator == null) {
-                    throw IllegalStateException(
-                        "Fixture ${property.name} has no default, and is read before set."
-                    )
-                } else {
-                    val nonNullCreator = creator
-                    // The Kotlin compiler seems to need this odd workaround
-                    thisRef.map[property.name] = thisRef.nonNullCreator()
-                }
-            }
-            return thisRef.map[property.name] as T
-        }
-
-        operator fun setValue(thisRef: Kosmos, property: KProperty<*>, value: T) {
-            check(!thisRef.gotten.contains(property.name)) {
-                "Tried to set fixture '${property.name}' after it's already been read."
-            }
-            thisRef.map[property.name] = value
-        }
-    }
-}
diff --git a/services/backup/flags.aconfig b/services/backup/flags.aconfig
index e9f959f..74adfe0 100644
--- a/services/backup/flags.aconfig
+++ b/services/backup/flags.aconfig
@@ -50,4 +50,12 @@
             "logger."
     bug: "296844513"
     is_fixed_read_only: true
+}
+
+flag {
+    name: "enable_increased_bmm_logging_for_restore_at_install"
+    namespace: "onboarding"
+    description: "Increase BMM logging coverage in restore at install flow."
+    bug: "331749778"
+    is_fixed_read_only: true
 }
\ No newline at end of file
diff --git a/services/backup/java/com/android/server/backup/UserBackupManagerService.java b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
index 976504a..dc1cfb9 100644
--- a/services/backup/java/com/android/server/backup/UserBackupManagerService.java
+++ b/services/backup/java/com/android/server/backup/UserBackupManagerService.java
@@ -3903,11 +3903,33 @@
             skip = true;
         }
 
+        BackupManagerMonitorEventSender  mBMMEventSender =
+                getBMMEventSender(/*monitor=*/ null);
+        PackageInfo packageInfo = getPackageInfoForBMMLogging(packageName);
         TransportConnection transportConnection =
                 mTransportManager.getCurrentTransportClient("BMS.restoreAtInstall()");
         if (transportConnection == null) {
             if (DEBUG) Slog.w(TAG, addUserIdToLogMessage(mUserId, "No transport client"));
             skip = true;
+        } else if (Flags.enableIncreasedBmmLoggingForRestoreAtInstall()) {
+            try {
+                BackupTransportClient transportClient = transportConnection.connectOrThrow(
+                        "BMS.restoreAtInstall");
+                mBMMEventSender.setMonitor(transportClient.getBackupManagerMonitor());
+            } catch (TransportNotAvailableException | RemoteException e) {
+                mBMMEventSender.monitorEvent(
+                        BackupManagerMonitor.LOG_EVENT_ID_TRANSPORT_IS_NULL, packageInfo,
+                        BackupManagerMonitor.LOG_EVENT_CATEGORY_TRANSPORT, null);
+            }
+        }
+
+        if (Flags.enableIncreasedBmmLoggingForRestoreAtInstall()) {
+            mBMMEventSender.monitorEvent(
+                    BackupManagerMonitor.LOG_EVENT_ID_RESTORE_AT_INSTALL_INVOKED, packageInfo,
+                    BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                    mBMMEventSender.putMonitoringExtra(/*extras=*/null,
+                            BackupManagerMonitor.EXTRA_LOG_OPERATION_TYPE,
+                            BackupAnnotations.OperationType.RESTORE));
         }
 
         if (!mAutoRestore) {
@@ -3943,7 +3965,7 @@
                         RestoreParams.createForRestoreAtInstall(
                                 transportConnection,
                                 /* observer */ null,
-                                /* monitor */ null,
+                                mBMMEventSender.getMonitor(),
                                 restoreSet,
                                 packageName,
                                 token,
@@ -3963,6 +3985,15 @@
         if (skip) {
             // Auto-restore disabled or no way to attempt a restore
 
+            if (Flags.enableIncreasedBmmLoggingForRestoreAtInstall()) {
+                mBMMEventSender.monitorEvent(
+                        BackupManagerMonitor.LOG_EVENT_ID_SKIP_RESTORE_AT_INSTALL, packageInfo,
+                        BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                        mBMMEventSender.putMonitoringExtra(/*extras=*/null,
+                                BackupManagerMonitor.EXTRA_LOG_OPERATION_TYPE,
+                                BackupAnnotations.OperationType.RESTORE));
+            }
+
             if (transportConnection != null) {
                 mTransportManager.disposeOfTransportClient(
                         transportConnection, "BMS.restoreAtInstall()");
@@ -3976,6 +4007,23 @@
         }
     }
 
+    private PackageInfo getPackageInfoForBMMLogging(String packageName) {
+        try {
+            return mPackageManager.getPackageInfoAsUser(packageName, 0, mUserId);
+        } catch (NameNotFoundException e) {
+            Slog.w(
+                    TAG,
+                    addUserIdToLogMessage(
+                            mUserId, "Asked to get PackageInfo for BMM logging of nonexistent pkg "
+                                    + packageName));
+
+            PackageInfo packageInfo = new PackageInfo();
+            packageInfo.packageName = packageName;
+
+            return packageInfo;
+        }
+    }
+
     /** Hand off a restore session. */
     public IRestoreSession beginRestoreSession(String packageName, String transport) {
         if (DEBUG) {
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 2c9eb51..b414b25 100644
--- a/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
+++ b/services/backup/java/com/android/server/backup/restore/FullRestoreEngine.java
@@ -29,6 +29,7 @@
 import android.app.IBackupAgent;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupAnnotations;
+import android.app.backup.BackupManagerMonitor;
 import android.app.backup.FullBackup;
 import android.app.backup.IBackupManagerMonitor;
 import android.app.backup.IFullBackupRestoreObserver;
@@ -57,6 +58,7 @@
 import com.android.server.backup.UserBackupManagerService;
 import com.android.server.backup.fullbackup.FullBackupObbConnection;
 import com.android.server.backup.utils.BackupEligibilityRules;
+import com.android.server.backup.utils.BackupManagerMonitorEventSender;
 import com.android.server.backup.utils.BytesReadListener;
 import com.android.server.backup.utils.FullBackupRestoreObserverUtils;
 import com.android.server.backup.utils.RestoreUtils;
@@ -143,6 +145,8 @@
 
     private FileMetadata mReadOnlyParent = null;
 
+    private BackupManagerMonitorEventSender mBackupManagerMonitorEventSender;
+
     public FullRestoreEngine(
             UserBackupManagerService backupManagerService, OperationStorage operationStorage,
             BackupRestoreTask monitorTask, IFullBackupRestoreObserver observer,
@@ -155,6 +159,7 @@
         mMonitorTask = monitorTask;
         mObserver = observer;
         mMonitor = monitor;
+        mBackupManagerMonitorEventSender = new BackupManagerMonitorEventSender(monitor);
         mOnlyPackage = onlyPackage;
         mAllowApks = allowApks;
         mAgentTimeoutParameters = Objects.requireNonNull(
@@ -225,6 +230,9 @@
                     // one app's data but see a different app's on the wire
                     if (onlyPackage != null) {
                         if (!pkg.equals(onlyPackage.packageName)) {
+                            logBMMEvent(
+                                    BackupManagerMonitor.LOG_EVENT_ID_RESTORE_DATA_DOES_NOT_BELONG_TO_PACKAGE,
+                                    onlyPackage);
                             Slog.w(TAG, "Expected data for " + onlyPackage + " but saw " + pkg);
                             setResult(RestoreEngine.TRANSPORT_FAILURE);
                             setRunning(false);
@@ -412,6 +420,9 @@
                         }
 
                         if (mAgent == null) {
+                            logBMMEvent(
+                                    BackupManagerMonitor.LOG_EVENT_ID_UNABLE_TO_CREATE_AGENT_FOR_RESTORE,
+                                    onlyPackage);
                             Slog.e(TAG, "Unable to create agent for " + pkg);
                             okay = false;
                             tearDownPipes();
@@ -501,6 +512,9 @@
                         } catch (RemoteException e) {
                             // whoops, remote entity went away.  We'll eat the content
                             // ourselves, then, and not copy it over.
+                            logBMMEvent(
+                                    BackupManagerMonitor.LOG_EVENT_ID_AGENT_CRASHED_BEFORE_RESTORE_DATA_IS_SENT,
+                                    onlyPackage);
                             Slog.e(TAG, "Agent crashed during full restore");
                             agentSuccess = false;
                             okay = false;
@@ -531,6 +545,9 @@
                                     } catch (IOException e) {
                                         Slog.e(TAG, "Failed to write to restore pipe: "
                                                 + e.getMessage());
+                                        logBMMEvent(
+                                                BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_SEND_DATA_TO_AGENT_DURING_RESTORE,
+                                                onlyPackage);
                                         pipeOkay = false;
                                     }
                                 }
@@ -548,6 +565,8 @@
                         // okay, if the remote end failed at any point, deal with
                         // it by ignoring the rest of the restore on it
                         if (!agentSuccess) {
+                            logBMMEvent(BackupManagerMonitor.LOG_EVENT_ID_AGENT_FAILURE_DURING_RESTORE,
+                                    onlyPackage);
                             Slog.w(TAG, "Agent failure restoring " + pkg + "; ending restore");
                             mBackupManagerService.getBackupHandler().removeMessages(
                                     MSG_RESTORE_OPERATION_TIMEOUT);
@@ -590,6 +609,8 @@
             if (DEBUG) {
                 Slog.w(TAG, "io exception on restore socket read: " + e.getMessage());
             }
+            logBMMEvent(BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_READ_DATA_FROM_TRANSPORT,
+                    onlyPackage);
             setResult(RestoreEngine.TRANSPORT_FAILURE);
             info = null;
         }
@@ -631,6 +652,16 @@
         return false;
     }
 
+    private void logBMMEvent(int eventId, PackageInfo pkgInfo) {
+        if (Flags.enableIncreasedBmmLoggingForRestoreAtInstall()) {
+            mBackupManagerMonitorEventSender.monitorEvent(eventId, pkgInfo,
+                    BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                    mBackupManagerMonitorEventSender.putMonitoringExtra(/*extras=*/
+                            null, BackupManagerMonitor.EXTRA_LOG_OPERATION_TYPE,
+                            BackupAnnotations.OperationType.RESTORE));
+        }
+    }
+
     private static boolean isValidParent(FileMetadata parentDir, @NonNull FileMetadata childDir) {
         return parentDir != null
                 && childDir.packageName.equals(parentDir.packageName)
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 8fece82..e536876 100644
--- a/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
+++ b/services/backup/java/com/android/server/backup/restore/PerformUnifiedRestoreTask.java
@@ -418,25 +418,6 @@
     private void startRestore() {
         sendStartRestore(mAcceptSet.size());
 
-        // If we're starting a full-system restore, set up to begin widget ID remapping
-        if (mIsSystemRestore) {
-            AppWidgetBackupBridge.systemRestoreStarting(mUserId);
-            Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
-            mBackupManagerMonitorEventSender.monitorEvent(
-                    BackupManagerMonitor.LOG_EVENT_ID_START_SYSTEM_RESTORE,
-                    null,
-                    BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
-                    monitoringExtras);
-        } else {
-            // We are either performing RestoreAtInstall or Bmgr.
-            Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
-            mBackupManagerMonitorEventSender.monitorEvent(
-                    BackupManagerMonitor.LOG_EVENT_ID_START_RESTORE_AT_INSTALL,
-                    null,
-                    BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
-                    monitoringExtras);
-        }
-
         try {
             String transportDirName =
                     mTransportManager.getTransportDirName(
@@ -459,6 +440,34 @@
                 mBackupManagerMonitorEventSender.setMonitor(transport.getBackupManagerMonitor());
             }
 
+            if (Flags.enableIncreasedBmmLoggingForRestoreAtInstall()) {
+                for (PackageInfo info : mAcceptSet) {
+                    mBackupManagerMonitorEventSender.monitorEvent(
+                            BackupManagerMonitor.LOG_EVENT_ID_PACKAGE_ACCEPTED_FOR_RESTORE, info,
+                            BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                            addRestoreOperationTypeToEvent(/* extras= */ null));
+                }
+            }
+
+            // If we're starting a full-system restore, set up to begin widget ID remapping
+            if (mIsSystemRestore) {
+                AppWidgetBackupBridge.systemRestoreStarting(mUserId);
+                Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
+                mBackupManagerMonitorEventSender.monitorEvent(
+                        BackupManagerMonitor.LOG_EVENT_ID_START_SYSTEM_RESTORE,
+                        null,
+                        BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                        monitoringExtras);
+            } else {
+                // We are either performing RestoreAtInstall or Bmgr.
+                Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
+                mBackupManagerMonitorEventSender.monitorEvent(
+                        BackupManagerMonitor.LOG_EVENT_ID_START_RESTORE_AT_INSTALL,
+                        null,
+                        BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
+                        monitoringExtras);
+            }
+
             mStatus = transport.startRestore(mToken, packages);
             if (mStatus != BackupTransport.TRANSPORT_OK) {
                 Slog.e(TAG, "Transport error " + mStatus + "; no restore possible");
@@ -638,7 +647,7 @@
                 Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
                 mBackupManagerMonitorEventSender.monitorEvent(
                         BackupManagerMonitor.LOG_EVENT_ID_PACKAGE_NOT_PRESENT,
-                        mCurrentPackage,
+                        createPackageInfoForBMMLogging(pkgName),
                         BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
                         monitoringExtras);
                 EventLog.writeEvent(
@@ -739,7 +748,7 @@
             Bundle monitoringExtras = addRestoreOperationTypeToEvent(/* extras= */ null);
             mBackupManagerMonitorEventSender.monitorEvent(
                     BackupManagerMonitor.LOG_EVENT_ID_NO_NEXT_RESTORE_TARGET,
-                    mCurrentPackage,
+                    null,
                     BackupManagerMonitor.LOG_EVENT_CATEGORY_BACKUP_MANAGER_POLICY,
                     monitoringExtras);
             EventLog.writeEvent(EventLogTags.RESTORE_TRANSPORT_FAILURE);
@@ -1804,4 +1813,10 @@
                 monitoringExtrasDenylist);
     }
 
+    private PackageInfo createPackageInfoForBMMLogging(String packageName) {
+        PackageInfo packageInfo = new PackageInfo();
+        packageInfo.packageName = packageName;
+
+        return packageInfo;
+    }
 }
diff --git a/services/backup/java/com/android/server/backup/utils/BackupManagerMonitorDumpsysUtils.java b/services/backup/java/com/android/server/backup/utils/BackupManagerMonitorDumpsysUtils.java
index 6d315ba..fad59d2 100644
--- a/services/backup/java/com/android/server/backup/utils/BackupManagerMonitorDumpsysUtils.java
+++ b/services/backup/java/com/android/server/backup/utils/BackupManagerMonitorDumpsysUtils.java
@@ -371,6 +371,24 @@
                     "V to U restore pkg not eligible";
             case BackupManagerMonitor.LOG_EVENT_ID_V_TO_U_RESTORE_SET_LIST ->
                     "V to U restore lists";
+            case BackupManagerMonitor.LOG_EVENT_ID_RESTORE_AT_INSTALL_INVOKED ->
+                    "Invoked restore at install";
+            case BackupManagerMonitor.LOG_EVENT_ID_SKIP_RESTORE_AT_INSTALL ->
+                    "Skip restore at install";
+            case BackupManagerMonitor.LOG_EVENT_ID_PACKAGE_ACCEPTED_FOR_RESTORE ->
+                    "Pkg accepted for restore";
+            case BackupManagerMonitor.LOG_EVENT_ID_RESTORE_DATA_DOES_NOT_BELONG_TO_PACKAGE ->
+                    "Restore data does not belong to package";
+            case BackupManagerMonitor.LOG_EVENT_ID_UNABLE_TO_CREATE_AGENT_FOR_RESTORE ->
+                    "Unable to create Agent";
+            case BackupManagerMonitor.LOG_EVENT_ID_AGENT_CRASHED_BEFORE_RESTORE_DATA_IS_SENT ->
+                    "Agent crashed before restore data is streamed";
+            case BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_SEND_DATA_TO_AGENT_DURING_RESTORE ->
+                    "Failed to send data to agent";
+            case BackupManagerMonitor.LOG_EVENT_ID_AGENT_FAILURE_DURING_RESTORE ->
+                    "Agent failure during restore";
+            case BackupManagerMonitor.LOG_EVENT_ID_FAILED_TO_READ_DATA_FROM_TRANSPORT ->
+                    "Failed to read data from Transport";
             default -> "Unknown log event ID: " + code;
         };
         return id;
diff --git a/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java b/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
index 9e3f5ce..f38d772b 100644
--- a/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
+++ b/services/companion/java/com/android/server/companion/virtual/GenericWindowPolicyController.java
@@ -334,7 +334,7 @@
             // The error dialog alerting users that streaming is blocked is always allowed.
             return true;
         }
-        if (!mAllowedUsers.contains(activityUser)) {
+        if (!activityUser.isSystem() && !mAllowedUsers.contains(activityUser)) {
             Slog.d(TAG, "Virtual device launch disallowed from user " + activityUser);
             return false;
         }
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 434985e..40f0362 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -150,6 +150,7 @@
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.display.BrightnessSynchronizer;
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.internal.foldables.FoldLockSettingAvailabilityProvider;
 import com.android.internal.os.BackgroundThread;
 import com.android.internal.util.ArrayUtils;
@@ -596,7 +597,7 @@
         mUiHandler = UiThread.getHandler();
         mDisplayDeviceRepo = new DisplayDeviceRepository(mSyncRoot, mPersistentDataStore);
         mLogicalDisplayMapper = new LogicalDisplayMapper(mContext,
-                foldSettingProvider,
+                foldSettingProvider, new FoldGracePeriodProvider(),
                 mDisplayDeviceRepo, new LogicalDisplayListener(), mSyncRoot, mHandler, mFlags);
         mDisplayModeDirector = new DisplayModeDirector(context, mHandler, mFlags);
         mBrightnessSynchronizer = new BrightnessSynchronizer(mContext,
diff --git a/services/core/java/com/android/server/display/LogicalDisplayMapper.java b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
index f727eac..6203a32 100644
--- a/services/core/java/com/android/server/display/LogicalDisplayMapper.java
+++ b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
@@ -40,6 +40,7 @@
 import android.view.DisplayInfo;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.server.display.feature.DisplayManagerFlags;
 import com.android.server.display.layout.DisplayIdProducer;
 import com.android.server.display.layout.Layout;
@@ -120,7 +121,7 @@
     /**
      * Sleep the device when transitioning into these device state.
      */
-    private final SparseBooleanArray mDeviceStatesOnWhichToSleep;
+    private final SparseBooleanArray mDeviceStatesOnWhichToSelectiveSleep;
 
     /**
      * Map of all logical displays indexed by logical display id.
@@ -153,6 +154,7 @@
     private final DisplayManagerService.SyncRoot mSyncRoot;
     private final LogicalDisplayMapperHandler mHandler;
     private final FoldSettingProvider mFoldSettingProvider;
+    private final FoldGracePeriodProvider mFoldGracePeriodProvider;
     private final PowerManager mPowerManager;
 
     /**
@@ -201,15 +203,18 @@
     private final DisplayManagerFlags mFlags;
 
     LogicalDisplayMapper(@NonNull Context context, FoldSettingProvider foldSettingProvider,
+            FoldGracePeriodProvider foldGracePeriodProvider,
             @NonNull DisplayDeviceRepository repo,
             @NonNull Listener listener, @NonNull DisplayManagerService.SyncRoot syncRoot,
             @NonNull Handler handler, DisplayManagerFlags flags) {
-        this(context, foldSettingProvider, repo, listener, syncRoot, handler,
+        this(context, foldSettingProvider, foldGracePeriodProvider, repo, listener, syncRoot,
+                handler,
                 new DeviceStateToLayoutMap((isDefault) -> isDefault ? DEFAULT_DISPLAY
                         : sNextNonDefaultDisplayId++, flags), flags);
     }
 
     LogicalDisplayMapper(@NonNull Context context, FoldSettingProvider foldSettingProvider,
+            FoldGracePeriodProvider foldGracePeriodProvider,
             @NonNull DisplayDeviceRepository repo,
             @NonNull Listener listener, @NonNull DisplayManagerService.SyncRoot syncRoot,
             @NonNull Handler handler, @NonNull DeviceStateToLayoutMap deviceStateToLayoutMap,
@@ -221,13 +226,15 @@
         mDisplayDeviceRepo = repo;
         mListener = listener;
         mFoldSettingProvider = foldSettingProvider;
+        mFoldGracePeriodProvider = foldGracePeriodProvider;
         mSingleDisplayDemoMode = SystemProperties.getBoolean("persist.demo.singledisplay", false);
         mSupportsConcurrentInternalDisplays = context.getResources().getBoolean(
                 com.android.internal.R.bool.config_supportsConcurrentInternalDisplays);
         mDeviceStatesOnWhichToWakeUp = toSparseBooleanArray(context.getResources().getIntArray(
                 com.android.internal.R.array.config_deviceStatesOnWhichToWakeUp));
-        mDeviceStatesOnWhichToSleep = toSparseBooleanArray(context.getResources().getIntArray(
-                com.android.internal.R.array.config_deviceStatesOnWhichToSleep));
+        mDeviceStatesOnWhichToSelectiveSleep = toSparseBooleanArray(
+                context.getResources().getIntArray(
+                        com.android.internal.R.array.config_deviceStatesOnWhichToSleep));
         mDisplayDeviceRepo.addListener(this);
         mDeviceStateToLayoutMap = deviceStateToLayoutMap;
         mFlags = flags;
@@ -404,7 +411,7 @@
         ipw.println("mSingleDisplayDemoMode=" + mSingleDisplayDemoMode);
         ipw.println("mCurrentLayout=" + mCurrentLayout);
         ipw.println("mDeviceStatesOnWhichToWakeUp=" + mDeviceStatesOnWhichToWakeUp);
-        ipw.println("mDeviceStatesOnWhichToSleep=" + mDeviceStatesOnWhichToSleep);
+        ipw.println("mDeviceStatesOnWhichSelectiveSleep=" + mDeviceStatesOnWhichToSelectiveSleep);
         ipw.println("mInteractive=" + mInteractive);
         ipw.println("mBootCompleted=" + mBootCompleted);
 
@@ -569,8 +576,8 @@
     boolean shouldDeviceBePutToSleep(int pendingState, int currentState, boolean isInteractive,
             boolean isBootCompleted) {
         return currentState != DeviceStateManager.INVALID_DEVICE_STATE_IDENTIFIER
-                && mDeviceStatesOnWhichToSleep.get(pendingState)
-                && !mDeviceStatesOnWhichToSleep.get(currentState)
+                && mDeviceStatesOnWhichToSelectiveSleep.get(pendingState)
+                && !mDeviceStatesOnWhichToSelectiveSleep.get(currentState)
                 && isInteractive
                 && isBootCompleted
                 && !mFoldSettingProvider.shouldStayAwakeOnFold();
@@ -611,9 +618,12 @@
         final boolean waitingToWakeDevice = mDeviceStatesOnWhichToWakeUp.get(mPendingDeviceState)
                 && !mDeviceStatesOnWhichToWakeUp.get(mDeviceState)
                 && !mInteractive && mBootCompleted;
-        final boolean waitingToSleepDevice = mDeviceStatesOnWhichToSleep.get(mPendingDeviceState)
-                && !mDeviceStatesOnWhichToSleep.get(mDeviceState)
-                && mInteractive && mBootCompleted;
+        // The device should only wait for sleep if #shouldStayAwakeOnFold method returns false.
+        // If not, device should be marked ready for transition immediately.
+        final boolean waitingToSleepDevice = mDeviceStatesOnWhichToSelectiveSleep.get(
+                mPendingDeviceState)
+                && !mDeviceStatesOnWhichToSelectiveSleep.get(mDeviceState)
+                && mInteractive && mBootCompleted && !shouldStayAwakeOnFold();
 
         final boolean displaysOff = areAllTransitioningDisplaysOffLocked();
         final boolean isReadyToTransition = displaysOff && !waitingToWakeDevice
@@ -1233,6 +1243,16 @@
         return retval;
     }
 
+    /**
+     * Returns true if the device would definitely have outer display ON/Stay Awake on fold based on
+     * the value of `Continue using app on fold` setting
+     */
+    private boolean shouldStayAwakeOnFold() {
+        return mFoldSettingProvider.shouldStayAwakeOnFold() || (
+                mFoldSettingProvider.shouldSelectiveStayAwakeOnFold()
+                        && mFoldGracePeriodProvider.isEnabled());
+    }
+
     private String displayEventToString(int msg) {
         switch(msg) {
             case LOGICAL_DISPLAY_EVENT_ADDED:
diff --git a/services/core/java/com/android/server/notification/NotificationChannelExtractor.java b/services/core/java/com/android/server/notification/NotificationChannelExtractor.java
index 2f60e42..bd73cb6 100644
--- a/services/core/java/com/android/server/notification/NotificationChannelExtractor.java
+++ b/services/core/java/com/android/server/notification/NotificationChannelExtractor.java
@@ -15,8 +15,16 @@
 */
 package com.android.server.notification;
 
+import static android.app.Flags.restrictAudioAttributesAlarm;
+import static android.app.Flags.restrictAudioAttributesCall;
+import static android.app.Flags.restrictAudioAttributesMedia;
+import static android.app.Notification.CATEGORY_ALARM;
+import static android.media.AudioAttributes.USAGE_NOTIFICATION;
+
+import android.app.Notification;
 import android.app.NotificationChannel;
 import android.content.Context;
+import android.media.AudioAttributes;
 import android.util.Slog;
 
 /**
@@ -50,6 +58,36 @@
                 record.getSbn().getShortcutId(), true, false);
         record.updateNotificationChannel(updatedChannel);
 
+        if (restrictAudioAttributesCall() || restrictAudioAttributesAlarm()
+                || restrictAudioAttributesMedia()) {
+            AudioAttributes attributes = record.getChannel().getAudioAttributes();
+            boolean updateAttributes =  false;
+            if (restrictAudioAttributesCall()
+                    && !record.getNotification().isStyle(Notification.CallStyle.class)
+                    && attributes.getUsage() == AudioAttributes.USAGE_NOTIFICATION_RINGTONE) {
+                updateAttributes = true;
+            }
+            if (restrictAudioAttributesAlarm()
+                    && record.getNotification().category != CATEGORY_ALARM
+                    && attributes.getUsage() == AudioAttributes.USAGE_ALARM) {
+                updateAttributes = true;
+            }
+
+            if (restrictAudioAttributesMedia()
+                    && (attributes.getUsage() == AudioAttributes.USAGE_UNKNOWN
+                    || attributes.getUsage() == AudioAttributes.USAGE_MEDIA)) {
+                updateAttributes = true;
+            }
+
+            if (updateAttributes) {
+                NotificationChannel clone = record.getChannel().copy();
+                clone.setSound(clone.getSound(), new AudioAttributes.Builder(attributes)
+                        .setUsage(USAGE_NOTIFICATION)
+                        .build());
+                record.updateNotificationChannel(clone);
+            }
+        }
+
         return null;
     }
 
diff --git a/services/core/java/com/android/server/notification/NotificationRecord.java b/services/core/java/com/android/server/notification/NotificationRecord.java
index a4464a1..97d2620 100644
--- a/services/core/java/com/android/server/notification/NotificationRecord.java
+++ b/services/core/java/com/android/server/notification/NotificationRecord.java
@@ -15,6 +15,9 @@
  */
 package com.android.server.notification;
 
+import static android.app.Flags.restrictAudioAttributesAlarm;
+import static android.app.Flags.restrictAudioAttributesCall;
+import static android.app.Flags.restrictAudioAttributesMedia;
 import static android.app.Flags.updateRankingTime;
 import static android.app.NotificationChannel.USER_LOCKED_IMPORTANCE;
 import static android.app.NotificationManager.IMPORTANCE_DEFAULT;
@@ -1159,6 +1162,11 @@
             mChannel = channel;
             calculateImportance();
             calculateUserSentiment();
+            mVibration = calculateVibration();
+            if (restrictAudioAttributesCall() || restrictAudioAttributesAlarm()
+                    || restrictAudioAttributesMedia()) {
+                mAttributes = channel.getAudioAttributes();
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
index 02f90f2..b5df30f 100644
--- a/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
+++ b/services/core/java/com/android/server/stats/pull/StatsPullAtomService.java
@@ -37,6 +37,7 @@
 import static android.net.NetworkTemplate.OEM_MANAGED_PAID;
 import static android.net.NetworkTemplate.OEM_MANAGED_PRIVATE;
 import static android.os.Debug.getIonHeapsSizeKb;
+import static android.os.Process.INVALID_UID;
 import static android.os.Process.LAST_SHARED_APPLICATION_GID;
 import static android.os.Process.SYSTEM_UID;
 import static android.os.Process.getUidForPid;
@@ -3537,17 +3538,23 @@
                     String roleName = roleEntry.getKey();
                     Set<String> packageNames = roleEntry.getValue();
 
-                    for (String packageName : packageNames) {
-                        PackageInfo pkg;
-                        try {
-                            pkg = pm.getPackageInfoAsUser(packageName, 0, userId);
-                        } catch (PackageManager.NameNotFoundException e) {
-                            Slog.w(TAG, "Role holder " + packageName + " not found");
-                            return StatsManager.PULL_SKIP;
-                        }
+                    if (!packageNames.isEmpty()) {
+                        for (String packageName : packageNames) {
+                            PackageInfo pkg;
+                            try {
+                                pkg = pm.getPackageInfoAsUser(packageName, 0, userId);
+                            } catch (PackageManager.NameNotFoundException e) {
+                                Slog.w(TAG, "Role holder " + packageName + " not found");
+                                return StatsManager.PULL_SKIP;
+                            }
 
+                            pulledData.add(FrameworkStatsLog.buildStatsEvent(
+                                    atomTag, pkg.applicationInfo.uid, packageName, roleName));
+                        }
+                    } else {
+                        // Ensure that roles set to None are logged with an empty state.
                         pulledData.add(FrameworkStatsLog.buildStatsEvent(
-                                atomTag, pkg.applicationInfo.uid, packageName, roleName));
+                                atomTag, INVALID_UID, "", roleName));
                     }
                 }
             }
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 8c4f9ef..2dbb370 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -67,8 +67,9 @@
     /** @see #setServerVisible(boolean) */
     private boolean mServerVisible;
 
-    ImeInsetsSourceProvider(InsetsSource source,
-            InsetsStateController stateController, DisplayContent displayContent) {
+    ImeInsetsSourceProvider(@NonNull InsetsSource source,
+            @NonNull InsetsStateController stateController,
+            @NonNull DisplayContent displayContent) {
         super(source, stateController, displayContent);
     }
 
@@ -230,7 +231,7 @@
         if (mImeRequesterStatsToken != null) {
             // Cancel the pre-existing stats token, if any.
             // Log state on pre-existing request cancel.
-            logShowImePostLayoutState();
+            logShowImePostLayoutState(false /* aborted */);
             ImeTracker.forLogging().onCancelled(
                     mImeRequesterStatsToken, ImeTracker.PHASE_WM_SHOW_IME_RUNNER);
         }
@@ -310,7 +311,7 @@
         ProtoLog.d(WM_DEBUG_IME, "abortShowImePostLayout");
         if (mImeRequesterStatsToken != null) {
             // Log state on abort.
-            logShowImePostLayoutState();
+            logShowImePostLayoutState(true /* aborted */);
             ImeTracker.forLogging().onFailed(
                     mImeRequesterStatsToken, ImeTracker.PHASE_WM_ABORT_SHOW_IME_POST_LAYOUT);
             mImeRequesterStatsToken = null;
@@ -333,11 +334,30 @@
         //  actual IME target.
         final InsetsControlTarget dcTarget = mDisplayContent.getImeTarget(IME_TARGET_LAYERING);
         if (dcTarget == null || mImeRequester == null) {
+            // Not ready to show if there is no IME layering target, or no IME requester.
             return false;
         }
-        // Not ready to show if there is no IME control target.
-        final InsetsControlTarget controlTarget = mDisplayContent.getImeTarget(IME_TARGET_CONTROL);
+        final InsetsControlTarget controlTarget = getControlTarget();
         if (controlTarget == null) {
+            // Not ready to show if there is no IME control target.
+            return false;
+        }
+        if (controlTarget != mDisplayContent.getImeTarget(IME_TARGET_CONTROL)) {
+            // Not ready to show if control target does not match the one in DisplayContent.
+            return false;
+        }
+        if (!mServerVisible || mFrozen) {
+            // Not ready to show if the window container is not available and considered visible.
+            // If frozen, the server visibility is not set until unfrozen.
+            return false;
+        }
+        if (mStateController.hasPendingControls(controlTarget)) {
+            // Not ready to show if control target has pending controls.
+            return false;
+        }
+        if (getLeash(controlTarget) == null) {
+            // Not ready to show if control target has no source control leash (or leash is not
+            // ready for dispatching).
             return false;
         }
 
@@ -353,35 +373,56 @@
     }
 
     /**
-     * Logs the current state required for scheduleShowImePostLayout's runnable to be triggered.
+     * Logs the current state required for showImePostLayout to be triggered.
+     *
+     * @param aborted whether the showImePostLayout was aborted or cancelled.
      */
-    private void logShowImePostLayoutState() {
+    private void logShowImePostLayoutState(boolean aborted) {
         final var windowState = mWindowContainer != null ? mWindowContainer.asWindowState() : null;
         final var dcTarget = mDisplayContent.getImeTarget(IME_TARGET_LAYERING);
-        final var controlTarget = mDisplayContent.getImeTarget(IME_TARGET_CONTROL);
+        final var controlTarget = getControlTarget();
         final var sb = new StringBuilder();
-        sb.append("mWindowContainer: ").append(mWindowContainer);
-        sb.append(" windowState: ").append(windowState);
+        sb.append("showImePostLayout ").append(aborted ? "aborted" : "cancelled");
+        sb.append(", mWindowContainer is: ");
+        sb.append(mWindowContainer != null ? "non-null" : "null");
+        sb.append(", windowState: ").append(windowState);
         if (windowState != null) {
-            sb.append(" windowState.isDrawn(): ").append(windowState.isDrawn());
-            sb.append(" windowState.mGivenInsetsPending: ").append(windowState.mGivenInsetsPending);
+            sb.append(", windowState.isDrawn(): ");
+            sb.append(windowState.isDrawn());
+            sb.append(", windowState.mGivenInsetsPending: ");
+            sb.append(windowState.mGivenInsetsPending);
         }
-        sb.append(" mIsImeLayoutDrawn: ").append(mIsImeLayoutDrawn);
-        sb.append(" mShowImeRunner: ").append(mShowImeRunner);
-        sb.append(" mImeRequester: ").append(mImeRequester);
-        sb.append(" dcTarget: ").append(dcTarget);
-        sb.append(" controlTarget: ").append(controlTarget);
-        sb.append(" isReadyToShowIme(): ").append(isReadyToShowIme());
+        sb.append(", mIsImeLayoutDrawn: ").append(mIsImeLayoutDrawn);
+        sb.append(", mShowImeRunner: ").append(mShowImeRunner);
+        sb.append(", mImeRequester: ").append(mImeRequester);
+        sb.append(", dcTarget: ").append(dcTarget);
+        sb.append(", controlTarget: ").append(controlTarget);
+        sb.append("\n");
+        sb.append("isReadyToShowIme(): ").append(isReadyToShowIme());
         if (mImeRequester != null && dcTarget != null && controlTarget != null) {
-            sb.append(" isImeLayeringTarget: ");
+            sb.append(", controlTarget == DisplayContent.controlTarget: ");
+            sb.append(controlTarget == mDisplayContent.getImeTarget(IME_TARGET_CONTROL));
+            sb.append(", hasPendingControls: ");
+            sb.append(mStateController.hasPendingControls(controlTarget));
+            sb.append(", serverVisible: ");
+            sb.append(mServerVisible);
+            sb.append(", frozen: ");
+            sb.append(mFrozen);
+            sb.append(", leash is: ");
+            sb.append(getLeash(controlTarget) != null ? "non-null" : "null");
+            sb.append(", control is: ");
+            sb.append(mControl != null ? "non-null" : "null");
+            sb.append(", mIsLeashReadyForDispatching: ");
+            sb.append(mIsLeashReadyForDispatching);
+            sb.append(", isImeLayeringTarget: ");
             sb.append(isImeLayeringTarget(mImeRequester, dcTarget));
-            sb.append(" isAboveImeLayeringTarget: ");
+            sb.append(", isAboveImeLayeringTarget: ");
             sb.append(isAboveImeLayeringTarget(mImeRequester, dcTarget));
-            sb.append(" isImeFallbackTarget: ");
+            sb.append(", isImeFallbackTarget: ");
             sb.append(isImeFallbackTarget(mImeRequester));
-            sb.append(" isImeInputTarget: ");
+            sb.append(", isImeInputTarget: ");
             sb.append(isImeInputTarget(mImeRequester));
-            sb.append(" sameAsImeControlTarget: ");
+            sb.append(", sameAsImeControlTarget: ");
             sb.append(sameAsImeControlTarget());
         }
         Slog.d(TAG, sb.toString());
diff --git a/services/core/java/com/android/server/wm/InsetsSourceProvider.java b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
index 83f44d2..a8cbc62 100644
--- a/services/core/java/com/android/server/wm/InsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/InsetsSourceProvider.java
@@ -64,15 +64,16 @@
 
     private static final Rect EMPTY_RECT = new Rect();
 
-    protected final DisplayContent mDisplayContent;
     protected final @NonNull InsetsSource mSource;
-    protected WindowContainer mWindowContainer;
+    protected final @NonNull DisplayContent mDisplayContent;
+    protected final @NonNull InsetsStateController mStateController;
+    protected @Nullable WindowContainer mWindowContainer;
+    protected @Nullable InsetsSourceControl mControl;
+    protected boolean mIsLeashReadyForDispatching;
 
     private final Rect mTmpRect = new Rect();
-    private final InsetsStateController mStateController;
     private final InsetsSourceControl mFakeControl;
     private final Consumer<Transaction> mSetLeashPositionConsumer;
-    private @Nullable InsetsSourceControl mControl;
     private @Nullable InsetsControlTarget mControlTarget;
     private @Nullable InsetsControlTarget mPendingControlTarget;
     private @Nullable InsetsControlTarget mFakeControlTarget;
@@ -82,7 +83,6 @@
     private SparseArray<TriFunction<DisplayFrames, WindowContainer, Rect, Integer>>
             mOverrideFrameProviders;
     private final SparseArray<Rect> mOverrideFrames = new SparseArray<Rect>();
-    private boolean mIsLeashReadyForDispatching;
     private final Rect mSourceFrame = new Rect();
     private final Rect mLastSourceFrame = new Rect();
     private @NonNull Insets mInsetsHint = Insets.NONE;
@@ -114,8 +114,9 @@
      */
     private boolean mCropToProvidingInsets = false;
 
-    InsetsSourceProvider(InsetsSource source, InsetsStateController stateController,
-            DisplayContent displayContent) {
+    InsetsSourceProvider(@NonNull InsetsSource source,
+            @NonNull InsetsStateController stateController,
+            @NonNull DisplayContent displayContent) {
         mClientVisible = (WindowInsets.Type.defaultVisible() & source.getType()) != 0;
         mSource = source;
         mDisplayContent = displayContent;
@@ -560,7 +561,7 @@
         mDisplayContent.mWmService.mWindowPlacerLocked.requestTraversal();
     }
 
-    @VisibleForTesting
+    @VisibleForTesting(visibility = VisibleForTesting.Visibility.PROTECTED)
     void setServerVisible(boolean serverVisible) {
         mServerVisible = serverVisible;
         updateSourceFrameForServerVisibility();
@@ -575,6 +576,14 @@
                 mServerVisible, mClientVisible);
     }
 
+    /**
+     * Gets the source control for the given control target. If this is the provider's control
+     * target, but the leash is not ready for dispatching, a new source control object with the
+     * leash set to {@code null} is returned.
+     *
+     * @param target the control target to get the source control for.
+     */
+    @Nullable
     InsetsSourceControl getControl(InsetsControlTarget target) {
         if (target == mControlTarget) {
             if (!mIsLeashReadyForDispatching && mControl != null) {
@@ -593,10 +602,25 @@
         return null;
     }
 
+    /**
+     * Gets the leash of the source control for the given control target. If this is not the
+     * provider's control target, or the leash is not ready for dispatching, this will
+     * return {@code null}.
+     *
+     * @param target the control target to get the source control leash for.
+     */
+    @Nullable
+    protected SurfaceControl getLeash(@NonNull InsetsControlTarget target) {
+        return target == mControlTarget && mIsLeashReadyForDispatching && mControl != null
+                ? mControl.getLeash() : null;
+    }
+
+    @Nullable
     InsetsControlTarget getControlTarget() {
         return mControlTarget;
     }
 
+    @Nullable
     InsetsControlTarget getFakeControlTarget() {
         return mFakeControlTarget;
     }
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index 6b9fcf4..3a04bcd 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -388,6 +388,9 @@
                 onRequestedVisibleTypesChanged(newControlTargets.valueAt(i));
             }
             newControlTargets.clear();
+            // Check for and try to run the scheduled show IME request (if it exists), as we
+            // now applied the surface transaction and notified the target of the new control.
+            getImeSourceProvider().checkShowImePostLayout();
         });
     }
 
@@ -395,6 +398,15 @@
         mDisplayContent.notifyInsetsChanged(mDispatchInsetsChanged);
     }
 
+    /**
+     * Checks if the control target has pending controls.
+     *
+     * @param target the control target to check.
+     */
+    boolean hasPendingControls(@NonNull InsetsControlTarget target) {
+        return mPendingControlChanged.contains(target);
+    }
+
     void dump(String prefix, PrintWriter pw) {
         pw.println(prefix + "WindowInsetsStateController");
         prefix = prefix + "  ";
diff --git a/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java b/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
index d0c7077..5a50510 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/LogicalDisplayMapperTest.java
@@ -20,6 +20,8 @@
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.DEFAULT_DISPLAY_GROUP;
 import static android.view.Display.FLAG_REAR;
+import static android.view.Display.STATE_OFF;
+import static android.view.Display.STATE_ON;
 import static android.view.Display.TYPE_EXTERNAL;
 import static android.view.Display.TYPE_INTERNAL;
 import static android.view.Display.TYPE_VIRTUAL;
@@ -28,6 +30,7 @@
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_ADDED;
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_CHANGED;
 import static com.android.server.display.DisplayAdapter.DISPLAY_DEVICE_EVENT_REMOVED;
+import static com.android.server.display.DisplayDeviceInfo.DIFF_EVERYTHING;
 import static com.android.server.display.DisplayDeviceInfo.FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY;
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_ADDED;
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_CONNECTED;
@@ -35,6 +38,9 @@
 import static com.android.server.display.LogicalDisplayMapper.LOGICAL_DISPLAY_EVENT_REMOVED;
 import static com.android.server.display.layout.Layout.Display.POSITION_REAR;
 import static com.android.server.display.layout.Layout.Display.POSITION_UNKNOWN;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_SELECTIVE_STAY_AWAKE;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_SLEEP_ON_FOLD;
+import static com.android.server.utils.FoldSettingProvider.SETTING_VALUE_STAY_AWAKE_ON_FOLD;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -72,6 +78,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
+import com.android.internal.foldables.FoldGracePeriodProvider;
 import com.android.server.display.feature.DisplayManagerFlags;
 import com.android.server.display.layout.DisplayIdProducer;
 import com.android.server.display.layout.Layout;
@@ -96,9 +103,13 @@
 @RunWith(AndroidJUnit4.class)
 public class LogicalDisplayMapperTest {
     private static int sUniqueTestDisplayId = 0;
+    private static final int TIMEOUT_STATE_TRANSITION_MILLIS = 500;
+    private static final int FOLD_SETTLE_DELAY = 1000;
     private static final int DEVICE_STATE_CLOSED = 0;
+    private static final int DEVICE_STATE_HALF_OPEN = 1;
     private static final int DEVICE_STATE_OPEN = 2;
     private static final int FLAG_GO_TO_SLEEP_ON_FOLD = 0;
+    private static final int FLAG_GO_TO_SLEEP_FLAG_SOFT_SLEEP = 2;
     private static int sNextNonDefaultDisplayId = DEFAULT_DISPLAY + 1;
     private static final File NON_EXISTING_FILE = new File("/non_existing_folder/should_not_exist");
 
@@ -116,6 +127,7 @@
     @Mock LogicalDisplayMapper.Listener mListenerMock;
     @Mock Context mContextMock;
     @Mock FoldSettingProvider mFoldSettingProviderMock;
+    @Mock FoldGracePeriodProvider mFoldGracePeriodProvider;
     @Mock Resources mResourcesMock;
     @Mock IPowerManager mIPowerManagerMock;
     @Mock IThermalService mIThermalServiceMock;
@@ -160,6 +172,7 @@
                 .thenReturn(Context.POWER_SERVICE);
         when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(false);
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
         when(mIPowerManagerMock.isInteractive()).thenReturn(true);
         when(mContextMock.getSystemService(PowerManager.class)).thenReturn(mPowerManager);
         when(mContextMock.getResources()).thenReturn(mResourcesMock);
@@ -177,6 +190,7 @@
         mLooper = new TestLooper();
         mHandler = new Handler(mLooper.getLooper());
         mLogicalDisplayMapper = new LogicalDisplayMapper(mContextMock, mFoldSettingProviderMock,
+                mFoldGracePeriodProvider,
                 mDisplayDeviceRepo,
                 mListenerMock, new DisplayManagerService.SyncRoot(), mHandler,
                 mDeviceStateToLayoutMapSpy, mFlagsMock);
@@ -681,22 +695,147 @@
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(true);
 
         finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
 
         verify(mIPowerManagerMock, atLeastOnce()).goToSleep(anyLong(), anyInt(),
                 eq(FLAG_GO_TO_SLEEP_ON_FOLD));
     }
 
     @Test
+    public void testDeviceShouldPutToSleepWhenFoldSettingSelective() throws RemoteException {
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
+
+        finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
+
+        verify(mIPowerManagerMock, atLeastOnce()).goToSleep(anyLong(), anyInt(),
+                eq(FLAG_GO_TO_SLEEP_FLAG_SOFT_SLEEP));
+    }
+
+    @Test
     public void testDeviceShouldNotBePutToSleepWhenSleepSettingFalse() throws RemoteException {
         when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
 
         finishBootAndFoldDevice();
+        advanceTime(FOLD_SETTLE_DELAY);
 
         verify(mIPowerManagerMock, never()).goToSleep(anyLong(), anyInt(),
                 eq(FLAG_GO_TO_SLEEP_ON_FOLD));
     }
 
     @Test
+    public void testWaitForSleepWhenFoldSettingSleep() {
+        // Test device should not be marked ready for transition immediately, when 'Continue
+        // using app on fold' set to 'Never'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testSwapDeviceStateWithDelayWhenFoldSettingSleep() {
+        // Test device should be marked ready for transition after a delay when 'Continue using
+        // app on fold' set to 'Never'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+        advanceTime(TIMEOUT_STATE_TRANSITION_MILLIS);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testDoNotWaitForSleepWhenFoldSettingStayAwake() {
+        // Test device should be marked ready for transition immediately when 'Continue using app
+        // on fold' set to 'Always'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_STAY_AWAKE_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testDoNotWaitForSleepWhenFoldSettingSelectiveStayAwake() {
+        // Test device should be marked ready for transition immediately when 'Continue using app
+        // on fold' set to 'Swipe up to continue'
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SELECTIVE_STAY_AWAKE);
+        setGracePeriodAvailability(true);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testWaitForSleepWhenGracePeriodSettingDisabled() {
+        // Test device should not be marked ready for transition immediately when 'Continue using
+        // app on fold' set to 'Swipe up to continue' but Grace Period flag is disabled
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SELECTIVE_STAY_AWAKE);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        foldableDisplayDevices.mInner.setState(STATE_OFF);
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testWaitForSleepWhenTransitionDisplayStaysOn() {
+        // Test device should not be marked ready for transition immediately, when 'Continue
+        // using app on fold' set to 'Always' but not all transitioning displays are OFF.
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_STAY_AWAKE_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+
+        assertDisplayDisabled(foldableDisplayDevices.mOuter);
+        assertDisplayEnabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
+    public void testSwapDeviceStateWithDelayWhenTransitionDisplayStaysOn() {
+        // Test device should be marked ready for transition after a delay, when 'Continue using
+        // app on fold' set to 'Never' but not all transitioning displays are OFF.
+        setFoldLockBehaviorSettingValue(SETTING_VALUE_SLEEP_ON_FOLD);
+        setGracePeriodAvailability(false);
+        FoldableDisplayDevices foldableDisplayDevices = createFoldableDeviceStateToLayoutMap();
+
+        finishBootAndFoldDevice();
+        notifyDisplayChanges(foldableDisplayDevices.mOuter);
+        advanceTime(TIMEOUT_STATE_TRANSITION_MILLIS);
+
+        assertDisplayEnabled(foldableDisplayDevices.mOuter);
+        assertDisplayDisabled(foldableDisplayDevices.mInner);
+    }
+
+    @Test
     public void testDeviceStateLocked() {
         DisplayDevice device1 = createDisplayDevice(TYPE_INTERNAL, 600, 800,
                 FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
@@ -950,13 +1089,77 @@
     // Helper Methods
     /////////////////
 
+    private void setGracePeriodAvailability(boolean isGracePeriodEnabled) {
+        when(mFoldGracePeriodProvider.isEnabled()).thenReturn(isGracePeriodEnabled);
+    }
+
+    private void setFoldLockBehaviorSettingValue(String foldLockBehaviorSettingValue) {
+        when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(false);
+        when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(false);
+
+        switch (foldLockBehaviorSettingValue) {
+            case SETTING_VALUE_STAY_AWAKE_ON_FOLD:
+                when(mFoldSettingProviderMock.shouldStayAwakeOnFold()).thenReturn(true);
+                break;
+
+            case SETTING_VALUE_SLEEP_ON_FOLD:
+                when(mFoldSettingProviderMock.shouldSleepOnFold()).thenReturn(true);
+                break;
+
+            default:
+                when(mFoldSettingProviderMock.shouldSelectiveStayAwakeOnFold()).thenReturn(true);
+                break;
+        }
+    }
+
+    private FoldableDisplayDevices createFoldableDeviceStateToLayoutMap() {
+        TestDisplayDevice outer = createDisplayDevice(TYPE_INTERNAL, 600, 800,
+                FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
+        TestDisplayDevice inner = createDisplayDevice(TYPE_INTERNAL, 600, 800,
+                FLAG_ALLOWED_TO_BE_DEFAULT_DISPLAY);
+        outer.setState(STATE_OFF);
+        inner.setState(STATE_ON);
+
+        Layout layout = new Layout();
+        createDefaultDisplay(layout, outer);
+        createNonDefaultDisplay(layout, inner, /* enabled= */ false, /* group= */ null);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_CLOSED)).thenReturn(layout);
+
+        layout = new Layout();
+        createNonDefaultDisplay(layout, outer, /* enabled= */ false, /* group= */ null);
+        createDefaultDisplay(layout, inner);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_HALF_OPEN)).thenReturn(layout);
+        when(mDeviceStateToLayoutMapSpy.get(DEVICE_STATE_OPEN)).thenReturn(layout);
+        when(mDeviceStateToLayoutMapSpy.size()).thenReturn(4);
+
+        add(outer);
+        add(inner);
+
+        return new FoldableDisplayDevices(outer, inner);
+    }
+
     private void finishBootAndFoldDevice() {
         mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_OPEN);
+        mLogicalDisplayMapper.onEarlyInteractivityChange(true);
         advanceTime(1000);
         mLogicalDisplayMapper.onBootCompleted();
         advanceTime(1000);
         mLogicalDisplayMapper.setDeviceStateLocked(DEVICE_STATE_CLOSED);
-        advanceTime(1000);
+    }
+
+    private void notifyDisplayChanges(TestDisplayDevice displayDevice) {
+        mLogicalDisplayMapper.onDisplayDeviceChangedLocked(displayDevice, DIFF_EVERYTHING);
+    }
+
+    private void assertDisplayEnabled(DisplayDevice displayDevice) {
+        assertThat(
+                mLogicalDisplayMapper.getDisplayLocked(displayDevice).isEnabledLocked()).isTrue();
+    }
+
+    private void assertDisplayDisabled(DisplayDevice displayDevice) {
+        assertThat(
+                mLogicalDisplayMapper.getDisplayLocked(displayDevice).isEnabledLocked()).isFalse();
     }
 
     private void createDefaultDisplay(Layout layout, DisplayDevice device) {
@@ -1058,6 +1261,16 @@
         assertNotEquals(DEFAULT_DISPLAY, id(displayRemoved));
     }
 
+    private final static class FoldableDisplayDevices {
+        final TestDisplayDevice mOuter;
+        final TestDisplayDevice mInner;
+
+        FoldableDisplayDevices(TestDisplayDevice outer, TestDisplayDevice inner) {
+            this.mOuter = outer;
+            this.mInner = inner;
+        }
+    }
+
     class TestDisplayDevice extends DisplayDevice {
         private DisplayDeviceInfo mInfo;
         private DisplayDeviceInfo mSentInfo;
@@ -1083,6 +1296,16 @@
             mSentInfo = null;
         }
 
+        public void setState(int state) {
+            mState = state;
+            if (mSentInfo == null) {
+                mInfo.state = state;
+            } else {
+                mInfo.state = state;
+                mSentInfo.state = state;
+            }
+        }
+
         @Override
         public boolean hasStableUniqueId() {
             return true;
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
index ec3e97b..0678140 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/GenericWindowPolicyControllerTest.java
@@ -123,6 +123,7 @@
     @Test
     public void containsUid() {
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         assertThat(gwpc.containsUid(TEST_UID)).isFalse();
 
@@ -136,6 +137,7 @@
     @Test
     public void isEnteringPipAllowed_falseByDefault() {
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         assertThat(gwpc.isEnteringPipAllowed(TEST_UID)).isFalse();
         verify(mPipBlockedCallback, timeout(TIMEOUT_MILLIS)).onEnteringPipBlocked(TEST_UID);
@@ -144,6 +146,7 @@
     @Test
     public void isEnteringPipAllowed_dpcSupportsPinned_allowed() {
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
         gwpc.setSupportedWindowingModes(new HashSet<>(
                 Arrays.asList(WindowConfiguration.WINDOWING_MODE_FULLSCREEN,
                         WindowConfiguration.WINDOWING_MODE_PINNED)));
@@ -160,11 +163,25 @@
                 NONBLOCKED_APP_PACKAGE_NAME,
                 NONBLOCKED_APP_PACKAGE_NAME,
                 /* displayOnRemoteDevices */ true,
-                /* targetDisplayCategory */ null);
+                /* targetDisplayCategory */ null,
+                /* uid */ UserHandle.PER_USER_RANGE + 1);
         assertActivityIsBlocked(gwpc, activityInfo);
     }
 
     @Test
+    public void userNotAllowlisted_systemUser_isNotBlocked() {
+        GenericWindowPolicyController gwpc = createGwpcWithNoAllowedUsers();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
+
+        ActivityInfo activityInfo = getActivityInfo(
+                NONBLOCKED_APP_PACKAGE_NAME,
+                NONBLOCKED_APP_PACKAGE_NAME,
+                /* displayOnRemoteDevices */ true,
+                /* targetDisplayCategory */ null);
+        assertActivityCanBeLaunched(gwpc, activityInfo);
+    }
+
+    @Test
     public void userNotAllowlisted_systemUserCanLaunchBlockedAppStreamingActivity() {
         GenericWindowPolicyController gwpc = createGwpcWithNoAllowedUsers();
         gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
@@ -521,6 +538,7 @@
     public void registerRunningAppsChangedListener_onRunningAppsChanged_listenersNotified() {
         ArraySet<Integer> uids = new ArraySet<>(Arrays.asList(TEST_UID));
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         gwpc.registerRunningAppsChangedListener(mRunningAppsChangedListener);
         gwpc.onRunningAppsChanged(uids);
@@ -545,6 +563,7 @@
     public void noRunningAppsChangedListener_onRunningAppsChanged_doesNotThrowException() {
         ArraySet<Integer> uids = new ArraySet<>(Arrays.asList(TEST_UID));
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         gwpc.onRunningAppsChanged(uids);
 
@@ -557,6 +576,7 @@
     public void registerUnregisterRunningAppsChangedListener_onRunningAppsChanged_doesNotThrowException() {
         ArraySet<Integer> uids = new ArraySet<>(Arrays.asList(TEST_UID));
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         gwpc.registerRunningAppsChangedListener(mRunningAppsChangedListener);
         gwpc.unregisterRunningAppsChangedListener(mRunningAppsChangedListener);
@@ -579,6 +599,7 @@
         doReturn(interceptor).when(interceptor).queryLocalInterface(anyString());
 
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
         ActivityInfo activityInfo = getActivityInfo(
                 NONBLOCKED_APP_PACKAGE_NAME,
                 NONBLOCKED_APP_PACKAGE_NAME,
@@ -603,6 +624,7 @@
         Intent intent = new Intent(Intent.ACTION_VIEW, Uri.parse("testing"));
 
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
         ActivityInfo activityInfo = getActivityInfo(
                 NONBLOCKED_APP_PACKAGE_NAME,
                 NONBLOCKED_APP_PACKAGE_NAME,
@@ -621,6 +643,7 @@
     @Test
     public void onTopActivitychanged_null_noCallback() {
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         gwpc.onTopActivityChanged(null, 0, 0);
         verify(mActivityListener, after(TIMEOUT_MILLIS).never())
@@ -697,6 +720,7 @@
     @Test
     public void getCustomHomeComponent_noneSet() {
         GenericWindowPolicyController gwpc = createGwpc();
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         assertThat(gwpc.getCustomHomeComponent()).isNull();
     }
@@ -705,6 +729,7 @@
     public void getCustomHomeComponent_returnsHomeComponent() {
         GenericWindowPolicyController gwpc = createGwpcWithCustomHomeComponent(
                 NONBLOCKED_COMPONENT);
+        gwpc.setDisplayId(DISPLAY_ID, /* isMirrorDisplay= */ false);
 
         assertThat(gwpc.getCustomHomeComponent()).isEqualTo(NONBLOCKED_COMPONENT);
     }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationChannelExtractorTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationChannelExtractorTest.java
index 77ce2f0..ad25d76 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationChannelExtractorTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationChannelExtractorTest.java
@@ -16,26 +16,43 @@
 
 package com.android.server.notification;
 
+import static android.app.Notification.CATEGORY_ALARM;
 import static android.app.NotificationManager.IMPORTANCE_HIGH;
 import static android.app.NotificationManager.IMPORTANCE_LOW;
 
+import static android.media.AudioAttributes.USAGE_ALARM;
+import static android.media.AudioAttributes.USAGE_MEDIA;
+import static android.media.AudioAttributes.USAGE_NOTIFICATION;
+import static android.media.AudioAttributes.USAGE_NOTIFICATION_RINGTONE;
+import static android.media.AudioAttributes.USAGE_UNKNOWN;
+import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
+import static com.google.common.truth.Truth.assertThat;
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertNull;
 
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import android.app.Flags;
 import android.app.Notification;
 import android.app.NotificationChannel;
+import android.app.PendingIntent;
+import android.app.Person;
+import android.media.AudioAttributes;
+import android.net.Uri;
 import android.os.UserHandle;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
 import android.service.notification.StatusBarNotification;
 
 import com.android.server.UiServiceTestCase;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
@@ -44,25 +61,34 @@
 
     @Mock RankingConfig mConfig;
 
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
+
+    NotificationChannelExtractor mExtractor;
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
+
+        mExtractor = new NotificationChannelExtractor();
+        mExtractor.setConfig(mConfig);
+        mExtractor.initialize(mContext, null);
+    }
+
+    private NotificationRecord getRecord(NotificationChannel channel, Notification n) {
+        StatusBarNotification sbn = new StatusBarNotification("", "", 0, "", 0,
+                0, n, UserHandle.ALL, null, System.currentTimeMillis());
+        return new NotificationRecord(getContext(), sbn, channel);
     }
 
     @Test
-    public void testExtractsUpdatedChannel() {
-        NotificationChannelExtractor extractor = new NotificationChannelExtractor();
-        extractor.setConfig(mConfig);
-        extractor.initialize(mContext, null);
-
+    public void testExtractsUpdatedConversationChannel() {
         NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_LOW);
-        final Notification.Builder builder = new Notification.Builder(getContext())
+        final Notification n = new Notification.Builder(getContext())
                 .setContentTitle("foo")
-                .setSmallIcon(android.R.drawable.sym_def_app_icon);
-        Notification n = builder.build();
-        StatusBarNotification sbn = new StatusBarNotification("", "", 0, "", 0,
-                0, n, UserHandle.ALL, null, System.currentTimeMillis());
-        NotificationRecord r = new NotificationRecord(getContext(), sbn, channel);
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
 
         NotificationChannel updatedChannel =
                 new NotificationChannel("a", "", IMPORTANCE_HIGH);
@@ -70,26 +96,19 @@
                 any(), anyInt(), eq("a"), eq(null), eq(true), eq(false)))
                 .thenReturn(updatedChannel);
 
-        assertNull(extractor.process(r));
+        assertNull(mExtractor.process(r));
         assertEquals(updatedChannel, r.getChannel());
     }
 
     @Test
-    public void testInvalidShortcutFlagEnabled_looksUpCorrectChannel() {
-
-        NotificationChannelExtractor extractor = new NotificationChannelExtractor();
-        extractor.setConfig(mConfig);
-        extractor.initialize(mContext, null);
-
+    public void testInvalidShortcutFlagEnabled_looksUpCorrectNonChannel() {
         NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_LOW);
-        final Notification.Builder builder = new Notification.Builder(getContext())
+        final Notification n = new Notification.Builder(getContext())
                 .setContentTitle("foo")
                 .setStyle(new Notification.MessagingStyle("name"))
-                .setSmallIcon(android.R.drawable.sym_def_app_icon);
-        Notification n = builder.build();
-        StatusBarNotification sbn = new StatusBarNotification("", "", 0, "tag", 0,
-                0, n, UserHandle.ALL, null, System.currentTimeMillis());
-        NotificationRecord r = new NotificationRecord(getContext(), sbn, channel);
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
 
         NotificationChannel updatedChannel =
                 new NotificationChannel("a", "", IMPORTANCE_HIGH);
@@ -98,26 +117,19 @@
                 eq(true), eq(false)))
                 .thenReturn(updatedChannel);
 
-        assertNull(extractor.process(r));
+        assertNull(mExtractor.process(r));
         assertEquals(updatedChannel, r.getChannel());
     }
 
     @Test
     public void testInvalidShortcutFlagDisabled_looksUpCorrectChannel() {
-
-        NotificationChannelExtractor extractor = new NotificationChannelExtractor();
-        extractor.setConfig(mConfig);
-        extractor.initialize(mContext, null);
-
         NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_LOW);
-        final Notification.Builder builder = new Notification.Builder(getContext())
+        final Notification n = new Notification.Builder(getContext())
                 .setContentTitle("foo")
                 .setStyle(new Notification.MessagingStyle("name"))
-                .setSmallIcon(android.R.drawable.sym_def_app_icon);
-        Notification n = builder.build();
-        StatusBarNotification sbn = new StatusBarNotification("", "", 0, "tag", 0,
-                0, n, UserHandle.ALL, null, System.currentTimeMillis());
-        NotificationRecord r = new NotificationRecord(getContext(), sbn, channel);
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
 
         NotificationChannel updatedChannel =
                 new NotificationChannel("a", "", IMPORTANCE_HIGH);
@@ -125,7 +137,129 @@
                 any(), anyInt(), eq("a"), eq(null), eq(true), eq(false)))
                 .thenReturn(updatedChannel);
 
-        assertNull(extractor.process(r));
+        assertNull(mExtractor.process(r));
         assertEquals(updatedChannel, r.getChannel());
     }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_CALL)
+    public void testAudioAttributes_callStyleCanUseCallUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_NOTIFICATION_RINGTONE)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setStyle(Notification.CallStyle.forIncomingCall(
+                        new Person.Builder().setName("A Caller").build(),
+                        mock(PendingIntent.class),
+                        mock(PendingIntent.class)
+                ))
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION_RINGTONE);
+        assertThat(r.getChannel()).isEqualTo(channel);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_CALL)
+    public void testAudioAttributes_nonCallStyleCannotUseCallUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_NOTIFICATION_RINGTONE)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        // instance updated
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION);
+        // in-memory channel unchanged
+        assertThat(channel.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION_RINGTONE);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_ALARM)
+    public void testAudioAttributes_alarmCategoryCanUseAlarmUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_ALARM)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setCategory(CATEGORY_ALARM)
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_ALARM);
+        assertThat(r.getChannel()).isEqualTo(channel);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_ALARM)
+    public void testAudioAttributes_nonAlarmCategoryCannotUseAlarmUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_ALARM)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        // instance updated
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION);
+        // in-memory channel unchanged
+        assertThat(channel.getAudioAttributes().getUsage()).isEqualTo(USAGE_ALARM);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_MEDIA)
+    public void testAudioAttributes_noMediaUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_MEDIA)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        // instance updated
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION);
+        // in-memory channel unchanged
+        assertThat(channel.getAudioAttributes().getUsage()).isEqualTo(USAGE_MEDIA);
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_MEDIA)
+    public void testAudioAttributes_noUnknownUsage() {
+        NotificationChannel channel = new NotificationChannel("a", "a", IMPORTANCE_HIGH);
+        channel.setSound(Uri.EMPTY, new AudioAttributes.Builder()
+                .setUsage(USAGE_UNKNOWN)
+                .build());
+        final Notification n = new Notification.Builder(getContext())
+                .setContentTitle("foo")
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        NotificationRecord r = getRecord(channel, n);
+
+        assertThat(mExtractor.process(r)).isNull();
+        // instance updated
+        assertThat(r.getAudioAttributes().getUsage()).isEqualTo(USAGE_NOTIFICATION);
+        // in-memory channel unchanged
+        assertThat(channel.getAudioAttributes().getUsage()).isEqualTo(USAGE_UNKNOWN);
+    }
 }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index 20d1e98..f6fa487 100755
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -72,6 +72,8 @@
 import static android.content.pm.PackageManager.FEATURE_WATCH;
 import static android.content.pm.PackageManager.PERMISSION_DENIED;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
+import static android.media.AudioAttributes.USAGE_MEDIA;
+import static android.media.AudioAttributes.USAGE_NOTIFICATION;
 import static android.os.Build.VERSION_CODES.O_MR1;
 import static android.os.Build.VERSION_CODES.P;
 import static android.os.Flags.FLAG_ALLOW_PRIVATE_PROFILE;
@@ -191,6 +193,7 @@
 import android.graphics.Bitmap;
 import android.graphics.Color;
 import android.graphics.drawable.Icon;
+import android.media.AudioAttributes;
 import android.media.AudioManager;
 import android.media.session.MediaSession;
 import android.net.Uri;
@@ -14868,6 +14871,33 @@
         assertThat(posted.getRankingTimeMs()).isEqualTo(posted.getSbn().getPostTime());
     }
 
+    @Test
+    @EnableFlags(android.app.Flags.FLAG_RESTRICT_AUDIO_ATTRIBUTES_MEDIA)
+    public void testRestrictAudioAttributes_listenersGetCorrectAttributes() throws Exception {
+        NotificationChannel sound = new NotificationChannel("a", "a", IMPORTANCE_DEFAULT);
+        sound.setSound(Uri.EMPTY, new AudioAttributes.Builder().setUsage(USAGE_MEDIA).build());
+        mBinderService.createNotificationChannels(mPkg, new ParceledListSlice(
+                Arrays.asList(sound)));
+
+        Notification n = new Notification.Builder(mContext, "a")
+                .setSmallIcon(android.R.drawable.sym_def_app_icon)
+                .build();
+        StatusBarNotification sbn = new StatusBarNotification(mPkg, mPkg, 9, null, mUid, 0,
+                n, UserHandle.getUserHandleForUid(mUid), null, 0);
+
+        mBinderService.enqueueNotificationWithTag(mPkg, mPkg, sbn.getTag(),
+                sbn.getId(), sbn.getNotification(), sbn.getUserId());
+        waitForIdle();
+
+        ArgumentCaptor<NotificationRecord> captor =
+                ArgumentCaptor.forClass(NotificationRecord.class);
+        verify(mListeners, times(1)).prepareNotifyPostedLocked(
+                captor.capture(), any(), anyBoolean());
+
+        assertThat(captor.getValue().getChannel().getAudioAttributes().getUsage())
+                .isEqualTo(USAGE_NOTIFICATION);
+    }
+
     private NotificationRecord createAndPostCallStyleNotification(String packageName,
             UserHandle userHandle, String testName) throws Exception {
         Person person = new Person.Builder().setName("caller").build();
diff --git a/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java b/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
index faa6d97..7380aec 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ImeInsetsSourceProviderTest.java
@@ -16,8 +16,6 @@
 
 package com.android.server.wm;
 
-import static android.view.InsetsSource.ID_IME;
-import static android.view.WindowInsets.Type.ime;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
 import static android.view.WindowManager.LayoutParams.TYPE_INPUT_METHOD;
 
@@ -26,7 +24,6 @@
 
 import android.graphics.PixelFormat;
 import android.platform.test.annotations.Presubmit;
-import android.view.InsetsSource;
 import android.view.inputmethod.ImeTracker;
 
 import androidx.test.filters.SmallTest;
@@ -46,57 +43,148 @@
 @RunWith(WindowTestRunner.class)
 public class ImeInsetsSourceProviderTest extends WindowTestsBase {
 
-    private InsetsSource mImeSource = new InsetsSource(ID_IME, ime());
     private ImeInsetsSourceProvider mImeProvider;
 
     @Before
     public void setUp() throws Exception {
-        mImeSource.setVisible(true);
-        mImeProvider = new ImeInsetsSourceProvider(mImeSource,
-                mDisplayContent.getInsetsStateController(), mDisplayContent);
+        mImeProvider = mDisplayContent.getInsetsStateController().getImeSourceProvider();
+        mImeProvider.getSource().setVisible(true);
     }
 
     @Test
     public void testTransparentControlTargetWindowCanShowIme() {
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+        makeWindowVisibleAndDrawn(ime);
+        mImeProvider.setWindowContainer(ime, null, null);
+
         final WindowState appWin = createWindow(null, TYPE_APPLICATION, "app");
         final WindowState popup = createWindow(appWin, TYPE_APPLICATION, "popup");
-        mDisplayContent.setImeControlTarget(popup);
-        mDisplayContent.setImeLayeringTarget(appWin);
         popup.mAttrs.format = PixelFormat.TRANSPARENT;
+        mDisplayContent.setImeLayeringTarget(appWin);
+        mDisplayContent.updateImeInputAndControlTarget(popup);
+        performSurfacePlacementAndWaitForWindowAnimator();
+
         mImeProvider.scheduleShowImePostLayout(appWin, ImeTracker.Token.empty());
         assertTrue(mImeProvider.isReadyToShowIme());
     }
 
+    /**
+     * Checks that scheduling with all the state set and manually triggering the show does succeed.
+     */
     @Test
-    public void testInputMethodInputTargetCanShowIme() {
-        WindowState target = createWindow(null, TYPE_APPLICATION, "app");
-        mDisplayContent.setImeLayeringTarget(target);
-        mDisplayContent.updateImeInputAndControlTarget(target);
-        mImeProvider.scheduleShowImePostLayout(target, ImeTracker.Token.empty());
-        assertTrue(mImeProvider.isReadyToShowIme());
-    }
-
-    @Test
-    public void testIsImeShowing() {
-        WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+    public void testScheduleShowIme() {
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
         makeWindowVisibleAndDrawn(ime);
         mImeProvider.setWindowContainer(ime, null, null);
 
-        WindowState target = createWindow(null, TYPE_APPLICATION, "app");
+        final WindowState target = createWindow(null, TYPE_APPLICATION, "app");
         mDisplayContent.setImeLayeringTarget(target);
-        mDisplayContent.setImeControlTarget(target);
+        mDisplayContent.updateImeInputAndControlTarget(target);
+        performSurfacePlacementAndWaitForWindowAnimator();
 
+        // Schedule (without triggering) after everything is ready.
         mImeProvider.scheduleShowImePostLayout(target, ImeTracker.Token.empty());
+        assertTrue(mImeProvider.isReadyToShowIme());
         assertFalse(mImeProvider.isImeShowing());
+
+        // Manually trigger the show.
         mImeProvider.checkShowImePostLayout();
+        // No longer ready as it was already shown.
+        assertFalse(mImeProvider.isReadyToShowIme());
         assertTrue(mImeProvider.isImeShowing());
-        mImeProvider.setImeShowing(false);
+    }
+
+    /**
+     * Checks that scheduling to show before any state is set does succeed when
+     * all the state becomes available.
+     */
+    @Test
+    public void testScheduleShowIme_noInitialState() {
+        final WindowState target = createWindow(null, TYPE_APPLICATION, "app");
+
+        // Schedule before anything is ready.
+        mImeProvider.scheduleShowImePostLayout(target, ImeTracker.Token.empty());
+        assertFalse(mImeProvider.isReadyToShowIme());
         assertFalse(mImeProvider.isImeShowing());
+
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+        makeWindowVisibleAndDrawn(ime);
+        mImeProvider.setWindowContainer(ime, null, null);
+
+        mDisplayContent.setImeLayeringTarget(target);
+        mDisplayContent.updateImeInputAndControlTarget(target);
+        // Performing surface placement picks up the show scheduled above.
+        performSurfacePlacementAndWaitForWindowAnimator();
+        // No longer ready as it was already shown.
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertTrue(mImeProvider.isImeShowing());
+    }
+
+    /**
+     * Checks that scheduling to show before starting the {@code afterPrepareSurfacesRunnable}
+     * from {@link InsetsStateController#notifyPendingInsetsControlChanged}
+     * does continue and succeed when the runnable is started.
+     */
+    @Test
+    public void testScheduleShowIme_delayedAfterPrepareSurfaces() {
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+        makeWindowVisibleAndDrawn(ime);
+        mImeProvider.setWindowContainer(ime, null, null);
+
+        final WindowState target = createWindow(null, TYPE_APPLICATION, "app");
+        mDisplayContent.setImeLayeringTarget(target);
+        mDisplayContent.updateImeInputAndControlTarget(target);
+
+        // Schedule before starting the afterPrepareSurfacesRunnable.
+        mImeProvider.scheduleShowImePostLayout(target, ImeTracker.Token.empty());
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertFalse(mImeProvider.isImeShowing());
+
+        // This tries to pick up the show scheduled above, but must fail as the
+        // afterPrepareSurfacesRunnable was not started yet.
+        mDisplayContent.applySurfaceChangesTransaction();
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertFalse(mImeProvider.isImeShowing());
+
+        // Starting the afterPrepareSurfacesRunnable picks up the show scheduled above.
+        mWm.mAnimator.executeAfterPrepareSurfacesRunnables();
+        // No longer ready as it was already shown.
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertTrue(mImeProvider.isImeShowing());
+    }
+
+    /**
+     * Checks that scheduling to show before the surface placement does continue and succeed
+     * when the surface placement happens.
+     */
+    @Test
+    public void testScheduleShowIme_delayedSurfacePlacement() {
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+        makeWindowVisibleAndDrawn(ime);
+        mImeProvider.setWindowContainer(ime, null, null);
+
+        final WindowState target = createWindow(null, TYPE_APPLICATION, "app");
+        mDisplayContent.setImeLayeringTarget(target);
+        mDisplayContent.updateImeInputAndControlTarget(target);
+
+        // Schedule before surface placement.
+        mImeProvider.scheduleShowImePostLayout(target, ImeTracker.Token.empty());
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertFalse(mImeProvider.isImeShowing());
+
+        // Performing surface placement picks up the show scheduled above, and succeeds.
+        // This first executes the afterPrepareSurfacesRunnable, and then
+        // applySurfaceChangesTransaction. Both of them try to trigger the show,
+        // but only the second one can succeed, as it comes after onPostLayout.
+        performSurfacePlacementAndWaitForWindowAnimator();
+        // No longer ready as it was already shown.
+        assertFalse(mImeProvider.isReadyToShowIme());
+        assertTrue(mImeProvider.isImeShowing());
     }
 
     @Test
     public void testSetFrozen() {
-        WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
+        final WindowState ime = createWindow(null, TYPE_INPUT_METHOD, "ime");
         makeWindowVisibleAndDrawn(ime);
         mImeProvider.setWindowContainer(ime, null, null);
         mImeProvider.setServerVisible(true);
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
index 4034dbc..2a025cd 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsSourceProviderTest.java
@@ -24,6 +24,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -163,6 +164,48 @@
     }
 
     @Test
+    public void testGetLeash() {
+        final WindowState statusBar = createWindow(null, TYPE_APPLICATION, "statusBar");
+        final WindowState target = createWindow(null, TYPE_APPLICATION, "target");
+        final WindowState fakeTarget = createWindow(null, TYPE_APPLICATION, "fakeTarget");
+        final WindowState otherTarget = createWindow(null, TYPE_APPLICATION, "otherTarget");
+        statusBar.getFrame().set(0, 0, 500, 100);
+
+        // We must not have control or control target before we have the insets source window,
+        // so also no leash.
+        mProvider.updateControlForTarget(target, true /* force */);
+        assertNull(mProvider.getControl(target));
+        assertNull(mProvider.getControlTarget());
+        assertNull(mProvider.getLeash(target));
+
+        // We can have the control or the control target after we have the insets source window,
+        // but no leash as this is not yet ready for dispatching.
+        mProvider.setWindowContainer(statusBar, null, null);
+        mProvider.updateControlForTarget(target, false /* force */);
+        assertNotNull(mProvider.getControl(target));
+        assertNotNull(mProvider.getControlTarget());
+        assertEquals(mProvider.getControlTarget(), target);
+        assertNull(mProvider.getLeash(target));
+
+        // After surface transactions are applied, the leash is ready for dispatching.
+        mProvider.onSurfaceTransactionApplied();
+        assertNotNull(mProvider.getLeash(target));
+
+        // We do have fake control for the fake control target, but that has no leash.
+        mProvider.updateFakeControlTarget(fakeTarget);
+        assertNotNull(mProvider.getControl(fakeTarget));
+        assertNotNull(mProvider.getFakeControlTarget());
+        assertNotEquals(mProvider.getControlTarget(), fakeTarget);
+        assertNull(mProvider.getLeash(fakeTarget));
+
+        // We don't have any control for a different (non-fake control target), so also no leash.
+        assertNull(mProvider.getControl(otherTarget));
+        assertNotNull(mProvider.getControlTarget());
+        assertNotEquals(mProvider.getControlTarget(), otherTarget);
+        assertNull(mProvider.getLeash(otherTarget));
+    }
+
+    @Test
     public void testUpdateSourceFrame() {
         final WindowState statusBar = createWindow(null, TYPE_APPLICATION, "statusBar");
         mProvider.setWindowContainer(statusBar, null, null);
diff --git a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
index 2085d61..0e1a1af 100644
--- a/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/InsetsStateControllerTest.java
@@ -552,6 +552,24 @@
                 control2.getInsetsHint().bottom);
     }
 
+    @Test
+    public void testHasPendingControls() {
+        final WindowState statusBar = createWindow(null, TYPE_APPLICATION, "statusBar");
+        final WindowState app = createWindow(null, TYPE_APPLICATION, "app");
+        getController().getOrCreateSourceProvider(ID_STATUS_BAR, statusBars())
+                .setWindowContainer(statusBar, null, null);
+        // No controls dispatched yet.
+        assertFalse(getController().hasPendingControls(app));
+
+        getController().onBarControlTargetChanged(app, null, null, null);
+        // Controls pending to be dispatched.
+        assertTrue(getController().hasPendingControls(app));
+
+        performSurfacePlacementAndWaitForWindowAnimator();
+        // Pending controls were dispatched.
+        assertFalse(getController().hasPendingControls(app));
+    }
+
     /** Creates a window which is associated with ActivityRecord. */
     private WindowState createTestWindow(String name) {
         final WindowState win = createWindow(null, TYPE_APPLICATION, name);