Merge "Trace coroutine section name generation" into main
diff --git a/cmds/uinput/src/com/android/commands/uinput/Device.java b/cmds/uinput/src/com/android/commands/uinput/Device.java
index 7e99c92..2cac6313 100644
--- a/cmds/uinput/src/com/android/commands/uinput/Device.java
+++ b/cmds/uinput/src/com/android/commands/uinput/Device.java
@@ -223,7 +223,7 @@
                         break;
                     }
                     long offsetMicros = args.argl1;
-                    if (mLastInjectTimestampMicros == -1 || offsetMicros == -1) {
+                    if (mLastInjectTimestampMicros == -1) {
                         // There's often a delay of a few milliseconds between the time specified to
                         // Handler.sendMessageAtTime and the handler actually being called, due to
                         // the way threads are scheduled. We don't take this into account when
@@ -239,6 +239,9 @@
                         // To prevent this, we need to use the time at which we scheduled this first
                         // batch, rather than the actual current time.
                         mLastInjectTimestampMicros = args.argl2 / 1000;
+                    } else if (offsetMicros == -1) {
+                        // No timestamp offset is specified for this event, so use the current time.
+                        mLastInjectTimestampMicros = SystemClock.uptimeNanos() / 1000;
                     } else {
                         mLastInjectTimestampMicros += offsetMicros;
                     }
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index f7e0e9f..436221f 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -765,11 +765,24 @@
             MessagingStyle.class, CallStyle.class);
 
     /** @hide */
-    @IntDef(flag = true, prefix = { "FLAG_" }, value = {FLAG_SHOW_LIGHTS, FLAG_ONGOING_EVENT,
-            FLAG_INSISTENT, FLAG_ONLY_ALERT_ONCE,
-            FLAG_AUTO_CANCEL, FLAG_NO_CLEAR, FLAG_FOREGROUND_SERVICE, FLAG_HIGH_PRIORITY,
-            FLAG_LOCAL_ONLY, FLAG_GROUP_SUMMARY, FLAG_AUTOGROUP_SUMMARY, FLAG_BUBBLE,
-            FLAG_USER_INITIATED_JOB})
+    @IntDef(flag = true, prefix = {"FLAG_"}, value = {
+            FLAG_SHOW_LIGHTS,
+            FLAG_ONGOING_EVENT,
+            FLAG_INSISTENT,
+            FLAG_ONLY_ALERT_ONCE,
+            FLAG_AUTO_CANCEL,
+            FLAG_NO_CLEAR,
+            FLAG_FOREGROUND_SERVICE,
+            FLAG_HIGH_PRIORITY,
+            FLAG_LOCAL_ONLY,
+            FLAG_GROUP_SUMMARY,
+            FLAG_AUTOGROUP_SUMMARY,
+            FLAG_CAN_COLORIZE,
+            FLAG_BUBBLE,
+            FLAG_NO_DISMISS,
+            FLAG_FSI_REQUESTED_BUT_DENIED,
+            FLAG_USER_INITIATED_JOB
+    })
     @Retention(RetentionPolicy.SOURCE)
     public @interface NotificationFlags{};
 
@@ -3779,10 +3792,10 @@
         if (this.tickerText != null) {
             sb.append(" tick");
         }
-        sb.append(" defaults=0x");
-        sb.append(Integer.toHexString(this.defaults));
-        sb.append(" flags=0x");
-        sb.append(Integer.toHexString(this.flags));
+        sb.append(" defaults=");
+        sb.append(defaultsToString(this.defaults));
+        sb.append(" flags=");
+        sb.append(flagsToString(this.flags));
         sb.append(String.format(" color=0x%08x", this.color));
         if (this.category != null) {
             sb.append(" category=");
@@ -3851,6 +3864,124 @@
     }
 
     /**
+     * {@hide}
+     */
+    public static String flagsToString(@NotificationFlags int flags) {
+        final List<String> flagStrings = new ArrayList<String>();
+        if ((flags & FLAG_SHOW_LIGHTS) != 0) {
+            flagStrings.add("SHOW_LIGHTS");
+            flags &= ~FLAG_SHOW_LIGHTS;
+        }
+        if ((flags & FLAG_ONGOING_EVENT) != 0) {
+            flagStrings.add("ONGOING_EVENT");
+            flags &= ~FLAG_ONGOING_EVENT;
+        }
+        if ((flags & FLAG_INSISTENT) != 0) {
+            flagStrings.add("INSISTENT");
+            flags &= ~FLAG_INSISTENT;
+        }
+        if ((flags & FLAG_ONLY_ALERT_ONCE) != 0) {
+            flagStrings.add("ONLY_ALERT_ONCE");
+            flags &= ~FLAG_ONLY_ALERT_ONCE;
+        }
+        if ((flags & FLAG_AUTO_CANCEL) != 0) {
+            flagStrings.add("AUTO_CANCEL");
+            flags &= ~FLAG_AUTO_CANCEL;
+        }
+        if ((flags & FLAG_NO_CLEAR) != 0) {
+            flagStrings.add("NO_CLEAR");
+            flags &= ~FLAG_NO_CLEAR;
+        }
+        if ((flags & FLAG_FOREGROUND_SERVICE) != 0) {
+            flagStrings.add("FOREGROUND_SERVICE");
+            flags &= ~FLAG_FOREGROUND_SERVICE;
+        }
+        if ((flags & FLAG_HIGH_PRIORITY) != 0) {
+            flagStrings.add("HIGH_PRIORITY");
+            flags &= ~FLAG_HIGH_PRIORITY;
+        }
+        if ((flags & FLAG_LOCAL_ONLY) != 0) {
+            flagStrings.add("LOCAL_ONLY");
+            flags &= ~FLAG_LOCAL_ONLY;
+        }
+        if ((flags & FLAG_GROUP_SUMMARY) != 0) {
+            flagStrings.add("GROUP_SUMMARY");
+            flags &= ~FLAG_GROUP_SUMMARY;
+        }
+        if ((flags & FLAG_AUTOGROUP_SUMMARY) != 0) {
+            flagStrings.add("AUTOGROUP_SUMMARY");
+            flags &= ~FLAG_AUTOGROUP_SUMMARY;
+        }
+        if ((flags & FLAG_CAN_COLORIZE) != 0) {
+            flagStrings.add("CAN_COLORIZE");
+            flags &= ~FLAG_CAN_COLORIZE;
+        }
+        if ((flags & FLAG_BUBBLE) != 0) {
+            flagStrings.add("BUBBLE");
+            flags &= ~FLAG_BUBBLE;
+        }
+        if ((flags & FLAG_NO_DISMISS) != 0) {
+            flagStrings.add("NO_DISMISS");
+            flags &= ~FLAG_NO_DISMISS;
+        }
+        if ((flags & FLAG_FSI_REQUESTED_BUT_DENIED) != 0) {
+            flagStrings.add("FSI_REQUESTED_BUT_DENIED");
+            flags &= ~FLAG_FSI_REQUESTED_BUT_DENIED;
+        }
+        if ((flags & FLAG_USER_INITIATED_JOB) != 0) {
+            flagStrings.add("USER_INITIATED_JOB");
+            flags &= ~FLAG_USER_INITIATED_JOB;
+        }
+        if (Flags.lifetimeExtensionRefactor()) {
+            if ((flags & FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY) != 0) {
+                flagStrings.add("LIFETIME_EXTENDED_BY_DIRECT_REPLY");
+                flags &= ~FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY;
+            }
+        }
+
+        if (flagStrings.isEmpty()) {
+            return "0";
+        }
+
+        if (flags != 0) {
+            flagStrings.add(String.format("UNKNOWN(0x%08x)", flags));
+        }
+
+        return String.join("|", flagStrings);
+    }
+
+    /** @hide */
+    public static String defaultsToString(int defaults) {
+        final List<String> defaultStrings = new ArrayList<String>();
+        if ((defaults & DEFAULT_ALL) == DEFAULT_ALL) {
+            defaultStrings.add("ALL");
+            defaults &= ~DEFAULT_ALL;
+        }
+        if ((defaults & DEFAULT_SOUND) != 0) {
+            defaultStrings.add("SOUND");
+            defaults &= ~DEFAULT_SOUND;
+        }
+        if ((defaults & DEFAULT_VIBRATE) != 0) {
+            defaultStrings.add("VIBRATE");
+            defaults &= ~DEFAULT_VIBRATE;
+        }
+        if ((defaults & DEFAULT_LIGHTS) != 0) {
+            defaultStrings.add("LIGHTS");
+            defaults &= ~DEFAULT_LIGHTS;
+        }
+
+        if (defaultStrings.isEmpty()) {
+            return "0";
+        }
+
+        if (defaults != 0) {
+            defaultStrings.add(String.format("UNKNOWN(0x%08x)", defaults));
+        }
+
+        return String.join("|", defaultStrings);
+    }
+
+    /**
      * @hide
      */
     public boolean hasCompletedProgress() {
diff --git a/core/java/android/app/contextualsearch/ContextualSearchManager.java b/core/java/android/app/contextualsearch/ContextualSearchManager.java
index c080a6b..cfbe741 100644
--- a/core/java/android/app/contextualsearch/ContextualSearchManager.java
+++ b/core/java/android/app/contextualsearch/ContextualSearchManager.java
@@ -27,6 +27,7 @@
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.SystemClock;
 import android.util.Log;
 
 import java.lang.annotation.Retention;
@@ -51,6 +52,7 @@
      */
     public static final String EXTRA_ENTRYPOINT =
             "android.app.contextualsearch.extra.ENTRYPOINT";
+
     /**
      * Key to get the flag_secure value from the extras of the activity launched by contextual
      * search. The value will be true if flag_secure is found in any of the visible activities.
@@ -58,12 +60,14 @@
      */
     public static final String EXTRA_FLAG_SECURE_FOUND =
             "android.app.contextualsearch.extra.FLAG_SECURE_FOUND";
+
     /**
      * Key to get the screenshot from the extras of the activity launched by contextual search.
      * Only supposed to be used with ACTON_LAUNCH_CONTEXTUAL_SEARCH.
      */
     public static final String EXTRA_SCREENSHOT =
             "android.app.contextualsearch.extra.SCREENSHOT";
+
     /**
      * Key to check whether managed profile is visible from the extras of the activity launched by
      * contextual search. The value will be true if any one of the visible apps is managed.
@@ -71,6 +75,7 @@
      */
     public static final String EXTRA_IS_MANAGED_PROFILE_VISIBLE =
             "android.app.contextualsearch.extra.IS_MANAGED_PROFILE_VISIBLE";
+
     /**
      * Key to get the list of visible packages from the extras of the activity launched by
      * contextual search.
@@ -80,6 +85,18 @@
             "android.app.contextualsearch.extra.VISIBLE_PACKAGE_NAMES";
 
     /**
+     * Key to get the time the user made the invocation request, based on
+     * {@link SystemClock#uptimeMillis()}.
+     * Only supposed to be used with ACTON_LAUNCH_CONTEXTUAL_SEARCH.
+     *
+     * TODO: un-hide in W
+     *
+     * @hide
+     */
+    public static final String EXTRA_INVOCATION_TIME_MS =
+            "android.app.contextualsearch.extra.INVOCATION_TIME_MS";
+
+    /**
      * Key to get the binder token from the extras of the activity launched by contextual search.
      * This token is needed to invoke {@link CallbackToken#getContextualSearchState} method.
      * Only supposed to be used with ACTON_LAUNCH_CONTEXTUAL_SEARCH.
diff --git a/core/java/android/content/pm/multiuser.aconfig b/core/java/android/content/pm/multiuser.aconfig
index 6d3a56d..e2dddef 100644
--- a/core/java/android/content/pm/multiuser.aconfig
+++ b/core/java/android/content/pm/multiuser.aconfig
@@ -113,6 +113,16 @@
     is_fixed_read_only: true
 }
 
+flag {
+    name: "fix_avatar_picker_read_back_order"
+    namespace: "multiuser"
+    description: "Talkback focus doesn't move to the 'If you change your Google Account picture…' after swiping next to move the focus from 'Choose a picture'"
+    bug: "330835921"
+    metadata {
+    	purpose: PURPOSE_BUGFIX
+  }
+}
+
 # This flag guards the private space feature and all its implementations excluding the APIs. APIs are guarded by android.os.Flags.allow_private_profile.
 flag {
     name: "enable_private_space_features"
diff --git a/core/java/android/os/connectivity/WifiActivityEnergyInfo.java b/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
index ad74a9f..2ceb1c7 100644
--- a/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
+++ b/core/java/android/os/connectivity/WifiActivityEnergyInfo.java
@@ -37,8 +37,10 @@
  * real-time.
  * @hide
  */
+@android.ravenwood.annotation.RavenwoodKeepWholeClass
 @SystemApi
 public final class WifiActivityEnergyInfo implements Parcelable {
+    private static final long DEFERRED_ENERGY_ESTIMATE = -1;
     @ElapsedRealtimeLong
     private final long mTimeSinceBootMillis;
     @StackState
@@ -52,7 +54,7 @@
     @IntRange(from = 0)
     private final long mControllerIdleDurationMillis;
     @IntRange(from = 0)
-    private final long mControllerEnergyUsedMicroJoules;
+    private long mControllerEnergyUsedMicroJoules;
 
     /** @hide */
     @Retention(RetentionPolicy.SOURCE)
@@ -99,9 +101,10 @@
                 rxDurationMillis,
                 scanDurationMillis,
                 idleDurationMillis,
-                calculateEnergyMicroJoules(txDurationMillis, rxDurationMillis, idleDurationMillis));
+                DEFERRED_ENERGY_ESTIMATE);
     }
 
+    @android.ravenwood.annotation.RavenwoodReplace
     private static long calculateEnergyMicroJoules(
             long txDurationMillis, long rxDurationMillis, long idleDurationMillis) {
         final Context context = ActivityThread.currentActivityThread().getSystemContext();
@@ -125,6 +128,11 @@
                 * voltage);
     }
 
+    private static long calculateEnergyMicroJoules$ravenwood(long txDurationMillis,
+            long rxDurationMillis, long idleDurationMillis) {
+        return 0;
+    }
+
     /** @hide */
     public WifiActivityEnergyInfo(
             @ElapsedRealtimeLong long timeSinceBootMillis,
@@ -152,7 +160,7 @@
                 + " mControllerRxDurationMillis=" + mControllerRxDurationMillis
                 + " mControllerScanDurationMillis=" + mControllerScanDurationMillis
                 + " mControllerIdleDurationMillis=" + mControllerIdleDurationMillis
-                + " mControllerEnergyUsedMicroJoules=" + mControllerEnergyUsedMicroJoules
+                + " mControllerEnergyUsedMicroJoules=" + getControllerEnergyUsedMicroJoules()
                 + " }";
     }
 
@@ -231,6 +239,11 @@
     /** Get the energy consumed by Wifi, in microjoules. */
     @IntRange(from = 0)
     public long getControllerEnergyUsedMicroJoules() {
+        if (mControllerEnergyUsedMicroJoules == DEFERRED_ENERGY_ESTIMATE) {
+            mControllerEnergyUsedMicroJoules = calculateEnergyMicroJoules(
+                    mControllerTxDurationMillis, mControllerRxDurationMillis,
+                    mControllerIdleDurationMillis);
+        }
         return mControllerEnergyUsedMicroJoules;
     }
 
diff --git a/core/java/android/permission/IPermissionManager.aidl b/core/java/android/permission/IPermissionManager.aidl
index 3b59901..55011e5 100644
--- a/core/java/android/permission/IPermissionManager.aidl
+++ b/core/java/android/permission/IPermissionManager.aidl
@@ -98,7 +98,7 @@
 
     IBinder registerAttributionSource(in AttributionSourceState source);
 
-    int getNumRegisteredAttributionSources(int uid);
+    int getRegisteredAttributionSourceCount(int uid);
 
     boolean isRegisteredAttributionSource(in AttributionSourceState source);
 
diff --git a/core/java/android/permission/PermissionManager.java b/core/java/android/permission/PermissionManager.java
index 2daf4ac..55bb430 100644
--- a/core/java/android/permission/PermissionManager.java
+++ b/core/java/android/permission/PermissionManager.java
@@ -1680,9 +1680,9 @@
      * @hide
      */
     @RequiresPermission(Manifest.permission.UPDATE_APP_OPS_STATS)
-    public int getNumRegisteredAttributionSourcesForTest(int uid) {
+    public int getRegisteredAttributionSourceCountForTest(int uid) {
         try {
-            return mPermissionManager.getNumRegisteredAttributionSources(uid);
+            return mPermissionManager.getRegisteredAttributionSourceCount(uid);
         } catch (RemoteException e) {
             e.rethrowFromSystemServer();
         }
diff --git a/core/java/android/service/dreams/DreamService.java b/core/java/android/service/dreams/DreamService.java
index 353828c..f896447 100644
--- a/core/java/android/service/dreams/DreamService.java
+++ b/core/java/android/service/dreams/DreamService.java
@@ -52,6 +52,7 @@
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.service.controls.flags.Flags;
+import android.service.dreams.utils.DreamAccessibility;
 import android.util.AttributeSet;
 import android.util.Log;
 import android.util.MathUtils;
@@ -273,6 +274,7 @@
     private boolean mDebug = false;
 
     private ComponentName mDreamComponent;
+    private DreamAccessibility mDreamAccessibility;
     private boolean mShouldShowComplications;
 
     private DreamServiceWrapper mDreamServiceWrapper;
@@ -664,6 +666,7 @@
      */
     public void setInteractive(boolean interactive) {
         mInteractive = interactive;
+        updateAccessibilityMessage();
     }
 
     /**
@@ -1430,7 +1433,7 @@
         // Hide all insets when the dream is showing
         mWindow.getDecorView().getWindowInsetsController().hide(WindowInsets.Type.systemBars());
         mWindow.setDecorFitsSystemWindows(false);
-
+        updateAccessibilityMessage();
         mWindow.getDecorView().addOnAttachStateChangeListener(
                 new View.OnAttachStateChangeListener() {
                     private Consumer<IDreamOverlayClient> mDreamStartOverlayConsumer;
@@ -1477,6 +1480,15 @@
                 });
     }
 
+    private void updateAccessibilityMessage() {
+        if (mWindow == null) return;
+        if (mDreamAccessibility == null) {
+            final View rootView = mWindow.getDecorView();
+            mDreamAccessibility = new DreamAccessibility(this, rootView);
+        }
+        mDreamAccessibility.updateAccessibilityConfiguration(isInteractive());
+    }
+
     private boolean getWindowFlagValue(int flag, boolean defaultValue) {
         return mWindow == null ? defaultValue : (mWindow.getAttributes().flags & flag) != 0;
     }
diff --git a/core/java/android/service/dreams/utils/DreamAccessibility.java b/core/java/android/service/dreams/utils/DreamAccessibility.java
new file mode 100644
index 0000000..c38f41b
--- /dev/null
+++ b/core/java/android/service/dreams/utils/DreamAccessibility.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.service.dreams.utils;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.view.View;
+import android.view.accessibility.AccessibilityNodeInfo;
+
+import com.android.internal.R;
+
+/**
+ * {@link DreamAccessibility} allows customization of accessibility
+ * actions for the root view of the dream overlay.
+ * @hide
+ */
+public class DreamAccessibility {
+    private final Context mContext;
+    private final View mView;
+    private final View.AccessibilityDelegate mAccessibilityDelegate;
+
+    public DreamAccessibility(@NonNull Context context, @NonNull View view) {
+        mContext = context;
+        mView = view;
+        mAccessibilityDelegate = createNewAccessibilityDelegate(mContext);
+    }
+
+    /**
+     * @param interactive
+     * Removes and add accessibility configuration depending if the dream is interactive or not
+     */
+    public void updateAccessibilityConfiguration(Boolean interactive) {
+        if (!interactive) {
+            addAccessibilityConfiguration();
+        } else {
+            removeCustomAccessibilityAction();
+        }
+    }
+
+    /**
+     * Configures the accessibility actions for the given root view.
+     */
+    private void addAccessibilityConfiguration() {
+        mView.setAccessibilityDelegate(mAccessibilityDelegate);
+    }
+
+    /**
+     * Removes Configured the accessibility actions for the given root view.
+     */
+    private void removeCustomAccessibilityAction() {
+        if (mView.getAccessibilityDelegate() == mAccessibilityDelegate) {
+            mView.setAccessibilityDelegate(null);
+        }
+    }
+
+    private View.AccessibilityDelegate createNewAccessibilityDelegate(Context context) {
+        return new View.AccessibilityDelegate() {
+            @Override
+            public void onInitializeAccessibilityNodeInfo(View host, AccessibilityNodeInfo info) {
+                super.onInitializeAccessibilityNodeInfo(host, info);
+                for (AccessibilityNodeInfo.AccessibilityAction action : info.getActionList()) {
+                    if (action.getId() == AccessibilityNodeInfo.ACTION_CLICK) {
+                        info.removeAction(action);
+                        break;
+                    }
+                }
+                info.addAction(new AccessibilityNodeInfo.AccessibilityAction(
+                        AccessibilityNodeInfo.ACTION_CLICK,
+                        context.getResources().getString(R.string.dream_accessibility_action_click)
+                ));
+            }
+        };
+    }
+}
diff --git a/core/java/android/telephony/PhoneStateListener.java b/core/java/android/telephony/PhoneStateListener.java
index fb57921..4281da1 100644
--- a/core/java/android/telephony/PhoneStateListener.java
+++ b/core/java/android/telephony/PhoneStateListener.java
@@ -1685,6 +1685,10 @@
         public final void onSimultaneousCallingStateChanged(int[] subIds) {
             // not supported on the deprecated interface - Use TelephonyCallback instead
         }
+
+        public final void onCarrierRoamingNtnModeChanged(boolean active) {
+            // not supported on the deprecated interface - Use TelephonyCallback instead
+        }
     }
 
     private void log(String s) {
diff --git a/core/java/android/telephony/TelephonyCallback.java b/core/java/android/telephony/TelephonyCallback.java
index dc6a035..b8b84d9 100644
--- a/core/java/android/telephony/TelephonyCallback.java
+++ b/core/java/android/telephony/TelephonyCallback.java
@@ -644,6 +644,15 @@
     public static final int EVENT_SIMULTANEOUS_CELLULAR_CALLING_SUBSCRIPTIONS_CHANGED = 41;
 
     /**
+     * Event for listening to changes in carrier roaming non-terrestrial network mode.
+     *
+     * @see CarrierRoamingNtnModeListener
+     *
+     * @hide
+     */
+    public static final int EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED = 42;
+
+    /**
      * @hide
      */
     @IntDef(prefix = {"EVENT_"}, value = {
@@ -687,7 +696,8 @@
             EVENT_TRIGGER_NOTIFY_ANBR,
             EVENT_MEDIA_QUALITY_STATUS_CHANGED,
             EVENT_EMERGENCY_CALLBACK_MODE_CHANGED,
-            EVENT_SIMULTANEOUS_CELLULAR_CALLING_SUBSCRIPTIONS_CHANGED
+            EVENT_SIMULTANEOUS_CELLULAR_CALLING_SUBSCRIPTIONS_CHANGED,
+            EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED
     })
     @Retention(RetentionPolicy.SOURCE)
     public @interface TelephonyEvent {
@@ -1686,6 +1696,24 @@
     }
 
     /**
+     * Interface for carrier roaming non-terrestrial network listener.
+     *
+     * @hide
+     */
+    public interface CarrierRoamingNtnModeListener {
+        /**
+         * Callback invoked when carrier roaming non-terrestrial network mode changes.
+         *
+         * @param active {@code true} If the device is connected to carrier roaming
+         *                           non-terrestrial network or was connected within the
+         *                           {CarrierConfigManager
+         *                           #KEY_SATELLITE_CONNECTION_HYSTERESIS_SEC_INT} duration,
+         *                           {code false} otherwise.
+         */
+        void onCarrierRoamingNtnModeChanged(boolean active);
+    }
+
+    /**
      * The callback methods need to be called on the handler thread where
      * this object was created.  If the binder did that for us it'd be nice.
      * <p>
@@ -2086,5 +2114,16 @@
             Binder.withCleanCallingIdentity(
                     () -> mExecutor.execute(() -> listener.onCallBackModeStopped(type, reason)));
         }
+
+        public void onCarrierRoamingNtnModeChanged(boolean active) {
+            if (!Flags.carrierEnabledSatelliteFlag()) return;
+
+            CarrierRoamingNtnModeListener listener =
+                    (CarrierRoamingNtnModeListener) mTelephonyCallbackWeakRef.get();
+            if (listener == null) return;
+
+            Binder.withCleanCallingIdentity(
+                    () -> mExecutor.execute(() -> listener.onCarrierRoamingNtnModeChanged(active)));
+        }
     }
 }
diff --git a/core/java/android/telephony/TelephonyRegistryManager.java b/core/java/android/telephony/TelephonyRegistryManager.java
index d39c4ce..6160fdb 100644
--- a/core/java/android/telephony/TelephonyRegistryManager.java
+++ b/core/java/android/telephony/TelephonyRegistryManager.java
@@ -1073,6 +1073,24 @@
     }
 
     /**
+     * Notify external listeners that carrier roaming non-terrestrial network mode changed.
+     * @param subId subscription ID.
+     * @param active {@code true} If the device is connected to carrier roaming
+     *                           non-terrestrial network or was connected within the
+     *                           {CarrierConfigManager#KEY_SATELLITE_CONNECTION_HYSTERESIS_SEC_INT}
+     *                           duration, {code false} otherwise.
+     * @hide
+     */
+    public void notifyCarrierRoamingNtnModeChanged(int subId, boolean active) {
+        try {
+            sRegistry.notifyCarrierRoamingNtnModeChanged(subId, active);
+        } catch (RemoteException ex) {
+            // system server crash
+            throw ex.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Processes potential event changes from the provided {@link TelephonyCallback}.
      *
      * @param telephonyCallback callback for monitoring callback changes to the telephony state.
@@ -1224,6 +1242,10 @@
             eventList.add(
                     TelephonyCallback.EVENT_SIMULTANEOUS_CELLULAR_CALLING_SUBSCRIPTIONS_CHANGED);
         }
+
+        if (telephonyCallback instanceof TelephonyCallback.CarrierRoamingNtnModeListener) {
+            eventList.add(TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED);
+        }
         return eventList;
     }
 
diff --git a/core/java/com/android/internal/telephony/IPhoneStateListener.aidl b/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
index 969f95d..792c223 100644
--- a/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
+++ b/core/java/com/android/internal/telephony/IPhoneStateListener.aidl
@@ -81,4 +81,5 @@
     void onCallBackModeStarted(int type);
     void onCallBackModeStopped(int type, int reason);
     void onSimultaneousCallingStateChanged(in int[] subIds);
+    void onCarrierRoamingNtnModeChanged(in boolean active);
 }
diff --git a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
index 0203ea4..04332cd 100644
--- a/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
+++ b/core/java/com/android/internal/telephony/ITelephonyRegistry.aidl
@@ -120,4 +120,5 @@
 
     void notifyCallbackModeStarted(int phoneId, int subId, int type);
     void notifyCallbackModeStopped(int phoneId, int subId, int type, int reason);
+    void notifyCarrierRoamingNtnModeChanged(int subId, in boolean active);
 }
diff --git a/core/res/res/drawable/ic_android_satellite_24px.xml b/core/res/res/drawable/ic_android_satellite_24px.xml
new file mode 100644
index 0000000..15f2884
--- /dev/null
+++ b/core/res/res/drawable/ic_android_satellite_24px.xml
@@ -0,0 +1,10 @@
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+  <path
+      android:fillColor="@android:color/white"
+      android:pathData="M487,600L392,505L346,551L376,580Q399,603 399,637Q399,671 376,694L333,737Q310,760 276.5,760Q243,760 220,737L103,620Q80,597 80,563.5Q80,530 103,507L146,464Q169,441 203,441Q237,441 260,464L289,494L335,448L160,273L273,160L448,335L493,290L464,260Q441,237 441,203Q441,169 464,146L507,103Q530,80 563.5,80Q597,80 620,103L737,220Q760,243 760,276.5Q760,310 737,333L694,376Q671,399 637,399Q603,399 580,376L550,347L505,392L600,487L487,600ZM520,880L520,800Q637,800 718.5,718.5Q800,637 800,520L880,520Q880,595 851.5,660.5Q823,726 774.5,774.5Q726,823 660.5,851.5Q595,880 520,880ZM520,720L520,640Q570,640 605,605Q640,570 640,520L720,520Q720,603 661.5,661.5Q603,720 520,720ZM520,203L550,233L593,190L563,160Q563,160 563,160Q563,160 563,160L520,203Q520,203 520,203Q520,203 520,203ZM160,563L190,593L233,550L203,520Q203,520 203,520Q203,520 203,520L160,563Q160,563 160,563Q160,563 160,563ZM637,320L680,277Q680,277 680,277Q680,277 680,277L650,247L607,290L637,320Q637,320 637,320Q637,320 637,320ZM277,680L320,637Q320,637 320,637Q320,637 320,637L290,607L247,650L277,680Q277,680 277,680Q277,680 277,680Z"/>
+</vector>
diff --git a/core/res/res/values/config_battery_stats.xml b/core/res/res/values/config_battery_stats.xml
index ae47899..8d97362 100644
--- a/core/res/res/values/config_battery_stats.xml
+++ b/core/res/res/values/config_battery_stats.xml
@@ -35,6 +35,9 @@
     <!-- Mobile Radio power stats collection throttle period in milliseconds. -->
     <integer name="config_defaultPowerStatsThrottlePeriodMobileRadio">3600000</integer>
 
+    <!-- Mobile Radio power stats collection throttle period in milliseconds. -->
+    <integer name="config_defaultPowerStatsThrottlePeriodWifi">3600000</integer>
+
     <!-- PowerStats aggregation period in milliseconds. This is the interval at which the power
     stats aggregation procedure is performed and the results stored in PowerStatsStore. -->
     <integer name="config_powerStatsAggregationPeriod">14400000</integer>
diff --git a/core/res/res/values/strings.xml b/core/res/res/values/strings.xml
index 60289c1..e96240d 100644
--- a/core/res/res/values/strings.xml
+++ b/core/res/res/values/strings.xml
@@ -1018,6 +1018,9 @@
     <!-- The title to use when a dream is opened in preview mode. [CHAR LIMIT=NONE] -->
     <string name="dream_preview_title">Preview, <xliff:g id="dream_name" example="Clock">%1$s</xliff:g></string>
 
+    <!-- The title to use when a dream is open in accessibility mode to let users know to double tap to dismiss the dream  [CHAR LIMIT=32] -->
+    <string name="dream_accessibility_action_click">dismiss</string>
+
     <!--  Permissions -->
 
     <!-- Title of an application permission, listed so the user can choose whether they want to allow the application to do this. -->
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index ead5827..0bf6347 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -4471,6 +4471,7 @@
   <java-symbol type="string" name="capability_title_canTakeScreenshot" />
 
   <java-symbol type="string" name="dream_preview_title" />
+  <java-symbol type="string" name="dream_accessibility_action_click" />
 
   <java-symbol type="string" name="config_servicesExtensionPackage" />
 
@@ -5218,6 +5219,7 @@
   <java-symbol type="bool" name="config_batteryStatsResetOnUnplugAfterSignificantCharge" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodCpu" />
   <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodMobileRadio" />
+  <java-symbol type="integer" name="config_defaultPowerStatsThrottlePeriodWifi" />
   <java-symbol type="integer" name="config_powerStatsAggregationPeriod" />
   <java-symbol type="integer" name="config_aggregatedPowerStatsSpanDuration" />
 
@@ -5380,6 +5382,7 @@
   <java-symbol type="string" name="satellite_notification_open_message" />
   <java-symbol type="string" name="satellite_notification_how_it_works" />
   <java-symbol type="drawable" name="ic_satellite_alt_24px" />
+  <java-symbol type="drawable" name="ic_android_satellite_24px" />
 
   <!-- DisplayManager configs. -->
   <java-symbol type="bool" name="config_evenDimmerEnabled" />
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipBoundsState.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipBoundsState.java
index b87c2f6..7ceaaea 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipBoundsState.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipBoundsState.java
@@ -125,6 +125,7 @@
     private @Nullable Runnable mOnMinimalSizeChangeCallback;
     private @Nullable TriConsumer<Boolean, Integer, Boolean> mOnShelfVisibilityChangeCallback;
     private List<Consumer<Rect>> mOnPipExclusionBoundsChangeCallbacks = new ArrayList<>();
+    private List<Consumer<Float>> mOnAspectRatioChangedCallbacks = new ArrayList<>();
 
     // the size of the current bounds relative to the max size spec
     private float mBoundsScale;
@@ -297,7 +298,12 @@
 
     /** Set the PIP aspect ratio. */
     public void setAspectRatio(float aspectRatio) {
-        mAspectRatio = aspectRatio;
+        if (Float.compare(mAspectRatio, aspectRatio) != 0) {
+            mAspectRatio = aspectRatio;
+            for (Consumer<Float> callback : mOnAspectRatioChangedCallbacks) {
+                callback.accept(mAspectRatio);
+            }
+        }
     }
 
     /** Get the PIP aspect ratio. */
@@ -527,6 +533,23 @@
         mOnPipExclusionBoundsChangeCallbacks.remove(onPipExclusionBoundsChangeCallback);
     }
 
+    /** Adds callback to listen on aspect ratio change. */
+    public void addOnAspectRatioChangedCallback(
+            @NonNull Consumer<Float> onAspectRatioChangedCallback) {
+        if (!mOnAspectRatioChangedCallbacks.contains(onAspectRatioChangedCallback)) {
+            mOnAspectRatioChangedCallbacks.add(onAspectRatioChangedCallback);
+            onAspectRatioChangedCallback.accept(mAspectRatio);
+        }
+    }
+
+    /** Removes callback to listen on aspect ratio change. */
+    public void removeOnAspectRatioChangedCallback(
+            @NonNull Consumer<Float> onAspectRatioChangedCallback) {
+        if (mOnAspectRatioChangedCallbacks.contains(onAspectRatioChangedCallback)) {
+            mOnAspectRatioChangedCallbacks.remove(onAspectRatioChangedCallback);
+        }
+    }
+
     public LauncherState getLauncherState() {
         return mLauncherState;
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index 73228de..6834e6d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -107,6 +107,7 @@
 import com.android.wm.shell.taskview.TaskViewFactoryController;
 import com.android.wm.shell.taskview.TaskViewTransitions;
 import com.android.wm.shell.transition.HomeTransitionObserver;
+import com.android.wm.shell.transition.MixedTransitionHandler;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
 import com.android.wm.shell.unfold.UnfoldAnimationController;
@@ -675,6 +676,22 @@
         return new TaskViewTransitions(transitions);
     }
 
+    // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
+    @BindsOptionalOf
+    @DynamicOverride
+    abstract MixedTransitionHandler optionalMixedTransitionHandler();
+
+    @WMSingleton
+    @Provides
+    static Optional<MixedTransitionHandler> provideMixedTransitionHandler(
+            @DynamicOverride Optional<MixedTransitionHandler> mixedTransitionHandler
+    ) {
+        if (mixedTransitionHandler.isPresent()) {
+            return mixedTransitionHandler;
+        }
+        return Optional.empty();
+    }
+
     //
     // Keyguard transitions (optional feature)
     //
@@ -934,6 +951,7 @@
             Optional<OneHandedController> oneHandedControllerOptional,
             Optional<HideDisplayCutoutController> hideDisplayCutoutControllerOptional,
             Optional<ActivityEmbeddingController> activityEmbeddingOptional,
+            Optional<MixedTransitionHandler> mixedTransitionHandler,
             Transitions transitions,
             StartingWindowController startingWindow,
             ProtoLogController protoLogController,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 26e7acb..b574b81 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -87,6 +87,7 @@
 import com.android.wm.shell.taskview.TaskViewTransitions;
 import com.android.wm.shell.transition.DefaultMixedHandler;
 import com.android.wm.shell.transition.HomeTransitionObserver;
+import com.android.wm.shell.transition.MixedTransitionHandler;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
 import com.android.wm.shell.unfold.UnfoldAnimationController;
@@ -373,8 +374,9 @@
     //
 
     @WMSingleton
+    @DynamicOverride
     @Provides
