Merge "Always allow the system to launch stuff on VDs." 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/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/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/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/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/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();