-    static DefaultMixedHandler provideDefaultMixedHandler(
+    static MixedTransitionHandler provideMixedTransitionHandler(
             ShellInit shellInit,
             Optional<SplitScreenController> splitScreenOptional,
             @Nullable PipTransitionController pipTransitionController,
@@ -655,7 +657,6 @@
     @Provides
     static Object provideIndependentShellComponentsToCreate(
             DragAndDropController dragAndDropController,
-            DefaultMixedHandler defaultMixedHandler,
             Optional<DesktopTasksTransitionObserver> desktopTasksTransitionObserverOptional) {
         return new Object();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
index c1adfff..d8ac8e9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
@@ -219,6 +219,7 @@
                 mMotionHelper, pipTaskOrganizer, mPipBoundsAlgorithm.getSnapAlgorithm(),
                 this::onAccessibilityShowMenu, this::updateMovementBounds,
                 this::animateToUnStashedState, mainExecutor);
+        mPipBoundsState.addOnAspectRatioChangedCallback(this::updateMinMaxSize);
 
         // TODO(b/181599115): This should really be initializes as part of the pip controller, but
         // until all PIP implementations derive from the controller, just initialize the touch handler
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
index cc8e3e0..472003c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip2/phone/PipTouchHandler.java
@@ -208,6 +208,7 @@
                 new PipResizeGestureHandler(context, pipBoundsAlgorithm, pipBoundsState,
                         mTouchState, this::updateMovementBounds, pipUiEventLogger,
                         menuController, mainExecutor, mPipPerfHintController);
+        mPipBoundsState.addOnAspectRatioChangedCallback(this::updateMinMaxSize);
 
         if (PipUtils.isPip2ExperimentEnabled()) {
             shellInit.addInitCallback(this::onInit, this);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
index 8746b8c..4bc0dc0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
@@ -61,9 +61,10 @@
 
 /**
  * A handler for dealing with transitions involving multiple other handlers. For example: an
- * activity in split-screen going into PiP.
+ * activity in split-screen going into PiP. Note this is provided as a handset-specific
+ * implementation of {@code MixedTransitionHandler}.
  */
-public class DefaultMixedHandler implements Transitions.TransitionHandler,
+public class DefaultMixedHandler implements MixedTransitionHandler,
         RecentsTransitionHandler.RecentsMixedHandler {
 
     private final Transitions mPlayer;
@@ -116,7 +117,7 @@
         final IBinder mTransition;
 
         protected final Transitions mPlayer;
-        protected final DefaultMixedHandler mMixedHandler;
+        protected final MixedTransitionHandler mMixedHandler;
         protected final PipTransitionController mPipHandler;
         protected final StageCoordinator mSplitHandler;
         protected final KeyguardTransitionHandler mKeyguardHandler;
@@ -142,7 +143,7 @@
         int mInFlightSubAnimations = 0;
 
         MixedTransition(int type, IBinder transition, Transitions player,
-                DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
+                MixedTransitionHandler mixedHandler, PipTransitionController pipHandler,
                 StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler) {
             mType = type;
             mTransition = transition;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
index e9cd73b..b028bd6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
@@ -41,7 +41,7 @@
     private final ActivityEmbeddingController mActivityEmbeddingController;
 
     DefaultMixedTransition(int type, IBinder transition, Transitions player,
-            DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
+            MixedTransitionHandler mixedHandler, PipTransitionController pipHandler,
             StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler,
             UnfoldTransitionHandler unfoldHandler,
             ActivityEmbeddingController activityEmbeddingController) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHandler.java
new file mode 100644
index 0000000..ff429fb
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHandler.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.transition;
+
+/**
+ * Interface for a {@link Transitions.TransitionHandler} that can take the subset of transitions
+ * that it handles and further decompose those transitions into sub-transitions which can be
+ * independently delegated to separate handlers.
+ */
+public interface MixedTransitionHandler extends Transitions.TransitionHandler {
+
+    // TODO(b/335685449) this currently exists purely as a marker interface for use in form-factor
+    // specific/sysui dagger modules. Going forward, we should define this in a meaningful
+    // way so as to provide a clear basis for expectations/behaviours associated with mixed
+    // transitions and their default handlers.
+
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
index 0974cd1..ffc0b76 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
@@ -44,7 +44,7 @@
             @NonNull SurfaceControl.Transaction startTransaction,
             @NonNull SurfaceControl.Transaction finishTransaction,
             @NonNull Transitions.TransitionFinishCallback finishCallback,
-            @NonNull Transitions player, @NonNull DefaultMixedHandler mixedHandler,
+            @NonNull Transitions player, @NonNull MixedTransitionHandler mixedHandler,
             @NonNull PipTransitionController pipHandler, @NonNull StageCoordinator splitHandler) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
                 + "entering PIP while Split-Screen is foreground.");
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
index 5b402a5..d6e64cf 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
@@ -43,7 +43,7 @@
     private final DesktopTasksController mDesktopTasksController;
 
     RecentsMixedTransition(int type, IBinder transition, Transitions player,
-            DefaultMixedHandler mixedHandler, PipTransitionController pipHandler,
+            MixedTransitionHandler mixedHandler, PipTransitionController pipHandler,
             StageCoordinator splitHandler, KeyguardTransitionHandler keyguardHandler,
             RecentsTransitionHandler recentsHandler,
             DesktopTasksController desktopTasksController) {
diff --git a/libs/hwui/pipeline/skia/ShaderCache.cpp b/libs/hwui/pipeline/skia/ShaderCache.cpp
index b870023..8e07a2f 100644
--- a/libs/hwui/pipeline/skia/ShaderCache.cpp
+++ b/libs/hwui/pipeline/skia/ShaderCache.cpp
@@ -15,14 +15,18 @@
  */
 
 #include "ShaderCache.h"
+
 #include <GrDirectContext.h>
 #include <SkData.h>
 #include <gui/TraceUtils.h>
 #include <log/log.h>
 #include <openssl/sha.h>
+
 #include <algorithm>
 #include <array>
+#include <mutex>
 #include <thread>
+
 #include "FileBlobCache.h"
 #include "Properties.h"
 
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index c664d3d..ffd6e16 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -545,16 +545,17 @@
                                     ? SCANNING_STATE_SCANNING_FULL
                                     : SCANNING_STATE_WHILE_INTERACTIVE);
 
-                    if (scanRequest.isScreenOffScan()) {
-                        mScreenOffScanRequestCount++;
-                    } else {
-                        mScreenOnScanRequestCount++;
-                    }
                 } catch (RemoteException ex) {
                     throw ex.rethrowFromSystemServer();
                 }
             }
 
+            if (scanRequest.isScreenOffScan()) {
+                mScreenOffScanRequestCount++;
+            } else {
+                mScreenOnScanRequestCount++;
+            }
+
             mScanRequestsMap.put(token.mId, scanRequest);
             return token;
         }
@@ -580,27 +581,29 @@
             }
 
             boolean shouldUpdate =
-                    mScreenOffScanRequestCount == 1
-                            && (request.isScreenOffScan() || mScreenOnScanRequestCount == 1);
+                    request.isScreenOffScan()
+                            ? mScreenOffScanRequestCount == 1
+                            : mScreenOnScanRequestCount == 1 && mScreenOffScanRequestCount == 0;
 
             if (shouldUpdate) {
                 try {
-                    if (request.isScreenOffScan() && mScreenOnScanRequestCount == 0) {
+                    if (!request.isScreenOffScan() || mScreenOnScanRequestCount == 0) {
                         mImpl.updateScanningState(SCANNING_STATE_NOT_SCANNING);
                     } else {
                         mImpl.updateScanningState(SCANNING_STATE_WHILE_INTERACTIVE);
                     }
 
-                    if (request.isScreenOffScan()) {
-                        mScreenOffScanRequestCount--;
-                    } else {
-                        mScreenOnScanRequestCount--;
-                    }
                 } catch (RemoteException ex) {
                     ex.rethrowFromSystemServer();
                 }
             }
 
+            if (request.isScreenOffScan()) {
+                mScreenOffScanRequestCount--;
+            } else {
+                mScreenOnScanRequestCount--;
+            }
+
             mScanRequestsMap.remove(token.mId);
         }
     }
diff --git a/media/java/android/media/RoutingSessionInfo.java b/media/java/android/media/RoutingSessionInfo.java
index 2202766..a3c8b68 100644
--- a/media/java/android/media/RoutingSessionInfo.java
+++ b/media/java/android/media/RoutingSessionInfo.java
@@ -84,7 +84,6 @@
      *
      * @hide
      */
-    @FlaggedApi(FLAG_ENABLE_BUILT_IN_SPEAKER_ROUTE_SUITABILITY_STATUSES)
     @IntDef(value = {TRANSFER_REASON_FALLBACK, TRANSFER_REASON_SYSTEM_REQUEST, TRANSFER_REASON_APP})
     @Retention(RetentionPolicy.SOURCE)
     public @interface TransferReason {}
@@ -382,14 +381,12 @@
     }
 
     /** @hide */
-    @FlaggedApi(FLAG_ENABLE_BUILT_IN_SPEAKER_ROUTE_SUITABILITY_STATUSES)
     @Nullable
     public UserHandle getTransferInitiatorUserHandle() {
         return mTransferInitiatorUserHandle;
     }
 
     /** @hide */
-    @FlaggedApi(FLAG_ENABLE_BUILT_IN_SPEAKER_ROUTE_SUITABILITY_STATUSES)
     @Nullable
     public String getTransferInitiatorPackageName() {
         return mTransferInitiatorPackageName;
diff --git a/packages/SettingsLib/src/com/android/settingslib/accounts/AuthenticatorHelper.java b/packages/SettingsLib/src/com/android/settingslib/accounts/AuthenticatorHelper.java
index 4af9e3c..2f6d839 100644
--- a/packages/SettingsLib/src/com/android/settingslib/accounts/AuthenticatorHelper.java
+++ b/packages/SettingsLib/src/com/android/settingslib/accounts/AuthenticatorHelper.java
@@ -106,8 +106,7 @@
                 AuthenticatorDescription desc = mTypeToAuthDescription.get(accountType);
                 Context authContext = context.createPackageContextAsUser(desc.packageName, 0,
                         mUserHandle);
-                icon = mContext.getPackageManager().getUserBadgedIcon(
-                        authContext.getDrawable(desc.iconId), mUserHandle);
+                icon = authContext.getDrawable(desc.iconId);
                 synchronized (mAccTypeIconCache) {
                     mAccTypeIconCache.put(accountType, icon);
                 }
diff --git a/packages/SystemUI/aconfig/systemui.aconfig b/packages/SystemUI/aconfig/systemui.aconfig
index 16bc7f2..b499a0e 100644
--- a/packages/SystemUI/aconfig/systemui.aconfig
+++ b/packages/SystemUI/aconfig/systemui.aconfig
@@ -814,3 +814,10 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enforce_brightness_base_user_restriction"
+    namespace: "systemui"
+    description: "Enforce BaseUserRestriction for DISALLOW_CONFIG_BRIGHTNESS."
+    bug: "329205638"
+}
diff --git a/packages/SystemUI/customization/src/com/android/systemui/shared/clocks/AnimatableClockView.kt b/packages/SystemUI/customization/src/com/android/systemui/shared/clocks/AnimatableClockView.kt
index f1f1b57..f539a23 100644
--- a/packages/SystemUI/customization/src/com/android/systemui/shared/clocks/AnimatableClockView.kt
+++ b/packages/SystemUI/customization/src/com/android/systemui/shared/clocks/AnimatableClockView.kt
@@ -97,6 +97,8 @@
     @VisibleForTesting var textAnimatorFactory: (Layout, () -> Unit) -> TextAnimator =
         { layout, invalidateCb ->
             TextAnimator(layout, NUM_CLOCK_FONT_ANIMATION_STEPS, invalidateCb) }
+
+    // Used by screenshot tests to provide stability
     @VisibleForTesting var isAnimationEnabled: Boolean = true
     @VisibleForTesting var timeOverrideInMillis: Long? = null
 
@@ -242,15 +244,8 @@
         }
 
         logger.d({ "onDraw($str1)"}) { str1 = text.toString() }
-        // Use textAnimator to render text if animation is enabled.
-        // Otherwise default to using standard draw functions.
-        if (isAnimationEnabled) {
-            // intentionally doesn't call super.onDraw here or else the text will be rendered twice
-            textAnimator?.draw(canvas)
-        } else {
-            super.onDraw(canvas)
-        }
-
+        // intentionally doesn't call super.onDraw here or else the text will be rendered twice
+        textAnimator?.draw(canvas)
         canvas.restore()
     }
 
@@ -318,7 +313,7 @@
             weight = lockScreenWeight,
             textSize = -1f,
             color = lockScreenColor,
-            animate = isAnimationEnabled,
+            animate = true,
             duration = APPEAR_ANIM_DURATION,
             interpolator = Interpolators.EMPHASIZED_DECELERATE,
             delay = 0,
@@ -327,7 +322,7 @@
     }
 
     fun animateFoldAppear(animate: Boolean = true) {
-        if (isAnimationEnabled && textAnimator == null) {
+        if (textAnimator == null) {
             return
         }
         logger.d("animateFoldAppear")
@@ -344,7 +339,7 @@
             weight = dozingWeightInternal,
             textSize = -1f,
             color = dozingColor,
-            animate = animate && isAnimationEnabled,
+            animate = animate,
             interpolator = Interpolators.EMPHASIZED_DECELERATE,
             duration = ANIMATION_DURATION_FOLD_TO_AOD.toLong(),
             delay = 0,
@@ -363,7 +358,7 @@
                 weight = if (isDozing()) dozingWeight else lockScreenWeight,
                 textSize = -1f,
                 color = null,
-                animate = isAnimationEnabled,
+                animate = true,
                 duration = CHARGE_ANIM_DURATION_PHASE_1,
                 delay = 0,
                 onAnimationEnd = null
@@ -373,7 +368,7 @@
             weight = if (isDozing()) lockScreenWeight else dozingWeight,
             textSize = -1f,
             color = null,
-            animate = isAnimationEnabled,
+            animate = true,
             duration = CHARGE_ANIM_DURATION_PHASE_0,
             delay = chargeAnimationDelay.toLong(),
             onAnimationEnd = startAnimPhase2
@@ -386,7 +381,7 @@
             weight = if (isDozing) dozingWeight else lockScreenWeight,
             textSize = -1f,
             color = if (isDozing) dozingColor else lockScreenColor,
-            animate = animate && isAnimationEnabled,
+            animate = animate,
             duration = DOZE_ANIM_DURATION,
             delay = 0,
             onAnimationEnd = null
@@ -445,9 +440,6 @@
                 onAnimationEnd = onAnimationEnd
             )
             textAnimator?.glyphFilter = glyphFilter
-            if (color != null && !isAnimationEnabled) {
-                setTextColor(color)
-            }
         } else {
             // when the text animator is set, update its start values
             onTextAnimatorInitialized = Runnable {
@@ -462,9 +454,6 @@
                     onAnimationEnd = onAnimationEnd
                 )
                 textAnimator?.glyphFilter = glyphFilter
-                if (color != null && !isAnimationEnabled) {
-                    setTextColor(color)
-                }
             }
         }
     }
@@ -482,7 +471,7 @@
             weight = weight,
             textSize = textSize,
             color = color,
-            animate = animate && isAnimationEnabled,
+            animate = animate,
             interpolator = null,
             duration = duration,
             delay = delay,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt
new file mode 100644
index 0000000..722eb2b
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt
@@ -0,0 +1,152 @@
+/*
+ * 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.communal
+
+import android.appwidget.AppWidgetManager
+import android.content.Context
+import android.content.Intent
+import android.content.mockedContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.FakeBroadcastDispatcher
+import com.android.systemui.broadcast.broadcastDispatcher
+import com.android.systemui.communal.domain.interactor.CommunalInteractor
+import com.android.systemui.communal.widgets.CommunalWidgetModule
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.log.logcatLogBuffer
+import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.kotlinArgumentCaptor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CommunalBackupRestoreStartableTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+
+    @Mock private lateinit var communalInteractor: CommunalInteractor
+
+    private val mapCaptor = kotlinArgumentCaptor<Map<Int, Int>>()
+
+    private lateinit var context: Context
+    private lateinit var broadcastDispatcher: FakeBroadcastDispatcher
+    private lateinit var underTest: CommunalBackupRestoreStartable
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        context = kosmos.mockedContext
+        broadcastDispatcher = kosmos.broadcastDispatcher
+
+        underTest =
+            CommunalBackupRestoreStartable(
+                broadcastDispatcher,
+                communalInteractor,
+                logcatLogBuffer("CommunalBackupRestoreStartable"),
+            )
+    }
+
+    @Test
+    fun testRestoreWidgetsUponHostRestored() =
+        testScope.runTest {
+            underTest.start()
+
+            // Verify restore widgets not called
+            verify(communalInteractor, never()).restoreWidgets(any())
+
+            // Trigger app widget host restored
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(
+                        AppWidgetManager.EXTRA_HOST_ID,
+                        CommunalWidgetModule.APP_WIDGET_HOST_ID
+                    )
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, intArrayOf(1, 2, 3))
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, intArrayOf(7, 8, 9))
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets called
+            verify(communalInteractor).restoreWidgets(mapCaptor.capture())
+            val oldToNewWidgetIdMap = mapCaptor.value
+            assertThat(oldToNewWidgetIdMap)
+                .containsExactlyEntriesIn(
+                    mapOf(
+                        Pair(1, 7),
+                        Pair(2, 8),
+                        Pair(3, 9),
+                    )
+                )
+        }
+
+    @Test
+    fun testDoNotRestoreWidgetsIfNotForCommunalWidgetHost() =
+        testScope.runTest {
+            underTest.start()
+
+            // Trigger app widget host restored, but for another host
+            val hostId = CommunalWidgetModule.APP_WIDGET_HOST_ID + 1
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(AppWidgetManager.EXTRA_HOST_ID, hostId)
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, intArrayOf(1, 2, 3))
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, intArrayOf(7, 8, 9))
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets not called
+            verify(communalInteractor, never()).restoreWidgets(any())
+        }
+
+    @Test
+    fun testAbortRestoreWidgetsIfOldToNewIdsMappingInvalid() =
+        testScope.runTest {
+            underTest.start()
+
+            // Trigger app widget host restored, but new ids list is one too many for old ids
+            val oldIds = intArrayOf(1, 2, 3)
+            val newIds = intArrayOf(6, 7, 8, 9)
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(
+                        AppWidgetManager.EXTRA_HOST_ID,
+                        CommunalWidgetModule.APP_WIDGET_HOST_ID
+                    )
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, oldIds)
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, newIds)
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets aborted
+            verify(communalInteractor).abortRestoreWidgets()
+            verify(communalInteractor, never()).restoreWidgets(any())
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
index a4c7abd..fe4d32d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
@@ -16,18 +16,23 @@
 
 package com.android.systemui.communal.data.repository
 
+import android.app.backup.BackupManager
 import android.appwidget.AppWidgetManager
 import android.appwidget.AppWidgetProviderInfo
 import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_CONFIGURATION_OPTIONAL
 import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_RECONFIGURABLE
 import android.content.ComponentName
+import android.content.applicationContext
 import android.os.UserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
+import com.android.systemui.communal.nano.CommunalHubState
+import com.android.systemui.communal.proto.toByteArray
 import com.android.systemui.communal.shared.model.CommunalWidgetContentModel
 import com.android.systemui.communal.widgets.CommunalAppWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetHost
@@ -42,6 +47,7 @@
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.mockito.withArgCaptor
 import com.google.common.truth.Truth.assertThat
 import java.util.Optional
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -68,7 +74,9 @@
     @Mock private lateinit var providerInfoA: AppWidgetProviderInfo
     @Mock private lateinit var communalWidgetHost: CommunalWidgetHost
     @Mock private lateinit var communalWidgetDao: CommunalWidgetDao
+    @Mock private lateinit var backupManager: BackupManager
 
+    private lateinit var backupUtils: CommunalBackupUtils
     private lateinit var logBuffer: LogBuffer
     private lateinit var fakeWidgets: MutableStateFlow<Map<CommunalItemRank, CommunalWidgetItem>>
 
@@ -89,6 +97,7 @@
         MockitoAnnotations.initMocks(this)
         fakeWidgets = MutableStateFlow(emptyMap())
         logBuffer = logcatLogBuffer(name = "CommunalWidgetRepoImplTest")
+        backupUtils = CommunalBackupUtils(kosmos.applicationContext)
 
         setAppWidgetIds(emptyList())
 
@@ -106,6 +115,8 @@
                 communalWidgetHost,
                 communalWidgetDao,
                 logBuffer,
+                backupManager,
+                backupUtils,
             )
     }
 
@@ -129,6 +140,9 @@
                         priority = communalItemRankEntry.rank,
                     )
                 )
+
+            // Verify backup not requested
+            verify(backupManager, never()).dataChanged()
         }
 
     @Test
@@ -152,6 +166,9 @@
 
             verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
             verify(communalWidgetDao).addWidget(id, provider, priority)
+
+            // Verify backup requested
+            verify(backupManager).dataChanged()
         }
 
     @Test
@@ -176,6 +193,9 @@
             verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
             verify(communalWidgetDao, never()).addWidget(id, provider, priority)
             verify(appWidgetHost).deleteAppWidgetId(id)
+
+            // Verify backup not requested
+            verify(backupManager, never()).dataChanged()
         }
 
     @Test
@@ -202,6 +222,9 @@
             verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
             verify(communalWidgetDao, never()).addWidget(id, provider, priority)
             verify(appWidgetHost).deleteAppWidgetId(id)
+
+            // Verify backup not requested
+            verify(backupManager, never()).dataChanged()
         }
 
     @Test
@@ -225,10 +248,13 @@
 
             verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
             verify(communalWidgetDao).addWidget(id, provider, priority)
+
+            // Verify backup requested
+            verify(backupManager).dataChanged()
         }
 
     @Test
-    fun deleteWidget_deletefromDbTrue_alsoDeleteFromHost() =
+    fun deleteWidget_deleteFromDbTrue_alsoDeleteFromHost() =
         testScope.runTest {
             val id = 1
             whenever(communalWidgetDao.deleteWidgetById(eq(id))).thenReturn(true)
@@ -237,10 +263,13 @@
 
             verify(communalWidgetDao).deleteWidgetById(id)
             verify(appWidgetHost).deleteAppWidgetId(id)
+
+            // Verify backup requested
+            verify(backupManager).dataChanged()
         }
 
     @Test
-    fun deleteWidget_deletefromDbFalse_doesNotDeleteFromHost() =
+    fun deleteWidget_deleteFromDbFalse_doesNotDeleteFromHost() =
         testScope.runTest {
             val id = 1
             whenever(communalWidgetDao.deleteWidgetById(eq(id))).thenReturn(false)
@@ -249,6 +278,9 @@
 
             verify(communalWidgetDao).deleteWidgetById(id)
             verify(appWidgetHost, never()).deleteAppWidgetId(id)
+
+            // Verify backup not requested
+            verify(backupManager, never()).dataChanged()
         }
 
     @Test
@@ -259,6 +291,147 @@
             runCurrent()
 
             verify(communalWidgetDao).updateWidgetOrder(widgetIdToPriorityMap)
+
+            // Verify backup requested
+            verify(backupManager).dataChanged()
+        }
+
+    @Test
+    fun restoreWidgets_deleteStateFileIfRestoreFails() =
+        testScope.runTest {
+            // Write a state file that is invalid, and verify it is written
+            backupUtils.writeBytesToDisk(byteArrayOf(1, 2, 3, 4, 5, 6))
+            assertThat(backupUtils.fileExists()).isTrue()
+
+            // Try to restore widgets
+            underTest.restoreWidgets(emptyMap())
+            runCurrent()
+
+            // The restore should fail, and verify that the file is deleted
+            assertThat(backupUtils.fileExists()).isFalse()
+        }
+
+    @Test
+    fun restoreWidgets_deleteStateFileAfterWidgetsRestored() =
+        testScope.runTest {
+            // Write a state file, and verify it is written
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+            assertThat(backupUtils.fileExists()).isTrue()
+
+            // Set up app widget host with widget ids
+            setAppWidgetIds(listOf(11, 12))
+
+            // Restore widgets
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify state restored
+            verify(communalWidgetDao).restoreCommunalHubState(any())
+
+            // Verify state file deleted
+            assertThat(backupUtils.fileExists()).isFalse()
+        }
+
+    @Test
+    fun restoreWidgets_restoredWidgetsNotRegisteredWithHostAreSkipped() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 12 (previously 2) is absent.
+            setAppWidgetIds(listOf(11))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify state restored, and widget 2 skipped
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(1)
+
+            val restoredWidget = restoredWidgets.first()
+            val expectedWidget = fakeState.widgets.first()
+
+            // Verify widget id is updated, and the rest remain the same as expected
+            assertThat(restoredWidget.widgetId).isEqualTo(11)
+            assertThat(restoredWidget.componentName).isEqualTo(expectedWidget.componentName)
+            assertThat(restoredWidget.rank).isEqualTo(expectedWidget.rank)
+        }
+
+    @Test
+    fun restoreWidgets_registeredWidgetsNotRestoredAreRemoved() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 13 will not be restored.
+            setAppWidgetIds(listOf(11, 12, 13))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify widget 1 and 2 are restored, and are now 11 and 12.
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(2)
+
+            val restoredWidget1 = restoredWidgets[0]
+            val expectedWidget1 = fakeState.widgets[0]
+            assertThat(restoredWidget1.widgetId).isEqualTo(11)
+            assertThat(restoredWidget1.componentName).isEqualTo(expectedWidget1.componentName)
+            assertThat(restoredWidget1.rank).isEqualTo(expectedWidget1.rank)
+
+            val restoredWidget2 = restoredWidgets[1]
+            val expectedWidget2 = fakeState.widgets[1]
+            assertThat(restoredWidget2.widgetId).isEqualTo(12)
+            assertThat(restoredWidget2.componentName).isEqualTo(expectedWidget2.componentName)
+            assertThat(restoredWidget2.rank).isEqualTo(expectedWidget2.rank)
+
+            // Verify widget 13 removed since it is not restored
+            verify(appWidgetHost).deleteAppWidgetId(13)
+        }
+
+    @Test
+    fun restoreWidgets_onlySomeWidgetsGotNewIds() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 2 gets a new id: 12, but widget 1 does
+            // not.
+            setAppWidgetIds(listOf(1, 12))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(2, 12)))
+            runCurrent()
+
+            // Verify widget 1 and 2 are restored, and are now 1 and 12.
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(2)
+
+            val restoredWidget1 = restoredWidgets[0]
+            val expectedWidget1 = fakeState.widgets[0]
+            assertThat(restoredWidget1.widgetId).isEqualTo(1)
+            assertThat(restoredWidget1.componentName).isEqualTo(expectedWidget1.componentName)
+            assertThat(restoredWidget1.rank).isEqualTo(expectedWidget1.rank)
+
+            val restoredWidget2 = restoredWidgets[1]
+            val expectedWidget2 = fakeState.widgets[1]
+            assertThat(restoredWidget2.widgetId).isEqualTo(12)
+            assertThat(restoredWidget2.componentName).isEqualTo(expectedWidget2.componentName)
+            assertThat(restoredWidget2.rank).isEqualTo(expectedWidget2.rank)
         }
 
     private fun installedProviders(providers: List<AppWidgetProviderInfo>) {
@@ -278,5 +451,22 @@
                 widgetFeatures =
                     WIDGET_FEATURE_CONFIGURATION_OPTIONAL or WIDGET_FEATURE_RECONFIGURABLE
             }
+        val fakeState =
+            CommunalHubState().apply {
+                widgets =
+                    listOf(
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 1
+                                componentName = "pk_name/fake_widget_1"
+                                rank = 1
+                            },
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 2
+                                componentName = "pk_name/fake_widget_2"
+                                rank = 2
+                            },
+                        )
+                        .toTypedArray()
+            }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/AvalancheControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/AvalancheControllerTest.kt
index 29f286f..7420ea0 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/AvalancheControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/AvalancheControllerTest.kt
@@ -248,6 +248,35 @@
     }
 
     @Test
+    fun testGetDurationMs_untrackedEntryEmptyAvalanche_useAutoDismissTime() {
+        val givenEntry = createHeadsUpEntry(id = 0)
+
+        // Nothing is showing
+        mAvalancheController.headsUpEntryShowing = null
+
+        // Nothing is next
+        mAvalancheController.clearNext()
+
+        val durationMs = mAvalancheController.getDurationMs(givenEntry, autoDismissMs = 5000)
+        Truth.assertThat(durationMs).isEqualTo(5000)
+    }
+
+    @Test
+    fun testGetDurationMs_untrackedEntryNonEmptyAvalanche_useAutoDismissTime() {
+        val givenEntry = createHeadsUpEntry(id = 0)
+
+        // Given entry not tracked
+        mAvalancheController.headsUpEntryShowing = createHeadsUpEntry(id = 1)
+
+        mAvalancheController.clearNext()
+        val nextEntry = createHeadsUpEntry(id = 2)
+        mAvalancheController.addToNext(nextEntry, runnableMock!!)
+
+        val durationMs = mAvalancheController.getDurationMs(givenEntry, autoDismissMs = 5000)
+        Truth.assertThat(durationMs).isEqualTo(5000)
+    }
+
+    @Test
     fun testGetDurationMs_lastEntry_useAutoDismissTime() {
         // Entry is showing
         val showingEntry = createHeadsUpEntry(id = 0)
@@ -261,7 +290,7 @@
     }
 
     @Test
-    fun testGetDurationMs_nextEntryLowerPriority_500() {
+    fun testGetDurationMs_nextEntryLowerPriority_5000() {
         // Entry is showing
         val showingEntry = createFsiHeadsUpEntry(id = 1)
         mAvalancheController.headsUpEntryShowing = showingEntry
diff --git a/packages/SystemUI/res/layout/combined_qs_header.xml b/packages/SystemUI/res/layout/combined_qs_header.xml
index beb481a..b9ef88e 100644
--- a/packages/SystemUI/res/layout/combined_qs_header.xml
+++ b/packages/SystemUI/res/layout/combined_qs_header.xml
@@ -129,7 +129,7 @@
             android:id="@+id/hover_system_icons_container"
             android:layout_width="wrap_content"
             android:layout_height="match_parent"
-            android:layout_gravity="right|center_vertical"
+            android:layout_gravity="end|center_vertical"
             android:gravity="center_vertical"
             android:paddingStart="@dimen/hover_system_icons_container_padding_start"
             android:paddingEnd="@dimen/hover_system_icons_container_padding_end"
diff --git a/packages/SystemUI/res/layout/screenshot_shelf.xml b/packages/SystemUI/res/layout/screenshot_shelf.xml
index 6a5b999f..26d3f43 100644
--- a/packages/SystemUI/res/layout/screenshot_shelf.xml
+++ b/packages/SystemUI/res/layout/screenshot_shelf.xml
@@ -147,4 +147,11 @@
         <include layout="@layout/screenshot_work_profile_first_run" />
         <include layout="@layout/screenshot_detection_notice" />
     </FrameLayout>
+    <ImageView
+        android:id="@+id/screenshot_flash"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent"
+        android:visibility="gone"
+        android:elevation="12dp"
+        android:src="@android:color/white"/>
 </com.android.systemui.screenshot.ui.ScreenshotShelfView>
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIApplication.java b/packages/SystemUI/src/com/android/systemui/SystemUIApplication.java
index 08d45fa..4b07402 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIApplication.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIApplication.java
@@ -299,6 +299,7 @@
         int serviceIndex = 0;
 
         do {
+            startedAny = false;
             queue = nextQueue;
             nextQueue = new ArrayDeque<>(startables.size());
 
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java b/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
index 7e94804..5c75a49 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
@@ -94,7 +94,7 @@
                     .setOneHanded(mWMComponent.getOneHanded())
                     .setBubbles(mWMComponent.getBubbles())
                     .setTaskViewFactory(mWMComponent.getTaskViewFactory())
-                    .setTransitions(mWMComponent.getTransitions())
+                    .setShellTransitions(mWMComponent.getShellTransitions())
                     .setKeyguardTransitions(mWMComponent.getKeyguardTransitions())
                     .setStartingSurface(mWMComponent.getStartingSurface())
                     .setDisplayAreaHelper(mWMComponent.getDisplayAreaHelper())
@@ -115,7 +115,7 @@
                     .setOneHanded(Optional.ofNullable(null))
                     .setBubbles(Optional.ofNullable(null))
                     .setTaskViewFactory(Optional.ofNullable(null))
-                    .setTransitions(new ShellTransitions() {})
+                    .setShellTransitions(new ShellTransitions() {})
                     .setKeyguardTransitions(new KeyguardTransitions() {})
                     .setDisplayAreaHelper(Optional.ofNullable(null))
                     .setStartingSurface(Optional.ofNullable(null))
diff --git a/packages/SystemUI/src/com/android/systemui/backup/BackupHelper.kt b/packages/SystemUI/src/com/android/systemui/backup/BackupHelper.kt
index a667de2..970699f 100644
--- a/packages/SystemUI/src/com/android/systemui/backup/BackupHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/backup/BackupHelper.kt
@@ -28,8 +28,9 @@
 import android.os.UserHandle
 import android.util.Log
 import com.android.app.tracing.traceSection
-import com.android.systemui.Flags.communalHub
 import com.android.systemui.backup.BackupHelper.Companion.ACTION_RESTORE_FINISHED
+import com.android.systemui.communal.data.backup.CommunalBackupHelper
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.domain.backup.CommunalPrefsBackupHelper
 import com.android.systemui.controls.controller.AuxiliaryPersistenceWrapper
 import com.android.systemui.controls.controller.ControlsFavoritePersistenceWrapper
@@ -59,6 +60,7 @@
             "systemui.keyguard.quickaffordance.shared_preferences"
         private const val COMMUNAL_PREFS_BACKUP_KEY =
             "systemui.communal.shared_preferences"
+        private const val COMMUNAL_STATE_BACKUP_KEY = "systemui.communal_state"
         val controlsDataLock = Any()
         const val ACTION_RESTORE_FINISHED = "com.android.systemui.backup.RESTORE_FINISHED"
         const val PERMISSION_SELF = "com.android.systemui.permission.SELF"
@@ -89,6 +91,10 @@
                     userId = userHandle.identifier,
                 )
             )
+            addHelper(
+                COMMUNAL_STATE_BACKUP_KEY,
+                CommunalBackupHelper(userHandle, CommunalBackupUtils(context = this)),
+            )
         }
     }
 
@@ -116,7 +122,7 @@
     }
 
     private fun communalEnabled(): Boolean {
-        return resources.getBoolean(R.bool.config_communalServiceEnabled) && communalHub()
+        return resources.getBoolean(R.bool.config_communalServiceEnabled)
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt b/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt
new file mode 100644
index 0000000..cdeeb6f
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt
@@ -0,0 +1,83 @@
+/*
+ * 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.communal
+
+import android.appwidget.AppWidgetManager
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import com.android.systemui.CoreStartable
+import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.communal.domain.interactor.CommunalInteractor
+import com.android.systemui.communal.widgets.CommunalWidgetModule
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.Logger
+import com.android.systemui.log.dagger.CommunalLog
+import javax.inject.Inject
+
+@SysUISingleton
+class CommunalBackupRestoreStartable
+@Inject
+constructor(
+    private val broadcastDispatcher: BroadcastDispatcher,
+    private val communalInteractor: CommunalInteractor,
+    @CommunalLog logBuffer: LogBuffer,
+) : CoreStartable, BroadcastReceiver() {
+
+    private val logger = Logger(logBuffer, TAG)
+
+    override fun start() {
+        broadcastDispatcher.registerReceiver(
+            receiver = this,
+            filter = IntentFilter(AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED),
+        )
+    }
+
+    override fun onReceive(context: Context?, intent: Intent?) {
+        if (intent == null) {
+            logger.w("On app widget host restored, but intent is null")
+            return
+        }
+
+        if (intent.action != AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED) {
+            return
+        }
+
+        val hostId = intent.getIntExtra(AppWidgetManager.EXTRA_HOST_ID, 0)
+        if (hostId != CommunalWidgetModule.APP_WIDGET_HOST_ID) {
+            return
+        }
+
+        val oldIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS)
+        val newIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS)
+
+        if (oldIds == null || newIds == null || oldIds.size != newIds.size) {
+            logger.w("On app widget host restored, but old to new ids mapping is invalid")
+            communalInteractor.abortRestoreWidgets()
+            return
+        }
+
+        val oldToNewWidgetIdMap = oldIds.zip(newIds).toMap()
+        communalInteractor.restoreWidgets(oldToNewWidgetIdMap)
+    }
+
+    companion object {
+        const val TAG = "CommunalBackupRestoreStartable"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
index 27af99e..7fa091a 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.communal.dagger
 
+import android.content.Context
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalDatabaseModule
 import com.android.systemui.communal.data.repository.CommunalMediaRepositoryModule
 import com.android.systemui.communal.data.repository.CommunalPrefsRepositoryModule
@@ -78,5 +80,13 @@
                 )
             return SceneDataSourceDelegator(applicationScope, config)
         }
+
+        @Provides
+        @SysUISingleton
+        fun providesCommunalBackupUtils(
+            @Application context: Context,
+        ): CommunalBackupUtils {
+            return CommunalBackupUtils(context)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupHelper.kt b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupHelper.kt
new file mode 100644
index 0000000..b95c966
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupHelper.kt
@@ -0,0 +1,114 @@
+/*
+ * 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.communal.data.backup
+
+import android.app.backup.BackupDataInputStream
+import android.app.backup.BackupDataOutput
+import android.app.backup.BackupHelper
+import android.os.ParcelFileDescriptor
+import android.os.UserHandle
+import android.util.Log
+import com.android.systemui.Flags.communalHub
+import com.android.systemui.communal.proto.toByteArray
+import java.io.IOException
+
+/** Helps with backup & restore of the communal hub widgets. */
+class CommunalBackupHelper(
+    private val userHandle: UserHandle,
+    private val communalBackupUtils: CommunalBackupUtils,
+) : BackupHelper {
+
+    override fun performBackup(
+        oldState: ParcelFileDescriptor?,
+        data: BackupDataOutput?,
+        newState: ParcelFileDescriptor?
+    ) {
+        if (!communalHub()) {
+            Log.d(TAG, "Skipping backup. Communal not enabled")
+            return
+        }
+
+        if (data == null) {
+            Log.e(TAG, "Backup failed. Data is null")
+            return
+        }
+
+        if (!userHandle.isSystem) {
+            Log.d(TAG, "Backup skipped for non-system user")
+            return
+        }
+
+        val state = communalBackupUtils.getCommunalHubState()
+        Log.i(TAG, "Backing up communal state: $state")
+
+        val bytes = state.toByteArray()
+        try {
+            data.writeEntityHeader(ENTITY_KEY, bytes.size)
+            data.writeEntityData(bytes, bytes.size)
+        } catch (e: IOException) {
+            Log.e(TAG, "Backup failed while writing data: ${e.localizedMessage}")
+            return
+        }
+
+        Log.i(TAG, "Backup complete")
+    }
+
+    override fun restoreEntity(data: BackupDataInputStream?) {
+        if (data == null) {
+            Log.e(TAG, "Restore failed. Data is null")
+            return
+        }
+
+        if (!userHandle.isSystem) {
+            Log.d(TAG, "Restore skipped for non-system user")
+            return
+        }
+
+        if (data.key != ENTITY_KEY) {
+            Log.d(TAG, "Restore skipped due to mismatching entity key")
+            return
+        }
+
+        val dataSize = data.size()
+        val bytes = ByteArray(dataSize)
+        try {
+            data.read(bytes, /* offset= */ 0, dataSize)
+        } catch (e: IOException) {
+            Log.e(TAG, "Restore failed while reading data: ${e.localizedMessage}")
+            return
+        }
+
+        try {
+            communalBackupUtils.writeBytesToDisk(bytes)
+        } catch (e: Exception) {
+            Log.e(TAG, "Restore failed while writing to disk: ${e.localizedMessage}")
+            return
+        }
+
+        Log.i(TAG, "Restore complete")
+    }
+
+    override fun writeNewStateDescription(newState: ParcelFileDescriptor?) {
+        // Do nothing because there is no partial backup
+    }
+
+    companion object {
+        private const val TAG = "CommunalBackupHelper"
+
+        const val ENTITY_KEY = "communal_hub_state"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt
new file mode 100644
index 0000000..a8e5174
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt
@@ -0,0 +1,113 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.communal.data.backup
+
+import android.content.Context
+import androidx.annotation.WorkerThread
+import com.android.systemui.communal.data.db.CommunalDatabase
+import com.android.systemui.communal.nano.CommunalHubState
+import java.io.File
+import java.io.FileInputStream
+import java.io.FileNotFoundException
+import java.io.FileOutputStream
+import java.io.IOException
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.runBlocking
+
+/** Utilities for communal backup and restore. */
+class CommunalBackupUtils(
+    private val context: Context,
+) {
+
+    /**
+     * Retrieves a communal hub state protobuf that represents the current state of the communal
+     * database.
+     */
+    @WorkerThread
+    fun getCommunalHubState(): CommunalHubState {
+        val database = CommunalDatabase.getInstance(context)
+        val widgetsFromDb = runBlocking { database.communalWidgetDao().getWidgets().first() }
+        val widgetsState = mutableListOf<CommunalHubState.CommunalWidgetItem>()
+        widgetsFromDb.keys.forEach { rankItem ->
+            widgetsState.add(
+                CommunalHubState.CommunalWidgetItem().apply {
+                    rank = rankItem.rank
+                    widgetId = widgetsFromDb[rankItem]!!.widgetId
+                    componentName = widgetsFromDb[rankItem]?.componentName
+                }
+            )
+        }
+        return CommunalHubState().apply { widgets = widgetsState.toTypedArray() }
+    }
+
+    /**
+     * Writes [data] to disk as a file as [FILE_NAME], overwriting existing content if any.
+     *
+     * @throws FileNotFoundException if the file exists but is a directory rather than a regular
+     *   file, does not exist but cannot be created, or cannot be opened for any other reason.
+     * @throws SecurityException if write access is denied.
+     * @throws IOException if writing fails.
+     */
+    @WorkerThread
+    fun writeBytesToDisk(data: ByteArray) {
+        val output = FileOutputStream(getFile())
+        output.write(data)
+        output.close()
+    }
+
+    /**
+     * Reads bytes from [FILE_NAME], and throws if file does not exist.
+     *
+     * @throws FileNotFoundException if file does not exist.
+     * @throws SecurityException if read access is denied.
+     * @throws IOException if reading fails.
+     */
+    @WorkerThread
+    fun readBytesFromDisk(): ByteArray {
+        val input = FileInputStream(getFile())
+        val bytes = input.readAllBytes()
+        input.close()
+
+        return bytes
+    }
+
+    /**
+     * Removes the bytes written to disk at [FILE_NAME].
+     *
+     * @return True if and only if the file is successfully deleted
+     * @throws SecurityException if permission is denied
+     */
+    @WorkerThread
+    fun clear(): Boolean {
+        return getFile().delete()
+    }
+
+    /** Whether [FILE_NAME] exists. */
+    @WorkerThread
+    fun fileExists(): Boolean {
+        return getFile().exists()
+    }
+
+    @WorkerThread
+    private fun getFile(): File {
+        return File(context.filesDir, FILE_NAME)
+    }
+
+    companion object {
+        private const val FILE_NAME = "communal_restore"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
index 595d320..3ce8109 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
@@ -16,10 +16,53 @@
 
 package com.android.systemui.communal.data.db
 
+import android.content.Context
+import androidx.annotation.VisibleForTesting
 import androidx.room.Database
+import androidx.room.Room
 import androidx.room.RoomDatabase
+import com.android.systemui.res.R
 
 @Database(entities = [CommunalWidgetItem::class, CommunalItemRank::class], version = 1)
 abstract class CommunalDatabase : RoomDatabase() {
     abstract fun communalWidgetDao(): CommunalWidgetDao
+
+    companion object {
+        private var instance: CommunalDatabase? = null
+
+        /**
+         * Gets a singleton instance of the communal database. If this is called for the first time
+         * globally, a new instance is created.
+         *
+         * @param context The context the database is created in. Only effective when a new instance
+         *   is created.
+         * @param callback An optional callback registered to the database. Only effective when a
+         *   new instance is created.
+         */
+        fun getInstance(
+            context: Context,
+            callback: Callback? = null,
+        ): CommunalDatabase {
+            if (instance == null) {
+                instance =
+                    Room.databaseBuilder(
+                            context,
+                            CommunalDatabase::class.java,
+                            context.resources.getString(R.string.config_communalDatabase)
+                        )
+                        .also { builder ->
+                            builder.fallbackToDestructiveMigration(dropAllTables = false)
+                            callback?.let { callback -> builder.addCallback(callback) }
+                        }
+                        .build()
+            }
+
+            return instance!!
+        }
+
+        @VisibleForTesting
+        fun setInstance(database: CommunalDatabase) {
+            instance = database
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabaseModule.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabaseModule.kt
index e766290..b8161fd 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabaseModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabaseModule.kt
@@ -17,10 +17,8 @@
 package com.android.systemui.communal.data.db
 
 import android.content.Context
-import androidx.room.Room
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.res.R
 import dagger.Module
 import dagger.Provides
 
@@ -33,14 +31,7 @@
             @Application context: Context,
             defaultWidgetPopulation: DefaultWidgetPopulation,
         ): CommunalDatabase {
-            return Room.databaseBuilder(
-                    context,
-                    CommunalDatabase::class.java,
-                    context.resources.getString(R.string.config_communalDatabase)
-                )
-                .fallbackToDestructiveMigration()
-                .addCallback(defaultWidgetPopulation)
-                .build()
+            return CommunalDatabase.getInstance(context, defaultWidgetPopulation)
         }
 
         @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
index 9cd77c4..d174fd1 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
@@ -23,6 +23,7 @@
 import androidx.room.RoomDatabase
 import androidx.room.Transaction
 import androidx.sqlite.db.SupportSQLiteDatabase
+import com.android.systemui.communal.nano.CommunalHubState
 import com.android.systemui.communal.widgets.CommunalWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetModule.Companion.DEFAULT_WIDGETS
 import com.android.systemui.dagger.qualifiers.Application
@@ -116,6 +117,10 @@
     @Query("UPDATE communal_item_rank_table SET rank = :order WHERE uid = :itemUid")
     fun updateItemRank(itemUid: Long, order: Int)
 
+    @Query("DELETE FROM communal_widget_table") fun clearCommunalWidgetsTable()
+
+    @Query("DELETE FROM communal_item_rank_table") fun clearCommunalItemRankTable()
+
     @Transaction
     fun updateWidgetOrder(widgetIdToPriorityMap: Map<Int, Int>) {
         widgetIdToPriorityMap.forEach { (id, priority) ->
@@ -128,9 +133,18 @@
 
     @Transaction
     fun addWidget(widgetId: Int, provider: ComponentName, priority: Int): Long {
-        return insertWidget(
+        return addWidget(
             widgetId = widgetId,
             componentName = provider.flattenToString(),
+            priority = priority,
+        )
+    }
+
+    @Transaction
+    fun addWidget(widgetId: Int, componentName: String, priority: Int): Long {
+        return insertWidget(
+            widgetId = widgetId,
+            componentName = componentName,
             itemId = insertItemRank(priority),
         )
     }
@@ -145,4 +159,13 @@
         deleteWidgets(widget)
         return true
     }
+
+    /** Wipes current database and restores the snapshot represented by [state]. */
+    @Transaction
+    fun restoreCommunalHubState(state: CommunalHubState) {
+        clearCommunalWidgetsTable()
+        clearCommunalItemRankTable()
+
+        state.widgets.forEach { addWidget(it.widgetId, it.componentName, it.rank) }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
index e395ca9..1f54e70 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
@@ -16,13 +16,17 @@
 
 package com.android.systemui.communal.data.repository
 
+import android.app.backup.BackupManager
 import android.appwidget.AppWidgetManager
 import android.content.ComponentName
 import android.os.UserHandle
 import androidx.annotation.WorkerThread
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
+import com.android.systemui.communal.nano.CommunalHubState
+import com.android.systemui.communal.proto.toCommunalHubState
 import com.android.systemui.communal.shared.model.CommunalWidgetContentModel
 import com.android.systemui.communal.widgets.CommunalAppWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetHost
@@ -69,6 +73,15 @@
      * @param widgetIdToPriorityMap mapping of the widget ids to the priority of the widget.
      */
     fun updateWidgetOrder(widgetIdToPriorityMap: Map<Int, Int>) {}
+
+    /**
+     * Restores the database by reading a state file from disk and updating the widget ids according
+     * to [oldToNewWidgetIdMap].
+     */
+    fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>)
+
+    /** Aborts the restore process and removes files from disk if necessary. */
+    fun abortRestoreWidgets()
 }
 
 @SysUISingleton
@@ -82,6 +95,8 @@
     private val communalWidgetHost: CommunalWidgetHost,
     private val communalWidgetDao: CommunalWidgetDao,
     @CommunalLog logBuffer: LogBuffer,
+    private val backupManager: BackupManager,
+    private val backupUtils: CommunalBackupUtils,
 ) : CommunalWidgetRepository {
     companion object {
         const val TAG = "CommunalWidgetRepository"
@@ -143,6 +158,7 @@
                     provider = provider,
                     priority = priority,
                 )
+                backupManager.dataChanged()
             } else {
                 appWidgetHost.deleteAppWidgetId(id)
             }
@@ -155,6 +171,7 @@
             if (communalWidgetDao.deleteWidgetById(widgetId)) {
                 appWidgetHost.deleteAppWidgetId(widgetId)
                 logger.i("Deleted widget with id $widgetId.")
+                backupManager.dataChanged()
             }
         }
     }
@@ -165,6 +182,76 @@
             logger.i({ "Updated the order of widget list with ids: $str1." }) {
                 str1 = widgetIdToPriorityMap.toString()
             }
+            backupManager.dataChanged()
+        }
+    }
+
+    override fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {
+        bgScope.launch {
+            // Read restored state file from disk
+            val state: CommunalHubState
+            try {
+                state = backupUtils.readBytesFromDisk().toCommunalHubState()
+            } catch (e: Exception) {
+                logger.e({ "Failed reading restore data from disk: $str1" }) {
+                    str1 = e.localizedMessage
+                }
+                abortRestoreWidgets()
+                return@launch
+            }
+
+            val widgetsWithHost = appWidgetHost.appWidgetIds.toList()
+            val widgetsToRemove = widgetsWithHost.toMutableList()
+
+            // Produce a new state to be restored, skipping invalid widgets
+            val newWidgets =
+                state.widgets.mapNotNull { restoredWidget ->
+                    val newWidgetId =
+                        oldToNewWidgetIdMap[restoredWidget.widgetId] ?: restoredWidget.widgetId
+
+                    // Skip if widget id is not registered with the host
+                    if (!widgetsWithHost.contains(newWidgetId)) {
+                        logger.d({
+                            "Skipped restoring widget (old:$int1 new:$int2) " +
+                                "because it is not registered with host"
+                        }) {
+                            int1 = restoredWidget.widgetId
+                            int2 = newWidgetId
+                        }
+                        return@mapNotNull null
+                    }
+
+                    widgetsToRemove.remove(newWidgetId)
+
+                    CommunalHubState.CommunalWidgetItem().apply {
+                        widgetId = newWidgetId
+                        componentName = restoredWidget.componentName
+                        rank = restoredWidget.rank
+                    }
+                }
+            val newState = CommunalHubState().apply { widgets = newWidgets.toTypedArray() }
+
+            // Restore database
+            logger.i("Restoring communal database $newState")
+            communalWidgetDao.restoreCommunalHubState(newState)
+
+            // Delete restored state file from disk
+            backupUtils.clear()
+
+            // Remove widgets from host that have not been restored
+            widgetsToRemove.forEach { widgetId ->
+                logger.i({ "Deleting widget $int1 from host since it has not been restored" }) {
+                    int1 = widgetId
+                }
+                appWidgetHost.deleteAppWidgetId(widgetId)
+            }
+        }
+    }
+
+    override fun abortRestoreWidgets() {
+        bgScope.launch {
+            logger.i("Restore widgets aborted")
+            backupUtils.clear()
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index 619e052..7448e14 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -170,6 +170,23 @@
     }
 
     /**
+     * Repopulates the communal widgets database by first reading a backed-up state from disk and
+     * updating the widget ids indicated by [oldToNewWidgetIdMap]. The backed-up state is removed
+     * from disk afterwards.
+     */
+    fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {
+        widgetRepository.restoreWidgets(oldToNewWidgetIdMap)
+    }
+
+    /**
+     * Aborts the task of restoring widgets from a backup. The backed up state stored on disk is
+     * removed.
+     */
+    fun abortRestoreWidgets() {
+        widgetRepository.abortRestoreWidgets()
+    }
+
+    /**
      * Updates the transition state of the hub [SceneTransitionLayout].
      *
      * Note that you must call is with `null` when the UI is done or risk a memory leak.
diff --git a/packages/SystemUI/src/com/android/systemui/communal/proto/CommunalHubStateExt.kt b/packages/SystemUI/src/com/android/systemui/communal/proto/CommunalHubStateExt.kt
new file mode 100644
index 0000000..2d661cd
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/proto/CommunalHubStateExt.kt
@@ -0,0 +1,34 @@
+/*
+ * 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.communal.proto
+
+import com.android.systemui.communal.nano.CommunalHubState
+import com.google.protobuf.nano.InvalidProtocolBufferNanoException
+
+/** Converts a [CommunalHubState] to bytes. */
+fun CommunalHubState.toByteArray(): ByteArray {
+    return CommunalHubState.toByteArray(this)
+}
+
+/**
+ * Converts bytes to a [CommunalHubState].
+ *
+ * @throws InvalidProtocolBufferNanoException if parsing fails.
+ */
+fun ByteArray.toCommunalHubState(): CommunalHubState {
+    return CommunalHubState.parseFrom(this)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto b/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto
new file mode 100644
index 0000000..0816259
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = "proto3";
+
+package com.android.systemui.communal;
+
+option java_multiple_files = true;
+
+// Represents the state of communal hub for backup & restore.
+message CommunalHubState {
+    // Widgets in the communal hub.
+    repeated CommunalWidgetItem widgets = 1;
+
+    // Represents a widget in the communal hub.
+    message CommunalWidgetItem {
+        // Id of the widget.
+        int32 widget_id = 1;
+
+        // Component name of the widget.
+        string component_name = 2;
+
+        // Rank or order of the widget in the communal hub.
+        int32 rank = 3;
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
index 60fb8d4..aa6516d 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
@@ -38,7 +38,7 @@
 @Module
 interface CommunalWidgetModule {
     companion object {
-        private const val APP_WIDGET_HOST_ID = 116
+        const val APP_WIDGET_HOST_ID = 116
         const val DEFAULT_WIDGETS = "default_widgets"
 
         @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
index 871ef1e..1dd3722 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SysUIComponent.java
@@ -94,7 +94,7 @@
         Builder setTaskViewFactory(Optional<TaskViewFactory> t);
 
         @BindsInstance
-        Builder setTransitions(ShellTransitions t);
+        Builder setShellTransitions(ShellTransitions t);
 
         @BindsInstance
         Builder setKeyguardTransitions(KeyguardTransitions k);
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
index 23fc8ac..593196c 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.biometrics.BiometricNotificationService
 import com.android.systemui.clipboardoverlay.ClipboardListener
 import com.android.systemui.communal.CommunalDreamStartable
+import com.android.systemui.communal.CommunalBackupRestoreStartable
 import com.android.systemui.communal.CommunalSceneStartable
 import com.android.systemui.communal.log.CommunalLoggerStartable
 import com.android.systemui.communal.widgets.CommunalAppWidgetHostStartable
@@ -342,6 +343,13 @@
 
     @Binds
     @IntoMap
+    @ClassKey(CommunalBackupRestoreStartable::class)
+    abstract fun bindCommunalBackupRestoreStartable(
+        impl: CommunalBackupRestoreStartable
+    ): CoreStartable
+
+    @Binds
+    @IntoMap
     @ClassKey(HomeControlsDreamStartable::class)
     abstract fun bindHomeControlsDreamStartable(impl: HomeControlsDreamStartable): CoreStartable
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index 93864799..3462164 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -18,6 +18,7 @@
 
 import android.app.INotificationManager;
 import android.app.Service;
+import android.app.backup.BackupManager;
 import android.content.Context;
 import android.service.dreams.IDreamManager;
 
@@ -318,6 +319,13 @@
         return new Monitor(executor, Collections.singleton(systemProcessCondition), logBuffer);
     }
 
+    /** Provides the package name for SystemUI. */
+    @SysUISingleton
+    @Provides
+    static BackupManager provideBackupManager(@Application Context context) {
+        return new BackupManager(context);
+    }
+
     @BindsOptionalOf
     abstract CommandQueue optionalCommandQueue();
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java b/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
index e04a0e5..a3cdb2e 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/WMComponent.java
@@ -94,7 +94,7 @@
     Optional<TaskViewFactory> getTaskViewFactory();
 
     @WMSingleton
-    ShellTransitions getTransitions();
+    ShellTransitions getShellTransitions();
 
     @WMSingleton
     KeyguardTransitions getKeyguardTransitions();
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaControlPanel.java b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaControlPanel.java
index bd3893b..e6c785e 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaControlPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/controller/MediaControlPanel.java
@@ -57,6 +57,7 @@
 import android.media.session.PlaybackState;
 import android.os.Process;
 import android.os.Trace;
+import android.os.UserHandle;
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.Log;
@@ -738,11 +739,11 @@
                                     mPackageName, mMediaViewHolder.getSeamlessButton());
                         } else {
                             mLogger.logOpenOutputSwitcher(mUid, mPackageName, mInstanceId);
-                            // TODO: b/321969740 - Populate the userHandle parameter. The user
-                            // handle is necessary to disambiguate the same package running on
-                            // different users.
                             mMediaOutputDialogManager.createAndShow(
-                                    mPackageName, true, mMediaViewHolder.getSeamlessButton(), null);
+                                    mPackageName,
+                                    /* aboveStatusBar */ true,
+                                    mMediaViewHolder.getSeamlessButton(),
+                                    UserHandle.getUserHandleForUid(mUid));
                         }
                     } else {
                         mLogger.logOpenOutputSwitcher(mUid, mPackageName, mInstanceId);
@@ -770,11 +771,11 @@
                                 Log.w(TAG, "Device pending intent is not an activity.");
                             }
                         } else {
-                            // TODO: b/321969740 - Populate the userHandle parameter. The user
-                            // handle is necessary to disambiguate the same package running on
-                            // different users.
                             mMediaOutputDialogManager.createAndShow(
-                                    mPackageName, true, mMediaViewHolder.getSeamlessButton(), null);
+                                    mPackageName,
+                                    /* aboveStatusBar */ true,
+                                    mMediaViewHolder.getSeamlessButton(),
+                                    UserHandle.getUserHandleForUid(mUid));
                         }
                     }
                 });
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
index 303eb78..12a3daa 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
@@ -31,6 +31,7 @@
 import android.view.WindowManager
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
+import androidx.core.animation.doOnEnd
 import com.android.internal.logging.UiEventLogger
 import com.android.systemui.log.DebugLogger.debugLog
 import com.android.systemui.res.R
@@ -109,7 +110,10 @@
     override fun updateOrientation(insets: WindowInsets) {}
 
     override fun createScreenshotDropInAnimation(screenRect: Rect, showFlash: Boolean): Animator {
-        return animationController.getEntranceAnimation()
+        val entrance = animationController.getEntranceAnimation(screenRect, showFlash)
+        // reset the timeout when animation finishes
+        entrance.doOnEnd { callbacks?.onUserInteraction() }
+        return entrance
     }
 
     override fun addQuickShareChip(quickShareAction: Notification.Action) {}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/ScreenshotAnimationController.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/ScreenshotAnimationController.kt
index 9cf8ca5..3f4f74b 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/ScreenshotAnimationController.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/ScreenshotAnimationController.kt
@@ -17,9 +17,16 @@
 package com.android.systemui.screenshot.ui
 
 import android.animation.Animator
-import android.animation.AnimatorListenerAdapter
+import android.animation.AnimatorSet
+import android.animation.ObjectAnimator
 import android.animation.ValueAnimator
+import android.graphics.PointF
+import android.graphics.Rect
+import android.util.MathUtils
 import android.view.View
+import android.view.animation.AnimationUtils
+import androidx.core.animation.doOnEnd
+import androidx.core.animation.doOnStart
 import com.android.systemui.res.R
 import kotlin.math.abs
 import kotlin.math.max
@@ -27,23 +34,57 @@
 
 class ScreenshotAnimationController(private val view: ScreenshotShelfView) {
     private var animator: Animator? = null
+    private val screenshotPreview = view.requireViewById<View>(R.id.screenshot_preview)
+    private val flashView = view.requireViewById<View>(R.id.screenshot_flash)
     private val actionContainer = view.requireViewById<View>(R.id.actions_container_background)
-
-    fun getEntranceAnimation(): Animator {
-        val animator = ValueAnimator.ofFloat(0f, 1f)
-        animator.addUpdateListener { view.alpha = it.animatedFraction }
-        animator.addListener(
-            object : AnimatorListenerAdapter() {
-                override fun onAnimationStart(animator: Animator) {
-                    view.alpha = 0f
-                }
-                override fun onAnimationEnd(animator: Animator) {
-                    view.alpha = 1f
-                }
-            }
+    private val fastOutSlowIn =
+        AnimationUtils.loadInterpolator(view.context, android.R.interpolator.fast_out_slow_in)
+    private val staticUI =
+        listOf<View>(
+            view.requireViewById(R.id.screenshot_preview_border),
+            view.requireViewById(R.id.actions_container_background),
+            view.requireViewById(R.id.screenshot_badge),
+            view.requireViewById(R.id.screenshot_dismiss_button)
         )
-        this.animator = animator
-        return animator
+
+    fun getEntranceAnimation(bounds: Rect, showFlash: Boolean): Animator {
+        val entranceAnimation = AnimatorSet()
+
+        val previewAnimator = getPreviewAnimator(bounds)
+
+        if (showFlash) {
+            val flashInAnimator =
+                ObjectAnimator.ofFloat(flashView, "alpha", 0f, 1f).apply {
+                    duration = FLASH_IN_DURATION_MS
+                    interpolator = fastOutSlowIn
+                }
+            val flashOutAnimator =
+                ObjectAnimator.ofFloat(flashView, "alpha", 1f, 0f).apply {
+                    duration = FLASH_OUT_DURATION_MS
+                    interpolator = fastOutSlowIn
+                }
+            flashInAnimator.doOnStart { flashView.visibility = View.VISIBLE }
+            flashOutAnimator.doOnEnd { flashView.visibility = View.GONE }
+            entranceAnimation.play(flashOutAnimator).after(flashInAnimator)
+            entranceAnimation.play(previewAnimator).with(flashOutAnimator)
+            entranceAnimation.doOnStart { screenshotPreview.visibility = View.INVISIBLE }
+        }
+
+        val fadeInAnimator = ValueAnimator.ofFloat(0f, 1f)
+        fadeInAnimator.addUpdateListener {
+            for (child in staticUI) {
+                child.alpha = it.animatedValue as Float
+            }
+        }
+        entranceAnimation.play(fadeInAnimator).after(previewAnimator)
+        entranceAnimation.doOnStart {
+            for (child in staticUI) {
+                child.alpha = 0f
+            }
+        }
+
+        this.animator = entranceAnimation
+        return entranceAnimation
     }
 
     fun getSwipeReturnAnimation(): Animator {
@@ -81,11 +122,49 @@
         animator?.cancel()
     }
 
+    private fun getPreviewAnimator(bounds: Rect): Animator {
+        val targetPosition = Rect()
+        screenshotPreview.getHitRect(targetPosition)
+        val startXScale = bounds.width() / targetPosition.width().toFloat()
+        val startYScale = bounds.height() / targetPosition.height().toFloat()
+        val startPos = PointF(bounds.exactCenterX(), bounds.exactCenterY())
+        val endPos = PointF(targetPosition.exactCenterX(), targetPosition.exactCenterY())
+
+        val previewYAnimator =
+            ValueAnimator.ofFloat(startPos.y, endPos.y).apply {
+                duration = PREVIEW_Y_ANIMATION_DURATION_MS
+                interpolator = fastOutSlowIn
+            }
+        previewYAnimator.addUpdateListener {
+            val progress = it.animatedValue as Float
+            screenshotPreview.y = progress - screenshotPreview.height / 2f
+        }
+        // scale animation starts/finishes at the same time as x placement
+        val previewXAndScaleAnimator =
+            ValueAnimator.ofFloat(0f, 1f).apply {
+                duration = PREVIEW_X_ANIMATION_DURATION_MS
+                interpolator = fastOutSlowIn
+            }
+        previewXAndScaleAnimator.addUpdateListener {
+            val t = it.animatedFraction
+            screenshotPreview.scaleX = MathUtils.lerp(startXScale, 1f, t)
+            screenshotPreview.scaleY = MathUtils.lerp(startYScale, 1f, t)
+            screenshotPreview.x =
+                MathUtils.lerp(startPos.x, endPos.x, t) - screenshotPreview.width / 2f
+        }
+
+        val previewAnimator = AnimatorSet()
+        previewAnimator.play(previewXAndScaleAnimator).with(previewYAnimator)
+
+        previewAnimator.doOnStart { screenshotPreview.visibility = View.VISIBLE }
+        return previewAnimator
+    }
+
     private fun getAdjustedVelocity(requestedVelocity: Float?): Float {
         return if (requestedVelocity == null) {
             val isLTR = view.resources.configuration.layoutDirection == View.LAYOUT_DIRECTION_LTR
             // dismiss to the left in LTR locales, to the right in RTL
-            if (isLTR) -1 * MINIMUM_VELOCITY else MINIMUM_VELOCITY
+            if (isLTR) -MINIMUM_VELOCITY else MINIMUM_VELOCITY
         } else {
             sign(requestedVelocity) * max(MINIMUM_VELOCITY, abs(requestedVelocity))
         }
@@ -93,5 +172,9 @@
 
     companion object {
         private const val MINIMUM_VELOCITY = 1.5f // pixels per second
+        private const val FLASH_IN_DURATION_MS: Long = 133
+        private const val FLASH_OUT_DURATION_MS: Long = 217
+        private const val PREVIEW_X_ANIMATION_DURATION_MS: Long = 234
+        private const val PREVIEW_Y_ANIMATION_DURATION_MS: Long = 500
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/SwipeGestureListener.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/SwipeGestureListener.kt
index 7274757..61d4489 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/SwipeGestureListener.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/SwipeGestureListener.kt
@@ -24,7 +24,7 @@
 
 class SwipeGestureListener(
     private val view: View,
-    private val onDismiss: (Float) -> Unit,
+    private val onDismiss: (Float?) -> Unit,
     private val onCancel: () -> Unit
 ) {
     private val velocityTracker = VelocityTracker.obtain()
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
index 3e35295..3376b8c 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.screenshot.ui.SwipeGestureListener
 import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
 import com.android.systemui.util.children
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 
 object ScreenshotShelfViewBinder {
@@ -60,7 +61,8 @@
             onDismissalRequested(ScreenshotEvent.SCREENSHOT_EXPLICIT_DISMISSAL, null)
         }
 
-        view.repeatWhenAttached {
+        // use immediate dispatcher to ensure screenshot bitmap is set before animation
+        view.repeatWhenAttached(Dispatchers.Main.immediate) {
             lifecycleScope.launch {
                 repeatOnLifecycle(Lifecycle.State.STARTED) {
                     launch {
@@ -96,9 +98,9 @@
                             // ID is unique.
                             val newIds = visibleActions.map { it.id }
 
-                            for (view in actionsContainer.children.toList()) {
-                                if (view.tag !in newIds) {
-                                    actionsContainer.removeView(view)
+                            for (child in actionsContainer.children.toList()) {
+                                if (child.tag !in newIds) {
+                                    actionsContainer.removeView(child)
                                 }
                             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/settings/brightness/BrightnessController.java b/packages/SystemUI/src/com/android/systemui/settings/brightness/BrightnessController.java
index 8397d9f..37f2a21 100644
--- a/packages/SystemUI/src/com/android/systemui/settings/brightness/BrightnessController.java
+++ b/packages/SystemUI/src/com/android/systemui/settings/brightness/BrightnessController.java
@@ -47,7 +47,9 @@
 import com.android.internal.display.BrightnessSynchronizer;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
+import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.RestrictedLockUtilsInternal;
+import com.android.systemui.Flags;
 import com.android.systemui.dagger.qualifiers.Background;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.settings.DisplayTracker;
@@ -370,10 +372,18 @@
         mBackgroundHandler.post(new Runnable() {
             @Override
             public void run() {
-                mControl.setEnforcedAdmin(
+                int userId = mUserTracker.getUserId();
+                RestrictedLockUtils.EnforcedAdmin enforcedAdmin =
                         RestrictedLockUtilsInternal.checkIfRestrictionEnforced(mContext,
                                 UserManager.DISALLOW_CONFIG_BRIGHTNESS,
-                                mUserTracker.getUserId()));
+                                userId);
+                if (Flags.enforceBrightnessBaseUserRestriction() && enforcedAdmin == null
+                        && RestrictedLockUtilsInternal.hasBaseUserRestriction(mContext,
+                        UserManager.DISALLOW_CONFIG_BRIGHTNESS,
+                        userId)) {
+                    enforcedAdmin = new RestrictedLockUtils.EnforcedAdmin();
+                }
+                mControl.setEnforcedAdmin(enforcedAdmin);
             }
         });
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
index 9bb860f..594c191 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/dagger/CentralSurfacesDependenciesModule.java
@@ -41,7 +41,6 @@
 import com.android.systemui.shade.ShadeSurface;
 import com.android.systemui.shade.ShadeSurfaceImpl;
 import com.android.systemui.shade.carrier.ShadeCarrierGroupController;
-import com.android.systemui.startable.Dependencies;
 import com.android.systemui.statusbar.CommandQueue;
 import com.android.systemui.statusbar.NotificationClickNotifier;
 import com.android.systemui.statusbar.NotificationMediaManager;
@@ -53,7 +52,6 @@
 import com.android.systemui.statusbar.notification.collection.NotifCollection;
 import com.android.systemui.statusbar.notification.collection.NotifPipeline;
 import com.android.systemui.statusbar.notification.collection.render.NotificationVisibilityProvider;
-import com.android.systemui.statusbar.phone.CentralSurfaces;
 import com.android.systemui.statusbar.phone.CentralSurfacesImpl;
 import com.android.systemui.statusbar.phone.ManagedProfileController;
 import com.android.systemui.statusbar.phone.ManagedProfileControllerImpl;
@@ -70,8 +68,6 @@
 import dagger.multibindings.ClassKey;
 import dagger.multibindings.IntoMap;
 
-import java.util.Set;
-
 import javax.inject.Provider;
 
 /**
@@ -161,15 +157,6 @@
     CoreStartable bindsStartStatusBarStateController(StatusBarStateControllerImpl sbsc);
 
     /** */
-    @Provides
-    @IntoMap
-    @Dependencies
-    @ClassKey(SysuiStatusBarStateController.class)
-    static Set<Class<? extends CoreStartable>> providesStatusBarStateControllerDeps() {
-        return Set.of(CentralSurfaces.class);
-    }
-
-    /** */
     @Binds
     StatusBarIconController provideStatusBarIconController(
             StatusBarIconControllerImpl controllerImpl);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
index f573530..36c23d3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/AvalancheController.kt
@@ -153,27 +153,45 @@
             // Use default duration, like we did before AvalancheController existed
             return autoDismissMs
         }
-
         val showingList: MutableList<HeadsUpEntry> = mutableListOf()
-        headsUpEntryShowing?.let { showingList.add(it) }
-
+        if (headsUpEntryShowing != null) {
+            showingList.add(headsUpEntryShowing!!)
+        }
         nextList.sort()
         val entryList = showingList + nextList
-        val thisEntryIndex = entryList.indexOf(entry)
+        if (entryList.isEmpty()) {
+            log { "No avalanche HUNs, use default ms: $autoDismissMs" }
+            return autoDismissMs
+        }
+        // entryList.indexOf(entry) returns -1 even when the entry is in entryList
+        var thisEntryIndex = -1
+        for ((i, e) in entryList.withIndex()) {
+            if (e == entry) {
+                thisEntryIndex = i
+            }
+        }
+        if (thisEntryIndex == -1) {
+            log { "Untracked entry, use default ms: $autoDismissMs" }
+            return autoDismissMs
+        }
         val nextEntryIndex = thisEntryIndex + 1
 
         // If last entry, use default duration
         if (nextEntryIndex >= entryList.size) {
+            log { "Last entry, use default ms: $autoDismissMs" }
             return autoDismissMs
         }
         val nextEntry = entryList[nextEntryIndex]
         if (nextEntry.compareNonTimeFields(entry) == -1) {
             // Next entry is higher priority
+            log { "Next entry is higher priority: 500ms" }
             return 500
         } else if (nextEntry.compareNonTimeFields(entry) == 0) {
             // Next entry is same priority
+            log { "Next entry is same priority: 1000ms" }
             return 1000
         } else {
+            log { "Next entry is lower priority, use default ms: $autoDismissMs" }
             return autoDismissMs
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputActionsInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputActionsInteractor.kt
index eebb6fb..22c0530 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputActionsInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputActionsInteractor.kt
@@ -21,7 +21,7 @@
 import com.android.systemui.animation.DialogTransitionAnimator
 import com.android.systemui.animation.Expandable
 import com.android.systemui.media.dialog.MediaOutputDialogManager
-import com.android.systemui.volume.panel.component.mediaoutput.shared.model.SessionWithPlayback
+import com.android.systemui.volume.panel.component.mediaoutput.shared.model.SessionWithPlaybackState
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
 import javax.inject.Inject
 
@@ -33,10 +33,10 @@
     private val mediaOutputDialogManager: MediaOutputDialogManager,
 ) {
 
-    fun onBarClick(sessionWithPlayback: SessionWithPlayback?, expandable: Expandable) {
-        if (sessionWithPlayback?.playback?.isActive == true) {
+    fun onBarClick(sessionWithPlaybackState: SessionWithPlaybackState?, expandable: Expandable) {
+        if (sessionWithPlaybackState?.isPlaybackActive == true) {
             mediaOutputDialogManager.createAndShowWithController(
-                sessionWithPlayback.session.packageName,
+                sessionWithPlaybackState.session.packageName,
                 false,
                 expandable.dialogController()
             )
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
index 41ad035..83b8029 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/domain/interactor/MediaOutputInteractor.kt
@@ -94,7 +94,7 @@
 
     /** Currently connected [MediaDevice]. */
     val currentConnectedDevice: Flow<MediaDevice?> =
-        localMediaRepository.flatMapLatest { it.currentConnectedDevice }
+        localMediaRepository.flatMapLatest { it.currentConnectedDevice }.distinctUntilChanged()
 
     private suspend fun getApplicationLabel(packageName: String): CharSequence? {
         return try {
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlayback.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlaybackState.kt
similarity index 86%
rename from packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlayback.kt
rename to packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlaybackState.kt
index c4476fc..bac969d 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlayback.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/shared/model/SessionWithPlaybackState.kt
@@ -16,9 +16,7 @@
 
 package com.android.systemui.volume.panel.component.mediaoutput.shared.model
 
-import android.media.session.PlaybackState
-
-data class SessionWithPlayback(
+data class SessionWithPlaybackState(
     val session: MediaDeviceSession,
-    val playback: PlaybackState,
+    val isPlaybackActive: Boolean,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
index f19fa20..d60d981 100644
--- a/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/volume/panel/component/mediaoutput/ui/viewmodel/MediaOutputViewModel.kt
@@ -25,7 +25,7 @@
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaDeviceSessionInteractor
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaOutputActionsInteractor
 import com.android.systemui.volume.panel.component.mediaoutput.domain.interactor.MediaOutputInteractor
-import com.android.systemui.volume.panel.component.mediaoutput.shared.model.SessionWithPlayback
+import com.android.systemui.volume.panel.component.mediaoutput.shared.model.SessionWithPlaybackState
 import com.android.systemui.volume.panel.dagger.scope.VolumePanelScope
 import com.android.systemui.volume.panel.shared.model.Result
 import com.android.systemui.volume.panel.ui.VolumePanelUiEvent
@@ -35,10 +35,9 @@
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.mapNotNull
 import kotlinx.coroutines.flow.stateIn
 
 /** Models the UI of the Media Output Volume Panel component. */
@@ -55,20 +54,19 @@
     private val uiEventLogger: UiEventLogger,
 ) {
 
-    private val sessionWithPlayback: StateFlow<Result<SessionWithPlayback?>> =
+    private val sessionWithPlaybackState: StateFlow<Result<SessionWithPlaybackState?>> =
         interactor.defaultActiveMediaSession
             .flatMapLatest { session ->
                 if (session == null) {
-                    flowOf(Result.Data<SessionWithPlayback?>(null))
+                    flowOf(Result.Data<SessionWithPlaybackState?>(null))
                 } else {
-                    mediaDeviceSessionInteractor
-                        .playbackState(session)
-                        .map { playback ->
-                            playback?.let {
-                                Result.Data<SessionWithPlayback?>(SessionWithPlayback(session, it))
-                            }
+                    mediaDeviceSessionInteractor.playbackState(session).mapNotNull { playback ->
+                        playback?.let {
+                            Result.Data<SessionWithPlaybackState?>(
+                                SessionWithPlaybackState(session, playback.isActive())
+                            )
                         }
-                        .filterNotNull()
+                    }
                 }
             }
             .stateIn(
@@ -78,14 +76,14 @@
             )
 
     val connectedDeviceViewModel: StateFlow<ConnectedDeviceViewModel?> =
-        combine(sessionWithPlayback, interactor.currentConnectedDevice) {
+        combine(sessionWithPlaybackState, interactor.currentConnectedDevice) {
                 mediaDeviceSession,
                 currentConnectedDevice ->
                 if (mediaDeviceSession !is Result.Data) {
                     return@combine null
                 }
                 ConnectedDeviceViewModel(
-                    if (mediaDeviceSession.data?.playback?.isActive == true) {
+                    if (mediaDeviceSession.data?.isPlaybackActive == true) {
                         context.getString(
                             R.string.media_output_label_title,
                             mediaDeviceSession.data.session.appLabel
@@ -103,19 +101,16 @@
             )
 
     val deviceIconViewModel: StateFlow<DeviceIconViewModel?> =
-        combine(sessionWithPlayback, interactor.currentConnectedDevice) {
+        combine(sessionWithPlaybackState, interactor.currentConnectedDevice) {
                 mediaDeviceSession,
                 currentConnectedDevice ->
                 if (mediaDeviceSession !is Result.Data) {
                     return@combine null
                 }
-                if (mediaDeviceSession.data?.playback?.isActive == true) {
-                    val icon =
-                        currentConnectedDevice?.icon?.let { Icon.Loaded(it, null) }
-                            ?: Icon.Resource(
-                                com.android.internal.R.drawable.ic_bt_headphones_a2dp,
-                                null
-                            )
+                val icon: Icon =
+                    currentConnectedDevice?.icon?.let { Icon.Loaded(it, null) }
+                        ?: Icon.Resource(R.drawable.ic_media_home_devices, null)
+                if (mediaDeviceSession.data?.isPlaybackActive == true) {
                     DeviceIconViewModel.IsPlaying(
                         icon = icon,
                         iconColor =
@@ -125,7 +120,7 @@
                     )
                 } else {
                     DeviceIconViewModel.IsNotPlaying(
-                        icon = Icon.Resource(R.drawable.ic_media_home_devices, null),
+                        icon = icon,
                         iconColor =
                             Color.Attribute(
                                 com.android.internal.R.attr.materialColorOnSurfaceVariant
@@ -143,7 +138,7 @@
 
     fun onBarClick(expandable: Expandable) {
         uiEventLogger.log(VolumePanelUiEvent.VOLUME_PANEL_MEDIA_OUTPUT_CLICKED)
-        val result = sessionWithPlayback.value
+        val result = sessionWithPlaybackState.value
         actionsInteractor.onBarClick((result as? Result.Data)?.data, expandable)
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
new file mode 100644
index 0000000..7094848
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.communal.data.backup
+
+import android.app.backup.BackupDataInput
+import android.app.backup.BackupDataInputStream
+import android.app.backup.BackupDataOutput
+import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import androidx.room.Room
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.Flags
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.backup.CommunalBackupUtilsTest.Companion.represents
+import com.android.systemui.communal.data.backup.CommunalBackupUtilsTest.FakeWidgetMetadata
+import com.android.systemui.communal.data.db.CommunalDatabase
+import com.android.systemui.communal.data.db.CommunalWidgetDao
+import com.android.systemui.communal.proto.toCommunalHubState
+import com.android.systemui.lifecycle.InstantTaskExecutorRule
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import java.io.FileInputStream
+import java.io.FileOutputStream
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CommunalBackupHelperTest : SysuiTestCase() {
+    @JvmField @Rule val instantTaskExecutor = InstantTaskExecutorRule()
+
+    private lateinit var database: CommunalDatabase
+    private lateinit var dao: CommunalWidgetDao
+    private lateinit var backupUtils: CommunalBackupUtils
+
+    // Temporary file used for storing backed-up data.
+    private lateinit var backupDataFile: File
+
+    private lateinit var underTest: CommunalBackupHelper
+
+    @Before
+    fun setup() {
+        database =
+            Room.inMemoryDatabaseBuilder(context, CommunalDatabase::class.java)
+                .allowMainThreadQueries()
+                .build()
+        CommunalDatabase.setInstance(database)
+
+        dao = database.communalWidgetDao()
+        backupUtils = CommunalBackupUtils(context)
+
+        backupDataFile = File(context.cacheDir, "backup_data_file")
+
+        underTest = CommunalBackupHelper(UserHandle.SYSTEM, backupUtils)
+    }
+
+    @After
+    fun teardown() {
+        backupDataFile.delete()
+        database.close()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_COMMUNAL_HUB)
+    fun backupAndRestoreCommunalHub() {
+        val expectedWidgets = setUpDatabase()
+
+        underTest.performBackup(oldState = null, data = getBackupDataOutput(), newState = null)
+        underTest.restoreEntity(getBackupDataInputStream())
+
+        // Verify restored state matches backed-up state
+        val restoredState = backupUtils.readBytesFromDisk().toCommunalHubState()
+        val restoredWidgets = restoredState.widgets.toList()
+        assertThat(restoredWidgets)
+            .comparingElementsUsing(represents)
+            .containsExactlyElementsIn(expectedWidgets)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_COMMUNAL_HUB)
+    fun backup_skippedWhenCommunalDisabled() {
+        setUpDatabase()
+
+        underTest.performBackup(oldState = null, data = getBackupDataOutput(), newState = null)
+
+        // Verify nothing written to the backup
+        assertThat(backupDataFile.length()).isEqualTo(0)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_COMMUNAL_HUB)
+    fun backup_skippedForNonSystemUser() {
+        setUpDatabase()
+
+        val helper = CommunalBackupHelper(UserHandle.CURRENT, backupUtils)
+        helper.performBackup(oldState = null, data = getBackupDataOutput(), newState = null)
+
+        // Verify nothing written to the backup
+        assertThat(backupDataFile.length()).isEqualTo(0)
+    }
+
+    private fun setUpDatabase(): List<FakeWidgetMetadata> {
+        return listOf(
+                FakeWidgetMetadata(11, "com.android.fakePackage1/fakeWidget1", 3),
+                FakeWidgetMetadata(12, "com.android.fakePackage2/fakeWidget2", 2),
+                FakeWidgetMetadata(13, "com.android.fakePackage3/fakeWidget3", 1),
+            )
+            .onEach { dao.addWidget(it.widgetId, it.componentName, it.rank) }
+    }
+
+    private fun getBackupDataInputStream(): BackupDataInputStream {
+        val input = BackupDataInput(FileInputStream(backupDataFile).fd).apply { readNextHeader() }
+
+        // Construct BackupDataInputStream using reflection because its constructor is package
+        // private
+        val inputStream = BackupDataInputStream::class.constructors.first().call(input)
+
+        // Set key
+        with(inputStream.javaClass.getDeclaredField("key")) {
+            isAccessible = true
+            set(inputStream, input.key)
+        }
+
+        // Set dataSize
+        with(inputStream.javaClass.getDeclaredField("dataSize")) {
+            isAccessible = true
+            set(inputStream, input.dataSize)
+        }
+
+        return inputStream
+    }
+
+    private fun getBackupDataOutput(): BackupDataOutput {
+        return BackupDataOutput(FileOutputStream(backupDataFile).fd)
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt
new file mode 100644
index 0000000..c3849f9
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt
@@ -0,0 +1,153 @@
+/*
+ * 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.communal.data.backup
+
+import androidx.room.Room
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.db.CommunalDatabase
+import com.android.systemui.communal.data.db.CommunalWidgetDao
+import com.android.systemui.communal.nano.CommunalHubState
+import com.android.systemui.lifecycle.InstantTaskExecutorRule
+import com.google.common.truth.Correspondence
+import com.google.common.truth.Truth.assertThat
+import java.io.FileNotFoundException
+import java.nio.charset.Charset
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CommunalBackupUtilsTest : SysuiTestCase() {
+    @JvmField @Rule val instantTaskExecutor = InstantTaskExecutorRule()
+
+    private lateinit var database: CommunalDatabase
+    private lateinit var dao: CommunalWidgetDao
+    private lateinit var underTest: CommunalBackupUtils
+
+    @Before
+    fun setup() {
+        database =
+            Room.inMemoryDatabaseBuilder(context, CommunalDatabase::class.java)
+                .allowMainThreadQueries()
+                .build()
+        CommunalDatabase.setInstance(database)
+
+        dao = database.communalWidgetDao()
+        underTest = CommunalBackupUtils(context)
+    }
+
+    @After
+    fun teardown() {
+        database.close()
+        underTest.clear()
+    }
+
+    @Test
+    fun getCommunalHubState_returnsExpectedWidgets() {
+        // Set up database
+        val expectedWidgets =
+            listOf(
+                FakeWidgetMetadata(11, "com.android.fakePackage1/fakeWidget1", 3),
+                FakeWidgetMetadata(12, "com.android.fakePackage2/fakeWidget2", 2),
+                FakeWidgetMetadata(13, "com.android.fakePackage3/fakeWidget3", 1),
+            )
+        expectedWidgets.forEach { dao.addWidget(it.widgetId, it.componentName, it.rank) }
+
+        // Get communal hub state
+        val state = underTest.getCommunalHubState()
+        val actualWidgets = state.widgets.toList()
+
+        // Verify the state contains widgets as expected
+        assertThat(actualWidgets)
+            .comparingElementsUsing(represents)
+            .containsExactlyElementsIn(expectedWidgets)
+    }
+
+    @Test
+    fun write_existingContentIsOverwritten() {
+        // Write old data
+        val dataToWrite = "I am old data. Erase me."
+        underTest.writeBytesToDisk(dataToWrite.toByteArray(Charset.defaultCharset()))
+
+        // Verify old data written
+        var dataRead = underTest.readBytesFromDisk().toString(Charset.defaultCharset())
+        assertThat(dataRead).isEqualTo(dataToWrite)
+
+        // Write new data
+        val newDataToWrite = "I am new data."
+        underTest.writeBytesToDisk(newDataToWrite.toByteArray(Charset.defaultCharset()))
+
+        // Verify new data overwrites old
+        dataRead = underTest.readBytesFromDisk().toString(Charset.defaultCharset())
+        assertThat(dataRead).isEqualTo(newDataToWrite)
+    }
+
+    @Test(expected = FileNotFoundException::class)
+    fun read_fileNotFoundException() {
+        underTest.readBytesFromDisk()
+    }
+
+    @Test(expected = FileNotFoundException::class)
+    fun clear_returnsTrueWhenFileDeleted() {
+        // Write bytes to disk
+        underTest.writeBytesToDisk(byteArrayOf(1, 2, 3))
+
+        assertThat(underTest.clear()).isTrue()
+
+        // Verify a read after that throws a FileNotFoundException
+        underTest.readBytesFromDisk()
+    }
+
+    @Test
+    fun clear_returnsFalseWhenFileDoesNotExist() {
+        assertThat(underTest.clear()).isFalse()
+    }
+
+    @Test
+    fun fileExists() {
+        assertThat(underTest.fileExists()).isFalse()
+
+        underTest.writeBytesToDisk(byteArrayOf(1, 2, 3))
+        assertThat(underTest.fileExists()).isTrue()
+
+        underTest.clear()
+        assertThat(underTest.fileExists()).isFalse()
+    }
+
+    data class FakeWidgetMetadata(val widgetId: Int, val componentName: String, val rank: Int)
+
+    companion object {
+        /**
+         * A comparator for whether a [CommunalHubState.CommunalWidgetItem] represents a
+         * [FakeWidgetMetadata]
+         */
+        val represents: Correspondence<CommunalHubState.CommunalWidgetItem, FakeWidgetMetadata> =
+            Correspondence.from(
+                { actual, expected ->
+                    actual?.widgetId == expected?.widgetId &&
+                        actual?.componentName == expected?.componentName &&
+                        actual?.rank == expected?.rank
+                },
+                "represents",
+            )
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
index 20dd913..f77c7a6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
@@ -21,6 +21,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.nano.CommunalHubState
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.lifecycle.InstantTaskExecutorRule
 import com.google.common.truth.Truth.assertThat
@@ -224,6 +225,42 @@
                 .inOrder()
         }
 
+    @Test
+    fun restoreCommunalHubState() =
+        testScope.runTest {
+            // Set up db
+            listOf(widgetInfo1, widgetInfo2, widgetInfo3).forEach { addWidget(it) }
+
+            // Restore db to fake state
+            communalWidgetDao.restoreCommunalHubState(fakeState)
+
+            // Verify db matches new state
+            val expected = mutableMapOf<CommunalItemRank, CommunalWidgetItem>()
+            fakeState.widgets.forEachIndexed { index, fakeWidget ->
+                // Auto-generated uid continues after the initial 3 widgets and starts at 4
+                val uid = index + 4L
+                val rank = CommunalItemRank(uid = uid, rank = fakeWidget.rank)
+                val widget =
+                    CommunalWidgetItem(
+                        uid = uid,
+                        widgetId = fakeWidget.widgetId,
+                        componentName = fakeWidget.componentName,
+                        itemId = rank.uid,
+                    )
+                expected[rank] = widget
+            }
+            val widgets by collectLastValue(communalWidgetDao.getWidgets())
+            assertThat(widgets).containsExactlyEntriesIn(expected)
+        }
+
+    private fun addWidget(metadata: FakeWidgetMetadata, priority: Int? = null) {
+        communalWidgetDao.addWidget(
+            widgetId = metadata.widgetId,
+            provider = metadata.provider,
+            priority = priority ?: metadata.priority,
+        )
+    }
+
     data class FakeWidgetMetadata(
         val widgetId: Int,
         val provider: ComponentName,
@@ -273,5 +310,22 @@
                 componentName = widgetInfo3.provider.flattenToString(),
                 itemId = communalItemRankEntry3.uid,
             )
+        val fakeState =
+            CommunalHubState().apply {
+                widgets =
+                    listOf(
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 1
+                                componentName = "pk_name/fake_widget_1"
+                                rank = 1
+                            },
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 2
+                                componentName = "pk_name/fake_widget_2"
+                                rank = 2
+                            },
+                        )
+                        .toTypedArray()
+            }
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
index 4ed6fe2..329c0f1 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
@@ -46,6 +46,10 @@
         _communalWidgets.value = _communalWidgets.value.filter { it.appWidgetId != widgetId }
     }
 
+    override fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {}
+
+    override fun abortRestoreWidgets() {}
+
     private fun onConfigured(id: Int, providerInfo: AppWidgetProviderInfo, priority: Int) {
         _communalWidgets.value += listOf(CommunalWidgetContentModel(id, providerInfo, priority))
     }
diff --git a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
index 9b4d378..243e224 100644
--- a/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
+++ b/ravenwood/texts/ravenwood-annotation-allowed-classes.txt
@@ -266,6 +266,8 @@
 android.telephony.ModemActivityInfo
 android.telephony.ServiceState
 
+android.os.connectivity.WifiActivityEnergyInfo
+
 com.android.server.LocalServices
 
 com.android.internal.util.BitUtils
diff --git a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
index 7f64786..9a73a2d 100644
--- a/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
+++ b/services/contextualsearch/java/com/android/server/contextualsearch/ContextualSearchManagerService.java
@@ -30,7 +30,6 @@
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 
 import static com.android.server.contextualsearch.flags.Flags.enableExcludePersistentUi;
-
 import static com.android.server.wm.ActivityTaskManagerInternal.ASSIST_KEY_CONTENT;
 import static com.android.server.wm.ActivityTaskManagerInternal.ASSIST_KEY_STRUCTURE;
 
@@ -65,6 +64,7 @@
 import android.os.ResultReceiver;
 import android.os.ServiceManager;
 import android.os.ShellCallback;
+import android.os.SystemClock;
 import android.util.Log;
 import android.util.Slog;
 import android.view.IWindowManager;
@@ -246,6 +246,9 @@
         if (DEBUG_USER) Log.d(TAG, "Launch component: " + launchIntent.getComponent());
         launchIntent.addFlags(FLAG_ACTIVITY_NEW_TASK | FLAG_ACTIVITY_NO_ANIMATION
                 | FLAG_ACTIVITY_NO_USER_ACTION);
+        launchIntent.putExtra(
+                ContextualSearchManager.EXTRA_INVOCATION_TIME_MS,
+                SystemClock.uptimeMillis());
         launchIntent.putExtra(ContextualSearchManager.EXTRA_ENTRYPOINT, entrypoint);
         launchIntent.putExtra(ContextualSearchManager.EXTRA_TOKEN, mToken);
         boolean isAssistDataAllowed = mAtmInternal.isAssistDataAllowed();
diff --git a/services/core/java/com/android/server/TelephonyRegistry.java b/services/core/java/com/android/server/TelephonyRegistry.java
index bd67cf420..e171064 100644
--- a/services/core/java/com/android/server/TelephonyRegistry.java
+++ b/services/core/java/com/android/server/TelephonyRegistry.java
@@ -425,6 +425,8 @@
     private int[] mSCBMReason;
     private boolean[] mSCBMStarted;
 
+    private boolean[] mCarrierRoamingNtnMode = null;
+
     /**
      * Per-phone map of precise data connection state. The key of the map is the pair of transport
      * type and APN setting. This is the cache to prevent redundant callbacks to the listeners.
@@ -723,6 +725,7 @@
             mECBMStarted = copyOf(mECBMStarted, mNumPhones);
             mSCBMReason = copyOf(mSCBMReason, mNumPhones);
             mSCBMStarted = copyOf(mSCBMStarted, mNumPhones);
+            mCarrierRoamingNtnMode = copyOf(mCarrierRoamingNtnMode, mNumPhones);
             // ds -> ss switch.
             if (mNumPhones < oldNumPhones) {
                 cutListToSize(mCellInfo, mNumPhones);
@@ -781,6 +784,7 @@
                 mECBMStarted[i] = false;
                 mSCBMReason[i] = TelephonyManager.STOP_REASON_UNKNOWN;
                 mSCBMStarted[i] = false;
+                mCarrierRoamingNtnMode[i] = false;
             }
         }
     }
@@ -854,6 +858,7 @@
         mECBMStarted = new boolean[numPhones];
         mSCBMReason = new int[numPhones];
         mSCBMStarted = new boolean[numPhones];
+        mCarrierRoamingNtnMode = new boolean[numPhones];
 
         for (int i = 0; i < numPhones; i++) {
             mCallState[i] =  TelephonyManager.CALL_STATE_IDLE;
@@ -897,6 +902,7 @@
             mECBMStarted[i] = false;
             mSCBMReason[i] = TelephonyManager.STOP_REASON_UNKNOWN;
             mSCBMStarted[i] = false;
+            mCarrierRoamingNtnMode[i] = false;
         }
 
         mAppOps = mContext.getSystemService(AppOpsManager.class);
@@ -1523,6 +1529,14 @@
                         remove(r.binder);
                     }
                 }
+                if (events.contains(TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED)) {
+                    try {
+                        r.callback.onCarrierRoamingNtnModeChanged(
+                                mCarrierRoamingNtnMode[r.phoneId]);
+                    } catch (RemoteException ex) {
+                        remove(r.binder);
+                    }
+                }
             }
         }
     }
@@ -3512,6 +3526,41 @@
         handleRemoveListLocked();
     }
 
+    /**
+     * Notify external listeners that carrier roaming non-terrestrial network mode changed.
+     * @param subId subscription ID.
+     * @param active {@code true} If the device is connected to carrier roaming
+     *                           non-terrestrial network or was connected within the
+     *                           {CarrierConfigManager#KEY_SATELLITE_CONNECTION_HYSTERESIS_SEC_INT}
+     *                           duration, {code false} otherwise.
+     */
+    public void notifyCarrierRoamingNtnModeChanged(int subId, boolean active) {
+        if (!checkNotifyPermission("notifyCarrierRoamingNtnModeChanged")) {
+            return;
+        }
+
+        if (VDBG) {
+            log("notifyCarrierRoamingNtnModeChanged: subId=" + subId + " active=" + active);
+        }
+
+        synchronized (mRecords) {
+            int phoneId = getPhoneIdFromSubId(subId);
+            mCarrierRoamingNtnMode[phoneId] = active;
+            for (Record r : mRecords) {
+                if (r.matchTelephonyCallbackEvent(
+                        TelephonyCallback.EVENT_CARRIER_ROAMING_NTN_MODE_CHANGED)
+                        && idMatch(r, subId, phoneId)) {
+                    try {
+                        r.callback.onCarrierRoamingNtnModeChanged(active);
+                    } catch (RemoteException ex) {
+                        mRemoveList.add(r.binder);
+                    }
+                }
+            }
+            handleRemoveListLocked();
+        }
+    }
+
     @NeverCompile // Avoid size overhead of debugging code.
     @Override
     public void dump(FileDescriptor fd, PrintWriter writer, String[] args) {
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index f98799d..4f84149 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -133,6 +133,7 @@
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.PowerStatsUidResolver;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
+import com.android.server.power.stats.WifiPowerStatsProcessor;
 import com.android.server.power.stats.wakeups.CpuWakeupStats;
 
 import java.io.File;
@@ -412,6 +413,8 @@
                 com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodCpu);
         final long powerStatsThrottlePeriodMobileRadio = context.getResources().getInteger(
                 com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodMobileRadio);
+        final long powerStatsThrottlePeriodWifi = context.getResources().getInteger(
+                com.android.internal.R.integer.config_defaultPowerStatsThrottlePeriodWifi);
         mBatteryStatsConfig =
                 new BatteryStatsImpl.BatteryStatsConfig.Builder()
                         .setResetOnUnplugHighBatteryLevel(resetOnUnplugHighBatteryLevel)
@@ -422,6 +425,9 @@
                         .setPowerStatsThrottlePeriodMillis(
                                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                                 powerStatsThrottlePeriodMobileRadio)
+                        .setPowerStatsThrottlePeriodMillis(
+                                BatteryConsumer.POWER_COMPONENT_WIFI,
+                                powerStatsThrottlePeriodWifi)
                         .build();
         mPowerStatsUidResolver = new PowerStatsUidResolver();
         mStats = new BatteryStatsImpl(mBatteryStatsConfig, Clock.SYSTEM_CLOCK, mMonotonicClock,
@@ -480,6 +486,7 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
                         new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -490,9 +497,21 @@
                         AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
                 .setProcessor(
                         new MobileRadioPowerStatsProcessor(mPowerProfile));
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_PHONE,
                         BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .setProcessor(new PhoneCallPowerStatsProcessor());
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN,
+                        AggregatedPowerStatsConfig.STATE_PROCESS_STATE)
+                .setProcessor(
+                        new WifiPowerStatsProcessor(mPowerProfile));
         return config;
     }
 
@@ -518,14 +537,22 @@
     public void systemServicesReady() {
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
-        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
-                Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
                 BatteryConsumer.POWER_COMPONENT_CPU,
                 Flags.streamlinedBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
+                Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_WIFI,
+                Flags.streamlinedConnectivityBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_WIFI,
+                Flags.streamlinedConnectivityBatteryStats());
+
         mWorker.systemServicesReady();
         mStats.systemServicesReady(mContext);
         mCpuWakeupStats.systemServicesReady();
diff --git a/services/core/java/com/android/server/am/CachedAppOptimizer.java b/services/core/java/com/android/server/am/CachedAppOptimizer.java
index fa6b54b..211f952 100644
--- a/services/core/java/com/android/server/am/CachedAppOptimizer.java
+++ b/services/core/java/com/android/server/am/CachedAppOptimizer.java
@@ -2371,6 +2371,14 @@
                     return;
                 }
 
+                if (opt.shouldNotFreeze()) {
+                    if (DEBUG_FREEZER) {
+                        Slog.d(TAG_AM, "Skipping freeze because process is marked "
+                                + "should not be frozen");
+                    }
+                    return;
+                }
+
                 if (pid == 0 || opt.isFrozen()) {
                     // Already frozen or not a real process, either one being
                     // launched or one being killed
diff --git a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClient.java b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClient.java
index acc086b..00a5911 100644
--- a/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClient.java
+++ b/services/core/java/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClient.java
@@ -35,6 +35,7 @@
 import android.hardware.biometrics.BiometricManager.Authenticators;
 import android.hardware.biometrics.BiometricSourceType;
 import android.hardware.biometrics.common.ICancellationSignal;
+import android.hardware.biometrics.common.OperationState;
 import android.hardware.biometrics.events.AuthenticationAcquiredInfo;
 import android.hardware.biometrics.events.AuthenticationErrorInfo;
 import android.hardware.biometrics.events.AuthenticationFailedInfo;
@@ -344,6 +345,12 @@
             if (session.hasContextMethods()) {
                 try {
                     session.getSession().onContextChanged(ctx);
+                    // TODO(b/317414324): Deprecate setIgnoreDisplayTouches
+                    if (ctx.operationState != null && ctx.operationState.getTag()
+                            == OperationState.fingerprintOperationState) {
+                        session.getSession().setIgnoreDisplayTouches(ctx.operationState
+                                .getFingerprintOperationState().isHardwareIgnoringTouches);
+                    }
                 } catch (RemoteException e) {
                     Slog.e(TAG, "Unable to notify context changed", e);
                 }
diff --git a/services/core/java/com/android/server/display/DisplayBrightnessState.java b/services/core/java/com/android/server/display/DisplayBrightnessState.java
index e3e3be2..baa154d 100644
--- a/services/core/java/com/android/server/display/DisplayBrightnessState.java
+++ b/services/core/java/com/android/server/display/DisplayBrightnessState.java
@@ -45,6 +45,7 @@
     private final float mCustomAnimationRate;
 
     private final BrightnessEvent mBrightnessEvent;
+    private final int mBrightnessAdjustmentFlag;
 
     private DisplayBrightnessState(Builder builder) {
         mBrightness = builder.getBrightness();
@@ -58,6 +59,7 @@
         mCustomAnimationRate = builder.getCustomAnimationRate();
         mShouldUpdateScreenBrightnessSetting = builder.shouldUpdateScreenBrightnessSetting();
         mBrightnessEvent = builder.getBrightnessEvent();
+        mBrightnessAdjustmentFlag = builder.getBrightnessAdjustmentFlag();
     }
 
     /**
@@ -138,6 +140,14 @@
         return mBrightnessEvent;
     }
 
+    /**
+     * Gets the flag representing the reason for the brightness adjustment. This can be
+     * automatic(e.g. because of the change in the lux), or user initiated(e.g. moving the slider)
+     */
+    public int getBrightnessAdjustmentFlag() {
+        return mBrightnessAdjustmentFlag;
+    }
+
     @Override
     public String toString() {
         StringBuilder stringBuilder = new StringBuilder("DisplayBrightnessState:");
@@ -157,6 +167,7 @@
                 .append(mShouldUpdateScreenBrightnessSetting);
         stringBuilder.append("\n    mBrightnessEvent:")
                 .append(Objects.toString(mBrightnessEvent, "null"));
+        stringBuilder.append("\n    mBrightnessAdjustmentFlag:").append(mBrightnessAdjustmentFlag);
         return stringBuilder.toString();
     }
 
@@ -187,7 +198,8 @@
                 && mCustomAnimationRate == otherState.getCustomAnimationRate()
                 && mShouldUpdateScreenBrightnessSetting
                     == otherState.shouldUpdateScreenBrightnessSetting()
-                && Objects.equals(mBrightnessEvent, otherState.getBrightnessEvent());
+                && Objects.equals(mBrightnessEvent, otherState.getBrightnessEvent())
+                && mBrightnessAdjustmentFlag == otherState.getBrightnessAdjustmentFlag();
     }
 
     @Override
@@ -195,7 +207,7 @@
         return Objects.hash(mBrightness, mSdrBrightness, mBrightnessReason,
                 mShouldUseAutoBrightness, mIsSlowChange, mMaxBrightness, mMinBrightness,
                 mCustomAnimationRate,
-                mShouldUpdateScreenBrightnessSetting, mBrightnessEvent);
+                mShouldUpdateScreenBrightnessSetting, mBrightnessEvent, mBrightnessAdjustmentFlag);
     }
 
     /**
@@ -222,6 +234,8 @@
 
         private BrightnessEvent mBrightnessEvent;
 
+        public int mBrightnessAdjustmentFlag = 0;
+
         /**
          * Create a builder starting with the values from the specified {@link
          * DisplayBrightnessState}.
@@ -242,6 +256,7 @@
             builder.setShouldUpdateScreenBrightnessSetting(
                     state.shouldUpdateScreenBrightnessSetting());
             builder.setBrightnessEvent(state.getBrightnessEvent());
+            builder.setBrightnessAdjustmentFlag(state.getBrightnessAdjustmentFlag());
             return builder;
         }
 
@@ -418,7 +433,6 @@
             return new DisplayBrightnessState(this);
         }
 
-
         /**
          * This is used to get the BrightnessEvent object from its builder
          */
@@ -434,5 +448,21 @@
             mBrightnessEvent = brightnessEvent;
             return this;
         }
+
+        /**
+         * This is used to get the brightness adjustment flag from its builder
+         */
+        public int getBrightnessAdjustmentFlag() {
+            return mBrightnessAdjustmentFlag;
+        }
+
+
+        /**
+         * This is used to set the brightness adjustment flag
+         */
+        public Builder setBrightnessAdjustmentFlag(int brightnessAdjustmentFlag) {
+            mBrightnessAdjustmentFlag = brightnessAdjustmentFlag;
+            return this;
+        }
     }
 }
diff --git a/services/core/java/com/android/server/display/DisplayPowerController.java b/services/core/java/com/android/server/display/DisplayPowerController.java
index e88ace1..c5d8686 100644
--- a/services/core/java/com/android/server/display/DisplayPowerController.java
+++ b/services/core/java/com/android/server/display/DisplayPowerController.java
@@ -1345,6 +1345,7 @@
         boolean slowChange = displayBrightnessState.isSlowChange();
         // custom transition duration
         float customAnimationRate = displayBrightnessState.getCustomAnimationRate();
+        int brightnessAdjustmentFlags = displayBrightnessState.getBrightnessAdjustmentFlag();
         final boolean userSetBrightnessChanged =
                 mDisplayBrightnessController.getIsUserSetScreenBrightnessUpdated();
         if (displayBrightnessState.getBrightnessEvent() != null) {
@@ -1392,15 +1393,10 @@
                 displayBrightnessState.shouldUpdateScreenBrightnessSetting();
         float currentBrightnessSetting = mDisplayBrightnessController.getCurrentBrightness();
         // Apply auto-brightness.
-        int brightnessAdjustmentFlags = 0;
         // All the conditions inside this if block will be moved to AutomaticBrightnessStrategy
         if (mFlags.isRefactorDisplayPowerControllerEnabled()
                 && displayBrightnessState.getBrightnessReason().getReason()
                         == BrightnessReason.REASON_AUTOMATIC) {
-            brightnessAdjustmentFlags =
-                    mAutomaticBrightnessStrategy.getAutoBrightnessAdjustmentReasonsFlags();
-            updateScreenBrightnessSetting = currentBrightnessSetting != brightnessState;
-            mBrightnessReasonTemp.setReason(BrightnessReason.REASON_AUTOMATIC);
             if (mScreenOffBrightnessSensorController != null) {
                 mScreenOffBrightnessSensorController.setLightSensorEnabled(false);
             }
diff --git a/services/core/java/com/android/server/display/brightness/DisplayBrightnessController.java b/services/core/java/com/android/server/display/brightness/DisplayBrightnessController.java
index 8b3e4a4..6a88a76 100644
--- a/services/core/java/com/android/server/display/brightness/DisplayBrightnessController.java
+++ b/services/core/java/com/android/server/display/brightness/DisplayBrightnessController.java
@@ -146,7 +146,8 @@
         synchronized (mLock) {
             mDisplayBrightnessStrategy = mDisplayBrightnessStrategySelector.selectStrategy(
                     constructStrategySelectionRequest(displayPowerRequest, targetDisplayState));
-            state = mDisplayBrightnessStrategy.updateBrightness(displayPowerRequest);
+            state = mDisplayBrightnessStrategy
+                        .updateBrightness(constructStrategyExecutionRequest(displayPowerRequest));
         }
 
         // This is a temporary measure until AutomaticBrightnessStrategy works as a traditional
@@ -550,4 +551,10 @@
         return new StrategySelectionRequest(displayPowerRequest, targetDisplayState,
                 lastUserSetScreenBrightness, userSetBrightnessChanged);
     }
+
+    private StrategyExecutionRequest constructStrategyExecutionRequest(
+            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+        float currentScreenBrightness = getCurrentBrightness();
+        return new StrategyExecutionRequest(displayPowerRequest, currentScreenBrightness);
+    }
 }
diff --git a/services/core/java/com/android/server/display/brightness/StrategyExecutionRequest.java b/services/core/java/com/android/server/display/brightness/StrategyExecutionRequest.java
new file mode 100644
index 0000000..82c0bbf
--- /dev/null
+++ b/services/core/java/com/android/server/display/brightness/StrategyExecutionRequest.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.display.brightness;
+
+import android.hardware.display.DisplayManagerInternal;
+
+import java.util.Objects;
+
+/**
+ * A wrapper class to encapsulate the request to execute the selected strategy
+ */
+public final class StrategyExecutionRequest {
+    // The request to change the associated display's state and brightness
+    private final DisplayManagerInternal.DisplayPowerRequest mDisplayPowerRequest;
+
+    private final float mCurrentScreenBrightness;
+
+    public StrategyExecutionRequest(DisplayManagerInternal.DisplayPowerRequest displayPowerRequest,
+            float currentScreenBrightness) {
+        mDisplayPowerRequest = displayPowerRequest;
+        mCurrentScreenBrightness = currentScreenBrightness;
+    }
+
+    public DisplayManagerInternal.DisplayPowerRequest getDisplayPowerRequest() {
+        return mDisplayPowerRequest;
+    }
+
+    public float getCurrentScreenBrightness() {
+        return mCurrentScreenBrightness;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (!(obj instanceof StrategyExecutionRequest)) {
+            return false;
+        }
+        StrategyExecutionRequest other = (StrategyExecutionRequest) obj;
+        return Objects.equals(mDisplayPowerRequest, other.getDisplayPowerRequest())
+                && mCurrentScreenBrightness == other.getCurrentScreenBrightness();
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(mDisplayPowerRequest, mCurrentScreenBrightness);
+    }
+}
diff --git a/services/core/java/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategy.java
index 4be7332..5c4fa842 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategy.java
@@ -20,7 +20,6 @@
 import android.annotation.Nullable;
 import android.content.Context;
 import android.hardware.display.BrightnessConfiguration;
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 import android.os.UserHandle;
 import android.provider.Settings;
@@ -32,6 +31,7 @@
 import com.android.server.display.brightness.BrightnessEvent;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -96,13 +96,21 @@
     // want to re-evaluate the auto-brightness state
     private boolean mIsConfigured;
 
-    public AutomaticBrightnessStrategy(Context context, int displayId) {
+    private Injector mInjector;
+
+    @VisibleForTesting
+    AutomaticBrightnessStrategy(Context context, int displayId, Injector injector) {
         super(context, displayId);
         mContext = context;
         mDisplayId = displayId;
         mAutoBrightnessAdjustment = getAutoBrightnessAdjustmentSetting();
         mPendingAutoBrightnessAdjustment = PowerManager.BRIGHTNESS_INVALID_FLOAT;
         mTemporaryAutoBrightnessAdjustment = PowerManager.BRIGHTNESS_INVALID_FLOAT;
+        mInjector = (injector == null) ? new RealInjector() : injector;
+    }
+
+    public AutomaticBrightnessStrategy(Context context, int displayId) {
+        this(context, displayId, null);
     }
 
     /**
@@ -252,10 +260,10 @@
 
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         BrightnessReason brightnessReason = new BrightnessReason();
         brightnessReason.setReason(BrightnessReason.REASON_AUTOMATIC);
-        BrightnessEvent brightnessEvent = new BrightnessEvent(mDisplayId);
+        BrightnessEvent brightnessEvent = mInjector.getBrightnessEvent(mDisplayId);
         float brightness = getAutomaticScreenBrightness(brightnessEvent);
         return new DisplayBrightnessState.Builder()
                 .setBrightness(brightness)
@@ -265,6 +273,9 @@
                 .setIsSlowChange(hasAppliedAutoBrightness()
                         && !getAutoBrightnessAdjustmentChanged())
                 .setBrightnessEvent(brightnessEvent)
+                .setBrightnessAdjustmentFlag(mAutoBrightnessAdjustmentReasonsFlags)
+                .setShouldUpdateScreenBrightnessSetting(
+                        brightness != strategyExecutionRequest.getCurrentScreenBrightness())
                 .build();
     }
 
@@ -360,13 +371,6 @@
     }
 
     /**
-     * Gets the auto-brightness adjustment flag change reason
-     */
-    public int getAutoBrightnessAdjustmentReasonsFlags() {
-        return mAutoBrightnessAdjustmentReasonsFlags;
-    }
-
-    /**
      * Returns if the auto brightness has been applied
      */
     public boolean hasAppliedAutoBrightness() {
@@ -497,4 +501,15 @@
                 Settings.System.SCREEN_AUTO_BRIGHTNESS_ADJ, 0.0f, UserHandle.USER_CURRENT);
         return Float.isNaN(adj) ? 0.0f : BrightnessUtils.clampBrightnessAdjustment(adj);
     }
+
+    @VisibleForTesting
+    interface Injector {
+        BrightnessEvent getBrightnessEvent(int displayId);
+    }
+
+    static class RealInjector implements Injector {
+        public BrightnessEvent getBrightnessEvent(int displayId) {
+            return new BrightnessEvent(displayId);
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/display/brightness/strategy/BoostBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/BoostBrightnessStrategy.java
index 9c1acea..009a47a 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/BoostBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/BoostBrightnessStrategy.java
@@ -16,12 +16,12 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -37,7 +37,7 @@
     // Set the brightness to the maximum value when display brightness boost is requested
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         DisplayBrightnessState displayBrightnessState =
diff --git a/services/core/java/com/android/server/display/brightness/strategy/DisplayBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/DisplayBrightnessStrategy.java
index 61dd6d5..e96b83a 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/DisplayBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/DisplayBrightnessStrategy.java
@@ -17,9 +17,9 @@
 package com.android.server.display.brightness.strategy;
 
 import android.annotation.NonNull;
-import android.hardware.display.DisplayManagerInternal;
 
 import com.android.server.display.DisplayBrightnessState;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -33,10 +33,10 @@
     /**
      * Decides the DisplayBrightnessState that the system should change to.
      *
-     * @param displayPowerRequest The request to evaluate the updated brightness
+     * @param strategyExecutionRequest The request to evaluate the updated brightness
      */
     DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest);
+            StrategyExecutionRequest strategyExecutionRequest);
 
     /**
      * Returns the name of the Strategy
diff --git a/services/core/java/com/android/server/display/brightness/strategy/DozeBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/DozeBrightnessStrategy.java
index 1f7efd1..2b493f3 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/DozeBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/DozeBrightnessStrategy.java
@@ -16,11 +16,10 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
-
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -32,11 +31,12 @@
 
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         return BrightnessUtils.constructDisplayBrightnessState(BrightnessReason.REASON_DOZE,
-                displayPowerRequest.dozeScreenBrightness, displayPowerRequest.dozeScreenBrightness,
+                strategyExecutionRequest.getDisplayPowerRequest().dozeScreenBrightness,
+                strategyExecutionRequest.getDisplayPowerRequest().dozeScreenBrightness,
                 getName());
     }
 
diff --git a/services/core/java/com/android/server/display/brightness/strategy/FollowerBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/FollowerBrightnessStrategy.java
index baac276..5a07ce2 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/FollowerBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/FollowerBrightnessStrategy.java
@@ -16,12 +16,12 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -49,7 +49,7 @@
 
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         return BrightnessUtils.constructDisplayBrightnessState(BrightnessReason.REASON_FOLLOWER,
diff --git a/services/core/java/com/android/server/display/brightness/strategy/InvalidBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/InvalidBrightnessStrategy.java
index dfd47a3..9dc6cff 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/InvalidBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/InvalidBrightnessStrategy.java
@@ -16,12 +16,12 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -32,7 +32,7 @@
 public class InvalidBrightnessStrategy implements DisplayBrightnessStrategy {
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         return BrightnessUtils.constructDisplayBrightnessState(BrightnessReason.REASON_UNKNOWN,
                 PowerManager.BRIGHTNESS_INVALID_FLOAT, PowerManager.BRIGHTNESS_INVALID_FLOAT,
                 getName());
diff --git a/services/core/java/com/android/server/display/brightness/strategy/OffloadBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/OffloadBrightnessStrategy.java
index 60e5eb0..b46873a 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/OffloadBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/OffloadBrightnessStrategy.java
@@ -16,11 +16,11 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 import com.android.server.display.feature.DisplayManagerFlags;
 
@@ -42,7 +42,7 @@
 
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         float offloadBrightness = mOffloadScreenBrightness;
         if (mDisplayManagerFlags.isRefactorDisplayPowerControllerEnabled()) {
             // We reset the offload brightness to invalid so that there is no stale value lingering
diff --git a/services/core/java/com/android/server/display/brightness/strategy/OverrideBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/OverrideBrightnessStrategy.java
index 9605a88..a2982b1 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/OverrideBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/OverrideBrightnessStrategy.java
@@ -16,11 +16,10 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
-
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -31,12 +30,13 @@
 public class OverrideBrightnessStrategy implements DisplayBrightnessStrategy {
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         return BrightnessUtils.constructDisplayBrightnessState(BrightnessReason.REASON_OVERRIDE,
-                displayPowerRequest.screenBrightnessOverride,
-                displayPowerRequest.screenBrightnessOverride, getName());
+                strategyExecutionRequest.getDisplayPowerRequest().screenBrightnessOverride,
+                strategyExecutionRequest.getDisplayPowerRequest()
+                        .screenBrightnessOverride, getName());
     }
 
     @Override
diff --git a/services/core/java/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategy.java
index c9dc298..6a3162c 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategy.java
@@ -16,12 +16,12 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -32,7 +32,7 @@
 public class ScreenOffBrightnessStrategy implements DisplayBrightnessStrategy {
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         return BrightnessUtils.constructDisplayBrightnessState(BrightnessReason.REASON_SCREEN_OFF,
diff --git a/services/core/java/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategy.java b/services/core/java/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategy.java
index 6a691d1..6b8817a 100644
--- a/services/core/java/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategy.java
+++ b/services/core/java/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategy.java
@@ -16,12 +16,12 @@
 
 package com.android.server.display.brightness.strategy;
 
-import android.hardware.display.DisplayManagerInternal;
 import android.os.PowerManager;
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
 import com.android.server.display.brightness.BrightnessUtils;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 
 import java.io.PrintWriter;
@@ -43,7 +43,7 @@
     // WindowManager or based on the display state.
     @Override
     public DisplayBrightnessState updateBrightness(
-            DisplayManagerInternal.DisplayPowerRequest displayPowerRequest) {
+            StrategyExecutionRequest strategyExecutionRequest) {
         // Todo(b/241308599): Introduce a validator class and add validations before setting
         // the brightness
         DisplayBrightnessState displayBrightnessState =
diff --git a/services/core/java/com/android/server/location/contexthub/ContextHubTransactionManager.java b/services/core/java/com/android/server/location/contexthub/ContextHubTransactionManager.java
index b18871c..a0dbfa0 100644
--- a/services/core/java/com/android/server/location/contexthub/ContextHubTransactionManager.java
+++ b/services/core/java/com/android/server/location/contexthub/ContextHubTransactionManager.java
@@ -28,6 +28,7 @@
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Random;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.ScheduledThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -77,9 +78,12 @@
     private final AtomicInteger mNextAvailableId = new AtomicInteger();
 
     /**
-     * The next available message sequence number
+     * The next available message sequence number. We choose a random
+     * number to start with to avoid collisions and limit the bound to
+     * half of the max value to avoid overflow.
      */
-    private final AtomicInteger mNextAvailableMessageSequenceNumber = new AtomicInteger();
+    private final AtomicInteger mNextAvailableMessageSequenceNumber =
+            new AtomicInteger(new Random().nextInt(Integer.MAX_VALUE / 2));
 
     /*
      * An executor and the future object for scheduling timeout timers
diff --git a/services/core/java/com/android/server/media/MediaRoute2ProviderServiceProxy.java b/services/core/java/com/android/server/media/MediaRoute2ProviderServiceProxy.java
index db83d4b..a7fd750 100644
--- a/services/core/java/com/android/server/media/MediaRoute2ProviderServiceProxy.java
+++ b/services/core/java/com/android/server/media/MediaRoute2ProviderServiceProxy.java
@@ -233,7 +233,7 @@
             return false;
         }
         boolean bindDueToManagerScan =
-                mIsManagerScanning && Flags.enablePreventionOfManagerScansWhenNoAppsScan();
+                mIsManagerScanning && !Flags.enablePreventionOfManagerScansWhenNoAppsScan();
         if (!getSessionInfos().isEmpty() || bindDueToManagerScan) {
             // We bind if any manager is scanning (regardless of whether an app is scanning) to give
             // the opportunity for providers to publish routing sessions that were established
diff --git a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
index e50189b..869b89a 100644
--- a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
+++ b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
@@ -3411,7 +3411,14 @@
             for (RouterRecord activeRouterRecord : activeRouterRecords) {
                 RouteDiscoveryPreference preference = activeRouterRecord.mDiscoveryPreference;
                 preferredFeatures.addAll(preference.getPreferredFeatures());
-                if (activeRouterRecord.isActivelyScanning()) {
+
+                boolean isRouterRecordActivelyScanning =
+                        Flags.enablePreventionOfManagerScansWhenNoAppsScan()
+                                ? (activeRouterRecord.isActivelyScanning() || shouldForceActiveScan)
+                                        && !preference.getPreferredFeatures().isEmpty()
+                                : activeRouterRecord.isActivelyScanning();
+
+                if (isRouterRecordActivelyScanning) {
                     activeScan = true;
                     activelyScanningPackages.add(activeRouterRecord.mPackageName);
                 }
diff --git a/services/core/java/com/android/server/media/MediaSessionRecord.java b/services/core/java/com/android/server/media/MediaSessionRecord.java
index 3d68555..a3c5d2d 100644
--- a/services/core/java/com/android/server/media/MediaSessionRecord.java
+++ b/services/core/java/com/android/server/media/MediaSessionRecord.java
@@ -66,7 +66,6 @@
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
-import android.os.DeadObjectException;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
@@ -90,7 +89,6 @@
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Objects;
@@ -191,6 +189,9 @@
     private final ForegroundServiceDelegationOptions mForegroundServiceDelegationOptions;
 
     private final Object mLock = new Object();
+    // This field is partially guarded by mLock. Writes and non-atomic iterations (for example:
+    // index-based-iterations) must be guarded by mLock. But it is safe to acquire an iterator
+    // without acquiring mLock.
     private final CopyOnWriteArrayList<ISessionControllerCallbackHolder>
             mControllerCallbackHolders = new CopyOnWriteArrayList<>();
 
@@ -886,24 +887,9 @@
             }
             playbackState = mPlaybackState;
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onPlaybackStateChanged(playbackState);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushPlaybackStateUpdate", holder,
-                        e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushPlaybackStateUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushPlaybackStateUpdate",
+                holder -> holder.mCallback.onPlaybackStateChanged(playbackState));
     }
 
     private void pushMetadataUpdate() {
@@ -914,23 +900,8 @@
             }
             metadata = mMetadata;
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onMetadataChanged(metadata);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushMetadataUpdate", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushMetadataUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushMetadataUpdate", holder -> holder.mCallback.onMetadataChanged(metadata));
     }
 
     private void pushQueueUpdate() {
@@ -941,31 +912,18 @@
             }
             toSend = mQueue == null ? null : new ArrayList<>(mQueue);
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            ParceledListSlice<QueueItem> parcelableQueue = null;
-            if (toSend != null) {
-                parcelableQueue = new ParceledListSlice<>(toSend);
-                // Limit the size of initial Parcel to prevent binder buffer overflow
-                // as onQueueChanged is an async binder call.
-                parcelableQueue.setInlineCountLimit(1);
-            }
-
-            try {
-                holder.mCallback.onQueueChanged(parcelableQueue);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushQueueUpdate", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushQueueUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushQueueUpdate",
+                holder -> {
+                    ParceledListSlice<QueueItem> parcelableQueue = null;
+                    if (toSend != null) {
+                        parcelableQueue = new ParceledListSlice<>(toSend);
+                        // Limit the size of initial Parcel to prevent binder buffer overflow
+                        // as onQueueChanged is an async binder call.
+                        parcelableQueue.setInlineCountLimit(1);
+                    }
+                    holder.mCallback.onQueueChanged(parcelableQueue);
+                });
     }
 
     private void pushQueueTitleUpdate() {
@@ -976,23 +934,8 @@
             }
             queueTitle = mQueueTitle;
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onQueueTitleChanged(queueTitle);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushQueueTitleUpdate", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushQueueTitleUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushQueueTitleUpdate", holder -> holder.mCallback.onQueueTitleChanged(queueTitle));
     }
 
     private void pushExtrasUpdate() {
@@ -1003,23 +946,8 @@
             }
             extras = mExtras;
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onExtrasChanged(extras);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushExtrasUpdate", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushExtrasUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushExtrasUpdate", holder -> holder.mCallback.onExtrasChanged(extras));
     }
 
     private void pushVolumeUpdate() {
@@ -1030,23 +958,8 @@
             }
             info = getVolumeAttributes();
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onVolumeInfoChanged(info);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushVolumeUpdate", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushVolumeUpdate", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders(
+                "pushVolumeUpdate", holder -> holder.mCallback.onVolumeInfoChanged(info));
     }
 
     private void pushEvent(String event, Bundle data) {
@@ -1055,23 +968,7 @@
                 return;
             }
         }
-        Collection<ISessionControllerCallbackHolder> deadCallbackHolders = null;
-        for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
-            try {
-                holder.mCallback.onEvent(event, data);
-            } catch (DeadObjectException e) {
-                if (deadCallbackHolders == null) {
-                    deadCallbackHolders = new ArrayList<>();
-                }
-                deadCallbackHolders.add(holder);
-                logCallbackException("Removing dead callback in pushEvent", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushEvent", holder, e);
-            }
-        }
-        if (deadCallbackHolders != null) {
-            removeControllerHoldersSafely(deadCallbackHolders);
-        }
+        performOnCallbackHolders("pushEvent", holder -> holder.mCallback.onEvent(event, data));
     }
 
     private void pushSessionDestroyed() {
@@ -1082,20 +979,37 @@
                 return;
             }
         }
+        performOnCallbackHolders(
+                "pushSessionDestroyed",
+                holder -> {
+                    holder.mCallback.asBinder().unlinkToDeath(holder.mDeathMonitor, 0);
+                    holder.mCallback.onSessionDestroyed();
+                });
+        // After notifying clear all listeners
+        synchronized (mLock) {
+            mControllerCallbackHolders.clear();
+        }
+    }
+
+    private interface ControllerCallbackCall {
+
+        void performOn(ISessionControllerCallbackHolder holder) throws RemoteException;
+    }
+
+    private void performOnCallbackHolders(String operationName, ControllerCallbackCall call) {
+        ArrayList<ISessionControllerCallbackHolder> deadCallbackHolders = new ArrayList<>();
         for (ISessionControllerCallbackHolder holder : mControllerCallbackHolders) {
             try {
-                holder.mCallback.asBinder().unlinkToDeath(holder.mDeathMonitor, 0);
-                holder.mCallback.onSessionDestroyed();
-            } catch (NoSuchElementException e) {
-                logCallbackException("error unlinking to binder death", holder, e);
-            } catch (DeadObjectException e) {
-                logCallbackException("Removing dead callback in pushSessionDestroyed", holder, e);
-            } catch (RemoteException e) {
-                logCallbackException("unexpected exception in pushSessionDestroyed", holder, e);
+                call.performOn(holder);
+            } catch (RemoteException | NoSuchElementException exception) {
+                deadCallbackHolders.add(holder);
+                logCallbackException(
+                        "Exception while executing: " + operationName, holder, exception);
             }
         }
-        // After notifying clear all listeners
-        removeControllerHoldersSafely(null);
+        synchronized (mLock) {
+            mControllerCallbackHolders.removeAll(deadCallbackHolders);
+        }
     }
 
     private PlaybackState getStateWithUpdatedPosition() {
@@ -1143,17 +1057,6 @@
         return -1;
     }
 
-    private void removeControllerHoldersSafely(
-            Collection<ISessionControllerCallbackHolder> holders) {
-        synchronized (mLock) {
-            if (holders == null) {
-                mControllerCallbackHolders.clear();
-            } else {
-                mControllerCallbackHolders.removeAll(holders);
-            }
-        }
-    }
-
     private PlaybackInfo getVolumeAttributes() {
         int volumeType;
         AudioAttributes attributes;
diff --git a/services/core/java/com/android/server/notification/NotificationRecord.java b/services/core/java/com/android/server/notification/NotificationRecord.java
index e75d0a3..38c95f7 100644
--- a/services/core/java/com/android/server/notification/NotificationRecord.java
+++ b/services/core/java/com/android/server/notification/NotificationRecord.java
@@ -502,8 +502,8 @@
         pw.println(prefix + "uid=" + getSbn().getUid() + " userId=" + getSbn().getUserId());
         pw.println(prefix + "opPkg=" + getSbn().getOpPkg());
         pw.println(prefix + "icon=" + notification.getSmallIcon());
-        pw.println(prefix + "flags=0x" + Integer.toHexString(notification.flags));
-        pw.println(prefix + "originalFlags=0x" + Integer.toHexString(mOriginalFlags));
+        pw.println(prefix + "flags=" + Notification.flagsToString(notification.flags));
+        pw.println(prefix + "originalFlags=" + Notification.flagsToString(mOriginalFlags));
         pw.println(prefix + "pri=" + notification.priority);
         pw.println(prefix + "key=" + getSbn().getKey());
         pw.println(prefix + "seen=" + mStats.hasSeen());
@@ -538,8 +538,7 @@
         pw.println(prefix + "mInterruptionTimeMs=" + mInterruptionTimeMs);
         pw.println(prefix + "mSuppressedVisualEffects= " + mSuppressedVisualEffects);
         if (mPreChannelsNotification) {
-            pw.println(prefix + String.format("defaults=0x%08x flags=0x%08x",
-                    notification.defaults, notification.flags));
+            pw.println(prefix + "defaults=" + Notification.defaultsToString(notification.defaults));
             pw.println(prefix + "n.sound=" + notification.sound);
             pw.println(prefix + "n.audioStreamType=" + notification.audioStreamType);
             pw.println(prefix + "n.audioAttributes=" + notification.audioAttributes);
diff --git a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
index 8f6aa95..99401a1 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/OnDeviceIntelligenceManagerService.java
@@ -207,8 +207,9 @@
                     return;
                 }
                 ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.run(
-                        service -> service.getFeature(Binder.getCallingUid(), id, featureCallback));
+                        service -> service.getFeature(callerUid, id, featureCallback));
             }
 
             @Override
@@ -227,8 +228,9 @@
                     return;
                 }
                 ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.run(
-                        service -> service.listFeatures(Binder.getCallingUid(),
+                        service -> service.listFeatures(callerUid,
                                 listFeaturesCallback));
             }
 
@@ -250,8 +252,9 @@
                     return;
                 }
                 ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.run(
-                        service -> service.getFeatureDetails(Binder.getCallingUid(), feature,
+                        service -> service.getFeatureDetails(callerUid, feature,
                                 featureDetailsCallback));
             }
 
@@ -272,8 +275,9 @@
                             PersistableBundle.EMPTY);
                 }
                 ensureRemoteIntelligenceServiceInitialized();
+                int callerUid = Binder.getCallingUid();
                 mRemoteOnDeviceIntelligenceService.run(
-                        service -> service.requestFeatureDownload(Binder.getCallingUid(), feature,
+                        service -> service.requestFeatureDownload(callerUid, feature,
                                 wrapCancellationFuture(cancellationSignalFuture),
                                 downloadCallback));
             }
@@ -301,9 +305,9 @@
                                 PersistableBundle.EMPTY);
                     }
                     ensureRemoteInferenceServiceInitialized();
-
+                    int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.post(
-                            service -> service.requestTokenInfo(Binder.getCallingUid(), feature,
+                            service -> service.requestTokenInfo(callerUid, feature,
                                     request,
                                     wrapCancellationFuture(cancellationSignalFuture),
                                     wrapWithValidation(tokenInfoCallback)));
@@ -340,8 +344,9 @@
                                 PersistableBundle.EMPTY);
                     }
                     ensureRemoteInferenceServiceInitialized();
+                    int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.post(
-                            service -> service.processRequest(Binder.getCallingUid(), feature,
+                            service -> service.processRequest(callerUid, feature,
                                     request,
                                     requestType,
                                     wrapCancellationFuture(cancellationSignalFuture),
@@ -379,8 +384,9 @@
                                 PersistableBundle.EMPTY);
                     }
                     ensureRemoteInferenceServiceInitialized();
+                    int callerUid = Binder.getCallingUid();
                     result = mRemoteInferenceService.post(
-                            service -> service.processRequestStreaming(Binder.getCallingUid(),
+                            service -> service.processRequestStreaming(callerUid,
                                     feature,
                                     request, requestType,
                                     wrapCancellationFuture(cancellationSignalFuture),
diff --git a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
index 20c5b5f..28254d0 100644
--- a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
+++ b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
@@ -462,8 +462,8 @@
     }
 
     @Override
-    public int getNumRegisteredAttributionSources(int uid) {
-        return mAttributionSourceRegistry.getNumRegisteredAttributionSources(uid);
+    public int getRegisteredAttributionSourceCount(int uid) {
+        return mAttributionSourceRegistry.getRegisteredAttributionSourceCount(uid);
     }
 
     @Override
@@ -943,7 +943,7 @@
             }
         }
 
-        public int getNumRegisteredAttributionSources(int uid) {
+        public int getRegisteredAttributionSourceCount(int uid) {
             mContext.enforceCallingOrSelfPermission(UPDATE_APP_OPS_STATS,
                     "getting the number of registered AttributionSources requires "
                             + "UPDATE_APP_OPS_STATS");
@@ -952,14 +952,13 @@
             System.gc();
             System.gc();
             synchronized (mLock) {
-                int[] numForUid = { 0 };
-                mAttributions.forEach((key, value) -> {
-                    if (value.getUid() == uid) {
-                        numForUid[0]++;
+                int numForUid = 0;
+                for (Map.Entry<IBinder, AttributionSource> entry : mAttributions.entrySet()) {
+                    if (entry.getValue().getUid() == uid) {
+                        numForUid++;
                     }
-
-                });
-                return numForUid[0];
+                }
+                return numForUid;
             }
         }
 
diff --git a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
index 54cb9c9..49c4000 100644
--- a/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
+++ b/services/core/java/com/android/server/power/stats/BatteryStatsImpl.java
@@ -19,6 +19,7 @@
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.os.BatteryStats.Uid.NUM_PROCESS_STATE;
+import static android.os.BatteryStats.Uid.NUM_WIFI_BATCHED_SCAN_BINS;
 import static android.os.BatteryStatsManager.NUM_WIFI_STATES;
 import static android.os.BatteryStatsManager.NUM_WIFI_SUPPL_STATES;
 
@@ -292,7 +293,25 @@
     private int[] mCpuPowerBracketMap;
     private final CpuPowerStatsCollector mCpuPowerStatsCollector;
     private final MobileRadioPowerStatsCollector mMobileRadioPowerStatsCollector;
+    private final WifiPowerStatsCollector mWifiPowerStatsCollector;
     private final SparseBooleanArray mPowerStatsCollectorEnabled = new SparseBooleanArray();
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+                @Override
+                public void retrieveWifiScanTimes(Callback callback) {
+                    synchronized (BatteryStatsImpl.this) {
+                        retrieveWifiScanTimesLocked(callback);
+                    }
+                }
+
+                @Override
+                public long getWifiActiveDuration() {
+                    synchronized (BatteryStatsImpl.this) {
+                        return getGlobalWifiRunningTime(mClock.elapsedRealtime() * 1000,
+                                STATS_SINCE_CHARGED) / 1000;
+                    }
+                }
+            };
 
     public LongSparseArray<SamplingTimer> getKernelMemoryStats() {
         return mKernelMemoryStats;
@@ -501,6 +520,8 @@
                         TimeUnit.MINUTES.toMillis(1));
                 setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                         TimeUnit.HOURS.toMillis(1));
+                setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_WIFI,
+                        TimeUnit.HOURS.toMillis(1));
             }
 
             /**
@@ -1885,11 +1906,12 @@
     }
 
     private class PowerStatsCollectorInjector implements CpuPowerStatsCollector.Injector,
-            MobileRadioPowerStatsCollector.Injector {
+            MobileRadioPowerStatsCollector.Injector, WifiPowerStatsCollector.Injector {
         private PackageManager mPackageManager;
         private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
         private NetworkStatsManager mNetworkStatsManager;
         private TelephonyManager mTelephonyManager;
+        private WifiManager mWifiManager;
 
         void setContext(Context context) {
             mPackageManager = context.getPackageManager();
@@ -1897,6 +1919,7 @@
                     LocalServices.getService(PowerStatsInternal.class));
             mNetworkStatsManager = context.getSystemService(NetworkStatsManager.class);
             mTelephonyManager = context.getSystemService(TelephonyManager.class);
+            mWifiManager = context.getSystemService(WifiManager.class);
         }
 
         @Override
@@ -1950,11 +1973,26 @@
         }
 
         @Override
+        public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+            return () -> readWifiNetworkStatsLocked(mNetworkStatsManager);
+        }
+
+        @Override
+        public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+            return mWifiStatsRetriever;
+        }
+
+        @Override
         public TelephonyManager getTelephonyManager() {
             return mTelephonyManager;
         }
 
         @Override
+        public WifiManager getWifiManager() {
+            return mWifiManager;
+        }
+
+        @Override
         public LongSupplier getCallDurationSupplier() {
             return () -> mPhoneOnTimer.getTotalTimeLocked(mClock.elapsedRealtime() * 1000,
                     STATS_SINCE_CHARGED);
@@ -6354,7 +6392,11 @@
                     HistoryItem.STATE2_WIFI_ON_FLAG);
             mWifiOn = true;
             mWifiOnTimer.startRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-off", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-off", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6365,7 +6407,11 @@
                     HistoryItem.STATE2_WIFI_ON_FLAG);
             mWifiOn = false;
             mWifiOnTimer.stopRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-on", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-on", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6757,8 +6803,11 @@
                             .noteWifiRunningLocked(elapsedRealtimeMs);
                 }
             }
-
-            scheduleSyncExternalStatsLocked("wifi-running", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-running", ExternalStatsSync.UPDATE_WIFI);
+            }
         } else {
             Log.w(TAG, "noteWifiRunningLocked -- called while WIFI running");
         }
@@ -6827,7 +6876,11 @@
                 }
             }
 
-            scheduleSyncExternalStatsLocked("wifi-stopped", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-stopped", ExternalStatsSync.UPDATE_WIFI);
+            }
         } else {
             Log.w(TAG, "noteWifiStoppedLocked -- called while WIFI not running");
         }
@@ -6842,7 +6895,11 @@
             }
             mWifiState = wifiState;
             mWifiStateTimer[wifiState].startRunningLocked(elapsedRealtimeMs);
-            scheduleSyncExternalStatsLocked("wifi-state", ExternalStatsSync.UPDATE_WIFI);
+            if (mWifiPowerStatsCollector.isEnabled()) {
+                mWifiPowerStatsCollector.schedule();
+            } else {
+                scheduleSyncExternalStatsLocked("wifi-state", ExternalStatsSync.UPDATE_WIFI);
+            }
         }
     }
 
@@ -6965,6 +7022,25 @@
                 .noteWifiBatchedScanStoppedLocked(elapsedRealtimeMs);
     }
 
+    private void retrieveWifiScanTimesLocked(
+            WifiPowerStatsCollector.WifiStatsRetriever.Callback callback) {
+        long elapsedTimeUs = mClock.elapsedRealtime() * 1000;
+        for (int i = mUidStats.size() - 1; i >= 0; i--) {
+            int uid = mUidStats.keyAt(i);
+            Uid uidStats = mUidStats.valueAt(i);
+            long scanTimeUs = uidStats.getWifiScanTime(elapsedTimeUs, STATS_SINCE_CHARGED);
+            long batchScanTimeUs = 0;
+            for (int bucket = 0; bucket < NUM_WIFI_BATCHED_SCAN_BINS; bucket++) {
+                batchScanTimeUs += uidStats.getWifiBatchedScanTime(bucket, elapsedTimeUs,
+                        STATS_SINCE_CHARGED);
+            }
+            if (scanTimeUs != 0 || batchScanTimeUs != 0) {
+                callback.onWifiScanTime(uid, (scanTimeUs + 500) / 1000,
+                        (batchScanTimeUs + 500) / 1000);
+            }
+        }
+    }
+
     private int mWifiMulticastNesting = 0;
 
     @GuardedBy("this")
@@ -11101,6 +11177,11 @@
                 BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
         mMobileRadioPowerStatsCollector.addConsumer(this::recordPowerStats);
 
+        mWifiPowerStatsCollector = new WifiPowerStatsCollector(
+                mPowerStatsCollectorInjector, mBatteryStatsConfig.getPowerStatsThrottlePeriod(
+                BatteryConsumer.POWER_COMPONENT_WIFI));
+        mWifiPowerStatsCollector.addConsumer(this::recordPowerStats);
+
         mStartCount++;
         initTimersAndCounters();
         mOnBattery = mOnBatteryInternal = false;
@@ -12095,10 +12176,10 @@
                 }
             }
             if (lastEntry != null) {
-                delta.mRxBytes = entry.getRxBytes() - lastEntry.getRxBytes();
-                delta.mRxPackets = entry.getRxPackets() - lastEntry.getRxPackets();
-                delta.mTxBytes = entry.getTxBytes() - lastEntry.getTxBytes();
-                delta.mTxPackets = entry.getTxPackets() - lastEntry.getTxPackets();
+                delta.mRxBytes = Math.max(0, entry.getRxBytes() - lastEntry.getRxBytes());
+                delta.mRxPackets = Math.max(0, entry.getRxPackets() - lastEntry.getRxPackets());
+                delta.mTxBytes = Math.max(0, entry.getTxBytes() - lastEntry.getTxBytes());
+                delta.mTxPackets = Math.max(0, entry.getTxPackets() - lastEntry.getTxPackets());
             } else {
                 delta.mRxBytes = entry.getRxBytes();
                 delta.mRxPackets = entry.getRxPackets();
@@ -12119,6 +12200,10 @@
     public void updateWifiState(@Nullable final WifiActivityEnergyInfo info,
             final long consumedChargeUC, long elapsedRealtimeMs, long uptimeMs,
             @NonNull NetworkStatsManager networkStatsManager) {
+        if (mWifiPowerStatsCollector.isEnabled()) {
+            return;
+        }
+
         if (DEBUG_ENERGY) {
             synchronized (mWifiNetworkLock) {
                 Slog.d(TAG, "Updating wifi stats: " + Arrays.toString(mWifiIfaces));
@@ -14507,6 +14592,10 @@
                 mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO));
         mMobileRadioPowerStatsCollector.schedule();
 
+        mWifiPowerStatsCollector.setEnabled(
+                mPowerStatsCollectorEnabled.get(BatteryConsumer.POWER_COMPONENT_WIFI));
+        mWifiPowerStatsCollector.schedule();
+
         mSystemReady = true;
     }
 
@@ -14521,6 +14610,8 @@
                 return mCpuPowerStatsCollector;
             case BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO:
                 return mMobileRadioPowerStatsCollector;
+            case BatteryConsumer.POWER_COMPONENT_WIFI:
+                return mWifiPowerStatsCollector;
         }
         return null;
     }
@@ -16056,6 +16147,7 @@
     public void schedulePowerStatsSampleCollection() {
         mCpuPowerStatsCollector.forceSchedule();
         mMobileRadioPowerStatsCollector.forceSchedule();
+        mWifiPowerStatsCollector.forceSchedule();
     }
 
     /**
@@ -16074,6 +16166,7 @@
     public void dumpStatsSample(PrintWriter pw) {
         mCpuPowerStatsCollector.collectAndDump(pw);
         mMobileRadioPowerStatsCollector.collectAndDump(pw);
+        mWifiPowerStatsCollector.collectAndDump(pw);
     }
 
     private final Runnable mWriteAsyncRunnable = () -> {
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index 97f0986..0d5eabc 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -87,7 +87,9 @@
                         mPowerCalculators.add(new PhonePowerCalculator(mPowerProfile));
                     }
                 }
-                mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_WIFI)) {
+                    mPowerCalculators.add(new WifiPowerCalculator(mPowerProfile));
+                }
                 mPowerCalculators.add(new BluetoothPowerCalculator(mPowerProfile));
                 mPowerCalculators.add(new SensorPowerCalculator(
                         mContext.getSystemService(SensorManager.class)));
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
new file mode 100644
index 0000000..6321053
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsCollector.java
@@ -0,0 +1,328 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.PersistableBundle;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.util.Slog;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.function.IntSupplier;
+import java.util.function.Supplier;
+
+public class WifiPowerStatsCollector extends PowerStatsCollector {
+    private static final String TAG = "WifiPowerStatsCollector";
+
+    private static final long WIFI_ACTIVITY_REQUEST_TIMEOUT = 20000;
+
+    private static final long ENERGY_UNSPECIFIED = -1;
+
+    interface WifiStatsRetriever {
+        interface Callback {
+            void onWifiScanTime(int uid, long scanTimeMs, long batchScanTimeMs);
+        }
+
+        void retrieveWifiScanTimes(Callback callback);
+        long getWifiActiveDuration();
+    }
+
+    interface Injector {
+        Handler getHandler();
+        Clock getClock();
+        PowerStatsUidResolver getUidResolver();
+        PackageManager getPackageManager();
+        ConsumedEnergyRetriever getConsumedEnergyRetriever();
+        IntSupplier getVoltageSupplier();
+        Supplier<NetworkStats> getWifiNetworkStatsSupplier();
+        WifiManager getWifiManager();
+        WifiStatsRetriever getWifiStatsRetriever();
+    }
+
+    private final Injector mInjector;
+
+    private WifiPowerStatsLayout mLayout;
+    private boolean mIsInitialized;
+    private boolean mPowerReportingSupported;
+
+    private PowerStats mPowerStats;
+    private long[] mDeviceStats;
+    private volatile WifiManager mWifiManager;
+    private volatile Supplier<NetworkStats> mNetworkStatsSupplier;
+    private volatile WifiStatsRetriever mWifiStatsRetriever;
+    private ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    private IntSupplier mVoltageSupplier;
+    private int[] mEnergyConsumerIds = new int[0];
+    private WifiActivityEnergyInfo mLastWifiActivityInfo =
+            new WifiActivityEnergyInfo(0, 0, 0, 0, 0, 0);
+    private NetworkStats mLastNetworkStats;
+    private long[] mLastConsumedEnergyUws;
+    private int mLastVoltageMv;
+
+    private static class WifiScanTimes {
+        public long basicScanTimeMs;
+        public long batchedScanTimeMs;
+    }
+    private final WifiScanTimes mScanTimes = new WifiScanTimes();
+    private final SparseArray<WifiScanTimes> mLastScanTimes = new SparseArray<>();
+    private long mLastWifiActiveDuration;
+
+    public WifiPowerStatsCollector(Injector injector, long throttlePeriodMs) {
+        super(injector.getHandler(), throttlePeriodMs, injector.getUidResolver(),
+                injector.getClock());
+        mInjector = injector;
+    }
+
+    @Override
+    public void setEnabled(boolean enabled) {
+        if (enabled) {
+            PackageManager packageManager = mInjector.getPackageManager();
+            super.setEnabled(packageManager != null
+                    && packageManager.hasSystemFeature(PackageManager.FEATURE_WIFI));
+        } else {
+            super.setEnabled(false);
+        }
+    }
+
+    private boolean ensureInitialized() {
+        if (mIsInitialized) {
+            return true;
+        }
+
+        if (!isEnabled()) {
+            return false;
+        }
+
+        mConsumedEnergyRetriever = mInjector.getConsumedEnergyRetriever();
+        mVoltageSupplier = mInjector.getVoltageSupplier();
+        mWifiManager = mInjector.getWifiManager();
+        mNetworkStatsSupplier = mInjector.getWifiNetworkStatsSupplier();
+        mWifiStatsRetriever = mInjector.getWifiStatsRetriever();
+        mPowerReportingSupported =
+                mWifiManager != null && mWifiManager.isEnhancedPowerReportingSupported();
+
+        mEnergyConsumerIds = mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI);
+        mLastConsumedEnergyUws = new long[mEnergyConsumerIds.length];
+        Arrays.fill(mLastConsumedEnergyUws, ENERGY_UNSPECIFIED);
+
+        mLayout = new WifiPowerStatsLayout();
+        mLayout.addDeviceWifiActivity(mPowerReportingSupported);
+        mLayout.addDeviceSectionEnergyConsumers(mEnergyConsumerIds.length);
+        mLayout.addUidNetworkStats();
+        mLayout.addDeviceSectionUsageDuration();
+        mLayout.addDeviceSectionPowerEstimate();
+        mLayout.addUidSectionPowerEstimate();
+
+        PersistableBundle extras = new PersistableBundle();
+        mLayout.toExtras(extras);
+        PowerStats.Descriptor powerStatsDescriptor = new PowerStats.Descriptor(
+                BatteryConsumer.POWER_COMPONENT_WIFI, mLayout.getDeviceStatsArrayLength(),
+                null, 0, mLayout.getUidStatsArrayLength(),
+                extras);
+        mPowerStats = new PowerStats(powerStatsDescriptor);
+        mDeviceStats = mPowerStats.stats;
+
+        mIsInitialized = true;
+        return true;
+    }
+
+    @Override
+    protected PowerStats collectStats() {
+        if (!ensureInitialized()) {
+            return null;
+        }
+
+        if (mPowerReportingSupported) {
+            collectWifiActivityInfo();
+        } else {
+            collectWifiActivityStats();
+        }
+        collectNetworkStats();
+        collectWifiScanTime();
+
+        if (mEnergyConsumerIds.length != 0) {
+            collectEnergyConsumers();
+        }
+
+        return mPowerStats;
+    }
+
+    private void collectWifiActivityInfo() {
+        CompletableFuture<WifiActivityEnergyInfo> immediateFuture = new CompletableFuture<>();
+        mWifiManager.getWifiActivityEnergyInfoAsync(Runnable::run,
+                immediateFuture::complete);
+
+        WifiActivityEnergyInfo activityInfo;
+        try {
+            activityInfo = immediateFuture.get(WIFI_ACTIVITY_REQUEST_TIMEOUT,
+                    TimeUnit.MILLISECONDS);
+        } catch (Exception e) {
+            Slog.e(TAG, "Cannot acquire WifiActivityEnergyInfo", e);
+            activityInfo = null;
+        }
+
+        if (activityInfo == null) {
+            return;
+        }
+
+        long rxDuration = activityInfo.getControllerRxDurationMillis()
+                - mLastWifiActivityInfo.getControllerRxDurationMillis();
+        long txDuration = activityInfo.getControllerTxDurationMillis()
+                - mLastWifiActivityInfo.getControllerTxDurationMillis();
+        long scanDuration = activityInfo.getControllerScanDurationMillis()
+                - mLastWifiActivityInfo.getControllerScanDurationMillis();
+        long idleDuration = activityInfo.getControllerIdleDurationMillis()
+                - mLastWifiActivityInfo.getControllerIdleDurationMillis();
+
+        mLayout.setDeviceRxTime(mDeviceStats, rxDuration);
+        mLayout.setDeviceTxTime(mDeviceStats, txDuration);
+        mLayout.setDeviceScanTime(mDeviceStats, scanDuration);
+        mLayout.setDeviceIdleTime(mDeviceStats, idleDuration);
+
+        mPowerStats.durationMs = rxDuration + txDuration + scanDuration + idleDuration;
+
+        mLastWifiActivityInfo = activityInfo;
+    }
+
+    private void collectWifiActivityStats() {
+        long duration = mWifiStatsRetriever.getWifiActiveDuration();
+        mLayout.setDeviceActiveTime(mDeviceStats, Math.max(0, duration - mLastWifiActiveDuration));
+        mLastWifiActiveDuration = duration;
+        mPowerStats.durationMs = duration;
+    }
+
+    private void collectNetworkStats() {
+        mPowerStats.uidStats.clear();
+
+        NetworkStats networkStats = mNetworkStatsSupplier.get();
+        if (networkStats == null) {
+            return;
+        }
+
+        List<BatteryStatsImpl.NetworkStatsDelta> delta =
+                BatteryStatsImpl.computeDelta(networkStats, mLastNetworkStats);
+        mLastNetworkStats = networkStats;
+        for (int i = delta.size() - 1; i >= 0; i--) {
+            BatteryStatsImpl.NetworkStatsDelta uidDelta = delta.get(i);
+            long rxBytes = uidDelta.getRxBytes();
+            long txBytes = uidDelta.getTxBytes();
+            long rxPackets = uidDelta.getRxPackets();
+            long txPackets = uidDelta.getTxPackets();
+            if (rxBytes == 0 && txBytes == 0 && rxPackets == 0 && txPackets == 0) {
+                continue;
+            }
+
+            int uid = mUidResolver.mapUid(uidDelta.getUid());
+            long[] stats = mPowerStats.uidStats.get(uid);
+            if (stats == null) {
+                stats = new long[mLayout.getUidStatsArrayLength()];
+                mPowerStats.uidStats.put(uid, stats);
+                mLayout.setUidRxBytes(stats, rxBytes);
+                mLayout.setUidTxBytes(stats, txBytes);
+                mLayout.setUidRxPackets(stats, rxPackets);
+                mLayout.setUidTxPackets(stats, txPackets);
+            } else {
+                mLayout.setUidRxBytes(stats, mLayout.getUidRxBytes(stats) + rxBytes);
+                mLayout.setUidTxBytes(stats, mLayout.getUidTxBytes(stats) + txBytes);
+                mLayout.setUidRxPackets(stats, mLayout.getUidRxPackets(stats) + rxPackets);
+                mLayout.setUidTxPackets(stats, mLayout.getUidTxPackets(stats) + txPackets);
+            }
+        }
+    }
+
+    private void collectWifiScanTime() {
+        mScanTimes.basicScanTimeMs = 0;
+        mScanTimes.batchedScanTimeMs = 0;
+        mWifiStatsRetriever.retrieveWifiScanTimes((uid, scanTimeMs, batchScanTimeMs) -> {
+            WifiScanTimes lastScanTimes = mLastScanTimes.get(uid);
+            if (lastScanTimes == null) {
+                lastScanTimes = new WifiScanTimes();
+                mLastScanTimes.put(uid, lastScanTimes);
+            }
+
+            long scanTimeDelta = Math.max(0, scanTimeMs - lastScanTimes.basicScanTimeMs);
+            long batchScanTimeDelta = Math.max(0,
+                    batchScanTimeMs - lastScanTimes.batchedScanTimeMs);
+            if (scanTimeDelta != 0 || batchScanTimeDelta != 0) {
+                mScanTimes.basicScanTimeMs += scanTimeDelta;
+                mScanTimes.batchedScanTimeMs += batchScanTimeDelta;
+                uid = mUidResolver.mapUid(uid);
+                long[] stats = mPowerStats.uidStats.get(uid);
+                if (stats == null) {
+                    stats = new long[mLayout.getUidStatsArrayLength()];
+                    mPowerStats.uidStats.put(uid, stats);
+                    mLayout.setUidScanTime(stats, scanTimeDelta);
+                    mLayout.setUidBatchScanTime(stats, batchScanTimeDelta);
+                } else {
+                    mLayout.setUidScanTime(stats, mLayout.getUidScanTime(stats) + scanTimeDelta);
+                    mLayout.setUidBatchScanTime(stats,
+                            mLayout.getUidBatchedScanTime(stats) + batchScanTimeDelta);
+                }
+            }
+            lastScanTimes.basicScanTimeMs = scanTimeMs;
+            lastScanTimes.batchedScanTimeMs = batchScanTimeMs;
+        });
+
+        mLayout.setDeviceBasicScanTime(mDeviceStats, mScanTimes.basicScanTimeMs);
+        mLayout.setDeviceBatchedScanTime(mDeviceStats, mScanTimes.batchedScanTimeMs);
+    }
+
+    private void collectEnergyConsumers() {
+        int voltageMv = mVoltageSupplier.getAsInt();
+        if (voltageMv <= 0) {
+            Slog.wtf(TAG, "Unexpected battery voltage (" + voltageMv
+                    + " mV) when querying energy consumers");
+            return;
+        }
+
+        int averageVoltage = mLastVoltageMv != 0 ? (mLastVoltageMv + voltageMv) / 2 : voltageMv;
+        mLastVoltageMv = voltageMv;
+
+        long[] energyUws = mConsumedEnergyRetriever.getConsumedEnergyUws(mEnergyConsumerIds);
+        if (energyUws == null) {
+            return;
+        }
+
+        for (int i = energyUws.length - 1; i >= 0; i--) {
+            long energyDelta = mLastConsumedEnergyUws[i] != ENERGY_UNSPECIFIED
+                    ? energyUws[i] - mLastConsumedEnergyUws[i] : 0;
+            if (energyDelta < 0) {
+                // Likely, restart of powerstats HAL
+                energyDelta = 0;
+            }
+            mLayout.setConsumedEnergy(mPowerStats.stats, i, uJtoUc(energyDelta, averageVoltage));
+            mLastConsumedEnergyUws[i] = energyUws[i];
+        }
+    }
+
+    @Override
+    protected void onUidRemoved(int uid) {
+        super.onUidRemoved(uid);
+        mLastScanTimes.remove(uid);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
new file mode 100644
index 0000000..0fa6ec6
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsLayout.java
@@ -0,0 +1,241 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.annotation.NonNull;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+public class WifiPowerStatsLayout extends PowerStatsLayout {
+    private static final String TAG = "WifiPowerStatsLayout";
+    private static final int UNSPECIFIED = -1;
+    private static final String EXTRA_POWER_REPORTING_SUPPORTED = "prs";
+    private static final String EXTRA_DEVICE_RX_TIME_POSITION = "dt-rx";
+    private static final String EXTRA_DEVICE_TX_TIME_POSITION = "dt-tx";
+    private static final String EXTRA_DEVICE_SCAN_TIME_POSITION = "dt-scan";
+    private static final String EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION = "dt-basic-scan";
+    private static final String EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION = "dt-batch-scan";
+    private static final String EXTRA_DEVICE_IDLE_TIME_POSITION = "dt-idle";
+    private static final String EXTRA_DEVICE_ACTIVE_TIME_POSITION = "dt-on";
+    private static final String EXTRA_UID_RX_BYTES_POSITION = "urxb";
+    private static final String EXTRA_UID_TX_BYTES_POSITION = "utxb";
+    private static final String EXTRA_UID_RX_PACKETS_POSITION = "urxp";
+    private static final String EXTRA_UID_TX_PACKETS_POSITION = "utxp";
+    private static final String EXTRA_UID_SCAN_TIME_POSITION = "ut-scan";
+    private static final String EXTRA_UID_BATCH_SCAN_TIME_POSITION = "ut-bscan";
+
+    private boolean mPowerReportingSupported;
+    private int mDeviceRxTimePosition;
+    private int mDeviceTxTimePosition;
+    private int mDeviceIdleTimePosition;
+    private int mDeviceScanTimePosition;
+    private int mDeviceBasicScanTimePosition;
+    private int mDeviceBatchedScanTimePosition;
+    private int mDeviceActiveTimePosition;
+    private int mUidRxBytesPosition;
+    private int mUidTxBytesPosition;
+    private int mUidRxPacketsPosition;
+    private int mUidTxPacketsPosition;
+    private int mUidScanTimePosition;
+    private int mUidBatchScanTimePosition;
+
+    WifiPowerStatsLayout() {
+    }
+
+    WifiPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
+
+    void addDeviceWifiActivity(boolean powerReportingSupported) {
+        mPowerReportingSupported = powerReportingSupported;
+        if (mPowerReportingSupported) {
+            mDeviceActiveTimePosition = UNSPECIFIED;
+            mDeviceRxTimePosition = addDeviceSection(1);
+            mDeviceTxTimePosition = addDeviceSection(1);
+            mDeviceIdleTimePosition = addDeviceSection(1);
+            mDeviceScanTimePosition = addDeviceSection(1);
+        } else {
+            mDeviceActiveTimePosition = addDeviceSection(1);
+            mDeviceRxTimePosition = UNSPECIFIED;
+            mDeviceTxTimePosition = UNSPECIFIED;
+            mDeviceIdleTimePosition = UNSPECIFIED;
+            mDeviceScanTimePosition = UNSPECIFIED;
+        }
+        mDeviceBasicScanTimePosition = addDeviceSection(1);
+        mDeviceBatchedScanTimePosition = addDeviceSection(1);
+    }
+
+    void addUidNetworkStats() {
+        mUidRxBytesPosition = addUidSection(1);
+        mUidTxBytesPosition = addUidSection(1);
+        mUidRxPacketsPosition = addUidSection(1);
+        mUidTxPacketsPosition = addUidSection(1);
+        mUidScanTimePosition = addUidSection(1);
+        mUidBatchScanTimePosition = addUidSection(1);
+    }
+
+    public boolean isPowerReportingSupported() {
+        return mPowerReportingSupported;
+    }
+
+    public void setDeviceRxTime(long[] stats, long durationMillis) {
+        stats[mDeviceRxTimePosition] = durationMillis;
+    }
+
+    public long getDeviceRxTime(long[] stats) {
+        return stats[mDeviceRxTimePosition];
+    }
+
+    public void setDeviceTxTime(long[] stats, long durationMillis) {
+        stats[mDeviceTxTimePosition] = durationMillis;
+    }
+
+    public long getDeviceTxTime(long[] stats) {
+        return stats[mDeviceTxTimePosition];
+    }
+
+    public void setDeviceScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceScanTime(long[] stats) {
+        return stats[mDeviceScanTimePosition];
+    }
+
+    public void setDeviceBasicScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceBasicScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceBasicScanTime(long[] stats) {
+        return stats[mDeviceBasicScanTimePosition];
+    }
+
+    public void setDeviceBatchedScanTime(long[] stats, long durationMillis) {
+        stats[mDeviceBatchedScanTimePosition] = durationMillis;
+    }
+
+    public long getDeviceBatchedScanTime(long[] stats) {
+        return stats[mDeviceBatchedScanTimePosition];
+    }
+
+    public void setDeviceIdleTime(long[] stats, long durationMillis) {
+        stats[mDeviceIdleTimePosition] = durationMillis;
+    }
+
+    public long getDeviceIdleTime(long[] stats) {
+        return stats[mDeviceIdleTimePosition];
+    }
+
+    public void setDeviceActiveTime(long[] stats, long durationMillis) {
+        stats[mDeviceActiveTimePosition] = durationMillis;
+    }
+
+    public long getDeviceActiveTime(long[] stats) {
+        return stats[mDeviceActiveTimePosition];
+    }
+
+    public void setUidRxBytes(long[] stats, long count) {
+        stats[mUidRxBytesPosition] = count;
+    }
+
+    public long getUidRxBytes(long[] stats) {
+        return stats[mUidRxBytesPosition];
+    }
+
+    public void setUidTxBytes(long[] stats, long count) {
+        stats[mUidTxBytesPosition] = count;
+    }
+
+    public long getUidTxBytes(long[] stats) {
+        return stats[mUidTxBytesPosition];
+    }
+
+    public void setUidRxPackets(long[] stats, long count) {
+        stats[mUidRxPacketsPosition] = count;
+    }
+
+    public long getUidRxPackets(long[] stats) {
+        return stats[mUidRxPacketsPosition];
+    }
+
+    public void setUidTxPackets(long[] stats, long count) {
+        stats[mUidTxPacketsPosition] = count;
+    }
+
+    public long getUidTxPackets(long[] stats) {
+        return stats[mUidTxPacketsPosition];
+    }
+
+    public void setUidScanTime(long[] stats, long count) {
+        stats[mUidScanTimePosition] = count;
+    }
+
+    public long getUidScanTime(long[] stats) {
+        return stats[mUidScanTimePosition];
+    }
+
+    public void setUidBatchScanTime(long[] stats, long count) {
+        stats[mUidBatchScanTimePosition] = count;
+    }
+
+    public long getUidBatchedScanTime(long[] stats) {
+        return stats[mUidBatchScanTimePosition];
+    }
+
+    /**
+     * Copies the elements of the stats array layout into <code>extras</code>
+     */
+    public void toExtras(PersistableBundle extras) {
+        super.toExtras(extras);
+        extras.putBoolean(EXTRA_POWER_REPORTING_SUPPORTED, mPowerReportingSupported);
+        extras.putInt(EXTRA_DEVICE_RX_TIME_POSITION, mDeviceRxTimePosition);
+        extras.putInt(EXTRA_DEVICE_TX_TIME_POSITION, mDeviceTxTimePosition);
+        extras.putInt(EXTRA_DEVICE_SCAN_TIME_POSITION, mDeviceScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION, mDeviceBasicScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION, mDeviceBatchedScanTimePosition);
+        extras.putInt(EXTRA_DEVICE_IDLE_TIME_POSITION, mDeviceIdleTimePosition);
+        extras.putInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION, mDeviceActiveTimePosition);
+        extras.putInt(EXTRA_UID_RX_BYTES_POSITION, mUidRxBytesPosition);
+        extras.putInt(EXTRA_UID_TX_BYTES_POSITION, mUidTxBytesPosition);
+        extras.putInt(EXTRA_UID_RX_PACKETS_POSITION, mUidRxPacketsPosition);
+        extras.putInt(EXTRA_UID_TX_PACKETS_POSITION, mUidTxPacketsPosition);
+        extras.putInt(EXTRA_UID_SCAN_TIME_POSITION, mUidScanTimePosition);
+        extras.putInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION, mUidBatchScanTimePosition);
+    }
+
+    /**
+     * Retrieves elements of the stats array layout from <code>extras</code>
+     */
+    public void fromExtras(PersistableBundle extras) {
+        super.fromExtras(extras);
+        mPowerReportingSupported = extras.getBoolean(EXTRA_POWER_REPORTING_SUPPORTED);
+        mDeviceRxTimePosition = extras.getInt(EXTRA_DEVICE_RX_TIME_POSITION);
+        mDeviceTxTimePosition = extras.getInt(EXTRA_DEVICE_TX_TIME_POSITION);
+        mDeviceScanTimePosition = extras.getInt(EXTRA_DEVICE_SCAN_TIME_POSITION);
+        mDeviceBasicScanTimePosition = extras.getInt(EXTRA_DEVICE_BASIC_SCAN_TIME_POSITION);
+        mDeviceBatchedScanTimePosition = extras.getInt(EXTRA_DEVICE_BATCHED_SCAN_TIME_POSITION);
+        mDeviceIdleTimePosition = extras.getInt(EXTRA_DEVICE_IDLE_TIME_POSITION);
+        mDeviceActiveTimePosition = extras.getInt(EXTRA_DEVICE_ACTIVE_TIME_POSITION);
+        mUidRxBytesPosition = extras.getInt(EXTRA_UID_RX_BYTES_POSITION);
+        mUidTxBytesPosition = extras.getInt(EXTRA_UID_TX_BYTES_POSITION);
+        mUidRxPacketsPosition = extras.getInt(EXTRA_UID_RX_PACKETS_POSITION);
+        mUidTxPacketsPosition = extras.getInt(EXTRA_UID_TX_PACKETS_POSITION);
+        mUidScanTimePosition = extras.getInt(EXTRA_UID_SCAN_TIME_POSITION);
+        mUidBatchScanTimePosition = extras.getInt(EXTRA_UID_BATCH_SCAN_TIME_POSITION);
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
new file mode 100644
index 0000000..5e9cc40
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/WifiPowerStatsProcessor.java
@@ -0,0 +1,425 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import android.util.Slog;
+
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+public class WifiPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "WifiPowerStatsProcessor";
+    private static final boolean DEBUG = false;
+
+    private final UsageBasedPowerEstimator mRxPowerEstimator;
+    private final UsageBasedPowerEstimator mTxPowerEstimator;
+    private final UsageBasedPowerEstimator mIdlePowerEstimator;
+
+    private final UsageBasedPowerEstimator mActivePowerEstimator;
+    private final UsageBasedPowerEstimator mScanPowerEstimator;
+    private final UsageBasedPowerEstimator mBatchedScanPowerEstimator;
+
+    private PowerStats.Descriptor mLastUsedDescriptor;
+    private WifiPowerStatsLayout mStatsLayout;
+    // Sequence of steps for power estimation and intermediate results.
+    private PowerEstimationPlan mPlan;
+
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpUidStatsArray;
+    private boolean mHasWifiPowerController;
+
+    public WifiPowerStatsProcessor(PowerProfile powerProfile) {
+        mRxPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX));
+        mTxPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX));
+        mIdlePowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE));
+        mActivePowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_ACTIVE));
+        mScanPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_SCAN));
+        mBatchedScanPowerEstimator = new UsageBasedPowerEstimator(
+                powerProfile.getAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN));
+    }
+
+    private static class Intermediates {
+        /**
+         * Estimated power for the RX state.
+         */
+        public double rxPower;
+        /**
+         * Estimated power for the TX state.
+         */
+        public double txPower;
+        /**
+         * Estimated power in the SCAN state
+         */
+        public double scanPower;
+        /**
+         * Estimated power for IDLE, SCAN states.
+         */
+        public double idlePower;
+        /**
+         * Number of received packets
+         */
+        public long rxPackets;
+        /**
+         * Number of transmitted packets
+         */
+        public long txPackets;
+        /**
+         * Total duration of unbatched scans across all UIDs.
+         */
+        public long basicScanDuration;
+        /**
+         * Estimated power in the unbatched SCAN state
+         */
+        public double basicScanPower;
+        /**
+         * Total duration of batched scans across all UIDs.
+         */
+        public long batchedScanDuration;
+        /**
+         * Estimated power in the BATCHED SCAN state
+         */
+        public double batchedScanPower;
+        /**
+         * Estimated total power when active; used only in the absence of WiFiManager power
+         * reporting.
+         */
+        public double activePower;
+        /**
+         * Measured consumed energy from power monitoring hardware (micro-coulombs)
+         */
+        public long consumedEnergy;
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats) {
+        if (stats.getPowerStatsDescriptor() == null) {
+            return;
+        }
+
+        unpackPowerStatsDescriptor(stats.getPowerStatsDescriptor());
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+            Intermediates intermediates = new Intermediates();
+            estimation.intermediates = intermediates;
+            computeDevicePowerEstimates(stats, estimation.stateValues, intermediates);
+        }
+
+        double ratio = 1.0;
+        if (mStatsLayout.getEnergyConsumerCount() != 0) {
+            ratio = computeEstimateAdjustmentRatioUsingConsumedEnergy();
+            if (ratio != 1) {
+                for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+                    DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+                    adjustDevicePowerEstimates(stats, estimation.stateValues,
+                            (Intermediates) estimation.intermediates, ratio);
+                }
+            }
+        }
+
+        combineDeviceStateEstimates();
+
+        ArrayList<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+        if (!uids.isEmpty()) {
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidActivityTotals(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+
+            for (int uid : uids) {
+                for (int i = 0; i < mPlan.uidStateEstimates.size(); i++) {
+                    computeUidPowerEstimates(stats, uid, mPlan.uidStateEstimates.get(i));
+                }
+            }
+        }
+        mPlan.resetIntermediates();
+    }
+
+    private void unpackPowerStatsDescriptor(PowerStats.Descriptor descriptor) {
+        if (descriptor.equals(mLastUsedDescriptor)) {
+            return;
+        }
+
+        mLastUsedDescriptor = descriptor;
+        mStatsLayout = new WifiPowerStatsLayout(descriptor);
+        mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
+        mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
+        mHasWifiPowerController = mStatsLayout.isPowerReportingSupported();
+    }
+
+    /**
+     * Compute power estimates using the power profile.
+     */
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates) {
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        for (int i = mStatsLayout.getEnergyConsumerCount() - 1; i >= 0; i--) {
+            intermediates.consumedEnergy += mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, i);
+        }
+
+        intermediates.basicScanDuration =
+                mStatsLayout.getDeviceBasicScanTime(mTmpDeviceStatsArray);
+        intermediates.batchedScanDuration =
+                mStatsLayout.getDeviceBatchedScanTime(mTmpDeviceStatsArray);
+        if (mHasWifiPowerController) {
+            intermediates.rxPower = mRxPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceRxTime(mTmpDeviceStatsArray));
+            intermediates.txPower = mTxPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceTxTime(mTmpDeviceStatsArray));
+            intermediates.scanPower = mScanPowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceScanTime(mTmpDeviceStatsArray));
+            intermediates.idlePower = mIdlePowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceIdleTime(mTmpDeviceStatsArray));
+            mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                    intermediates.rxPower + intermediates.txPower + intermediates.scanPower
+                            + intermediates.idlePower);
+        } else {
+            intermediates.activePower = mActivePowerEstimator.calculatePower(
+                    mStatsLayout.getDeviceActiveTime(mTmpDeviceStatsArray));
+            intermediates.basicScanPower =
+                    mScanPowerEstimator.calculatePower(intermediates.basicScanDuration);
+            intermediates.batchedScanPower =
+                    mBatchedScanPowerEstimator.calculatePower(intermediates.batchedScanDuration);
+            mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray,
+                    intermediates.activePower + intermediates.basicScanPower
+                            + intermediates.batchedScanPower);
+        }
+
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * Compute an adjustment ratio using the total power estimated using the power profile
+     * and the total power measured by hardware.
+     */
+    private double computeEstimateAdjustmentRatioUsingConsumedEnergy() {
+        long totalConsumedEnergy = 0;
+        double totalPower = 0;
+
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            Intermediates intermediates =
+                    (Intermediates) mPlan.deviceStateEstimations.get(i).intermediates;
+            if (mHasWifiPowerController) {
+                totalPower += intermediates.rxPower + intermediates.txPower
+                        + intermediates.scanPower + intermediates.idlePower;
+            } else {
+                totalPower += intermediates.activePower + intermediates.basicScanPower
+                        + intermediates.batchedScanPower;
+            }
+            totalConsumedEnergy += intermediates.consumedEnergy;
+        }
+
+        if (totalPower == 0) {
+            return 1;
+        }
+
+        return uCtoMah(totalConsumedEnergy) / totalPower;
+    }
+
+    /**
+     * Uniformly apply the same adjustment to all power estimates in order to ensure that the total
+     * estimated power matches the measured consumed power.  We are not claiming that all
+     * averages captured in the power profile have to be off by the same percentage in reality.
+     */
+    private void adjustDevicePowerEstimates(PowerComponentAggregatedPowerStats stats,
+            int[] deviceStates, Intermediates intermediates, double ratio) {
+        double adjutedPower;
+        if (mHasWifiPowerController) {
+            intermediates.rxPower *= ratio;
+            intermediates.txPower *= ratio;
+            intermediates.scanPower *= ratio;
+            intermediates.idlePower *= ratio;
+            adjutedPower = intermediates.rxPower + intermediates.txPower + intermediates.scanPower
+                    + intermediates.idlePower;
+        } else {
+            intermediates.activePower *= ratio;
+            intermediates.basicScanPower *= ratio;
+            intermediates.batchedScanPower *= ratio;
+            adjutedPower = intermediates.activePower + intermediates.basicScanPower
+                    + intermediates.batchedScanPower;
+        }
+
+        if (!stats.getDeviceStats(mTmpDeviceStatsArray, deviceStates)) {
+            return;
+        }
+
+        mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray, adjutedPower);
+        stats.setDeviceStats(deviceStates, mTmpDeviceStatsArray);
+    }
+
+    /**
+     * Combine power estimates before distributing them proportionally to UIDs.
+     */
+    private void combineDeviceStateEstimates() {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate cdse = mPlan.combinedDeviceStateEstimations.get(i);
+            Intermediates
+                    cdseIntermediates = new Intermediates();
+            cdse.intermediates = cdseIntermediates;
+            List<DeviceStateEstimation> deviceStateEstimations = cdse.deviceStateEstimations;
+            for (int j = deviceStateEstimations.size() - 1; j >= 0; j--) {
+                DeviceStateEstimation dse = deviceStateEstimations.get(j);
+                Intermediates intermediates = (Intermediates) dse.intermediates;
+                if (mHasWifiPowerController) {
+                    cdseIntermediates.rxPower += intermediates.rxPower;
+                    cdseIntermediates.txPower += intermediates.txPower;
+                    cdseIntermediates.scanPower += intermediates.scanPower;
+                    cdseIntermediates.idlePower += intermediates.idlePower;
+                } else {
+                    cdseIntermediates.activePower += intermediates.activePower;
+                    cdseIntermediates.basicScanPower += intermediates.basicScanPower;
+                    cdseIntermediates.batchedScanPower += intermediates.batchedScanPower;
+                }
+                cdseIntermediates.basicScanDuration += intermediates.basicScanDuration;
+                cdseIntermediates.batchedScanDuration += intermediates.batchedScanDuration;
+                cdseIntermediates.consumedEnergy += intermediates.consumedEnergy;
+            }
+        }
+    }
+
+    private void computeUidActivityTotals(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            intermediates.rxPackets += mStatsLayout.getUidRxPackets(mTmpUidStatsArray);
+            intermediates.txPackets += mStatsLayout.getUidTxPackets(mTmpUidStatsArray);
+        }
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats, int uid,
+            UidStateEstimate uidStateEstimate) {
+        Intermediates intermediates =
+                (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+        for (UidStateProportionalEstimate proportionalEstimate :
+                uidStateEstimate.proportionalEstimates) {
+            if (!stats.getUidStats(mTmpUidStatsArray, uid, proportionalEstimate.stateValues)) {
+                continue;
+            }
+
+            double power = 0;
+            if (mHasWifiPowerController) {
+                if (intermediates.rxPackets != 0) {
+                    power += intermediates.rxPower * mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                            / intermediates.rxPackets;
+                }
+                if (intermediates.txPackets != 0) {
+                    power += intermediates.txPower * mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                            / intermediates.txPackets;
+                }
+                long totalScanDuration =
+                        intermediates.basicScanDuration + intermediates.batchedScanDuration;
+                if (totalScanDuration != 0) {
+                    long scanDuration = mStatsLayout.getUidScanTime(mTmpUidStatsArray)
+                            + mStatsLayout.getUidBatchedScanTime(mTmpUidStatsArray);
+                    power += intermediates.scanPower * scanDuration / totalScanDuration;
+                }
+            } else {
+                long totalPackets = intermediates.rxPackets + intermediates.txPackets;
+                if (totalPackets != 0) {
+                    long packets = mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                            + mStatsLayout.getUidTxPackets(mTmpUidStatsArray);
+                    power += intermediates.activePower * packets / totalPackets;
+                }
+
+                if (intermediates.basicScanDuration != 0) {
+                    long scanDuration = mStatsLayout.getUidScanTime(mTmpUidStatsArray);
+                    power += intermediates.basicScanPower * scanDuration
+                            / intermediates.basicScanDuration;
+                }
+
+                if (intermediates.batchedScanDuration != 0) {
+                    long batchedScanDuration = mStatsLayout.getUidBatchedScanTime(
+                            mTmpUidStatsArray);
+                    power += intermediates.batchedScanPower * batchedScanDuration
+                            / intermediates.batchedScanDuration;
+                }
+            }
+            mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+            stats.setUidStats(uid, proportionalEstimate.stateValues, mTmpUidStatsArray);
+
+            if (DEBUG) {
+                Slog.d(TAG, "UID: " + uid
+                        + " states: " + Arrays.toString(proportionalEstimate.stateValues)
+                        + " stats: " + Arrays.toString(mTmpUidStatsArray)
+                        + " rx: " + mStatsLayout.getUidRxPackets(mTmpUidStatsArray)
+                        + " rx-power: " + intermediates.rxPower
+                        + " rx-packets: " + intermediates.rxPackets
+                        + " tx: " + mStatsLayout.getUidTxPackets(mTmpUidStatsArray)
+                        + " tx-power: " + intermediates.txPower
+                        + " tx-packets: " + intermediates.txPackets
+                        + " power: " + power);
+            }
+        }
+    }
+
+    @Override
+    String deviceStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        if (mHasWifiPowerController) {
+            return "rx: " + mStatsLayout.getDeviceRxTime(stats)
+                    + " tx: " + mStatsLayout.getDeviceTxTime(stats)
+                    + " scan: " + mStatsLayout.getDeviceScanTime(stats)
+                    + " idle: " + mStatsLayout.getDeviceIdleTime(stats)
+                    + " power: " + mStatsLayout.getDevicePowerEstimate(stats);
+        } else {
+            return "active: " + mStatsLayout.getDeviceActiveTime(stats)
+                    + " scan: " + mStatsLayout.getDeviceBasicScanTime(stats)
+                    + " batched-scan: " + mStatsLayout.getDeviceBatchedScanTime(stats)
+                    + " power: " + mStatsLayout.getDevicePowerEstimate(stats);
+        }
+    }
+
+    @Override
+    String stateStatsToString(PowerStats.Descriptor descriptor, int key, long[] stats) {
+        // Unsupported for this power component
+        return null;
+    }
+
+    @Override
+    String uidStatsToString(PowerStats.Descriptor descriptor, long[] stats) {
+        unpackPowerStatsDescriptor(descriptor);
+        return "rx: " + mStatsLayout.getUidRxPackets(stats)
+                + " tx: " + mStatsLayout.getUidTxPackets(stats)
+                + " scan: " + mStatsLayout.getUidScanTime(stats)
+                + " batched-scan: " + mStatsLayout.getUidBatchedScanTime(stats)
+                + " power: " + mStatsLayout.getUidPowerEstimate(stats);
+    }
+}
diff --git a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
index 430232c..c74284e 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskSupervisor.java
@@ -2311,8 +2311,12 @@
             mService.setLastResumedActivityUncheckLocked(mTopResumedActivity, reason);
         }
         scheduleTopResumedActivityStateIfNeeded();
-
-        mService.updateTopApp(mTopResumedActivity);
+        // If the device is not sleeping and there is no top resumed, do not update top app because
+        // it may be an intermediate state while moving a task to front. The actual top will be set
+        // when TaskFragment#setResumedActivity is called.
+        if (mTopResumedActivity != null || mService.isSleepingLocked()) {
+            mService.updateTopApp(mTopResumedActivity);
+        }
 
         return mTopResumedActivity;
     }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 1e88fe4..622a809 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -174,6 +174,7 @@
     private static final int SET_SCREEN_BRIGHTNESS_OVERRIDE = 1;
     private static final int SET_USER_ACTIVITY_TIMEOUT = 2;
     private static final int MSG_SEND_SLEEP_TRANSITION = 3;
+    private static final int PINNED_TASK_ABORT_TIMEOUT = 1000;
 
     static final String TAG_TASKS = TAG + POSTFIX_TASKS;
     static final String TAG_STATES = TAG + POSTFIX_STATES;
@@ -295,6 +296,9 @@
 
     };
 
+    // TODO: b/335866033 Remove the abort PiP on timeout once PiP2 flag is on.
+    @Nullable private Runnable mMaybeAbortPipEnterRunnable = null;
+
     private final FindTaskResult mTmpFindTaskResult = new FindTaskResult();
 
     static class FindTaskResult implements Predicate<Task> {
@@ -2272,8 +2276,67 @@
         resumeFocusedTasksTopActivities();
 
         notifyActivityPipModeChanged(r.getTask(), r);
+
+        if (!isPip2ExperimentEnabled()) {
+            // TODO: b/335866033 Remove the abort PiP on timeout once PiP2 flag is on.
+            // Set up a timeout callback to potentially abort PiP enter if in an inconsistent state.
+            scheduleTimeoutAbortPipEnter(rootTask);
+        }
     }
 
+    private void scheduleTimeoutAbortPipEnter(Task rootTask) {
+        if (mMaybeAbortPipEnterRunnable != null) {
+            // If there is an abort enter PiP check pending already remove it and abort
+            // immediately since we are trying to enter PiP in an inconsistent state
+            mHandler.removeCallbacks(mMaybeAbortPipEnterRunnable);
+            mMaybeAbortPipEnterRunnable.run();
+        }
+        // Snapshot a throwable early on to display the callstack upon abort later on timeout.
+        final Throwable enterPipThrowable = new Throwable();
+        // Set up a timeout to potentially roll back the task change to PINNED mode
+        // by aborting PiP.
+        mMaybeAbortPipEnterRunnable = new Runnable() {
+            @Override
+            public void run() {
+                synchronized (mService.mGlobalLock) {
+                    if (mTransitionController.inTransition()) {
+                        // If this task is a part an active transition aborting PiP might break
+                        // it; so run the timeout callback directly once idle.
+
+                        final Runnable expectedMaybeAbortAtTimeout = mMaybeAbortPipEnterRunnable;
+                        mTransitionController.mStateValidators.add(() -> {
+                            // If a second PiP transition comes in, it runs the abort runnable for
+                            // the first transition pre-emptively, so we need to avoid calling
+                            // the same runnable twice when validating states.
+                            if (expectedMaybeAbortAtTimeout != mMaybeAbortPipEnterRunnable) return;
+                            mMaybeAbortPipEnterRunnable = null;
+                            run();
+                        });
+                        return;
+                    } else {
+                        mMaybeAbortPipEnterRunnable = null;
+                    }
+                    mService.deferWindowLayout();
+                    final ActivityRecord top = rootTask.getTopMostActivity();
+                    final ActivityManager.RunningTaskInfo beforeTaskInfo =
+                            rootTask.getTaskInfo();
+                    if (top != null && !top.inPinnedWindowingMode()
+                            && rootTask.abortPipEnter(top)) {
+                        Slog.wtf(TAG, "Enter PiP was aborted via a scheduled timeout"
+                                        + "task_state_before=" + beforeTaskInfo
+                                        + "task_state_after=" + rootTask.getTaskInfo(),
+                                enterPipThrowable);
+                    }
+                    mService.continueWindowLayout();
+                }
+            }
+        };
+        mHandler.postDelayed(mMaybeAbortPipEnterRunnable, PINNED_TASK_ABORT_TIMEOUT);
+        Slog.d(TAG, "a delayed check for potentially aborting PiP if "
+                + "in a wrong state is scheduled.");
+    }
+
+
     /**
      * Notifies when an activity enters or leaves PIP mode.
      *
@@ -2403,14 +2466,6 @@
             return false;
         }
 
-        return resumeFocusedTasksTopActivitiesUnchecked(targetRootTask, target, targetOptions,
-                deferPause);
-    }
-
-    @VisibleForTesting
-    boolean resumeFocusedTasksTopActivitiesUnchecked(
-            Task targetRootTask, ActivityRecord target, ActivityOptions targetOptions,
-            boolean deferPause) {
         boolean result = false;
         if (targetRootTask != null && (targetRootTask.isTopRootTaskInDisplayArea()
                 || getTopDisplayFocusedRootTask() == targetRootTask)) {
@@ -2900,6 +2955,14 @@
         mService.mH.post(mDestroyAllActivitiesRunnable);
     }
 
+    void removeAllMaybeAbortPipEnterRunnable() {
+        if (mMaybeAbortPipEnterRunnable == null) {
+            return;
+        }
+        mHandler.removeCallbacks(mMaybeAbortPipEnterRunnable);
+        mMaybeAbortPipEnterRunnable = null;
+    }
+
     // Tries to put all activity tasks to sleep. Returns true if all tasks were
     // successfully put to sleep.
     boolean putTasksToSleep(boolean allowDelay, boolean shuttingDown) {
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 9d3ffa9..f23a440 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -4852,11 +4852,13 @@
     /**
      * Abort an incomplete pip-entry. If left in this state, it will cover everything but remain
      * paused. If this is needed, there is a bug -- this should only be used for recovery.
+     *
+     * @return true if there is an inconsistency in the task and activity state.
      */
-    void abortPipEnter(ActivityRecord top) {
+    boolean abortPipEnter(ActivityRecord top) {
         // an incomplete state has the task PINNED but the activity not.
         if (!inPinnedWindowingMode() || top.inPinnedWindowingMode() || !canMoveTaskToBack(this)) {
-            return;
+            return false;
         }
         final Transition transition = new Transition(TRANSIT_TO_BACK, 0 /* flags */,
                 mTransitionController, mWmService.mSyncEngine);
@@ -4878,6 +4880,7 @@
             top.setWindowingMode(WINDOWING_MODE_UNDEFINED);
             top.mWaitForEnteringPinnedMode = false;
         }
+        return true;
     }
 
     void resumeNextFocusAfterReparent() {
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 664019a..1543263 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1429,20 +1429,17 @@
                 }
             }
             if (mTransientLaunches != null) {
-                InsetsControlTarget prevImeTarget = dc.getImeTarget(
-                        DisplayContent.IME_TARGET_CONTROL);
-                InsetsControlTarget newImeTarget = null;
                 TaskDisplayArea transientTDA = null;
-                // Transient-launch activities cannot be IME target (WindowState#canBeImeTarget),
-                // so re-compute in case the IME target is changed after transition.
                 for (int t = 0; t < mTransientLaunches.size(); ++t) {
                     if (mTransientLaunches.keyAt(t).getDisplayContent() == dc) {
-                        newImeTarget = dc.computeImeTarget(true /* updateImeTarget */);
+                        if (hasVisibleTransientLaunch) {
+                            updateImeForVisibleTransientLaunch(dc);
+                        }
                         transientTDA = mTransientLaunches.keyAt(i).getTaskDisplayArea();
                         break;
                     }
                 }
-                if (mRecentsDisplayId != INVALID_DISPLAY && prevImeTarget == newImeTarget) {
+                if (!hasVisibleTransientLaunch && mRecentsDisplayId == dc.mDisplayId) {
                     // Restore IME icon only when moving the original app task to front from
                     // recents, in case IME icon may missing if the moving task has already been
                     // the current focused task.
@@ -1540,6 +1537,35 @@
         return null;
     }
 
+    /**
+     * Transient-launch activities cannot be IME target (see {@link WindowState#canBeImeTarget}),
+     * so re-compute in case the IME target is changed after transition.
+     */
+    private void updateImeForVisibleTransientLaunch(@NonNull DisplayContent dc) {
+        final WindowState imeTarget = dc.computeImeTarget(true /* updateImeTarget */);
+        final WindowState imeWindow = dc.mInputMethodWindow;
+        if (imeWindow == null || imeTarget == null
+                || !mController.hasCollectingRotationChange(dc, dc.getRotation())) {
+            return;
+        }
+        // Drop the insets leash if it is still controlled by previous (invisible) app. This avoids
+        // showing IME with old rotation on an app with new rotation if IME parent is updated
+        // but insets leash hasn't been refreshed, i.e. DisplayContent#updateImeParent is called
+        // but InsetsStateController#notifyControlTargetChanged still waits for IME to redraw.
+        final InsetsSourceProvider sourceProvider = imeWindow.getControllableInsetProvider();
+        if (sourceProvider == null || sourceProvider.mControl == null
+                || !sourceProvider.isClientVisible()
+                || imeTarget == sourceProvider.getControlTarget()) {
+            return;
+        }
+        final SurfaceControl imeInsetsLeash = sourceProvider.mControl.getLeash();
+        final InsetsControlTarget controlTarget = sourceProvider.getControlTarget();
+        if (imeInsetsLeash != null && controlTarget != null && controlTarget.getWindow() != null
+                && !controlTarget.getWindow().mToken.isVisible()) {
+            dc.getSyncTransaction().reparent(imeInsetsLeash, null);
+        }
+    }
+
     void abort() {
         // This calls back into itself via controller.abort, so just early return here.
         if (mState == STATE_ABORT) return;
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index 909c1b3..1573d09 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -19,6 +19,7 @@
 import static android.Manifest.permission.START_TASKS_FROM_RECENTS;
 import static android.app.ActivityManager.isStartResultSuccessful;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
+import static android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED;
 import static android.app.WindowConfiguration.WINDOW_CONFIG_BOUNDS;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.window.TaskFragmentOperation.OP_TYPE_CLEAR_ADJACENT_TASK_FRAGMENTS;
@@ -69,7 +70,6 @@
 import static com.android.server.wm.ActivityRecord.State.PAUSING;
 import static com.android.server.wm.ActivityRecord.State.RESUMED;
 import static com.android.server.wm.ActivityTaskManagerService.enforceTaskPermission;
-import static com.android.server.wm.ActivityTaskManagerService.isPip2ExperimentEnabled;
 import static com.android.server.wm.ActivityTaskSupervisor.REMOVE_FROM_RECENTS;
 import static com.android.server.wm.Task.FLAG_FORCE_HIDDEN_FOR_PINNED_TASK;
 import static com.android.server.wm.Task.FLAG_FORCE_HIDDEN_FOR_TASK_ORG;
@@ -838,8 +838,6 @@
     }
 
     private int applyTaskChanges(Task tr, WindowContainerTransaction.Change c) {
-        final boolean wasPrevFocusableAndVisible = tr.isFocusableAndVisible();
-
         int effects = applyChanges(tr, c);
         final SurfaceControl.Transaction t = c.getBoundsChangeTransaction();
 
@@ -860,6 +858,17 @@
         }
 
         final int childWindowingMode = c.getActivityWindowingMode();
+        if (!ActivityTaskManagerService.isPip2ExperimentEnabled()
+                && tr.getWindowingMode() == WINDOWING_MODE_PINNED
+                && (childWindowingMode == WINDOWING_MODE_PINNED
+                || childWindowingMode == WINDOWING_MODE_UNDEFINED)) {
+            // If setActivityWindowingMode requested to match its pinned task's windowing mode,
+            // remove any inconsistency checking timeout callbacks for PiP.
+            Slog.d(TAG, "Task and activity windowing modes match, so remove any timeout "
+                    + "abort PiP callbacks scheduled if needed; task_win_mode="
+                    + tr.getWindowingMode() + ", activity_win_mode=" + childWindowingMode);
+            mService.mRootWindowContainer.removeAllMaybeAbortPipEnterRunnable();
+        }
         if (childWindowingMode > -1) {
             tr.forAllActivities(a -> { a.setWindowingMode(childWindowingMode); });
         }
@@ -885,17 +894,8 @@
                 boolean canEnterPip = activity.checkEnterPictureInPictureState(
                         "applyTaskChanges", true /* beforeStopping */);
                 if (canEnterPip) {
-                    mService.mTaskSupervisor.beginDeferResume();
-                    try {
-                        canEnterPip = mService.mActivityClientController
-                                .requestPictureInPictureMode(activity);
-                    } finally {
-                        mService.mTaskSupervisor.endDeferResume();
-                        if (canEnterPip && !isPip2ExperimentEnabled()) {
-                            // Wait until the transaction is applied to only resume once.
-                            effects |= TRANSACT_EFFECTS_LIFECYCLE;
-                        }
-                    }
+                    canEnterPip = mService.mActivityClientController
+                            .requestPictureInPictureMode(activity);
                 }
                 if (!canEnterPip) {
                     // Restore the flag to its previous state when the activity cannot enter PIP.
@@ -904,11 +904,6 @@
             }
         }
 
-        // Activity in this Task may resume/pause when enter/exit pip.
-        if (wasPrevFocusableAndVisible != tr.isFocusableAndVisible()) {
-            effects |= TRANSACT_EFFECTS_LIFECYCLE;
-        }
-
         return effects;
     }
 
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index be235b3b..7761311 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -76,6 +76,7 @@
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_SECURITY_LOGGING;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_SMS;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_STATUS_BAR;
+import static android.Manifest.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_SUPPORT_MESSAGE;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_SYSTEM_DIALOGS;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_SYSTEM_UPDATES;
@@ -88,7 +89,6 @@
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WIFI;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WINDOWS;
 import static android.Manifest.permission.MANAGE_DEVICE_POLICY_WIPE_DATA;
-import static android.Manifest.permission.MANAGE_DEVICE_POLICY_STORAGE_LIMIT;
 import static android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS;
 import static android.Manifest.permission.MASTER_CLEAR;
 import static android.Manifest.permission.NOTIFY_PENDING_SYSTEM_UPDATE;
@@ -116,7 +116,6 @@
 import static android.app.admin.DeviceAdminInfo.USES_POLICY_WIPE_DATA;
 import static android.app.admin.DeviceAdminReceiver.ACTION_COMPLIANCE_ACKNOWLEDGEMENT_REQUIRED;
 import static android.app.admin.DeviceAdminReceiver.EXTRA_TRANSFER_OWNERSHIP_ADMIN_EXTRAS_BUNDLE;
-import static android.app.admin.DevicePolicyIdentifiers.AUTO_TIMEZONE_POLICY;
 import static android.app.admin.DevicePolicyManager.ACTION_CHECK_POLICY_COMPLIANCE;
 import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_FINANCING_STATE_CHANGED;
 import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_RESOURCE_UPDATED;
@@ -407,7 +406,6 @@
 import android.location.Location;
 import android.location.LocationManager;
 import android.media.AudioManager;
-import android.media.IAudioService;
 import android.net.ConnectivityManager;
 import android.net.ConnectivitySettingsManager;
 import android.net.IIpConnectivityMetrics;
@@ -1770,10 +1768,6 @@
                     ServiceManager.getService(Context.BACKUP_SERVICE));
         }
 
-        IAudioService getIAudioService() {
-            return IAudioService.Stub.asInterface(ServiceManager.getService(Context.AUDIO_SERVICE));
-        }
-
         PersistentDataBlockManagerInternal getPersistentDataBlockManagerInternal() {
             return LocalServices.getService(PersistentDataBlockManagerInternal.class);
         }
@@ -1946,10 +1940,6 @@
                     name, value, userHandle);
         }
 
-        void settingsSecurePutInt(String name, int value) {
-            Settings.Secure.putInt(mContext.getContentResolver(), name, value);
-        }
-
         int settingsGlobalGetInt(String name, int def) {
             return Settings.Global.getInt(mContext.getContentResolver(), name, def);
         }
@@ -1963,10 +1953,6 @@
             Settings.Global.putInt(mContext.getContentResolver(), name, value);
         }
 
-        void settingsSecurePutString(String name, String value) {
-            Settings.Secure.putString(mContext.getContentResolver(), name, value);
-        }
-
         void settingsGlobalPutString(String name, String value) {
             Settings.Global.putString(mContext.getContentResolver(), name, value);
         }
@@ -2907,16 +2893,6 @@
         return poAdmin;
     }
 
-    @NonNull ActiveAdmin getOrganizationOwnedProfileOwnerLocked(final CallerIdentity caller) {
-        Preconditions.checkCallAuthorization(
-                mOwners.isProfileOwnerOfOrganizationOwnedDevice(caller.getUserId()),
-                "Caller %s is not an admin of an org-owned device",
-                caller.getComponentName());
-        final ActiveAdmin profileOwner = getProfileOwnerLocked(caller.getUserId());
-
-        return profileOwner;
-    }
-
     ActiveAdmin getProfileOwnerOrDeviceOwnerLocked(@UserIdInt int userId) {
         ensureLocked();
         // Try to find an admin which can use reqPolicy
@@ -2929,18 +2905,6 @@
         return getDeviceOwnerLocked(userId);
     }
 
-    ActiveAdmin getProfileOwnerOrDefaultDeviceOwnerLocked(@UserIdInt int userId) {
-        ensureLocked();
-        // Try to find an admin which can use reqPolicy
-        final ComponentName poAdminComponent = mOwners.getProfileOwnerComponent(userId);
-
-        if (poAdminComponent != null) {
-            return getProfileOwnerLocked(userId);
-        }
-
-        return getDefaultDeviceOwnerLocked(userId);
-    }
-
     @NonNull ActiveAdmin getParentOfAdminIfRequired(ActiveAdmin admin, boolean parent) {
         Objects.requireNonNull(admin);
         return parent ? admin.getParentActiveAdmin() : admin;
@@ -4453,25 +4417,6 @@
      * <ul>
      * <li>The active admins associated with the userHandle itself</li>
      * <li>The parent active admins for each managed profile associated with the userHandle</li>
-     * </ul>
-     *
-     * @param userHandle the affected user for whom to get the active admins
-     * @return the list of active admins for the affected user
-     */
-    @GuardedBy("getLockObject()")
-    private List<ActiveAdmin> getActiveAdminsForAffectedUserLocked(int userHandle) {
-        if (isManagedProfile(userHandle)) {
-            return getUserDataUnchecked(userHandle).mAdminList;
-        }
-        return getActiveAdminsForUserAndItsManagedProfilesLocked(userHandle,
-                /* shouldIncludeProfileAdmins */ (user) -> false);
-    }
-
-    /**
-     * Get the list of active admins for an affected user:
-     * <ul>
-     * <li>The active admins associated with the userHandle itself</li>
-     * <li>The parent active admins for each managed profile associated with the userHandle</li>
      * <li>The permission based admin associated with the userHandle itself</li>
      * </ul>
      *
@@ -8688,17 +8633,6 @@
                 .write();
     }
 
-    // Set the latest screen capture policy, overriding any existing ones.
-    // userHandle can be one of USER_ALL, USER_NULL or a concrete userId.
-    private void setScreenCaptureDisabled(int userHandle) {
-        int current = mPolicyCache.getScreenCaptureDisallowedUser();
-        if (userHandle == current) {
-            return;
-        }
-        mPolicyCache.setScreenCaptureDisallowedUser(userHandle);
-        updateScreenCaptureDisabled();
-    }
-
     /**
      * Returns whether or not screen capture is disabled for  any active admin.
      */
@@ -9915,16 +9849,6 @@
         return doOrPo;
     }
 
-    ActiveAdmin getDeviceOwnerOrProfileOwnerOfOrganizationOwnedDeviceParentLocked(int userId) {
-        ensureLocked();
-        ActiveAdmin admin = getDeviceOwnerAdminLocked();
-        if (admin != null) {
-            return admin;
-        }
-        admin = getProfileOwnerOfOrganizationOwnedDeviceLocked(userId);
-        return admin != null ? admin.getParentActiveAdmin() : null;
-    }
-
     @Override
     public void clearDeviceOwner(String packageName) {
         Objects.requireNonNull(packageName, "packageName is null");
@@ -10253,13 +10177,6 @@
         return mInjector.hasUserSetupCompleted(getUserData(userHandle));
     }
 
-    private boolean hasPaired(int userHandle) {
-        if (!mHasFeature) {
-            return true;
-        }
-        return getUserData(userHandle).mPaired;
-    }
-
     @Override
     public int getUserProvisioningState(int userHandle) {
         if (!mHasFeature) {
@@ -11054,16 +10971,6 @@
         return enforcingAdmin;
     }
 
-    private void enforceCanCallLockTaskLocked(CallerIdentity caller) {
-        Preconditions.checkCallAuthorization(isProfileOwner(caller)
-                || isDefaultDeviceOwner(caller) || isFinancedDeviceOwner(caller));
-
-        final int userId =  caller.getUserId();
-        if (!canDPCManagedUserUseLockTaskLocked(userId)) {
-            throw new SecurityException("User " + userId + " is not allowed to use lock task");
-        }
-    }
-
     private boolean isSystemUid(CallerIdentity caller) {
         return UserHandle.isSameApp(caller.getUid(), Process.SYSTEM_UID);
     }
@@ -14595,15 +14502,6 @@
         }
     }
 
-    private void setLockTaskPackagesLocked(int userHandle, List<String> packages) {
-        DevicePolicyData policy = getUserData(userHandle);
-        policy.mLockTaskPackages = packages;
-
-        // Store the settings persistently.
-        saveSettingsLocked(userHandle);
-        updateLockTaskPackagesLocked(mContext, packages, userHandle);
-    }
-
     @Override
     public String[] getLockTaskPackages(ComponentName who, String callerPackageName) {
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
@@ -14653,7 +14551,6 @@
             "Cannot use LOCK_TASK_FEATURE_NOTIFICATIONS without LOCK_TASK_FEATURE_HOME");
 
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
-        final int userHandle = caller.getUserId();
         synchronized (getLockObject()) {
             checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_SET_LOCK_TASK_FEATURES);
         }
@@ -14689,13 +14586,6 @@
         }
     }
 
-    private void setLockTaskFeaturesLocked(int userHandle, int flags) {
-        DevicePolicyData policy = getUserData(userHandle);
-        policy.mLockTaskFeatures = flags;
-        saveSettingsLocked(userHandle);
-        updateLockTaskFeaturesLocked(flags, userHandle);
-    }
-
     @Override
     public int getLockTaskFeatures(ComponentName who, String callerPackageName) {
         CallerIdentity caller = getCallerIdentity(who, callerPackageName);
@@ -16171,29 +16061,6 @@
         }
     }
 
-
-    /**
-     *  Excludes restrictions imposed by UserManager.
-     */
-    private List<UserManager.EnforcingUser> getDevicePolicySources(
-            List<UserManager.EnforcingUser> sources) {
-        int sizeBefore = sources.size();
-        List<UserManager.EnforcingUser> realSources = new ArrayList<>(sizeBefore);
-        for (int i = 0; i < sizeBefore; i++) {
-            UserManager.EnforcingUser source = sources.get(i);
-            int type = source.getUserRestrictionSource();
-            if (type != UserManager.RESTRICTION_SOURCE_PROFILE_OWNER
-                    && type != UserManager.RESTRICTION_SOURCE_DEVICE_OWNER) {
-                // TODO(b/128928355): add unit test
-                Slogf.d(LOG_TAG, "excluding source of type %s at index %d",
-                        userRestrictionSourceToString(type), i);
-                continue;
-            }
-            realSources.add(source);
-        }
-        return realSources;
-    }
-
     private static String userRestrictionSourceToString(@UserRestrictionSource int source) {
         return DebugUtils.flagsToString(UserManager.class, "RESTRICTION_", source);
     }
@@ -20870,17 +20737,6 @@
         }
     }
 
-    private void suspendPersonalAppsInPackageManager(int userId) {
-        mInjector.binderWithCleanCallingIdentity(() -> {
-            final String[] appsToSuspend = mInjector.getPersonalAppsForSuspension(userId);
-            final String[] failedApps = mInjector.getPackageManagerInternal()
-                    .setPackagesSuspendedByAdmin(userId, appsToSuspend, true);
-            if (!ArrayUtils.isEmpty(failedApps)) {
-                Slogf.wtf(LOG_TAG, "Failed to suspend apps: " + String.join(",", failedApps));
-            }
-        });
-    }
-
     private void notifyIfManagedSubscriptionsAreUnavailable(
             UserHandle managedProfile, boolean managedProfileAvailable) {
         if (!isManagedProfile(managedProfile.getIdentifier())) {
@@ -23329,14 +23185,6 @@
         return getEnforcingAdminForCaller(admin, callerPackageName);
     }
 
-    private static final HashMap<String, String> POLICY_IDENTIFIER_TO_PERMISSION = new HashMap<>();
-    {
-        POLICY_IDENTIFIER_TO_PERMISSION.put(AUTO_TIMEZONE_POLICY, SET_TIME_ZONE);
-    }
-
-    private static final HashMap<String, Integer> POLICY_IDENTIFIER_TO_ACTIVE_ADMIN_POLICY =
-            new HashMap<>();
-
     /**
      * Checks if the calling process has been granted permission to apply a device policy.
      *
@@ -23355,7 +23203,6 @@
         }
     }
 
-
     /**
      * Checks if the calling process has been granted permission to apply a device policy on a
      * specific user. Only one permission provided in the list needs to be granted to pass this
@@ -24455,52 +24302,6 @@
         }
     }
 
-    // We need to add a mapping of policyId to permission in POLICY_IDENTIFIER_TO_PERMISSION
-    // for each migrated permission.
-    private List<ActiveAdmin> getNonDPCActiveAdminsForPolicyLocked(String policyIdentifier) {
-        Integer activeAdminPolicy = POLICY_IDENTIFIER_TO_ACTIVE_ADMIN_POLICY.get(policyIdentifier);
-        if (activeAdminPolicy == null) {
-            Slogf.e(LOG_TAG,
-                    "Can't find a active admin policy for %s in POLICY_IDENTIFIER_TO_PERMISSION",
-                    policyIdentifier);
-            return new ArrayList<>();
-        }
-
-        List<ActiveAdmin> admins = new ArrayList<>();
-        for (UserInfo userInfo : mUserManager.getUsers()) {
-            List<ComponentName> activeAdmins = getActiveAdmins(userInfo.id);
-            for (ComponentName admin : activeAdmins) {
-                if (isDeviceOwner(admin, userInfo.id) || isProfileOwner(admin, userInfo.id)) {
-                    continue;
-                }
-                DevicePolicyData policy = getUserDataUnchecked(userInfo.id);
-                if (isActiveAdminWithPolicyForUserLocked(
-                        policy.mAdminMap.get(admin), activeAdminPolicy,
-                        userInfo.id)) {
-                    admins.add(policy.mAdminMap.get(admin));
-                }
-            }
-        }
-        return admins;
-    }
-
-    // TODO: this can actually accept an EnforcingAdmin that gets created in the permission
-    //  check method.
-    private boolean isCallerActiveAdminOrDelegate(
-            CallerIdentity caller, @Nullable String delegateScope) {
-        return mInjector.binderWithCleanCallingIdentity(() -> {
-            List<ComponentName> activeAdmins = getActiveAdmins(caller.getUserId());
-            if (activeAdmins != null) {
-                for (ComponentName admin : activeAdmins) {
-                    if (admin.getPackageName().equals(caller.getPackageName())) {
-                        return true;
-                    }
-                }
-            }
-            return delegateScope != null && isCallerDelegate(caller, delegateScope);
-        });
-    }
-
     private ActiveAdmin getActiveAdminForCaller(@Nullable ComponentName who,
             CallerIdentity caller) {
         synchronized (getLockObject()) {
diff --git a/services/java/com/android/server/SystemServer.java b/services/java/com/android/server/SystemServer.java
index 648b810..8caf5ae 100644
--- a/services/java/com/android/server/SystemServer.java
+++ b/services/java/com/android/server/SystemServer.java
@@ -91,6 +91,7 @@
 import android.util.Dumpable;
 import android.util.EventLog;
 import android.util.IndentingPrintWriter;
+import android.util.Log;
 import android.util.Pair;
 import android.util.Slog;
 import android.util.TimeUtils;
@@ -2005,9 +2006,11 @@
             mSystemServiceManager.startService(new FontManagerService.Lifecycle(context, safeMode));
             t.traceEnd();
 
-            t.traceBegin("StartTextServicesManager");
-            mSystemServiceManager.startService(TextServicesManagerService.Lifecycle.class);
-            t.traceEnd();
+            if (!isWatch || !android.server.Flags.removeTextService()) {
+                t.traceBegin("StartTextServicesManager");
+                mSystemServiceManager.startService(TextServicesManagerService.Lifecycle.class);
+                t.traceEnd();
+            }
 
             if (!disableSystemTextClassifier) {
                 t.traceBegin("StartTextClassificationManagerService");
diff --git a/services/java/com/android/server/flags.aconfig b/services/java/com/android/server/flags.aconfig
index 854bc0f..38354e8 100644
--- a/services/java/com/android/server/flags.aconfig
+++ b/services/java/com/android/server/flags.aconfig
@@ -8,3 +8,10 @@
      is_fixed_read_only: true
      bug: "324153471"
 }
+
+flag {
+     name: "remove_text_service"
+     namespace: "wear_frameworks"
+     description: "Remove TextServiceManagerService on Wear"
+     bug: "323720705"
+}
\ No newline at end of file
diff --git a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
index 2134278..1535298 100644
--- a/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
+++ b/services/tests/InputMethodSystemServerTests/src/com/android/inputmethodservice/InputMethodServiceTest.java
@@ -38,7 +38,6 @@
 import android.support.test.uiautomator.UiObject2;
 import android.support.test.uiautomator.Until;
 import android.util.Log;
-import android.view.KeyEvent;
 import android.view.WindowManagerGlobal;
 import android.view.inputmethod.EditorInfo;
 import android.view.inputmethod.InputMethodManager;
@@ -601,28 +600,6 @@
                 false /* orientationPortrait */);
     }
 
-    @Test
-    public void switchesKeyboardLayout_withShortcut_onlyIfImeVisible() throws Exception {
-        setShowImeWithHardKeyboard(true /* enabled */);
-
-        assertThat(mInputMethodService.isInputViewShown()).isFalse();
-        assertThat(mInputMethodService.onKeyDown(KeyEvent.KEYCODE_SPACE,
-                new KeyEvent(0 /* downTime */, 0 /* eventTime */, KeyEvent.ACTION_DOWN,
-                        KeyEvent.KEYCODE_SPACE, 0 /* repeat */,
-                        KeyEvent.META_CTRL_LEFT_ON | KeyEvent.META_CTRL_ON))).isFalse();
-
-        verifyInputViewStatusOnMainSync(
-                () -> assertThat(mActivity.showImeWithInputMethodManager(0 /* flags */)).isTrue(),
-                true /* expected */,
-                true /* inputViewStarted */);
-
-        assertThat(mInputMethodService.isInputViewShown()).isTrue();
-        assertThat(mInputMethodService.onKeyDown(KeyEvent.KEYCODE_SPACE,
-                new KeyEvent(0 /* downTime */, 0 /* eventTime */, KeyEvent.ACTION_DOWN,
-                        KeyEvent.KEYCODE_SPACE, 0 /* repeat */,
-                        KeyEvent.META_CTRL_LEFT_ON | KeyEvent.META_CTRL_ON))).isTrue();
-    }
-
     /**
      * This checks that when the system navigation bar is not created (e.g. emulator),
      * then the IME caption bar is also not created.
diff --git a/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java b/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
index 6d89e80..8db896b 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/DisplayBrightnessStateTest.java
@@ -47,6 +47,7 @@
         float sdrBrightness = 0.2f;
         boolean shouldUseAutoBrightness = true;
         boolean shouldUpdateScreenBrightnessSetting = true;
+        int brightnessAdjustmentFlag = 2;
         BrightnessReason brightnessReason = new BrightnessReason();
         brightnessReason.setReason(BrightnessReason.REASON_AUTOMATIC);
         brightnessReason.setModifier(BrightnessReason.MODIFIER_DIMMED);
@@ -56,6 +57,7 @@
                 .setBrightnessReason(brightnessReason)
                 .setShouldUseAutoBrightness(shouldUseAutoBrightness)
                 .setShouldUpdateScreenBrightnessSetting(shouldUpdateScreenBrightnessSetting)
+                .setBrightnessAdjustmentFlag(brightnessAdjustmentFlag)
                 .build();
 
         assertEquals(displayBrightnessState.getBrightness(), brightness, FLOAT_DELTA);
@@ -105,7 +107,9 @@
                 .append("\n    shouldUpdateScreenBrightnessSetting:")
                 .append(displayBrightnessState.shouldUpdateScreenBrightnessSetting())
                 .append("\n    mBrightnessEvent:")
-                .append(Objects.toString(displayBrightnessState.getBrightnessEvent(), "null"));
+                .append(Objects.toString(displayBrightnessState.getBrightnessEvent(), "null"))
+                .append("\n    mBrightnessAdjustmentFlag:")
+                .append(displayBrightnessState.getBrightnessAdjustmentFlag());
         return sb.toString();
     }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/DisplayBrightnessControllerTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/DisplayBrightnessControllerTest.java
index 1ae4099..13a1445 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/DisplayBrightnessControllerTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/DisplayBrightnessControllerTest.java
@@ -20,6 +20,7 @@
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -113,7 +114,8 @@
         when(mDisplayBrightnessStrategySelector.selectStrategy(
                 any(StrategySelectionRequest.class))).thenReturn(displayBrightnessStrategy);
         mDisplayBrightnessController.updateBrightness(displayPowerRequest, targetDisplayState);
-        verify(displayBrightnessStrategy).updateBrightness(displayPowerRequest);
+        verify(displayBrightnessStrategy).updateBrightness(
+                eq(new StrategyExecutionRequest(displayPowerRequest, DEFAULT_BRIGHTNESS)));
         assertEquals(mDisplayBrightnessController.getCurrentDisplayBrightnessStrategy(),
                 displayBrightnessStrategy);
     }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategyTest.java
index 6e163ca..54f2268 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/AutomaticBrightnessStrategyTest.java
@@ -41,8 +41,10 @@
 import com.android.internal.util.test.FakeSettingsProvider;
 import com.android.internal.util.test.FakeSettingsProviderRule;
 import com.android.server.display.AutomaticBrightnessController;
+import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessEvent;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.After;
 import org.junit.Before;
@@ -295,15 +297,11 @@
                 mContext.getContentResolver(),
                 Settings.System.SCREEN_AUTO_BRIGHTNESS_ADJ,
                 UserHandle.USER_CURRENT), 0.0f);
-        assertEquals(BrightnessReason.ADJUSTMENT_AUTO,
-                mAutomaticBrightnessStrategy.getAutoBrightnessAdjustmentReasonsFlags());
         float invalidBrightness = -0.5f;
         mAutomaticBrightnessStrategy
                 .adjustAutomaticBrightnessStateIfValid(invalidBrightness);
         assertEquals(autoBrightnessAdjustment,
                 mAutomaticBrightnessStrategy.getAutoBrightnessAdjustment(), 0.0f);
-        assertEquals(0,
-                mAutomaticBrightnessStrategy.getAutoBrightnessAdjustmentReasonsFlags());
     }
 
     @Test
@@ -426,6 +424,80 @@
         assertTrue(mAutomaticBrightnessStrategy.isAutoBrightnessValid());
     }
 
+    @Test
+    public void
+            updateBrightness_constructsDisplayBrightnessState_withAdjustmentAutoAdjustmentFlag() {
+        BrightnessEvent brightnessEvent = new BrightnessEvent(DISPLAY_ID);
+        mAutomaticBrightnessStrategy = new AutomaticBrightnessStrategy(
+                mContext, DISPLAY_ID, displayId -> brightnessEvent);
+        new AutomaticBrightnessStrategy(mContext, DISPLAY_ID);
+        mAutomaticBrightnessStrategy.setAutomaticBrightnessController(
+                mAutomaticBrightnessController);
+        float brightness = 0.4f;
+        BrightnessReason brightnessReason = new BrightnessReason();
+        brightnessReason.setReason(BrightnessReason.REASON_AUTOMATIC);
+        when(mAutomaticBrightnessController.getAutomaticScreenBrightness(brightnessEvent))
+                .thenReturn(brightness);
+
+        DisplayManagerInternal.DisplayPowerRequest displayPowerRequest =
+                mock(DisplayManagerInternal.DisplayPowerRequest.class);
+        DisplayBrightnessState expectedDisplayBrightnessState = new DisplayBrightnessState.Builder()
+                .setBrightness(brightness)
+                .setSdrBrightness(brightness)
+                .setBrightnessReason(brightnessReason)
+                .setDisplayBrightnessStrategyName(mAutomaticBrightnessStrategy.getName())
+                .setIsSlowChange(false)
+                .setBrightnessEvent(brightnessEvent)
+                .setBrightnessAdjustmentFlag(BrightnessReason.ADJUSTMENT_AUTO)
+                .setShouldUpdateScreenBrightnessSetting(true)
+                .build();
+        DisplayBrightnessState actualDisplayBrightnessState = mAutomaticBrightnessStrategy
+                .updateBrightness(new StrategyExecutionRequest(displayPowerRequest, 0.6f));
+        assertEquals(expectedDisplayBrightnessState, actualDisplayBrightnessState);
+    }
+
+    @Test
+    public void
+            updateBrightness_constructsDisplayBrightnessState_withAdjustmentTempAdjustmentFlag() {
+        BrightnessEvent brightnessEvent = new BrightnessEvent(DISPLAY_ID);
+        mAutomaticBrightnessStrategy = new AutomaticBrightnessStrategy(
+                mContext, DISPLAY_ID, displayId -> brightnessEvent);
+        new AutomaticBrightnessStrategy(mContext, DISPLAY_ID);
+        mAutomaticBrightnessStrategy.setAutomaticBrightnessController(
+                mAutomaticBrightnessController);
+        float brightness = 0.4f;
+        BrightnessReason brightnessReason = new BrightnessReason();
+        brightnessReason.setReason(BrightnessReason.REASON_AUTOMATIC);
+        when(mAutomaticBrightnessController.getAutomaticScreenBrightness(brightnessEvent))
+                .thenReturn(brightness);
+
+        DisplayManagerInternal.DisplayPowerRequest displayPowerRequest =
+                mock(DisplayManagerInternal.DisplayPowerRequest.class);
+        float temporaryBrightness = 0.3f;
+        float autoBrightnessAdjustment = 0.1f;
+        mAutomaticBrightnessStrategy.setTemporaryAutoBrightnessAdjustment(temporaryBrightness);
+        mAutomaticBrightnessStrategy.accommodateUserBrightnessChanges(true,
+                brightness, DisplayManagerInternal.DisplayPowerRequest.POLICY_BRIGHT,
+                Display.STATE_ON, mock(BrightnessConfiguration.class),
+                AutomaticBrightnessController.AUTO_BRIGHTNESS_ENABLED);
+        when(mAutomaticBrightnessController.getAutomaticScreenBrightnessAdjustment()).thenReturn(
+                autoBrightnessAdjustment);
+
+        DisplayBrightnessState expectedDisplayBrightnessState = new DisplayBrightnessState.Builder()
+                .setBrightness(brightness)
+                .setSdrBrightness(brightness)
+                .setBrightnessReason(brightnessReason)
+                .setDisplayBrightnessStrategyName(mAutomaticBrightnessStrategy.getName())
+                .setIsSlowChange(false)
+                .setBrightnessEvent(brightnessEvent)
+                .setBrightnessAdjustmentFlag(BrightnessReason.ADJUSTMENT_AUTO_TEMP)
+                .setShouldUpdateScreenBrightnessSetting(true)
+                .build();
+        DisplayBrightnessState actualDisplayBrightnessState = mAutomaticBrightnessStrategy
+                .updateBrightness(new StrategyExecutionRequest(displayPowerRequest, 0.6f));
+        assertEquals(expectedDisplayBrightnessState, actualDisplayBrightnessState);
+    }
+
     private void setPendingAutoBrightnessAdjustment(float pendingAutoBrightnessAdjustment) {
         Settings.System.putFloat(mContext.getContentResolver(),
                 Settings.System.SCREEN_AUTO_BRIGHTNESS_ADJ, pendingAutoBrightnessAdjustment);
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/BoostBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/BoostBrightnessStrategyTest.java
index c434631..47f1746 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/BoostBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/BoostBrightnessStrategyTest.java
@@ -27,6 +27,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -58,7 +59,8 @@
                         .setDisplayBrightnessStrategyName(mBoostBrightnessStrategy.getName())
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mBoostBrightnessStrategy.updateBrightness(displayPowerRequest);
+                mBoostBrightnessStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
     }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/DozeBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/DozeBrightnessStrategyTest.java
index d60caf6..9246780 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/DozeBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/DozeBrightnessStrategyTest.java
@@ -25,6 +25,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -55,7 +56,8 @@
                         .setDisplayBrightnessStrategyName(mDozeBrightnessModeStrategy.getName())
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mDozeBrightnessModeStrategy.updateBrightness(displayPowerRequest);
+                mDozeBrightnessModeStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
     }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/FollowerBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/FollowerBrightnessStrategyTest.java
index d8569f7..682c9cc 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/FollowerBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/FollowerBrightnessStrategyTest.java
@@ -26,6 +26,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -59,7 +60,8 @@
                         .setIsSlowChange(slowChange)
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mFollowerBrightnessStrategy.updateBrightness(displayPowerRequest);
+                mFollowerBrightnessStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(expectedDisplayBrightnessState, updatedDisplayBrightnessState);
     }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OffloadBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OffloadBrightnessStrategyTest.java
index c2fa4eb..ccf6309 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OffloadBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OffloadBrightnessStrategyTest.java
@@ -28,6 +28,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 import com.android.server.display.brightness.StrategySelectionNotifyRequest;
 import com.android.server.display.feature.DisplayManagerFlags;
 
@@ -70,7 +71,8 @@
                         .setShouldUpdateScreenBrightnessSetting(true)
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mOffloadBrightnessStrategy.updateBrightness(displayPowerRequest);
+                mOffloadBrightnessStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
         assertEquals(PowerManager.BRIGHTNESS_INVALID_FLOAT, mOffloadBrightnessStrategy
                 .getOffloadScreenBrightness(), 0.0f);
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OverrideBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OverrideBrightnessStrategyTest.java
index 530245d..8e7b463 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OverrideBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/OverrideBrightnessStrategyTest.java
@@ -26,6 +26,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -58,7 +59,8 @@
                         .setDisplayBrightnessStrategyName(mOverrideBrightnessStrategy.getName())
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mOverrideBrightnessStrategy.updateBrightness(displayPowerRequest);
+                mOverrideBrightnessStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
     }
 
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategyTest.java
index 7147aa8..e799d0e 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/ScreenOffBrightnessStrategyTest.java
@@ -26,6 +26,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -56,7 +57,8 @@
                                 .getName())
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mScreenOffBrightnessModeStrategy.updateBrightness(displayPowerRequest);
+                mScreenOffBrightnessModeStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
     }
 }
diff --git a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategyTest.java b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategyTest.java
index 9830edb..aaada41 100644
--- a/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategyTest.java
+++ b/services/tests/displayservicetests/src/com/android/server/display/brightness/strategy/TemporaryBrightnessStrategyTest.java
@@ -26,6 +26,7 @@
 
 import com.android.server.display.DisplayBrightnessState;
 import com.android.server.display.brightness.BrightnessReason;
+import com.android.server.display.brightness.StrategyExecutionRequest;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -58,7 +59,8 @@
                         .setDisplayBrightnessStrategyName(mTemporaryBrightnessStrategy.getName())
                         .build();
         DisplayBrightnessState updatedDisplayBrightnessState =
-                mTemporaryBrightnessStrategy.updateBrightness(displayPowerRequest);
+                mTemporaryBrightnessStrategy.updateBrightness(
+                        new StrategyExecutionRequest(displayPowerRequest, 0.2f));
         assertEquals(updatedDisplayBrightnessState, expectedDisplayBrightnessState);
     }
 
diff --git a/services/tests/dreamservicetests/src/com/android/server/dreams/DreamAccessibilityTest.java b/services/tests/dreamservicetests/src/com/android/server/dreams/DreamAccessibilityTest.java
new file mode 100644
index 0000000..99968d5
--- /dev/null
+++ b/services/tests/dreamservicetests/src/com/android/server/dreams/DreamAccessibilityTest.java
@@ -0,0 +1,163 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.dreams;
+
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.res.Resources;
+import android.service.dreams.utils.DreamAccessibility;
+import android.text.TextUtils;
+import android.view.View;
+import android.view.accessibility.AccessibilityNodeInfo;
+
+import androidx.test.filters.SmallTest;
+import androidx.test.runner.AndroidJUnit4;
+
+import com.android.internal.R;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.ArrayList;
+import java.util.Collections;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DreamAccessibilityTest {
+
+    @Mock
+    private View mView;
+
+    @Mock
+    private Context mContext;
+
+    @Mock
+    private Resources mResources;
+
+    @Mock
+    private AccessibilityNodeInfo mAccessibilityNodeInfo;
+
+    @Captor
+    private ArgumentCaptor<View.AccessibilityDelegate> mAccessibilityDelegateArgumentCaptor;
+
+    private DreamAccessibility mDreamAccessibility;
+    private static final String CUSTOM_ACTION = "Custom Action";
+    private static final String EXISTING_ACTION = "Existing Action";
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mDreamAccessibility = new DreamAccessibility(mContext, mView);
+
+        when(mContext.getResources()).thenReturn(mResources);
+        when(mResources.getString(R.string.dream_accessibility_action_click))
+                .thenReturn(CUSTOM_ACTION);
+    }
+    /**
+     * Test to verify the configuration of accessibility actions within a view delegate.
+     */
+    @Test
+    public void testConfigureAccessibilityActions() {
+        when(mAccessibilityNodeInfo.getActionList()).thenReturn(new ArrayList<>());
+
+        mDreamAccessibility.updateAccessibilityConfiguration(false);
+
+        verify(mView).setAccessibilityDelegate(mAccessibilityDelegateArgumentCaptor.capture());
+        View.AccessibilityDelegate capturedDelegate =
+                mAccessibilityDelegateArgumentCaptor.getValue();
+
+        capturedDelegate.onInitializeAccessibilityNodeInfo(mView, mAccessibilityNodeInfo);
+
+        verify(mAccessibilityNodeInfo).addAction(argThat(action ->
+                action.getId() == AccessibilityNodeInfo.ACTION_CLICK
+                        && TextUtils.equals(action.getLabel(), CUSTOM_ACTION)));
+    }
+
+    /**
+     * Test to verify the configuration of accessibility actions within a view delegate,
+     * specifically checking the removal of an existing click action and addition
+     * of a new custom action.
+     */
+    @Test
+    public void testConfigureAccessibilityActions_RemovesExistingClickAction() {
+        AccessibilityNodeInfo.AccessibilityAction existingAction =
+                new AccessibilityNodeInfo.AccessibilityAction(AccessibilityNodeInfo.ACTION_CLICK,
+                        EXISTING_ACTION);
+        when(mAccessibilityNodeInfo.getActionList())
+                .thenReturn(Collections.singletonList(existingAction));
+
+        mDreamAccessibility.updateAccessibilityConfiguration(false);
+
+        verify(mView).setAccessibilityDelegate(mAccessibilityDelegateArgumentCaptor.capture());
+        View.AccessibilityDelegate capturedDelegate =
+                mAccessibilityDelegateArgumentCaptor.getValue();
+
+        capturedDelegate.onInitializeAccessibilityNodeInfo(mView, mAccessibilityNodeInfo);
+
+        verify(mAccessibilityNodeInfo).removeAction(existingAction);
+        verify(mAccessibilityNodeInfo).addAction(argThat(action ->
+                action.getId() == AccessibilityNodeInfo.ACTION_CLICK
+                        && TextUtils.equals(action.getLabel(), CUSTOM_ACTION)));
+
+    }
+
+    /**
+     * Test to verify the removal of a custom accessibility action within a view delegate.
+     */
+    @Test
+    public void testRemoveCustomAccessibilityAction() {
+
+        AccessibilityNodeInfo.AccessibilityAction existingAction =
+                new AccessibilityNodeInfo.AccessibilityAction(AccessibilityNodeInfo.ACTION_CLICK,
+                        EXISTING_ACTION);
+        when(mAccessibilityNodeInfo.getActionList())
+                .thenReturn(Collections.singletonList(existingAction));
+
+        mDreamAccessibility.updateAccessibilityConfiguration(false);
+        verify(mView).setAccessibilityDelegate(mAccessibilityDelegateArgumentCaptor.capture());
+        View.AccessibilityDelegate capturedDelegate =
+                mAccessibilityDelegateArgumentCaptor.getValue();
+        when(mView.getAccessibilityDelegate()).thenReturn(capturedDelegate);
+        clearInvocations(mView);
+
+        mDreamAccessibility.updateAccessibilityConfiguration(true);
+        verify(mView).setAccessibilityDelegate(null);
+    }
+
+    /**
+     * Test to verify the removal of custom accessibility action is not called if delegate is not
+     * set by the dreamService.
+     */
+    @Test
+    public void testRemoveCustomAccessibility_DoesNotRemoveDelegateNotSetByDreamAccessibility() {
+        mDreamAccessibility.updateAccessibilityConfiguration(true);
+        verify(mView, never()).setAccessibilityDelegate(any());
+    }
+}
+
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
new file mode 100644
index 0000000..8b1d423
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsCollectorTest.java
@@ -0,0 +1,416 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.when;
+
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.BatteryStatsManager;
+import android.os.Handler;
+import android.os.WorkSource;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.IndentingPrintWriter;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerStats;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.StringWriter;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.IntSupplier;
+import java.util.function.Supplier;
+
+public class WifiPowerStatsCollectorTest {
+    private static final int APP_UID1 = 42;
+    private static final int APP_UID2 = 24;
+    private static final int APP_UID3 = 44;
+    private static final int ISOLATED_UID = 99123;
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setPowerStatsThrottlePeriodMillis(BatteryConsumer.POWER_COMPONENT_WIFI, 1000);
+
+    private MockBatteryStatsImpl mBatteryStats;
+
+    private final MockClock mClock = mStatsRule.getMockClock();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private WifiManager mWifiManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+
+    private NetworkStats mNetworkStats;
+    private List<NetworkStats.Entry> mNetworkStatsEntries;
+
+    private static class ScanTimes {
+        public long scanTimeMs;
+        public long batchScanTimeMs;
+    }
+
+    private final SparseArray<ScanTimes> mScanTimes = new SparseArray<>();
+    private long mWifiActiveDuration;
+
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+        @Override
+        public void retrieveWifiScanTimes(Callback callback) {
+            for (int i = 0; i < mScanTimes.size(); i++) {
+                int uid = mScanTimes.keyAt(i);
+                ScanTimes scanTimes = mScanTimes.valueAt(i);
+                callback.onWifiScanTime(uid, scanTimes.scanTimeMs, scanTimes.batchScanTimeMs);
+            }
+        }
+
+        @Override
+        public long getWifiActiveDuration() {
+            return mWifiActiveDuration;
+        }
+    };
+
+    private final List<PowerStats> mRecordedPowerStats = new ArrayList<>();
+
+    private WifiPowerStatsCollector.Injector mInjector = new WifiPowerStatsCollector.Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return mPowerStatsUidResolver;
+        }
+
+        @Override
+        public PackageManager getPackageManager() {
+            return mPackageManager;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> 3500;
+        }
+
+        @Override
+        public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+            return mNetworkStatsSupplier;
+        }
+
+        @Override
+        public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+            return mWifiStatsRetriever;
+        }
+
+        @Override
+        public WifiManager getWifiManager() {
+            return mWifiManager;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt())).thenAnswer(invocation -> {
+            int uid = invocation.getArgument(0);
+            if (uid == ISOLATED_UID) {
+                return APP_UID2;
+            } else {
+                return uid;
+            }
+        });
+        mBatteryStats = mStatsRule.getBatteryStats();
+    }
+
+    @SuppressWarnings("GuardedBy")
+    @Test
+    public void triggering() throws Throwable {
+        PowerStatsCollector collector = mBatteryStats.getPowerStatsCollector(
+                BatteryConsumer.POWER_COMPONENT_WIFI);
+        collector.addConsumer(mRecordedPowerStats::add);
+
+        mBatteryStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_WIFI, true);
+
+        mockWifiActivityInfo(1000, 2000, 3000, 600, 100);
+
+        // This should trigger a sample collection to establish a baseline
+        mBatteryStats.onSystemReady(mContext);
+
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(20000, 20000);
+        mBatteryStats.noteWifiOnLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(40000, 40000);
+        mBatteryStats.noteWifiOffLocked(mClock.realtime, mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(50000, 50000);
+        mBatteryStats.noteWifiRunningLocked(new WorkSource(APP_UID1), mClock.realtime,
+                mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(60000, 60000);
+        mBatteryStats.noteWifiStoppedLocked(new WorkSource(APP_UID1), mClock.realtime,
+                mClock.uptime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+
+        mRecordedPowerStats.clear();
+        mStatsRule.setTime(70000, 70000);
+        mBatteryStats.noteWifiStateLocked(BatteryStatsManager.WIFI_STATE_ON_CONNECTED_STA,
+                "mywyfy", mClock.realtime);
+        mStatsRule.waitForBackgroundThread();
+        assertThat(mRecordedPowerStats).hasSize(1);
+    }
+
+    @Test
+    public void collectStats_powerReportingSupported() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        assertThat(powerStats.durationMs).isEqualTo(7500);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(descriptor);
+        assertThat(layout.isPowerReportingSupported()).isTrue();
+        assertThat(layout.getDeviceRxTime(powerStats.stats)).isEqualTo(6000);
+        assertThat(layout.getDeviceTxTime(powerStats.stats)).isEqualTo(1000);
+        assertThat(layout.getDeviceScanTime(powerStats.stats)).isEqualTo(200);
+        assertThat(layout.getDeviceIdleTime(powerStats.stats)).isEqualTo(300);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        verifyUidStats(powerStats);
+    }
+
+    @Test
+    public void collectStats_powerReportingUnsupported() {
+        PowerStats powerStats = collectPowerStats(false);
+        assertThat(powerStats.durationMs).isEqualTo(13200);
+
+        PowerStats.Descriptor descriptor = powerStats.descriptor;
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(descriptor);
+        assertThat(layout.isPowerReportingSupported()).isFalse();
+        assertThat(layout.getDeviceActiveTime(powerStats.stats)).isEqualTo(7500);
+        assertThat(layout.getDeviceBasicScanTime(powerStats.stats)).isEqualTo(234 + 100 + 300);
+        assertThat(layout.getDeviceBatchedScanTime(powerStats.stats)).isEqualTo(345 + 200 + 400);
+        assertThat(layout.getConsumedEnergy(powerStats.stats, 0))
+                .isEqualTo((64321 - 10000) * 1000 / 3500);
+
+        verifyUidStats(powerStats);
+    }
+
+    private void verifyUidStats(PowerStats powerStats) {
+        WifiPowerStatsLayout layout = new WifiPowerStatsLayout(powerStats.descriptor);
+        assertThat(powerStats.uidStats.size()).isEqualTo(2);
+        long[] actual1 = powerStats.uidStats.get(APP_UID1);
+        assertThat(layout.getUidRxBytes(actual1)).isEqualTo(1000);
+        assertThat(layout.getUidTxBytes(actual1)).isEqualTo(2000);
+        assertThat(layout.getUidRxPackets(actual1)).isEqualTo(100);
+        assertThat(layout.getUidTxPackets(actual1)).isEqualTo(200);
+        assertThat(layout.getUidScanTime(actual1)).isEqualTo(234);
+        assertThat(layout.getUidBatchedScanTime(actual1)).isEqualTo(345);
+
+        // Combines APP_UID2 and ISOLATED_UID
+        long[] actual2 = powerStats.uidStats.get(APP_UID2);
+        assertThat(layout.getUidRxBytes(actual2)).isEqualTo(6000);
+        assertThat(layout.getUidTxBytes(actual2)).isEqualTo(3000);
+        assertThat(layout.getUidRxPackets(actual2)).isEqualTo(60);
+        assertThat(layout.getUidTxPackets(actual2)).isEqualTo(30);
+        assertThat(layout.getUidScanTime(actual2)).isEqualTo(100 + 300);
+        assertThat(layout.getUidBatchedScanTime(actual2)).isEqualTo(200 + 400);
+
+        assertThat(powerStats.uidStats.get(ISOLATED_UID)).isNull();
+        assertThat(powerStats.uidStats.get(APP_UID3)).isNull();
+    }
+
+    @Test
+    public void dump() throws Throwable {
+        PowerStats powerStats = collectPowerStats(true);
+        StringWriter sw = new StringWriter();
+        IndentingPrintWriter pw = new IndentingPrintWriter(sw);
+        powerStats.dump(pw);
+        pw.flush();
+        String dump = sw.toString();
+        assertThat(dump).contains("duration=7500");
+        assertThat(dump).contains(
+                "stats=[6000, 1000, 300, 200, 634, 945, " + ((64321 - 10000) * 1000 / 3500)
+                        + ", 0, 0]");
+        assertThat(dump).contains("UID 24: [6000, 3000, 60, 30, 400, 600, 0]");
+        assertThat(dump).contains("UID 42: [1000, 2000, 100, 200, 234, 345, 0]");
+    }
+
+    private PowerStats collectPowerStats(boolean hasPowerReporting) {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(hasPowerReporting);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[]{777});
+
+        if (hasPowerReporting) {
+            mockWifiActivityInfo(1000, 600, 100, 2000, 3000);
+        } else {
+            mWifiActiveDuration = 5700;
+        }
+        mockNetworkStats(1000);
+        mockNetworkStatsEntry(APP_UID1, 4321, 321, 1234, 23);
+        mockNetworkStatsEntry(APP_UID2, 4000, 40, 2000, 20);
+        mockNetworkStatsEntry(ISOLATED_UID, 2000, 20, 1000, 10);
+        mockNetworkStatsEntry(APP_UID3, 314, 281, 314, 281);
+        mockWifiScanTimes(APP_UID1, 1000, 2000);
+        mockWifiScanTimes(APP_UID2, 3000, 4000);
+        mockWifiScanTimes(ISOLATED_UID, 5000, 6000);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{10000});
+
+        collector.collectStats();
+
+        if (hasPowerReporting) {
+            mockWifiActivityInfo(1100, 6600, 1100, 2200, 3300);
+        } else {
+            mWifiActiveDuration = 13200;
+        }
+        mockNetworkStats(1100);
+        mockNetworkStatsEntry(APP_UID1, 5321, 421, 3234, 223);
+        mockNetworkStatsEntry(APP_UID2, 8000, 80, 4000, 40);
+        mockNetworkStatsEntry(ISOLATED_UID, 4000, 40, 2000, 20);
+        mockNetworkStatsEntry(APP_UID3, 314, 281, 314, 281);    // Unchanged
+        mockWifiScanTimes(APP_UID1, 1234, 2345);
+        mockWifiScanTimes(APP_UID2, 3100, 4200);
+        mockWifiScanTimes(ISOLATED_UID, 5300, 6400);
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(eq(new int[]{777})))
+                .thenReturn(new long[]{64321});
+
+        mStatsRule.setTime(20000, 20000);
+        return collector.collectStats();
+    }
+
+    private void mockWifiActivityInfo(long timestamp, long rxTimeMs, long txTimeMs, int scanTimeMs,
+            int idleTimeMs) {
+        int stackState = 0;
+        WifiActivityEnergyInfo info = new WifiActivityEnergyInfo(timestamp, stackState, txTimeMs,
+                rxTimeMs, scanTimeMs, idleTimeMs);
+        doAnswer(invocation -> {
+            WifiManager.OnWifiActivityEnergyInfoListener listener = invocation.getArgument(1);
+            listener.onWifiActivityEnergyInfo(info);
+            return null;
+        }).when(mWifiManager).getWifiActivityEnergyInfoAsync(any(), any());
+    }
+
+    private void mockNetworkStats(long elapsedRealtime) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            mNetworkStats = mock(NetworkStats.class);
+            ArrayList<NetworkStats.Entry> networkStatsEntries = new ArrayList<>();
+            when(mNetworkStats.iterator()).thenAnswer(inv -> networkStatsEntries.iterator());
+            mNetworkStatsEntries = networkStatsEntries;
+        } else {
+            mNetworkStats = new NetworkStats(elapsedRealtime, 1);
+        }
+        when(mNetworkStatsSupplier.get()).thenReturn(mNetworkStats);
+    }
+
+    private void mockNetworkStatsEntry(int uid, long rxBytes, long rxPackets, long txBytes,
+            long txPackets) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(METERED_NO);
+            when(entry.getRoaming()).thenReturn(ROAMING_NO);
+            when(entry.getDefaultNetwork()).thenReturn(DEFAULT_NETWORK_NO);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(100L);
+            mNetworkStatsEntries.add(entry);
+        } else {
+            mNetworkStats = mNetworkStats
+                    .addEntry(new NetworkStats.Entry("wifi", uid, 0, 0,
+                            METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, rxBytes, rxPackets,
+                            txBytes, txPackets, 100));
+            reset(mNetworkStatsSupplier);
+            when(mNetworkStatsSupplier.get()).thenReturn(mNetworkStats);
+        }
+    }
+
+    private void mockWifiScanTimes(int uid, long scanTimeMs, long batchScanTimeMs) {
+        ScanTimes scanTimes = new ScanTimes();
+        scanTimes.scanTimeMs = scanTimeMs;
+        scanTimes.batchScanTimeMs = batchScanTimeMs;
+        mScanTimes.put(uid, scanTimes);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
new file mode 100644
index 0000000..257a1a6
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/WifiPowerStatsProcessorTest.java
@@ -0,0 +1,592 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import static android.net.NetworkStats.DEFAULT_NETWORK_NO;
+import static android.net.NetworkStats.METERED_NO;
+import static android.net.NetworkStats.ROAMING_NO;
+import static android.os.BatteryConsumer.PROCESS_STATE_BACKGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_CACHED;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND;
+import static android.os.BatteryConsumer.PROCESS_STATE_FOREGROUND_SERVICE;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.annotation.Nullable;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.hardware.power.stats.EnergyConsumerType;
+import android.net.NetworkStats;
+import android.net.wifi.WifiManager;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.os.Process;
+import android.os.connectivity.WifiActivityEnergyInfo;
+import android.platform.test.ravenwood.RavenwoodRule;
+import android.util.SparseArray;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.List;
+import java.util.function.IntSupplier;
+import java.util.function.Supplier;
+
+public class WifiPowerStatsProcessorTest {
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    private static final double PRECISION = 0.00001;
+    private static final int APP_UID1 = Process.FIRST_APPLICATION_UID + 42;
+    private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
+    private static final int WIFI_ENERGY_CONSUMER_ID = 1;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_IDLE, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_RX, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_CONTROLLER_TX, 720.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_ACTIVE, 360.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_SCAN, 480.0)
+            .setAveragePower(PowerProfile.POWER_WIFI_BATCHED_SCAN, 720.0)
+            .initMeasuredEnergyStatsLocked();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PowerStatsUidResolver mPowerStatsUidResolver;
+    @Mock
+    private PackageManager mPackageManager;
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private Supplier<NetworkStats> mNetworkStatsSupplier;
+    @Mock
+    private WifiManager mWifiManager;
+
+    private static class ScanTimes {
+        public long scanTimeMs;
+        public long batchScanTimeMs;
+    }
+
+    private final SparseArray<ScanTimes> mScanTimes = new SparseArray<>();
+    private long mWifiActiveDuration;
+
+    private final WifiPowerStatsCollector.WifiStatsRetriever mWifiStatsRetriever =
+            new WifiPowerStatsCollector.WifiStatsRetriever() {
+        @Override
+        public void retrieveWifiScanTimes(Callback callback) {
+            for (int i = 0; i < mScanTimes.size(); i++) {
+                int uid = mScanTimes.keyAt(i);
+                ScanTimes scanTimes = mScanTimes.valueAt(i);
+                callback.onWifiScanTime(uid, scanTimes.scanTimeMs, scanTimes.batchScanTimeMs);
+            }
+        }
+
+        @Override
+        public long getWifiActiveDuration() {
+            return mWifiActiveDuration;
+        }
+    };
+
+    private final WifiPowerStatsCollector.Injector mInjector =
+            new WifiPowerStatsCollector.Injector() {
+                @Override
+                public Handler getHandler() {
+                    return mStatsRule.getHandler();
+                }
+
+                @Override
+                public Clock getClock() {
+                    return mStatsRule.getMockClock();
+                }
+
+                @Override
+                public PowerStatsUidResolver getUidResolver() {
+                    return mPowerStatsUidResolver;
+                }
+
+                @Override
+                public PackageManager getPackageManager() {
+                    return mPackageManager;
+                }
+
+                @Override
+                public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+                    return mConsumedEnergyRetriever;
+                }
+
+                @Override
+                public IntSupplier getVoltageSupplier() {
+                    return () -> VOLTAGE_MV;
+                }
+
+                @Override
+                public Supplier<NetworkStats> getWifiNetworkStatsSupplier() {
+                    return mNetworkStatsSupplier;
+                }
+
+                @Override
+                public WifiManager getWifiManager() {
+                    return mWifiManager;
+                }
+
+                @Override
+                public WifiPowerStatsCollector.WifiStatsRetriever getWifiStatsRetriever() {
+                    return mWifiStatsRetriever;
+                }
+            };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        when(mContext.getPackageManager()).thenReturn(mPackageManager);
+        when(mPackageManager.hasSystemFeature(PackageManager.FEATURE_WIFI)).thenReturn(true);
+        when(mPowerStatsUidResolver.mapUid(anyInt()))
+                .thenAnswer(invocation -> invocation.getArgument(0));
+    }
+
+    @Test
+    public void powerProfileModel_powerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
+
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[0]);
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty WifiActivityEnergyInfo.
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
+                WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(10000,
+                WifiActivityEnergyInfo.STACK_STATE_STATE_ACTIVE, 2000, 3000, 100, 600));
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // RX power = 'rx-duration * PowerProfile[wifi.controller.rx]`
+        //        RX power = 3000 * 480 = 1440000 mA-ms = 0.4 mAh
+        // TX power = 'tx-duration * PowerProfile[wifi.controller.tx]`
+        //        TX power = 2000 * 720 = 1440000 mA-ms = 0.4 mAh
+        // Scan power = 'scan-duration * PowerProfile[wifi.scan]`
+        //        Scan power = 100 * 480 = 48000 mA-ms = 0.013333 mAh
+        // Idle power = 'idle-duration * PowerProfile[wifi.idle]`
+        //        Idle power = 600 * 360 = 216000 mA-ms = 0.06 mAh
+        // Total power = RX + TX + Scan + Idle = 0.873333
+        // Screen-on  - 25%
+        // Screen-off - 75%
+        double expectedPower = 0.873333;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1 =
+        //     (1500 / 2000) * 0.4        // rx
+        //     + (300 / 400) * 0.4        // tx
+        //     + (700 / 1000) * 0.013333  // scan (basic + batched)
+        //   = 0.609333 mAh
+        double expectedPower1 = 0.609333;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2 =
+        //     (500 / 2000) * 0.4         // rx
+        //     + (100 / 400) * 0.4        // tx
+        //     + (300 / 1000) * 0.013333  // scan (basic + batched)
+        //   = 0.204 mAh
+        double expectedPower2 = 0.204;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    @Test
+    public void consumedEnergyModel_powerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(true);
+
+        // PowerStats hardware is available
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[] {WIFI_ENERGY_CONSUMER_ID});
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Initial empty WifiActivityEnergyInfo.
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(0L,
+                WifiActivityEnergyInfo.STACK_STATE_INVALID, 0L, 0L, 0L, 0L));
+
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{WIFI_ENERGY_CONSUMER_ID}))
+                .thenReturn(new long[]{0});
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mockWifiActivityEnergyInfo(new WifiActivityEnergyInfo(10000,
+                WifiActivityEnergyInfo.STACK_STATE_STATE_ACTIVE, 2000, 3000, 100, 600));
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        // 10 mAh represented as microWattSeconds
+        long energyUws = 10 * 3600 * VOLTAGE_MV;
+        when(mConsumedEnergyRetriever.getConsumedEnergyUws(
+                new int[]{WIFI_ENERGY_CONSUMER_ID})).thenReturn(new long[]{energyUws});
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // All estimates are computed as in the #powerProfileModel_powerController test,
+        // except they are all scaled by the same ratio to ensure that the total estimated
+        // energy is equal to the measured energy
+        double expectedPower = 10;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1
+        //   0.609333           // power profile model estimate
+        //   0.873333           // power profile model estimate for total power
+        //   10                 // total consumed energy
+        //   = 0.609333 * (10 / 0.873333) = 6.9771
+        double expectedPower1 = 6.9771;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2
+        //   0.204              // power profile model estimate
+        //   0.873333           // power profile model estimate for total power
+        //   10                 // total consumed energy
+        //   = 0.204 * (10 / 0.873333) = 2.33588
+        double expectedPower2 = 2.33588;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    @Test
+    public void powerProfileModel_noPowerController() {
+        when(mWifiManager.isEnhancedPowerReportingSupported()).thenReturn(false);
+
+        // No power monitoring hardware
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.WIFI))
+                .thenReturn(new int[0]);
+
+        WifiPowerStatsProcessor processor =
+                new WifiPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        WifiPowerStatsCollector collector = new WifiPowerStatsCollector(mInjector, 0);
+        collector.setEnabled(true);
+
+        // Establish a baseline
+        aggregatedStats.addPowerStats(collector.collectStats(), 0);
+
+        // Turn the screen off after 2.5 seconds
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_BACKGROUND, 2500);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND_SERVICE,
+                5000);
+
+        // Note application network activity
+        NetworkStats networkStats = mockNetworkStats(10000, 1,
+                mockNetworkStatsEntry("wifi", APP_UID1, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 10000, 1500, 20000, 300, 100),
+                mockNetworkStatsEntry("wifi", APP_UID2, 0, 0,
+                        METERED_NO, ROAMING_NO, DEFAULT_NETWORK_NO, 5000, 500, 3000, 100, 111));
+        when(mNetworkStatsSupplier.get()).thenReturn(networkStats);
+
+        mScanTimes.clear();
+        mWifiActiveDuration = 8000;
+        mockWifiScanTimes(APP_UID1, 300, 400);
+        mockWifiScanTimes(APP_UID2, 100, 200);
+
+        mStatsRule.setTime(10_000, 10_000);
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 10_000);
+
+        processor.finish(aggregatedStats);
+
+        WifiPowerStatsLayout statsLayout =
+                new WifiPowerStatsLayout(aggregatedStats.getPowerStatsDescriptor());
+
+        // Total active power = 'active-duration * PowerProfile[wifi.on]`
+        //        active = 8000 * 360 = 2880000 mA-ms = 0.8 mAh
+        // UID1 rxPackets + txPackets = 1800
+        // UID2 rxPackets + txPackets = 600
+        // Total rx+tx packets = 2400
+        // Total scan power = `scan-duration * PowerProfile[wifi.scan]`
+        //        scan = (100 + 300) * 480 = 192000 mA-ms = 0.05333 mAh
+        // Total batch scan power = `(200 + 400) * PowerProfile[wifi.batchedscan]`
+        //        bscan = (200 + 400) * 720 = 432000 mA-ms = 0.12 mAh
+        //
+        // Expected power = active + scan + bscan = 0.97333
+        double expectedPower = 0.97333;
+        long[] deviceStats = new long[aggregatedStats.getPowerStatsDescriptor().statsArrayLength];
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_ON));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.25);
+
+        aggregatedStats.getDeviceStats(deviceStats, states(POWER_STATE_OTHER, SCREEN_STATE_OTHER));
+        assertThat(statsLayout.getDevicePowerEstimate(deviceStats))
+                .isWithin(PRECISION).of(expectedPower * 0.75);
+
+        // UID1 =
+        //     (1800 / 2400) * 0.8      // active
+        //     + (300 / 400) * 0.05333  // scan
+        //     + (400 / 600) * 0.12     // batched scan
+        //   = 0.72 mAh
+        double expectedPower1 = 0.72;
+        long[] uidStats = new long[aggregatedStats.getPowerStatsDescriptor().uidStatsArrayLength];
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_FOREGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_BACKGROUND));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID1,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_FOREGROUND_SERVICE));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower1 * 0.5);
+
+        // UID2 =
+        //     (600 / 2400) * 0.8       // active
+        //     + (100 / 400) * 0.05333  // scan
+        //     + (200 / 600) * 0.12     // batched scan
+        //   = 0.253333 mAh
+        double expectedPower2 = 0.25333;
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_ON, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.25);
+
+        aggregatedStats.getUidStats(uidStats, APP_UID2,
+                states(POWER_STATE_OTHER, SCREEN_STATE_OTHER, PROCESS_STATE_CACHED));
+        assertThat(statsLayout.getUidPowerEstimate(uidStats))
+                .isWithin(PRECISION).of(expectedPower2 * 0.75);
+    }
+
+    private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
+            WifiPowerStatsProcessor processor) {
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(BatteryConsumer.POWER_COMPONENT_WIFI)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN, STATE_PROCESS_STATE)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+        aggregatedStats.setUidState(APP_UID1, STATE_PROCESS_STATE, PROCESS_STATE_FOREGROUND, 0);
+        aggregatedStats.setUidState(APP_UID2, STATE_PROCESS_STATE, PROCESS_STATE_CACHED, 0);
+
+        return aggregatedStats;
+    }
+
+    private int[] states(int... states) {
+        return states;
+    }
+
+    private void mockWifiActivityEnergyInfo(WifiActivityEnergyInfo waei) {
+        doAnswer(invocation -> {
+            WifiManager.OnWifiActivityEnergyInfoListener
+                    listener = invocation.getArgument(1);
+            listener.onWifiActivityEnergyInfo(waei);
+            return null;
+        }).when(mWifiManager).getWifiActivityEnergyInfoAsync(any(), any());
+    }
+
+    private NetworkStats mockNetworkStats(int elapsedTime, int initialSize,
+            NetworkStats.Entry... entries) {
+        NetworkStats stats;
+        if (RavenwoodRule.isOnRavenwood()) {
+            stats = mock(NetworkStats.class);
+            when(stats.iterator()).thenAnswer(inv -> List.of(entries).iterator());
+        } else {
+            stats = new NetworkStats(elapsedTime, initialSize);
+            for (NetworkStats.Entry entry : entries) {
+                stats = stats.addEntry(entry);
+            }
+        }
+        return stats;
+    }
+
+    private static NetworkStats.Entry mockNetworkStatsEntry(@Nullable String iface, int uid,
+            int set, int tag, int metered, int roaming, int defaultNetwork, long rxBytes,
+            long rxPackets, long txBytes, long txPackets, long operations) {
+        if (RavenwoodRule.isOnRavenwood()) {
+            NetworkStats.Entry entry = mock(NetworkStats.Entry.class);
+            when(entry.getUid()).thenReturn(uid);
+            when(entry.getMetered()).thenReturn(metered);
+            when(entry.getRoaming()).thenReturn(roaming);
+            when(entry.getDefaultNetwork()).thenReturn(defaultNetwork);
+            when(entry.getRxBytes()).thenReturn(rxBytes);
+            when(entry.getRxPackets()).thenReturn(rxPackets);
+            when(entry.getTxBytes()).thenReturn(txBytes);
+            when(entry.getTxPackets()).thenReturn(txPackets);
+            when(entry.getOperations()).thenReturn(operations);
+            return entry;
+        } else {
+            return new NetworkStats.Entry(iface, uid, set, tag, metered,
+                    roaming, defaultNetwork, rxBytes, rxPackets, txBytes, txPackets, operations);
+        }
+    }
+
+    private void mockWifiScanTimes(int uid, long scanTimeMs, long batchScanTimeMs) {
+        ScanTimes scanTimes = new ScanTimes();
+        scanTimes.scanTimeMs = scanTimeMs;
+        scanTimes.batchScanTimeMs = batchScanTimeMs;
+        mScanTimes.put(uid, scanTimes);
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java b/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
index a8e0c26..bd16813 100644
--- a/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
+++ b/services/tests/servicestests/src/com/android/server/biometrics/sensors/fingerprint/aidl/FingerprintAuthenticationClientTest.java
@@ -49,6 +49,7 @@
 import android.hardware.biometrics.BiometricSourceType;
 import android.hardware.biometrics.common.ICancellationSignal;
 import android.hardware.biometrics.common.OperationContext;
+import android.hardware.biometrics.common.OperationState;
 import android.hardware.biometrics.events.AuthenticationAcquiredInfo;
 import android.hardware.biometrics.events.AuthenticationErrorInfo;
 import android.hardware.biometrics.events.AuthenticationFailedInfo;
@@ -391,8 +392,8 @@
         verify(mBiometricContext).subscribe(mOperationContextCaptor.capture(),
                 mStartHalConsumerCaptor.capture(), mContextInjector.capture(), any());
 
-        mStartHalConsumerCaptor.getValue().accept(mOperationContextCaptor
-                .getValue().toAidlContext());
+        final OperationContextExt operationContext = mOperationContextCaptor.getValue();
+        mStartHalConsumerCaptor.getValue().accept(operationContext.toAidlContext());
         final ArgumentCaptor<OperationContext> captor =
                 ArgumentCaptor.forClass(OperationContext.class);
 
@@ -400,16 +401,20 @@
 
         OperationContext opContext = captor.getValue();
 
-        assertThat(opContext).isSameInstanceAs(
-                mOperationContextCaptor.getValue().toAidlContext());
+        assertThat(opContext).isSameInstanceAs(operationContext.toAidlContext());
 
+        opContext.operationState = new OperationState();
+        opContext.operationState.setFingerprintOperationState(
+                new OperationState.FingerprintOperationState());
         mContextInjector.getValue().accept(opContext);
 
         verify(mHal).onContextChanged(same(opContext));
+        verify(mHal, times(2)).setIgnoreDisplayTouches(
+                opContext.operationState.getFingerprintOperationState().isHardwareIgnoringTouches);
 
         client.stopHalOperation();
 
-        verify(mBiometricContext).unsubscribe(same(mOperationContextCaptor.getValue()));
+        verify(mBiometricContext).unsubscribe(same(operationContext));
     }
 
     @Test
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
index de3cfbf..855c658 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
@@ -30,8 +30,6 @@
 import android.content.pm.IPackageManager;
 import android.content.pm.PackageManagerInternal;
 import android.database.ContentObserver;
-import android.hardware.usb.UsbManager;
-import android.media.IAudioService;
 import android.net.IIpConnectivityMetrics;
 import android.net.Uri;
 import android.os.Bundle;
@@ -216,11 +214,6 @@
         }
 
         @Override
-        IAudioService getIAudioService() {
-            return services.iaudioService;
-        }
-
-        @Override
         PersistentDataBlockManagerInternal getPersistentDataBlockManagerInternal() {
             return services.persistentDataBlockManagerInternal;
         }
@@ -244,11 +237,6 @@
         }
 
         @Override
-        UsbManager getUsbManager() {
-            return services.usbManager;
-        }
-
-        @Override
         boolean storageManagerIsFileBasedEncryptionEnabled() {
             return services.storageManager.isFileBasedEncryptionEnabled();
         }
@@ -380,21 +368,11 @@
         }
 
         @Override
-        void settingsSecurePutInt(String name, int value) {
-            services.settings.settingsSecurePutInt(name, value);
-        }
-
-        @Override
         void settingsGlobalPutInt(String name, int value) {
             services.settings.settingsGlobalPutInt(name, value);
         }
 
         @Override
-        void settingsSecurePutString(String name, String value) {
-            services.settings.settingsSecurePutString(name, value);
-        }
-
-        @Override
         void settingsGlobalPutString(String name, String value) {
             services.settings.settingsGlobalPutString(name, value);
         }
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskSupervisorTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskSupervisorTests.java
index 6b614fa..27a4a2b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskSupervisorTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskSupervisorTests.java
@@ -343,6 +343,20 @@
         verify(mAtm).setLastResumedActivityUncheckLocked(any(), eq("test"));
     }
 
+    @Test
+    public void testUpdateTopResumed_moveToFront() {
+        final ActivityRecord activity1 = new ActivityBuilder(mAtm).setCreateTask(true).build();
+        final ActivityRecord activity2 = new ActivityBuilder(mAtm).setCreateTask(true).build();
+        activity2.setState(ActivityRecord.State.RESUMED, "test");
+        assertEquals(activity2.app, mAtm.mTopApp);
+        activity1.getTask().moveToFront("test");
+        // If the device is not sleeping, the app should be only set with resumed state.
+        assertEquals(activity2.app, mAtm.mTopApp);
+        activity2.setState(ActivityRecord.State.PAUSED, "test");
+        activity1.setState(ActivityRecord.State.RESUMED, "test");
+        assertEquals(activity1.app, mAtm.mTopApp);
+    }
+
     /**
      * We need to launch home again after user unlocked for those displays that do not have
      * encryption aware home app.
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
index 69b5c37..43b424f 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowOrganizerTests.java
@@ -1681,8 +1681,7 @@
         WindowContainerToken wct = rootTask.mRemoteToken.toWindowContainerToken();
         t.setWindowingMode(wct, WINDOWING_MODE_PINNED);
         mWm.mAtmService.mWindowOrganizerController.applyTransaction(t);
-        verify(mWm.mAtmService.mRootWindowContainer).resumeFocusedTasksTopActivitiesUnchecked(any(),
-                any(), any(), anyBoolean());
+        verify(mWm.mAtmService.mRootWindowContainer).resumeFocusedTasksTopActivities();
 
         clearInvocations(mWm.mAtmService.mRootWindowContainer);
         // The token for the PIP root task may have changed when the task entered PIP mode, so do
@@ -1691,8 +1690,7 @@
                 record.getRootTask().mRemoteToken.toWindowContainerToken();
         t.setWindowingMode(newToken, WINDOWING_MODE_FULLSCREEN);
         mWm.mAtmService.mWindowOrganizerController.applyTransaction(t);
-        verify(mWm.mAtmService.mRootWindowContainer).resumeFocusedTasksTopActivitiesUnchecked(any(),
-                any(), any(), anyBoolean());
+        verify(mWm.mAtmService.mRootWindowContainer).resumeFocusedTasksTopActivities();
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowProcessControllerTests.java b/services/tests/wmtests/src/com/android/server/wm/WindowProcessControllerTests.java
index 400e4b6..059fed20 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowProcessControllerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowProcessControllerTests.java
@@ -406,6 +406,7 @@
         verify(tracker).onActivityResumedWhileVisible(mWpc);
         assertTrue(tracker.hasResumedActivity(mWpc.mUid));
 
+        mAtm.mTopApp = null;
         activity.makeFinishingLocked();
         activity.setState(PAUSING, "test");