Merge "Add OWNERS file for bubbles sub-module in WM shell" into main
diff --git a/apex/jobscheduler/service/aconfig/job.aconfig b/apex/jobscheduler/service/aconfig/job.aconfig
index e20f525..e489c1a 100644
--- a/apex/jobscheduler/service/aconfig/job.aconfig
+++ b/apex/jobscheduler/service/aconfig/job.aconfig
@@ -38,3 +38,13 @@
         purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+   name: "thermal_restrictions_to_fgs_jobs"
+   namespace: "backstage_power"
+   description: "Apply thermal restrictions to FGS jobs."
+   bug: "315157163"
+   metadata {
+       purpose: PURPOSE_BUGFIX
+   }
+}
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
index 3bb395f..ba8e3e8 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobConcurrencyManager.java
@@ -1375,8 +1375,10 @@
             final JobServiceContext jsc = mActiveServices.get(i);
             final JobStatus jobStatus = jsc.getRunningJobLocked();
 
-            if (jobStatus != null && !jsc.isWithinExecutionGuaranteeTime()
-                    && restriction.isJobRestricted(jobStatus)) {
+            if (jobStatus != null
+                    && !jsc.isWithinExecutionGuaranteeTime()
+                    && restriction.isJobRestricted(
+                            jobStatus, mService.evaluateJobBiasLocked(jobStatus))) {
                 jsc.cancelExecutingJobLocked(restriction.getStopReason(),
                         restriction.getInternalReason(),
                         JobParameters.getInternalReasonCodeDescription(
diff --git a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
index 5d1433c..384d786 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/JobSchedulerService.java
@@ -310,7 +310,8 @@
      * Note: do not add to or remove from this list at runtime except in the constructor, because we
      * do not synchronize access to this list.
      */
-    private final List<JobRestriction> mJobRestrictions;
+    @VisibleForTesting
+    final List<JobRestriction> mJobRestrictions;
 
     @GuardedBy("mLock")
     @VisibleForTesting
@@ -3498,8 +3499,6 @@
 
     /**
      * Check if a job is restricted by any of the declared {@link JobRestriction JobRestrictions}.
-     * Note, that the jobs with {@link JobInfo#BIAS_FOREGROUND_SERVICE} bias or higher may not
-     * be restricted, thus we won't even perform the check, but simply return null early.
      *
      * @param job to be checked
      * @return the first {@link JobRestriction} restricting the given job that has been found; null
@@ -3508,13 +3507,9 @@
      */
     @GuardedBy("mLock")
     JobRestriction checkIfRestricted(JobStatus job) {
-        if (evaluateJobBiasLocked(job) >= JobInfo.BIAS_FOREGROUND_SERVICE) {
-            // Jobs with BIAS_FOREGROUND_SERVICE or higher should not be restricted
-            return null;
-        }
         for (int i = mJobRestrictions.size() - 1; i >= 0; i--) {
             final JobRestriction restriction = mJobRestrictions.get(i);
-            if (restriction.isJobRestricted(job)) {
+            if (restriction.isJobRestricted(job, evaluateJobBiasLocked(job))) {
                 return restriction;
             }
         }
@@ -4221,6 +4216,7 @@
         return curBias;
     }
 
+    /** Gets and returns the adjusted Job Bias **/
     int evaluateJobBiasLocked(JobStatus job) {
         int bias = job.getBias();
         if (bias >= JobInfo.BIAS_BOUND_FOREGROUND_SERVICE) {
@@ -5907,7 +5903,7 @@
                     if (isRestricted) {
                         for (int i = mJobRestrictions.size() - 1; i >= 0; i--) {
                             final JobRestriction restriction = mJobRestrictions.get(i);
-                            if (restriction.isJobRestricted(job)) {
+                            if (restriction.isJobRestricted(job, evaluateJobBiasLocked(job))) {
                                 final int reason = restriction.getInternalReason();
                                 pw.print(" ");
                                 pw.print(JobParameters.getInternalReasonCodeDescription(reason));
@@ -6240,7 +6236,7 @@
                         proto.write(JobSchedulerServiceDumpProto.JobRestriction.REASON,
                                 restriction.getInternalReason());
                         proto.write(JobSchedulerServiceDumpProto.JobRestriction.IS_RESTRICTING,
-                                restriction.isJobRestricted(job));
+                                restriction.isJobRestricted(job, evaluateJobBiasLocked(job)));
                         proto.end(restrictionsToken);
                     }
 
diff --git a/apex/jobscheduler/service/java/com/android/server/job/restrictions/JobRestriction.java b/apex/jobscheduler/service/java/com/android/server/job/restrictions/JobRestriction.java
index 7aab67a..555a118 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/restrictions/JobRestriction.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/restrictions/JobRestriction.java
@@ -62,10 +62,11 @@
      * fine with it).
      *
      * @param job to be checked
+     * @param bias job bias to be checked
      * @return false if the {@link JobSchedulerService} should not schedule this job at the moment,
      * true - otherwise
      */
-    public abstract boolean isJobRestricted(JobStatus job);
+    public abstract boolean isJobRestricted(JobStatus job, int bias);
 
     /** Dump any internal constants the Restriction may have. */
     public abstract void dumpConstants(IndentingPrintWriter pw);
diff --git a/apex/jobscheduler/service/java/com/android/server/job/restrictions/ThermalStatusRestriction.java b/apex/jobscheduler/service/java/com/android/server/job/restrictions/ThermalStatusRestriction.java
index ef634b5..ba01113 100644
--- a/apex/jobscheduler/service/java/com/android/server/job/restrictions/ThermalStatusRestriction.java
+++ b/apex/jobscheduler/service/java/com/android/server/job/restrictions/ThermalStatusRestriction.java
@@ -24,6 +24,7 @@
 import android.util.IndentingPrintWriter;
 
 import com.android.internal.annotations.VisibleForTesting;
+import com.android.server.job.Flags;
 import com.android.server.job.JobSchedulerService;
 import com.android.server.job.controllers.JobStatus;
 
@@ -85,7 +86,18 @@
     }
 
     @Override
-    public boolean isJobRestricted(JobStatus job) {
+    public boolean isJobRestricted(JobStatus job, int bias) {
+        if (Flags.thermalRestrictionsToFgsJobs()) {
+            if (bias >= JobInfo.BIAS_TOP_APP) {
+                // Jobs with BIAS_TOP_APP should not be restricted
+                return false;
+            }
+        } else {
+            if (bias >= JobInfo.BIAS_FOREGROUND_SERVICE) {
+                // Jobs with BIAS_FOREGROUND_SERVICE or higher should not be restricted
+                return false;
+            }
+        }
         if (mThermalStatus >= UPPER_THRESHOLD) {
             return true;
         }
@@ -107,6 +119,17 @@
                         || (mService.isCurrentlyRunningLocked(job)
                                 && mService.isJobInOvertimeLocked(job));
             }
+            if (Flags.thermalRestrictionsToFgsJobs()) {
+                // Only let foreground jobs run if:
+                // 1. They haven't previously run
+                // 2. They're already running and aren't yet in overtime
+                if (bias >= JobInfo.BIAS_FOREGROUND_SERVICE
+                        && job.getJob().isImportantWhileForeground()) {
+                    return job.getNumPreviousAttempts() > 0
+                            || (mService.isCurrentlyRunningLocked(job)
+                                    && mService.isJobInOvertimeLocked(job));
+                }
+            }
             if (priority == JobInfo.PRIORITY_HIGH) {
                 return !mService.isCurrentlyRunningLocked(job)
                         || mService.isJobInOvertimeLocked(job);
@@ -114,6 +137,13 @@
             return true;
         }
         if (mThermalStatus >= LOW_PRIORITY_THRESHOLD) {
+            if (Flags.thermalRestrictionsToFgsJobs()) {
+                if (bias >= JobInfo.BIAS_FOREGROUND_SERVICE) {
+                    // No restrictions on foreground jobs
+                    // on LOW_PRIORITY_THRESHOLD and below
+                    return false;
+                }
+            }
             // For light throttling, throttle all min priority jobs and all low priority jobs that
             // aren't already running or have been running for long enough.
             return priority == JobInfo.PRIORITY_MIN
diff --git a/cmds/bootanimation/FORMAT.md b/cmds/bootanimation/FORMAT.md
index 01e8fe1..da8331a 100644
--- a/cmds/bootanimation/FORMAT.md
+++ b/cmds/bootanimation/FORMAT.md
@@ -126,7 +126,7 @@
 Use `zopflipng` if you have it, otherwise `pngcrush` will do. e.g.:
 
     for fn in *.png ; do
-        zopflipng -m ${fn}s ${fn}s.new && mv -f ${fn}s.new ${fn}
+        zopflipng -m ${fn} ${fn}.new && mv -f ${fn}.new ${fn}
         # or: pngcrush -q ....
     done
 
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 3c402ca..4ac6bac 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -6490,7 +6490,8 @@
         // visual regressions.
         @SuppressWarnings("AndroidFrameworkCompatChange")
         private boolean bigContentViewRequired() {
-            if (mContext.getApplicationInfo().targetSdkVersion >= Build.VERSION_CODES.S) {
+            if (!Flags.notificationExpansionOptional()
+                    && mContext.getApplicationInfo().targetSdkVersion >= Build.VERSION_CODES.S) {
                 return true;
             }
             // Notifications with contentView and without a bigContentView, style, or actions would
diff --git a/core/java/android/app/notification.aconfig b/core/java/android/app/notification.aconfig
index 63ffaa0..50c7b7f 100644
--- a/core/java/android/app/notification.aconfig
+++ b/core/java/android/app/notification.aconfig
@@ -60,6 +60,13 @@
 }
 
 flag {
+  name: "notification_expansion_optional"
+  namespace: "systemui"
+  description: "Experiment to restore the pre-S behavior where standard notifications are not expandable unless they have actions."
+  bug: "339523906"
+}
+
+flag {
   name: "keyguard_private_notifications"
   namespace: "systemui"
   description: "Fixes the behavior of KeyguardManager#setPrivateNotificationsAllowed()"
diff --git a/core/java/android/content/pm/flags.aconfig b/core/java/android/content/pm/flags.aconfig
index 45591d7..cee8d96 100644
--- a/core/java/android/content/pm/flags.aconfig
+++ b/core/java/android/content/pm/flags.aconfig
@@ -137,6 +137,18 @@
 }
 
 flag {
+    name: "get_package_storage_stats"
+    namespace: "system_performance"
+    is_exported: true
+    description: "Add dumpsys entry point for package StorageStats"
+    bug: "332905331"
+    is_fixed_read_only: true
+    metadata {
+        purpose: PURPOSE_BUGFIX
+    }
+}
+
+flag {
     name: "provide_info_of_apk_in_apex"
     is_exported: true
     namespace: "package_manager_service"
diff --git a/core/java/android/view/SurfaceView.java b/core/java/android/view/SurfaceView.java
index a23df79..1d70d18 100644
--- a/core/java/android/view/SurfaceView.java
+++ b/core/java/android/view/SurfaceView.java
@@ -986,7 +986,7 @@
 
             updateBackgroundVisibility(surfaceUpdateTransaction);
             updateBackgroundColor(surfaceUpdateTransaction);
-            if (mLimitedHdrEnabled && hdrHeadroomChanged) {
+            if (mLimitedHdrEnabled && (hdrHeadroomChanged || creating)) {
                 surfaceUpdateTransaction.setDesiredHdrHeadroom(
                         mBlastSurfaceControl, mHdrHeadroom);
             }
diff --git a/core/java/android/webkit/WebChromeClient.java b/core/java/android/webkit/WebChromeClient.java
index a07141b..b7ee0b8 100644
--- a/core/java/android/webkit/WebChromeClient.java
+++ b/core/java/android/webkit/WebChromeClient.java
@@ -520,6 +520,13 @@
      * To cancel the request, call <code>filePathCallback.onReceiveValue(null)</code> and
      * return {@code true}.
      *
+     * <p class="note"><b>Note:</b> WebView does not enforce any restrictions on
+     * the chosen file(s). WebView can access all files that your app can access.
+     * In case the file(s) are chosen through an untrusted source such as a third-party
+     * app, it is your own app's responsibility to check what the returned Uris
+     * refer to before calling the <code>filePathCallback</code>. See
+     * {@link #createIntent} and {@link #parseResult} for more details.</p>
+     *
      * @param webView The WebView instance that is initiating the request.
      * @param filePathCallback Invoke this callback to supply the list of paths to files to upload,
      *                         or {@code null} to cancel. Must only be called if the
@@ -556,6 +563,15 @@
          * Parse the result returned by the file picker activity. This method should be used with
          * {@link #createIntent}. Refer to {@link #createIntent} for how to use it.
          *
+         * <p class="note"><b>Note:</b> The intent returned by the file picker activity
+         * should be treated as untrusted. A third-party app handling the implicit
+         * intent created by {@link #createIntent} might return Uris that the third-party
+         * app itself does not have access to, such as your own app's sensitive data files.
+         * WebView does not enforce any restrictions on the returned Uris. It is the
+         * app's responsibility to ensure that the untrusted source (such as a third-party
+         * app) has access the Uris it has returned and that the Uris are not pointing
+         * to any sensitive data files.</p>
+         *
          * @param resultCode the integer result code returned by the file picker activity.
          * @param data the intent returned by the file picker activity.
          * @return the Uris of selected file(s) or {@code null} if the resultCode indicates
@@ -618,6 +634,12 @@
          *   WebChromeClient#onShowFileChooser}</li>
          * </ol>
          *
+         * <p class="note"><b>Note:</b> The created intent may be handled by
+         * third-party applications on device. The received result must be treated
+         * as untrusted as it can contain Uris pointing to your own app's sensitive
+         * data files. Your app should check the resultant Uris in {@link #parseResult}
+         * before calling the <code>filePathCallback</code>.</p>
+         *
          * @return an Intent that supports basic file chooser sources.
          */
         public abstract Intent createIntent();
diff --git a/core/jni/com_android_internal_os_Zygote.cpp b/core/jni/com_android_internal_os_Zygote.cpp
index 12d62cc..062fab3 100644
--- a/core/jni/com_android_internal_os_Zygote.cpp
+++ b/core/jni/com_android_internal_os_Zygote.cpp
@@ -116,7 +116,7 @@
 
 using android::zygote::ZygoteFailure;
 
-using Action = android_mallopt_gwp_asan_options_t::Action;
+using Mode = android_mallopt_gwp_asan_options_t::Mode;
 
 // This type is duplicated in fd_utils.h
 typedef const std::function<void(std::string)>& fail_fn_t;
@@ -2101,21 +2101,21 @@
     switch (runtime_flags & RuntimeFlags::GWP_ASAN_LEVEL_MASK) {
         default:
         case RuntimeFlags::GWP_ASAN_LEVEL_DEFAULT:
-            gwp_asan_options.desire = GetBoolProperty(kGwpAsanAppRecoverableSysprop, true)
-                    ? Action::TURN_ON_FOR_APP_SAMPLED_NON_CRASHING
-                    : Action::DONT_TURN_ON_UNLESS_OVERRIDDEN;
+            gwp_asan_options.mode = GetBoolProperty(kGwpAsanAppRecoverableSysprop, true)
+                    ? Mode::APP_MANIFEST_DEFAULT
+                    : Mode::APP_MANIFEST_NEVER;
             android_mallopt(M_INITIALIZE_GWP_ASAN, &gwp_asan_options, sizeof(gwp_asan_options));
             break;
         case RuntimeFlags::GWP_ASAN_LEVEL_NEVER:
-            gwp_asan_options.desire = Action::DONT_TURN_ON_UNLESS_OVERRIDDEN;
+            gwp_asan_options.mode = Mode::APP_MANIFEST_NEVER;
             android_mallopt(M_INITIALIZE_GWP_ASAN, &gwp_asan_options, sizeof(gwp_asan_options));
             break;
         case RuntimeFlags::GWP_ASAN_LEVEL_ALWAYS:
-            gwp_asan_options.desire = Action::TURN_ON_FOR_APP;
+            gwp_asan_options.mode = Mode::APP_MANIFEST_ALWAYS;
             android_mallopt(M_INITIALIZE_GWP_ASAN, &gwp_asan_options, sizeof(gwp_asan_options));
             break;
         case RuntimeFlags::GWP_ASAN_LEVEL_LOTTERY:
-            gwp_asan_options.desire = Action::TURN_ON_WITH_SAMPLING;
+            gwp_asan_options.mode = Mode::APP_MANIFEST_DEFAULT;
             android_mallopt(M_INITIALIZE_GWP_ASAN, &gwp_asan_options, sizeof(gwp_asan_options));
             break;
     }
diff --git a/core/res/AndroidManifest.xml b/core/res/AndroidManifest.xml
index bfbfb3a..70d923b 100644
--- a/core/res/AndroidManifest.xml
+++ b/core/res/AndroidManifest.xml
@@ -8771,6 +8771,7 @@
 
         <service android:name="com.android.server.companion.datatransfer.contextsync.CallMetadataSyncInCallService"
                  android:permission="android.permission.BIND_INCALL_SERVICE"
+                 android:enabled="@bool/config_enableContextSyncInCall"
                  android:exported="true">
             <meta-data android:name="android.telecom.INCLUDE_SELF_MANAGED_CALLS"
                        android:value="true" />
diff --git a/core/res/res/values/config.xml b/core/res/res/values/config.xml
index edaf8b5..cefc648 100644
--- a/core/res/res/values/config.xml
+++ b/core/res/res/values/config.xml
@@ -7051,6 +7051,9 @@
          event gets ignored. -->
     <integer name="config_defaultMinEmergencyGestureTapDurationMillis">200</integer>
 
+    <!-- Control whether to enable CallMetadataSyncInCallService. -->
+    <bool name="config_enableContextSyncInCall">false</bool>
+
     <!-- Whether the system uses auto-suspend mode. -->
     <bool name="config_useAutoSuspend">true</bool>
 </resources>
diff --git a/media/java/android/media/MediaRouter2.java b/media/java/android/media/MediaRouter2.java
index b2838c8..7ddf11e 100644
--- a/media/java/android/media/MediaRouter2.java
+++ b/media/java/android/media/MediaRouter2.java
@@ -50,6 +50,7 @@
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.media.flags.Flags;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -984,37 +985,7 @@
     @SystemApi
     @RequiresPermission(android.Manifest.permission.MEDIA_CONTENT_CONTROL)
     public void transfer(@NonNull RoutingController controller, @NonNull MediaRoute2Info route) {
-        mImpl.transfer(
-                controller.getRoutingSessionInfo(),
-                route,
-                Process.myUserHandle(),
-                mContext.getPackageName());
-    }
-
-    /**
-     * Transfers the media of a routing controller to the given route.
-     *
-     * <p>This will be no-op for non-system media routers.
-     *
-     * @param controller a routing controller controlling media routing.
-     * @param route the route you want to transfer the media to.
-     * @param transferInitiatorUserHandle the user handle of the app that initiated the transfer
-     *     request.
-     * @param transferInitiatorPackageName the package name of the app that initiated the transfer.
-     *     This value is used with the user handle to populate {@link
-     *     RoutingController#wasTransferInitiatedBySelf()}.
-     * @hide
-     */
-    public void transfer(
-            @NonNull RoutingController controller,
-            @NonNull MediaRoute2Info route,
-            @NonNull UserHandle transferInitiatorUserHandle,
-            @NonNull String transferInitiatorPackageName) {
-        mImpl.transfer(
-                controller.getRoutingSessionInfo(),
-                route,
-                transferInitiatorUserHandle,
-                transferInitiatorPackageName);
+        mImpl.transfer(controller.getRoutingSessionInfo(), route);
     }
 
     void requestCreateController(
@@ -1913,13 +1884,7 @@
          */
         @FlaggedApi(FLAG_ENABLE_BUILT_IN_SPEAKER_ROUTE_SUITABILITY_STATUSES)
         public boolean wasTransferInitiatedBySelf() {
-            RoutingSessionInfo sessionInfo = getRoutingSessionInfo();
-
-            UserHandle transferInitiatorUserHandle = sessionInfo.getTransferInitiatorUserHandle();
-            String transferInitiatorPackageName = sessionInfo.getTransferInitiatorPackageName();
-
-            return Objects.equals(Process.myUserHandle(), transferInitiatorUserHandle)
-                    && Objects.equals(mContext.getPackageName(), transferInitiatorPackageName);
+            return mImpl.wasTransferredBySelf(getRoutingSessionInfo());
         }
 
         /**
@@ -2082,12 +2047,28 @@
             Objects.requireNonNull(route, "route must not be null");
             synchronized (mControllerLock) {
                 if (isReleased()) {
-                    Log.w(TAG, "transferToRoute: Called on released controller. Ignoring.");
+                    Log.w(
+                            TAG,
+                            "tryTransferWithinProvider: Called on released controller. Ignoring.");
                     return true;
                 }
 
-                if (!mSessionInfo.getTransferableRoutes().contains(route.getId())) {
-                    Log.w(TAG, "Ignoring transferring to a non-transferable route=" + route);
+                // If this call is trying to transfer to a selected system route, we let them
+                // through as a provider driven transfer in order to update the transfer reason and
+                // initiator data.
+                boolean isSystemRouteReselection =
+                        Flags.enableBuiltInSpeakerRouteSuitabilityStatuses()
+                                && mSessionInfo.isSystemSession()
+                                && route.isSystemRoute()
+                                && mSessionInfo.getSelectedRoutes().contains(route.getId());
+                if (!isSystemRouteReselection
+                        && !mSessionInfo.getTransferableRoutes().contains(route.getId())) {
+                    Log.i(
+                            TAG,
+                            "Transferring to a non-transferable route="
+                                    + route
+                                    + " session= "
+                                    + mSessionInfo.getId());
                     return false;
                 }
             }
@@ -2498,11 +2479,7 @@
 
         void stop();
 
-        void transfer(
-                @NonNull RoutingSessionInfo sessionInfo,
-                @NonNull MediaRoute2Info route,
-                @NonNull UserHandle transferInitiatorUserHandle,
-                @NonNull String transferInitiatorPackageName);
+        void transfer(@NonNull RoutingSessionInfo sessionInfo, @NonNull MediaRoute2Info route);
 
         List<RoutingController> getControllers();
 
@@ -2523,6 +2500,11 @@
                 boolean shouldNotifyStop,
                 RoutingController controller);
 
+        /**
+         * Returns the value of {@link RoutingController#wasTransferInitiatedBySelf()} for the app
+         * associated with this router.
+         */
+        boolean wasTransferredBySelf(RoutingSessionInfo sessionInfo);
     }
 
     /**
@@ -2723,7 +2705,7 @@
 
             List<RoutingSessionInfo> sessionInfos = getRoutingSessions();
             RoutingSessionInfo targetSession = sessionInfos.get(sessionInfos.size() - 1);
-            transfer(targetSession, route, mClientUser, mContext.getPackageName());
+            transfer(targetSession, route);
         }
 
         @Override
@@ -2746,24 +2728,15 @@
          *
          * @param sessionInfo The {@link RoutingSessionInfo routing session} to transfer.
          * @param route The {@link MediaRoute2Info route} to transfer to.
-         * @param transferInitiatorUserHandle The user handle of the app that initiated the
-         *     transfer.
-         * @param transferInitiatorPackageName The package name if of the app that initiated the
-         *     transfer.
          * @see #transferToRoute(RoutingSessionInfo, MediaRoute2Info, UserHandle, String)
          * @see #requestCreateSession(RoutingSessionInfo, MediaRoute2Info)
          */
         @Override
         @SuppressWarnings("AndroidFrameworkRequiresPermission")
         public void transfer(
-                @NonNull RoutingSessionInfo sessionInfo,
-                @NonNull MediaRoute2Info route,
-                @NonNull UserHandle transferInitiatorUserHandle,
-                @NonNull String transferInitiatorPackageName) {
+                @NonNull RoutingSessionInfo sessionInfo, @NonNull MediaRoute2Info route) {
             Objects.requireNonNull(sessionInfo, "sessionInfo must not be null");
             Objects.requireNonNull(route, "route must not be null");
-            Objects.requireNonNull(transferInitiatorUserHandle);
-            Objects.requireNonNull(transferInitiatorPackageName);
 
             Log.v(
                     TAG,
@@ -2780,15 +2753,19 @@
                 return;
             }
 
-            if (sessionInfo.getTransferableRoutes().contains(route.getId())) {
-                transferToRoute(
-                        sessionInfo,
-                        route,
-                        transferInitiatorUserHandle,
-                        transferInitiatorPackageName);
+            // If this call is trying to transfer to a selected system route, we let them
+            // through as a provider driven transfer in order to update the transfer reason and
+            // initiator data.
+            boolean isSystemRouteReselection =
+                    Flags.enableBuiltInSpeakerRouteSuitabilityStatuses()
+                            && sessionInfo.isSystemSession()
+                            && route.isSystemRoute()
+                            && sessionInfo.getSelectedRoutes().contains(route.getId());
+            if (sessionInfo.getTransferableRoutes().contains(route.getId())
+                    || isSystemRouteReselection) {
+                transferToRoute(sessionInfo, route, mClientUser, mClientPackageName);
             } else {
-                requestCreateSession(sessionInfo, route, transferInitiatorUserHandle,
-                        transferInitiatorPackageName);
+                requestCreateSession(sessionInfo, route, mClientUser, mClientPackageName);
             }
         }
 
@@ -3043,6 +3020,14 @@
             releaseSession(controller.getRoutingSessionInfo());
         }
 
+        @Override
+        public boolean wasTransferredBySelf(RoutingSessionInfo sessionInfo) {
+            UserHandle transferInitiatorUserHandle = sessionInfo.getTransferInitiatorUserHandle();
+            String transferInitiatorPackageName = sessionInfo.getTransferInitiatorPackageName();
+            return Objects.equals(mClientUser, transferInitiatorUserHandle)
+                    && Objects.equals(mClientPackageName, transferInitiatorPackageName);
+        }
+
         /**
          * Retrieves the system session info for the given package.
          *
@@ -3619,10 +3604,7 @@
          */
         @Override
         public void transfer(
-                @NonNull RoutingSessionInfo sessionInfo,
-                @NonNull MediaRoute2Info route,
-                @NonNull UserHandle transferInitiatorUserHandle,
-                @NonNull String transferInitiatorPackageName) {
+                @NonNull RoutingSessionInfo sessionInfo, @NonNull MediaRoute2Info route) {
             // Do nothing.
         }
 
@@ -3741,6 +3723,14 @@
             }
         }
 
+        @Override
+        public boolean wasTransferredBySelf(RoutingSessionInfo sessionInfo) {
+            UserHandle transferInitiatorUserHandle = sessionInfo.getTransferInitiatorUserHandle();
+            String transferInitiatorPackageName = sessionInfo.getTransferInitiatorPackageName();
+            return Objects.equals(Process.myUserHandle(), transferInitiatorUserHandle)
+                    && Objects.equals(mContext.getPackageName(), transferInitiatorPackageName);
+        }
+
         @GuardedBy("mLock")
         private void registerRouterStubIfNeededLocked() throws RemoteException {
             if (mStub == null) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryTest.kt
new file mode 100644
index 0000000..1e5599b
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryTest.kt
@@ -0,0 +1,165 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.data.repository
+
+import android.content.ComponentName
+import android.content.packageManager
+import android.content.pm.PackageManager
+import android.content.pm.ServiceInfo
+import android.content.pm.UserInfo
+import android.graphics.drawable.TestStubDrawable
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.kosmos.mainCoroutineContext
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.qs.panels.shared.model.EditTileData
+import com.android.systemui.qs.pipeline.data.repository.FakeInstalledTilesComponentRepository
+import com.android.systemui.qs.pipeline.data.repository.fakeInstalledTilesRepository
+import com.android.systemui.qs.pipeline.data.repository.installedTilesRepository
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.settings.FakeUserTracker
+import com.android.systemui.settings.fakeUserTracker
+import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.whenever
+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
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class IconAndNameCustomRepositoryTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
+    private val packageManager: PackageManager = kosmos.packageManager
+    private val userTracker: FakeUserTracker =
+        kosmos.fakeUserTracker.apply {
+            whenever(userContext.packageManager).thenReturn(packageManager)
+        }
+
+    private val service1 =
+        FakeInstalledTilesComponentRepository.ServiceInfo(
+            component1,
+            tileService1,
+            drawable1,
+            appName1,
+        )
+
+    private val service2 =
+        FakeInstalledTilesComponentRepository.ServiceInfo(
+            component2,
+            tileService2,
+            drawable2,
+            appName2,
+        )
+
+    private val underTest =
+        with(kosmos) {
+            IconAndNameCustomRepository(
+                installedTilesRepository,
+                userTracker,
+                mainCoroutineContext,
+            )
+        }
+
+    @Before
+    fun setUp() {
+        kosmos.fakeInstalledTilesRepository.setInstalledServicesForUser(
+            userTracker.userId,
+            listOf(service1, service2)
+        )
+    }
+
+    @Test
+    fun loadDataForCurrentServices() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTileDataList = underTest.getCustomTileData()
+                val expectedData1 =
+                    EditTileData(
+                        TileSpec.create(component1),
+                        Icon.Loaded(drawable1, ContentDescription.Loaded(tileService1)),
+                        Text.Loaded(tileService1),
+                        Text.Loaded(appName1),
+                    )
+                val expectedData2 =
+                    EditTileData(
+                        TileSpec.create(component2),
+                        Icon.Loaded(drawable2, ContentDescription.Loaded(tileService2)),
+                        Text.Loaded(tileService2),
+                        Text.Loaded(appName2),
+                    )
+
+                assertThat(editTileDataList).containsExactly(expectedData1, expectedData2)
+            }
+        }
+
+    @Test
+    fun loadDataForCurrentServices_otherCurrentUser_empty() =
+        with(kosmos) {
+            testScope.runTest {
+                userTracker.set(listOf(UserInfo(11, "", 0)), 0)
+                val editTileDataList = underTest.getCustomTileData()
+
+                assertThat(editTileDataList).isEmpty()
+            }
+        }
+
+    @Test
+    fun loadDataForCurrentServices_serviceInfoWithNullIcon_notInList() =
+        with(kosmos) {
+            testScope.runTest {
+                val serviceNullIcon =
+                    FakeInstalledTilesComponentRepository.ServiceInfo(
+                        component2,
+                        tileService2,
+                    )
+                fakeInstalledTilesRepository.setInstalledServicesForUser(
+                    userTracker.userId,
+                    listOf(service1, serviceNullIcon)
+                )
+
+                val expectedData1 =
+                    EditTileData(
+                        TileSpec.create(component1),
+                        Icon.Loaded(drawable1, ContentDescription.Loaded(tileService1)),
+                        Text.Loaded(tileService1),
+                        Text.Loaded(appName1),
+                    )
+
+                val editTileDataList = underTest.getCustomTileData()
+                assertThat(editTileDataList).containsExactly(expectedData1)
+            }
+        }
+
+    private companion object {
+        val drawable1 = TestStubDrawable("drawable1")
+        val appName1 = "App1"
+        val tileService1 = "Tile Service 1"
+        val component1 = ComponentName("pkg1", "srv1")
+
+        val drawable2 = TestStubDrawable("drawable2")
+        val appName2 = "App2"
+        val tileService2 = "Tile Service 2"
+        val component2 = ComponentName("pkg2", "srv2")
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryTest.kt
new file mode 100644
index 0000000..56cead1
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryTest.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.data.repository
+
+import android.content.res.mainResources
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.res.R
+import com.android.systemui.testKosmos
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class StockTilesRepositoryTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
+    private val underTest = StockTilesRepository(kosmos.mainResources)
+
+    @Test
+    fun stockTilesMatchesResources() {
+        val expected =
+            kosmos.mainResources
+                .getString(R.string.quick_settings_tiles_stock)
+                .split(",")
+                .map(TileSpec::create)
+        assertThat(underTest.stockTiles).isEqualTo(expected)
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorTest.kt
new file mode 100644
index 0000000..deefbf5
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorTest.kt
@@ -0,0 +1,198 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.domain.interactor
+
+import android.content.ComponentName
+import android.graphics.drawable.TestStubDrawable
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.qs.panels.data.repository.iconAndNameCustomRepository
+import com.android.systemui.qs.panels.data.repository.stockTilesRepository
+import com.android.systemui.qs.panels.shared.model.EditTileData
+import com.android.systemui.qs.pipeline.data.repository.FakeInstalledTilesComponentRepository
+import com.android.systemui.qs.pipeline.data.repository.fakeInstalledTilesRepository
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.tiles.impl.battery.qsBatterySaverTileConfig
+import com.android.systemui.qs.tiles.impl.flashlight.qsFlashlightTileConfig
+import com.android.systemui.qs.tiles.impl.internet.qsInternetTileConfig
+import com.android.systemui.qs.tiles.viewmodel.QSTileConfig
+import com.android.systemui.qs.tiles.viewmodel.fakeQSTileConfigProvider
+import com.android.systemui.qs.tiles.viewmodel.qSTileConfigProvider
+import com.android.systemui.settings.userTracker
+import com.android.systemui.testKosmos
+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
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class EditTilesListInteractorTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
+    // Only have some configurations so we can test the effect of missing configurations.
+    // As the configurations are injected by dagger, we'll have all the existing configurations
+    private val internetTileConfig = kosmos.qsInternetTileConfig
+    private val flashlightTileConfig = kosmos.qsFlashlightTileConfig
+    private val batteryTileConfig = kosmos.qsBatterySaverTileConfig
+
+    private val serviceInfo =
+        FakeInstalledTilesComponentRepository.ServiceInfo(
+            component,
+            tileName,
+            icon,
+            appName,
+        )
+
+    private val underTest =
+        with(kosmos) {
+            EditTilesListInteractor(
+                stockTilesRepository,
+                qSTileConfigProvider,
+                iconAndNameCustomRepository,
+            )
+        }
+
+    @Before
+    fun setUp() {
+        with(kosmos) {
+            fakeInstalledTilesRepository.setInstalledServicesForUser(
+                userTracker.userId,
+                listOf(serviceInfo)
+            )
+
+            with(fakeQSTileConfigProvider) {
+                putConfig(internetTileConfig.tileSpec, internetTileConfig)
+                putConfig(flashlightTileConfig.tileSpec, flashlightTileConfig)
+                putConfig(batteryTileConfig.tileSpec, batteryTileConfig)
+            }
+        }
+    }
+
+    @Test
+    fun getTilesToEdit_stockTilesHaveNoAppName() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTiles = underTest.getTilesToEdit()
+
+                assertThat(editTiles.stockTiles.all { it.appName == null }).isTrue()
+            }
+        }
+
+    @Test
+    fun getTilesToEdit_stockTilesAreAllPlatformSpecs() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTiles = underTest.getTilesToEdit()
+
+                assertThat(editTiles.stockTiles.all { it.tileSpec is TileSpec.PlatformTileSpec })
+                    .isTrue()
+            }
+        }
+
+    @Test
+    fun getTilesToEdit_stockTiles_sameOrderAsRepository() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTiles = underTest.getTilesToEdit()
+
+                assertThat(editTiles.stockTiles.map { it.tileSpec })
+                    .isEqualTo(stockTilesRepository.stockTiles)
+            }
+        }
+
+    @Test
+    fun getTilesToEdit_customTileData_matchesService() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTiles = underTest.getTilesToEdit()
+                val expected =
+                    EditTileData(
+                        tileSpec = TileSpec.create(component),
+                        icon = Icon.Loaded(icon, ContentDescription.Loaded(tileName)),
+                        label = Text.Loaded(tileName),
+                        appName = Text.Loaded(appName),
+                    )
+
+                assertThat(editTiles.customTiles).hasSize(1)
+                assertThat(editTiles.customTiles[0]).isEqualTo(expected)
+            }
+        }
+
+    @Test
+    fun getTilesToEdit_tilesInConfigProvider_correctData() =
+        with(kosmos) {
+            testScope.runTest {
+                val editTiles = underTest.getTilesToEdit()
+
+                assertThat(
+                        editTiles.stockTiles.first { it.tileSpec == internetTileConfig.tileSpec }
+                    )
+                    .isEqualTo(internetTileConfig.toEditTileData())
+                assertThat(
+                        editTiles.stockTiles.first { it.tileSpec == flashlightTileConfig.tileSpec }
+                    )
+                    .isEqualTo(flashlightTileConfig.toEditTileData())
+                assertThat(editTiles.stockTiles.first { it.tileSpec == batteryTileConfig.tileSpec })
+                    .isEqualTo(batteryTileConfig.toEditTileData())
+            }
+        }
+
+    @Test
+    fun getTilesToEdit_tilesNotInConfigProvider_useDefaultData() =
+        with(kosmos) {
+            testScope.runTest {
+                underTest
+                    .getTilesToEdit()
+                    .stockTiles
+                    .filterNot { qSTileConfigProvider.hasConfig(it.tileSpec.spec) }
+                    .forEach { assertThat(it).isEqualTo(it.tileSpec.missingConfigEditTileData()) }
+            }
+        }
+
+    private companion object {
+        val component = ComponentName("pkg", "srv")
+        const val tileName = "Tile Service"
+        const val appName = "App"
+        val icon = TestStubDrawable("icon")
+
+        fun TileSpec.missingConfigEditTileData(): EditTileData {
+            return EditTileData(
+                tileSpec = this,
+                icon = Icon.Resource(android.R.drawable.star_on, ContentDescription.Loaded(spec)),
+                label = Text.Loaded(spec),
+                appName = null
+            )
+        }
+
+        fun QSTileConfig.toEditTileData(): EditTileData {
+            return EditTileData(
+                tileSpec = tileSpec,
+                icon =
+                    Icon.Resource(uiConfig.iconRes, ContentDescription.Resource(uiConfig.labelRes)),
+                label = Text.Resource(uiConfig.labelRes),
+                appName = null,
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelTest.kt
new file mode 100644
index 0000000..9fb25a2
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelTest.kt
@@ -0,0 +1,507 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.ui.viewmodel
+
+import android.R
+import android.content.ComponentName
+import android.graphics.drawable.TestStubDrawable
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.qs.FakeQSFactory
+import com.android.systemui.qs.FakeQSTile
+import com.android.systemui.qs.panels.data.repository.stockTilesRepository
+import com.android.systemui.qs.panels.domain.interactor.editTilesListInteractor
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutMap
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutTypeInteractor
+import com.android.systemui.qs.panels.domain.interactor.infiniteGridLayout
+import com.android.systemui.qs.panels.shared.model.EditTileData
+import com.android.systemui.qs.pipeline.data.repository.FakeInstalledTilesComponentRepository
+import com.android.systemui.qs.pipeline.data.repository.MinimumTilesFixedRepository
+import com.android.systemui.qs.pipeline.data.repository.fakeInstalledTilesRepository
+import com.android.systemui.qs.pipeline.data.repository.fakeMinimumTilesRepository
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
+import com.android.systemui.qs.pipeline.domain.interactor.minimumTilesInteractor
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.qs.qsTileFactory
+import com.android.systemui.qs.tiles.impl.alarm.qsAlarmTileConfig
+import com.android.systemui.qs.tiles.impl.battery.qsBatterySaverTileConfig
+import com.android.systemui.qs.tiles.impl.flashlight.qsFlashlightTileConfig
+import com.android.systemui.qs.tiles.impl.internet.qsInternetTileConfig
+import com.android.systemui.qs.tiles.impl.sensorprivacy.qsCameraSensorPrivacyToggleTileConfig
+import com.android.systemui.qs.tiles.impl.sensorprivacy.qsMicrophoneSensorPrivacyToggleTileConfig
+import com.android.systemui.qs.tiles.viewmodel.QSTileConfig
+import com.android.systemui.qs.tiles.viewmodel.fakeQSTileConfigProvider
+import com.android.systemui.qs.tiles.viewmodel.qSTileConfigProvider
+import com.android.systemui.settings.userTracker
+import com.android.systemui.testKosmos
+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
+
+@RunWith(AndroidJUnit4::class)
+@SmallTest
+class EditModeViewModelTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+
+    // Only have some configurations so we can test the effect of missing configurations.
+    // As the configurations are injected by dagger, we'll have all the existing configurations
+    private val configs =
+        with(kosmos) {
+            setOf(
+                qsInternetTileConfig,
+                qsFlashlightTileConfig,
+                qsBatterySaverTileConfig,
+                qsAlarmTileConfig,
+                qsCameraSensorPrivacyToggleTileConfig,
+                qsMicrophoneSensorPrivacyToggleTileConfig,
+            )
+        }
+
+    private val serviceInfo1 =
+        FakeInstalledTilesComponentRepository.ServiceInfo(
+            component1,
+            tileService1,
+            drawable1,
+            appName1,
+        )
+
+    private val serviceInfo2 =
+        FakeInstalledTilesComponentRepository.ServiceInfo(
+            component2,
+            tileService2,
+            drawable2,
+            appName2,
+        )
+
+    private val underTest: EditModeViewModel by lazy {
+        with(kosmos) {
+            EditModeViewModel(
+                editTilesListInteractor,
+                currentTilesInteractor,
+                minimumTilesInteractor,
+                infiniteGridLayout,
+                applicationCoroutineScope,
+                gridLayoutTypeInteractor,
+                gridLayoutMap,
+            )
+        }
+    }
+
+    @Before
+    fun setUp() {
+        with(kosmos) {
+            fakeMinimumTilesRepository = MinimumTilesFixedRepository(minNumberOfTiles)
+
+            fakeInstalledTilesRepository.setInstalledServicesForUser(
+                userTracker.userId,
+                listOf(serviceInfo1, serviceInfo2)
+            )
+
+            with(fakeQSTileConfigProvider) { configs.forEach { putConfig(it.tileSpec, it) } }
+            qsTileFactory = FakeQSFactory { FakeQSTile(userTracker.userId, available = true) }
+        }
+    }
+
+    @Test
+    fun isEditing() =
+        with(kosmos) {
+            testScope.runTest {
+                val isEditing by collectLastValue(underTest.isEditing)
+
+                assertThat(isEditing).isFalse()
+
+                underTest.startEditing()
+                assertThat(isEditing).isTrue()
+
+                underTest.stopEditing()
+                assertThat(isEditing).isFalse()
+            }
+        }
+
+    @Test
+    fun editing_false_emptyFlowOfTiles() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                assertThat(tiles).isNull()
+            }
+        }
+
+    @Test
+    fun editing_true_notEmptyTileData() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                underTest.startEditing()
+
+                assertThat(tiles).isNotEmpty()
+            }
+        }
+
+    @Test
+    fun tilesData_hasAllStockTiles() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                underTest.startEditing()
+
+                assertThat(
+                        tiles!!
+                            .filter { it.tileSpec is TileSpec.PlatformTileSpec }
+                            .map { it.tileSpec }
+                    )
+                    .containsExactlyElementsIn(stockTilesRepository.stockTiles)
+            }
+        }
+
+    @Test
+    fun tilesData_stockTiles_haveCorrectUiValues() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filter { it.tileSpec is TileSpec.PlatformTileSpec }
+                    .forEach {
+                        val data = getEditTileData(it.tileSpec)
+
+                        assertThat(it.label).isEqualTo(data.label)
+                        assertThat(it.icon).isEqualTo(data.icon)
+                        assertThat(it.appName).isNull()
+                    }
+            }
+        }
+
+    @Test
+    fun tilesData_hasAllCustomTiles() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                underTest.startEditing()
+
+                assertThat(
+                        tiles!!
+                            .filter { it.tileSpec is TileSpec.CustomTileSpec }
+                            .map { it.tileSpec }
+                    )
+                    .containsExactly(TileSpec.create(component1), TileSpec.create(component2))
+            }
+        }
+
+    @Test
+    fun tilesData_customTiles_haveCorrectUiValues() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+
+                underTest.startEditing()
+
+                // service1
+                val model1 = tiles!!.first { it.tileSpec == TileSpec.create(component1) }
+                assertThat(model1.label).isEqualTo(Text.Loaded(tileService1))
+                assertThat(model1.appName).isEqualTo(Text.Loaded(appName1))
+                assertThat(model1.icon)
+                    .isEqualTo(Icon.Loaded(drawable1, ContentDescription.Loaded(tileService1)))
+
+                // service2
+                val model2 = tiles!!.first { it.tileSpec == TileSpec.create(component2) }
+                assertThat(model2.label).isEqualTo(Text.Loaded(tileService2))
+                assertThat(model2.appName).isEqualTo(Text.Loaded(appName2))
+                assertThat(model2.icon)
+                    .isEqualTo(Icon.Loaded(drawable2, ContentDescription.Loaded(tileService2)))
+            }
+        }
+
+    @Test
+    fun currentTiles_inCorrectOrder_markedAsCurrent() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    listOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                assertThat(tiles!!.filter { it.isCurrent }.map { it.tileSpec })
+                    .containsExactlyElementsIn(currentTiles)
+                    .inOrder()
+            }
+        }
+
+    @Test
+    fun notCurrentTiles() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    listOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                val remainingTiles =
+                    stockTilesRepository.stockTiles.filterNot { it in currentTiles } +
+                        listOf(TileSpec.create(component1))
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                assertThat(tiles!!.filterNot { it.isCurrent }.map { it.tileSpec })
+                    .containsExactlyElementsIn(remainingTiles)
+            }
+        }
+
+    @Test
+    fun currentTilesChange_trackingChange() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                val newTile = TileSpec.create("internet")
+                val position = 1
+                currentTilesInteractor.addTile(newTile, position)
+                currentTiles.add(position, newTile)
+
+                assertThat(tiles!!.filter { it.isCurrent }.map { it.tileSpec })
+                    .containsExactlyElementsIn(currentTiles)
+                    .inOrder()
+            }
+        }
+
+    @Test
+    fun nonCurrentTiles_orderPreservedWhenCurrentTilesChange() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                val nonCurrentSpecs = tiles!!.filterNot { it.isCurrent }.map { it.tileSpec }
+                val newTile = TileSpec.create("internet")
+                currentTilesInteractor.addTile(newTile)
+
+                assertThat(tiles!!.filterNot { it.isCurrent }.map { it.tileSpec })
+                    .containsExactlyElementsIn(nonCurrentSpecs - listOf(newTile))
+                    .inOrder()
+            }
+        }
+
+    @Test
+    fun nonCurrentTiles_haveOnlyAddAction() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filterNot { it.isCurrent }
+                    .forEach {
+                        assertThat(it.availableEditActions)
+                            .containsExactly(AvailableEditActions.ADD)
+                    }
+            }
+        }
+
+    @Test
+    fun currentTiles_moreThanMinimumTiles_haveRemoveAction() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+                assertThat(currentTiles.size).isGreaterThan(minNumberOfTiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filter { it.isCurrent }
+                    .forEach {
+                        assertThat(it.availableEditActions).contains(AvailableEditActions.REMOVE)
+                    }
+            }
+        }
+
+    @Test
+    fun currentTiles_minimumTiles_dontHaveRemoveAction() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+                assertThat(currentTiles.size).isEqualTo(minNumberOfTiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filter { it.isCurrent }
+                    .forEach {
+                        assertThat(it.availableEditActions)
+                            .doesNotContain(AvailableEditActions.REMOVE)
+                    }
+            }
+        }
+
+    @Test
+    fun currentTiles_lessThanMinimumTiles_dontHaveRemoveAction() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+                assertThat(currentTiles.size).isLessThan(minNumberOfTiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filter { it.isCurrent }
+                    .forEach {
+                        assertThat(it.availableEditActions)
+                            .doesNotContain(AvailableEditActions.REMOVE)
+                    }
+            }
+        }
+
+    @Test
+    fun currentTiles_haveMoveAction() =
+        with(kosmos) {
+            testScope.runTest {
+                val tiles by collectLastValue(underTest.tiles)
+                val currentTiles =
+                    mutableListOf(
+                        TileSpec.create("flashlight"),
+                        TileSpec.create("airplane"),
+                        TileSpec.create(component2),
+                        TileSpec.create("alarm"),
+                    )
+                currentTilesInteractor.setTiles(currentTiles)
+
+                underTest.startEditing()
+
+                tiles!!
+                    .filter { it.isCurrent }
+                    .forEach {
+                        assertThat(it.availableEditActions).contains(AvailableEditActions.MOVE)
+                    }
+            }
+        }
+
+    private companion object {
+        val drawable1 = TestStubDrawable("drawable1")
+        val appName1 = "App1"
+        val tileService1 = "Tile Service 1"
+        val component1 = ComponentName("pkg1", "srv1")
+
+        val drawable2 = TestStubDrawable("drawable2")
+        val appName2 = "App2"
+        val tileService2 = "Tile Service 2"
+        val component2 = ComponentName("pkg2", "srv2")
+
+        fun TileSpec.missingConfigEditTileData(): EditTileData {
+            return EditTileData(
+                tileSpec = this,
+                icon = Icon.Resource(R.drawable.star_on, ContentDescription.Loaded(spec)),
+                label = Text.Loaded(spec),
+                appName = null
+            )
+        }
+
+        fun QSTileConfig.toEditTileData(): EditTileData {
+            return EditTileData(
+                tileSpec = tileSpec,
+                icon =
+                    Icon.Resource(uiConfig.iconRes, ContentDescription.Resource(uiConfig.labelRes)),
+                label = Text.Resource(uiConfig.labelRes),
+                appName = null,
+            )
+        }
+
+        fun Kosmos.getEditTileData(tileSpec: TileSpec): EditTileData {
+            return if (qSTileConfigProvider.hasConfig(tileSpec.spec)) {
+                qSTileConfigProvider.getConfig(tileSpec.spec).toEditTileData()
+            } else {
+                tileSpec.missingConfigEditTileData()
+            }
+        }
+
+        val minNumberOfTiles = 3
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepositoryImplTest.kt
index bc57ce6..a0dec8c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepositoryImplTest.kt
@@ -90,7 +90,7 @@
     }
 
     @Test
-    fun componentsLoadedOnStart() =
+    fun servicesLoadedOnStart() =
         testScope.runTest {
             val userId = 0
             val resolveInfo =
@@ -106,12 +106,14 @@
 
             val componentNames by collectLastValue(underTest.getInstalledTilesComponents(userId))
             runCurrent()
+            val services = underTest.getInstalledTilesServiceInfos(userId)
 
             assertThat(componentNames).containsExactly(TEST_COMPONENT)
+            assertThat(services).containsExactly(resolveInfo.serviceInfo)
         }
 
     @Test
-    fun componentAdded_foundAfterPackageChange() =
+    fun serviceAdded_foundAfterPackageChange() =
         testScope.runTest {
             val userId = 0
             val resolveInfo =
@@ -132,12 +134,14 @@
                 .thenReturn(listOf(resolveInfo))
             kosmos.fakePackageChangeRepository.notifyChange(PackageChangeModel.Empty)
             runCurrent()
+            val services = underTest.getInstalledTilesServiceInfos(userId)
 
             assertThat(componentNames).containsExactly(TEST_COMPONENT)
+            assertThat(services).containsExactly(resolveInfo.serviceInfo)
         }
 
     @Test
-    fun componentWithoutPermission_notValid() =
+    fun serviceWithoutPermission_notValid() =
         testScope.runTest {
             val userId = 0
             val resolveInfo =
@@ -152,13 +156,15 @@
                 .thenReturn(listOf(resolveInfo))
 
             val componentNames by collectLastValue(underTest.getInstalledTilesComponents(userId))
+            val services = underTest.getInstalledTilesServiceInfos(userId)
             runCurrent()
 
             assertThat(componentNames).isEmpty()
+            assertThat(services).isEmpty()
         }
 
     @Test
-    fun componentNotEnabled_notValid() =
+    fun serviceNotEnabled_notValid() =
         testScope.runTest {
             val userId = 0
             val resolveInfo =
@@ -173,9 +179,11 @@
                 .thenReturn(listOf(resolveInfo))
 
             val componentNames by collectLastValue(underTest.getInstalledTilesComponents(userId))
+            val services = underTest.getInstalledTilesServiceInfos(userId)
             runCurrent()
 
             assertThat(componentNames).isEmpty()
+            assertThat(services).isEmpty()
         }
 
     @Test
@@ -221,30 +229,22 @@
 
             val componentNames by collectLastValue(underTest.getInstalledTilesComponents(userId))
             runCurrent()
+            val service = underTest.getInstalledTilesServiceInfos(userId)
 
             assertThat(componentNames).containsExactly(TEST_COMPONENT)
+            assertThat(service).containsExactly(resolveInfo.serviceInfo)
         }
 
     @Test
-    fun loadComponentsForSameUserTwice_returnsSameFlow() =
+    fun loadServicesForSameUserTwice_returnsSameFlow() =
         testScope.runTest {
-            val flowForUser1 = underTest.getInstalledTilesComponents(1)
-            val flowForUser1TheSecondTime = underTest.getInstalledTilesComponents(1)
+            val flowForUser1 = underTest.getInstalledTilesServiceInfos(1)
+            val flowForUser1TheSecondTime = underTest.getInstalledTilesServiceInfos(1)
             runCurrent()
 
             assertThat(flowForUser1TheSecondTime).isEqualTo(flowForUser1)
         }
 
-    @Test
-    fun loadComponentsForDifferentUsers_returnsDifferentFlow() =
-        testScope.runTest {
-            val flowForUser1 = underTest.getInstalledTilesComponents(1)
-            val flowForUser2 = underTest.getInstalledTilesComponents(2)
-            runCurrent()
-
-            assertThat(flowForUser2).isNotEqualTo(flowForUser1)
-        }
-
     // Tests that a ServiceInfo that is returned by queryIntentServicesAsUser but shortly
     // after uninstalled, doesn't crash SystemUI.
     @Test
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AccessibilityTilesInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AccessibilityTilesInteractorTest.kt
index 61e4774..3faab50 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AccessibilityTilesInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AccessibilityTilesInteractorTest.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.accessibility.data.repository.FakeAccessibilityQsShortcutsRepository
 import com.android.systemui.qs.FakeQSFactory
+import com.android.systemui.qs.FakeQSTile
 import com.android.systemui.qs.pipeline.domain.model.TileModel
 import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.qs.tiles.ColorCorrectionTile
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AutoAddInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AutoAddInteractorTest.kt
index 8ae9172..167eff1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AutoAddInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/AutoAddInteractorTest.kt
@@ -22,6 +22,7 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.dump.DumpManager
+import com.android.systemui.qs.FakeQSTile
 import com.android.systemui.qs.pipeline.data.repository.FakeAutoAddRepository
 import com.android.systemui.qs.pipeline.domain.autoaddable.FakeAutoAddable
 import com.android.systemui.qs.pipeline.domain.model.AutoAddTracking
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractorImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractorImplTest.kt
index 634c5fa..1c73fe2b 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractorImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractorImplTest.kt
@@ -32,6 +32,7 @@
 import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.plugins.qs.QSTile.BooleanState
 import com.android.systemui.qs.FakeQSFactory
+import com.android.systemui.qs.FakeQSTile
 import com.android.systemui.qs.external.CustomTile
 import com.android.systemui.qs.external.CustomTileStatePersister
 import com.android.systemui.qs.external.TileLifecycleManager
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/NoLowNumberOfTilesTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/NoLowNumberOfTilesTest.kt
index 90c8304..260189d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/NoLowNumberOfTilesTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/NoLowNumberOfTilesTest.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.qs.FakeQSFactory
+import com.android.systemui.qs.FakeQSTile
 import com.android.systemui.qs.pipeline.data.model.RestoreData
 import com.android.systemui.qs.pipeline.data.repository.FakeDefaultTilesRepository
 import com.android.systemui.qs.pipeline.data.repository.MinimumTilesFixedRepository
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/WorkProfileAutoAddedAfterRestoreTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/WorkProfileAutoAddedAfterRestoreTest.kt
index 4207a9c..dffd0d7 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/WorkProfileAutoAddedAfterRestoreTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/WorkProfileAutoAddedAfterRestoreTest.kt
@@ -27,6 +27,7 @@
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.plugins.qs.QSTile
 import com.android.systemui.qs.FakeQSFactory
+import com.android.systemui.qs.FakeQSTile
 import com.android.systemui.qs.pipeline.data.model.RestoreData
 import com.android.systemui.qs.pipeline.data.repository.fakeRestoreRepository
 import com.android.systemui.qs.pipeline.data.repository.fakeTileSpecRepository
@@ -54,9 +55,7 @@
 @OptIn(ExperimentalCoroutinesApi::class)
 class WorkProfileAutoAddedAfterRestoreTest : SysuiTestCase() {
 
-    private val kosmos by lazy {
-        Kosmos().apply { fakeUserTracker.set(listOf(USER_0_INFO), 0) }
-    }
+    private val kosmos by lazy { Kosmos().apply { fakeUserTracker.set(listOf(USER_0_INFO), 0) } }
     // Getter here so it can change when there is a managed profile.
     private val workTileAvailable: Boolean
         get() = hasManagedProfile()
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepository.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepository.kt
index a49b3ae..c11c49c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepository.kt
@@ -19,6 +19,7 @@
 
 import android.os.Handler
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.shared.model.KeyguardBlueprint
@@ -57,21 +58,7 @@
         TreeMap<String, KeyguardBlueprint>().apply { putAll(blueprints.associateBy { it.id }) }
     val blueprint: MutableStateFlow<KeyguardBlueprint> = MutableStateFlow(blueprintIdMap[DEFAULT]!!)
     val refreshTransition = MutableSharedFlow<Config>(extraBufferCapacity = 1)
-    private var targetTransitionConfig: Config? = null
-
-    /**
-     * Emits the blueprint value to the collectors.
-     *
-     * @param blueprintId
-     * @return whether the transition has succeeded.
-     */
-    fun applyBlueprint(index: Int): Boolean {
-        ArrayList(blueprintIdMap.values)[index]?.let {
-            applyBlueprint(it)
-            return true
-        }
-        return false
-    }
+    @VisibleForTesting var targetTransitionConfig: Config? = null
 
     /**
      * Emits the blueprint value to the collectors.
@@ -81,27 +68,21 @@
      */
     fun applyBlueprint(blueprintId: String?): Boolean {
         val blueprint = blueprintIdMap[blueprintId]
-        return if (blueprint != null) {
-            applyBlueprint(blueprint)
-            true
-        } else {
+        if (blueprint == null) {
             Log.e(
                 TAG,
                 "Could not find blueprint with id: $blueprintId. " +
                     "Perhaps it was not added to KeyguardBlueprintModule?"
             )
-            false
+            return false
         }
-    }
 
-    /** Emits the blueprint value to the collectors. */
-    fun applyBlueprint(blueprint: KeyguardBlueprint?) {
         if (blueprint == this.blueprint.value) {
-            refreshBlueprint()
-            return
+            return true
         }
 
-        blueprint?.let { this.blueprint.value = it }
+        this.blueprint.value = blueprint
+        return true
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractor.kt b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractor.kt
index da4f85e..857096e 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractor.kt
@@ -36,9 +36,12 @@
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.launch
 
 @SysUISingleton
@@ -64,12 +67,7 @@
 
     /** Current BlueprintId */
     val blueprintId =
-        combine(
-            configurationInteractor.onAnyConfigurationChange,
-            fingerprintPropertyInteractor.propertiesInitialized.filter { it },
-            clockInteractor.currentClock,
-            shadeInteractor.shadeMode,
-        ) { _, _, _, shadeMode ->
+        shadeInteractor.shadeMode.map { shadeMode ->
             val useSplitShade = shadeMode == ShadeMode.Split && !ComposeLockscreen.isEnabled
             when {
                 useSplitShade -> SplitShadeKeyguardBlueprint.ID
@@ -77,17 +75,29 @@
             }
         }
 
+    private val refreshEvents: Flow<Unit> =
+        merge(
+            configurationInteractor.onAnyConfigurationChange,
+            fingerprintPropertyInteractor.propertiesInitialized.filter { it }.map { Unit },
+        )
+
     init {
         applicationScope.launch { blueprintId.collect { transitionToBlueprint(it) } }
+        applicationScope.launch { refreshEvents.collect { refreshBlueprint() } }
     }
 
     /**
-     * Transitions to a blueprint.
+     * Transitions to a blueprint, or refreshes it if already applied.
      *
      * @param blueprintId
      * @return whether the transition has succeeded.
      */
-    fun transitionToBlueprint(blueprintId: String): Boolean {
+    fun transitionOrRefreshBlueprint(blueprintId: String): Boolean {
+        if (blueprintId == blueprint.value.id) {
+            refreshBlueprint()
+            return true
+        }
+
         return keyguardBlueprintRepository.applyBlueprint(blueprintId)
     }
 
@@ -97,7 +107,7 @@
      * @param blueprintId
      * @return whether the transition has succeeded.
      */
-    fun transitionToBlueprint(blueprintId: Int): Boolean {
+    fun transitionToBlueprint(blueprintId: String): Boolean {
         return keyguardBlueprintRepository.applyBlueprint(blueprintId)
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
index ccc48b5..bda6438 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/binder/KeyguardRootViewBinder.kt
@@ -36,7 +36,6 @@
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.repeatOnLifecycle
 import com.android.app.animation.Interpolators
-import com.android.app.tracing.coroutines.launch
 import com.android.internal.jank.InteractionJankMonitor
 import com.android.internal.jank.InteractionJankMonitor.CUJ_SCREEN_OFF_SHOW_AOD
 import com.android.systemui.Flags.newAodTransition
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListener.kt b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListener.kt
index ce7ec0e..962cdf1 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListener.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListener.kt
@@ -46,15 +46,14 @@
                 return
             }
 
-            if (
-                arg.isDigitsOnly() && keyguardBlueprintInteractor.transitionToBlueprint(arg.toInt())
-            ) {
-                pw.println("Transition succeeded!")
-            } else if (keyguardBlueprintInteractor.transitionToBlueprint(arg)) {
-                pw.println("Transition succeeded!")
-            } else {
-                pw.println("Invalid argument! To see available blueprint ids, run:")
-                pw.println("$ adb shell cmd statusbar blueprint help")
+            when {
+                arg.isDigitsOnly() -> pw.println("Invalid argument! Use string ids.")
+                keyguardBlueprintInteractor.transitionOrRefreshBlueprint(arg) ->
+                    pw.println("Transition succeeded!")
+                else -> {
+                    pw.println("Invalid argument! To see available blueprint ids, run:")
+                    pw.println("$ adb shell cmd statusbar blueprint help")
+                }
             }
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepository.kt
new file mode 100644
index 0000000..28c1fbf
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepository.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.data.repository
+
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.qs.panels.shared.model.EditTileData
+import com.android.systemui.qs.pipeline.data.repository.InstalledTilesComponentRepository
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.settings.UserTracker
+import javax.inject.Inject
+import kotlin.coroutines.CoroutineContext
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.withContext
+
+@SysUISingleton
+class IconAndNameCustomRepository
+@Inject
+constructor(
+    private val installedTilesComponentRepository: InstalledTilesComponentRepository,
+    private val userTracker: UserTracker,
+    @Background private val backgroundContext: CoroutineContext,
+) {
+    /**
+     * Returns a list of the icon/labels for all available (installed and enabled) tile services.
+     *
+     * No order is guaranteed.
+     */
+    suspend fun getCustomTileData(): List<EditTileData> {
+        return withContext(backgroundContext) {
+            val installedTiles =
+                installedTilesComponentRepository.getInstalledTilesServiceInfos(userTracker.userId)
+            val packageManager = userTracker.userContext.packageManager
+            installedTiles
+                .map {
+                    val tileSpec = TileSpec.create(it.componentName)
+                    val label = it.loadLabel(packageManager)
+                    val icon = it.loadIcon(packageManager)
+                    val appName = it.applicationInfo.loadLabel(packageManager)
+                    if (icon != null) {
+                        EditTileData(
+                            tileSpec,
+                            Icon.Loaded(icon, ContentDescription.Loaded(label.toString())),
+                            Text.Loaded(label.toString()),
+                            Text.Loaded(appName.toString()),
+                        )
+                    } else {
+                        null
+                    }
+                }
+                .filterNotNull()
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/StockTilesRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/StockTilesRepository.kt
new file mode 100644
index 0000000..ec9d151
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/data/repository/StockTilesRepository.kt
@@ -0,0 +1,41 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.data.repository
+
+import android.content.res.Resources
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Main
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import com.android.systemui.res.R
+import javax.inject.Inject
+
+@SysUISingleton
+class StockTilesRepository
+@Inject
+constructor(
+    @Main private val resources: Resources,
+) {
+    /**
+     * List of stock platform tiles. All of the specs will be of type [TileSpec.PlatformTileSpec].
+     */
+    val stockTiles =
+        resources
+            .getString(R.string.quick_settings_tiles_stock)
+            .split(",")
+            .map(TileSpec::create)
+            .filterNot { it is TileSpec.Invalid }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractor.kt
new file mode 100644
index 0000000..3b29422
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractor.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.domain.interactor
+
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.qs.panels.data.repository.IconAndNameCustomRepository
+import com.android.systemui.qs.panels.data.repository.StockTilesRepository
+import com.android.systemui.qs.panels.domain.model.EditTilesModel
+import com.android.systemui.qs.panels.shared.model.EditTileData
+import com.android.systemui.qs.tiles.viewmodel.QSTileConfigProvider
+import javax.inject.Inject
+
+@SysUISingleton
+class EditTilesListInteractor
+@Inject
+constructor(
+    private val stockTilesRepository: StockTilesRepository,
+    private val qsTileConfigProvider: QSTileConfigProvider,
+    private val iconAndNameCustomRepository: IconAndNameCustomRepository,
+) {
+    /**
+     * Provides a list of the tiles to edit, with their UI information (icon, labels).
+     *
+     * The icons have the label as their content description.
+     */
+    suspend fun getTilesToEdit(): EditTilesModel {
+        val stockTiles =
+            stockTilesRepository.stockTiles.map {
+                if (qsTileConfigProvider.hasConfig(it.spec)) {
+                    val config = qsTileConfigProvider.getConfig(it.spec)
+                    EditTileData(
+                        it,
+                        Icon.Resource(
+                            config.uiConfig.iconRes,
+                            ContentDescription.Resource(config.uiConfig.labelRes)
+                        ),
+                        Text.Resource(config.uiConfig.labelRes),
+                        null,
+                    )
+                } else {
+                    EditTileData(
+                        it,
+                        Icon.Resource(
+                            android.R.drawable.star_on,
+                            ContentDescription.Loaded(it.spec)
+                        ),
+                        Text.Loaded(it.spec),
+                        null
+                    )
+                }
+            }
+        return EditTilesModel(stockTiles, iconAndNameCustomRepository.getCustomTileData())
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/domain/model/EditTilesModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/model/EditTilesModel.kt
new file mode 100644
index 0000000..b573b9a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/domain/model/EditTilesModel.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.domain.model
+
+import com.android.systemui.qs.panels.shared.model.EditTileData
+
+data class EditTilesModel(
+    val stockTiles: List<EditTileData>,
+    val customTiles: List<EditTileData>,
+)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/EditTileData.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/EditTileData.kt
new file mode 100644
index 0000000..8b70bb9
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/shared/model/EditTileData.kt
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.shared.model
+
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.qs.pipeline.shared.TileSpec
+
+data class EditTileData(
+    val tileSpec: TileSpec,
+    val icon: Icon,
+    val label: Text,
+    val appName: Text?,
+) {
+    init {
+        check(
+            (tileSpec is TileSpec.PlatformTileSpec && appName == null) ||
+                (tileSpec is TileSpec.CustomTileSpec && appName != null)
+        ) {
+            "tileSpec: $tileSpec - appName: $appName. " +
+                "appName must be non-null for custom tiles and only for custom tiles."
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt
new file mode 100644
index 0000000..5c17fd1
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/EditMode.kt
@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.ui.compose
+
+import androidx.activity.compose.BackHandler
+import androidx.compose.foundation.layout.Column
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.DisposableEffect
+import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.getValue
+import androidx.compose.ui.Modifier
+import com.android.systemui.qs.panels.ui.viewmodel.EditModeViewModel
+
+@Composable
+fun EditMode(
+    viewModel: EditModeViewModel,
+    modifier: Modifier = Modifier,
+) {
+    val gridLayout by viewModel.gridLayout.collectAsState()
+    val tiles by viewModel.tiles.collectAsState(emptyList())
+
+    BackHandler { viewModel.stopEditing() }
+
+    DisposableEffect(Unit) { onDispose { viewModel.stopEditing() } }
+
+    Column(modifier) {
+        gridLayout.EditTileGrid(
+            tiles,
+            Modifier,
+            viewModel::addTile,
+            viewModel::removeTile,
+        )
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
index 68ce5d8..8806931 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/GridLayout.kt
@@ -18,7 +18,9 @@
 
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
+import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
+import com.android.systemui.qs.pipeline.shared.TileSpec
 
 interface GridLayout {
     @Composable
@@ -26,4 +28,12 @@
         tiles: List<TileViewModel>,
         modifier: Modifier,
     )
+
+    @Composable
+    fun EditTileGrid(
+        tiles: List<EditTileViewModel>,
+        modifier: Modifier,
+        onAddTile: (TileSpec, Int) -> Unit,
+        onRemoveTile: (TileSpec) -> Unit,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
index e2143e0..6539cf3 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/compose/InfiniteGridLayout.kt
@@ -28,6 +28,8 @@
 import androidx.compose.foundation.basicMarquee
 import androidx.compose.foundation.clickable
 import androidx.compose.foundation.layout.Arrangement
+import androidx.compose.foundation.layout.Arrangement.spacedBy
+import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.fillMaxHeight
@@ -37,8 +39,14 @@
 import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.lazy.grid.GridCells
 import androidx.compose.foundation.lazy.grid.GridItemSpan
+import androidx.compose.foundation.lazy.grid.LazyGridScope
 import androidx.compose.foundation.lazy.grid.LazyVerticalGrid
+import androidx.compose.foundation.shape.CircleShape
 import androidx.compose.foundation.shape.RoundedCornerShape
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.Add
+import androidx.compose.material.icons.filled.Remove
+import androidx.compose.material3.Icon
 import androidx.compose.material3.Text
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
@@ -47,6 +55,7 @@
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.rememberUpdatedState
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
@@ -56,14 +65,28 @@
 import androidx.compose.ui.platform.LocalContext
 import androidx.compose.ui.res.dimensionResource
 import androidx.compose.ui.res.integerResource
+import androidx.compose.ui.res.stringResource
+import androidx.compose.ui.semantics.contentDescription
+import androidx.compose.ui.semantics.onClick
+import androidx.compose.ui.semantics.semantics
+import androidx.compose.ui.semantics.stateDescription
+import androidx.compose.ui.unit.dp
+import com.android.compose.modifiers.background
 import com.android.compose.theme.colorAttr
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.ui.compose.Icon
+import com.android.systemui.common.ui.compose.load
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.qs.panels.domain.interactor.IconTilesInteractor
+import com.android.systemui.qs.panels.ui.viewmodel.ActiveTileColorAttributes
+import com.android.systemui.qs.panels.ui.viewmodel.AvailableEditActions
+import com.android.systemui.qs.panels.ui.viewmodel.EditTileViewModel
+import com.android.systemui.qs.panels.ui.viewmodel.TileColorAttributes
 import com.android.systemui.qs.panels.ui.viewmodel.TileUiState
 import com.android.systemui.qs.panels.ui.viewmodel.TileViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.toUiState
+import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor.Companion.POSITION_AT_END
+import com.android.systemui.qs.pipeline.shared.TileSpec
 import com.android.systemui.qs.tileimpl.QSTileImpl
 import com.android.systemui.res.R
 import javax.inject.Inject
@@ -75,6 +98,8 @@
 class InfiniteGridLayout @Inject constructor(private val iconTilesInteractor: IconTilesInteractor) :
     GridLayout {
 
+    private object TileType
+
     @Composable
     override fun TileGrid(
         tiles: List<TileViewModel>,
@@ -88,17 +113,7 @@
         val iconTilesSpecs by
             iconTilesInteractor.iconTilesSpecs.collectAsState(initial = emptySet())
 
-        LazyVerticalGrid(
-            columns =
-                GridCells.Fixed(
-                    integerResource(R.integer.quick_settings_infinite_grid_num_columns)
-                ),
-            verticalArrangement =
-                Arrangement.spacedBy(dimensionResource(R.dimen.qs_tile_margin_vertical)),
-            horizontalArrangement =
-                Arrangement.spacedBy(dimensionResource(R.dimen.qs_tile_margin_horizontal)),
-            modifier = modifier
-        ) {
+        TileLazyGrid(modifier) {
             items(
                 tiles.size,
                 span = { index ->
@@ -131,29 +146,11 @@
                 .mapLatest { it.toUiState() }
                 .collectAsState(initial = tile.currentState.toUiState())
         val context = LocalContext.current
-        val horizontalAlignment =
-            if (iconOnly) {
-                Alignment.CenterHorizontally
-            } else {
-                Alignment.Start
-            }
 
         Row(
-            modifier =
-                modifier
-                    .fillMaxWidth()
-                    .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
-                    .clickable { tile.onClick(null) }
-                    .background(colorAttr(state.colors.background))
-                    .padding(
-                        horizontal = dimensionResource(id = R.dimen.qs_label_container_margin)
-                    ),
+            modifier = modifier.clickable { tile.onClick(null) }.tileModifier(state.colors),
             verticalAlignment = Alignment.CenterVertically,
-            horizontalArrangement =
-                Arrangement.spacedBy(
-                    space = dimensionResource(id = R.dimen.qs_label_container_margin),
-                    alignment = horizontalAlignment
-                )
+            horizontalArrangement = tileHorizontalArrangement(iconOnly)
         ) {
             val icon =
                 remember(state.icon) {
@@ -165,62 +162,275 @@
                         }
                     }
                 }
-            TileIcon(icon, colorAttr(state.colors.icon))
+            TileContent(
+                label = state.label.toString(),
+                secondaryLabel = state.secondaryLabel.toString(),
+                icon = icon,
+                colors = state.colors,
+                iconOnly = iconOnly
+            )
+        }
+    }
 
-            if (!iconOnly) {
-                Column(
-                    verticalArrangement = Arrangement.Center,
-                    modifier = Modifier.fillMaxHeight()
-                ) {
-                    Text(
-                        state.label.toString(),
-                        color = colorAttr(state.colors.label),
-                        modifier = Modifier.basicMarquee(),
-                    )
-                    if (!TextUtils.isEmpty(state.secondaryLabel)) {
-                        Text(
-                            state.secondaryLabel.toString(),
-                            color = colorAttr(state.colors.secondaryLabel),
-                            modifier = Modifier.basicMarquee(),
-                        )
-                    }
+    @Composable
+    override fun EditTileGrid(
+        tiles: List<EditTileViewModel>,
+        modifier: Modifier,
+        onAddTile: (TileSpec, Int) -> Unit,
+        onRemoveTile: (TileSpec) -> Unit,
+    ) {
+        val (currentTiles, otherTiles) = tiles.partition { it.isCurrent }
+        val (otherTilesStock, otherTilesCustom) = otherTiles.partition { it.appName == null }
+        val addTileToEnd: (TileSpec) -> Unit by rememberUpdatedState {
+            onAddTile(it, POSITION_AT_END)
+        }
+        val iconOnlySpecs by iconTilesInteractor.iconTilesSpecs.collectAsState(initial = emptySet())
+        val isIconOnly: (TileSpec) -> Boolean =
+            remember(iconOnlySpecs) { { tileSpec: TileSpec -> tileSpec in iconOnlySpecs } }
+
+        TileLazyGrid(modifier = modifier) {
+            // These Text are just placeholders to see the different sections. Not final UI.
+            item(span = { GridItemSpan(maxLineSpan) }) {
+                Text("Current tiles", color = Color.White)
+            }
+
+            editTiles(
+                currentTiles,
+                ClickAction.REMOVE,
+                onRemoveTile,
+                isIconOnly,
+                indicatePosition = true,
+            )
+
+            item(span = { GridItemSpan(maxLineSpan) }) { Text("Tiles to add", color = Color.White) }
+
+            editTiles(
+                otherTilesStock,
+                ClickAction.ADD,
+                addTileToEnd,
+                isIconOnly,
+            )
+
+            item(span = { GridItemSpan(maxLineSpan) }) {
+                Text("Custom tiles to add", color = Color.White)
+            }
+
+            editTiles(
+                otherTilesCustom,
+                ClickAction.ADD,
+                addTileToEnd,
+                isIconOnly,
+            )
+        }
+    }
+
+    private fun LazyGridScope.editTiles(
+        tiles: List<EditTileViewModel>,
+        clickAction: ClickAction,
+        onClick: (TileSpec) -> Unit,
+        isIconOnly: (TileSpec) -> Boolean,
+        indicatePosition: Boolean = false,
+    ) {
+        items(
+            count = tiles.size,
+            key = { tiles[it].tileSpec.spec },
+            span = { GridItemSpan(if (isIconOnly(tiles[it].tileSpec)) 1 else 2) },
+            contentType = { TileType }
+        ) {
+            val viewModel = tiles[it]
+            val canClick =
+                when (clickAction) {
+                    ClickAction.ADD -> AvailableEditActions.ADD in viewModel.availableEditActions
+                    ClickAction.REMOVE ->
+                        AvailableEditActions.REMOVE in viewModel.availableEditActions
+                }
+            val onClickActionName =
+                when (clickAction) {
+                    ClickAction.ADD ->
+                        stringResource(id = R.string.accessibility_qs_edit_tile_add_action)
+                    ClickAction.REMOVE ->
+                        stringResource(id = R.string.accessibility_qs_edit_remove_tile_action)
+                }
+            val stateDescription =
+                if (indicatePosition) {
+                    stringResource(id = R.string.accessibility_qs_edit_position, it + 1)
+                } else {
+                    ""
+                }
+
+            Box(
+                modifier =
+                    Modifier.clickable(enabled = canClick) { onClick.invoke(viewModel.tileSpec) }
+                        .animateItem()
+                        .semantics {
+                            onClick(onClickActionName) { false }
+                            this.stateDescription = stateDescription
+                        }
+            ) {
+                EditTile(
+                    tileViewModel = viewModel,
+                    isIconOnly(viewModel.tileSpec),
+                    modifier = Modifier.height(dimensionResource(id = R.dimen.qs_tile_height))
+                )
+                if (canClick) {
+                    Badge(clickAction, Modifier.align(Alignment.TopEnd))
                 }
             }
         }
     }
 
-    @OptIn(ExperimentalAnimationGraphicsApi::class)
     @Composable
-    private fun TileIcon(icon: Icon, color: Color) {
-        val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
-        val context = LocalContext.current
-        val loadedDrawable =
-            remember(icon, context) {
-                when (icon) {
-                    is Icon.Loaded -> icon.drawable
-                    is Icon.Resource -> AppCompatResources.getDrawable(context, icon.res)
-                }
-            }
-        if (loadedDrawable !is Animatable) {
+    private fun Badge(action: ClickAction, modifier: Modifier = Modifier) {
+        Box(modifier = modifier.size(16.dp).background(Color.Cyan, shape = CircleShape)) {
             Icon(
-                icon = icon,
-                tint = color,
-                modifier = modifier,
+                imageVector =
+                    when (action) {
+                        ClickAction.ADD -> Icons.Filled.Add
+                        ClickAction.REMOVE -> Icons.Filled.Remove
+                    },
+                "",
+                tint = Color.Black,
             )
-        } else if (icon is Icon.Resource) {
-            val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
-            var atEnd by remember(icon.res) { mutableStateOf(false) }
-            LaunchedEffect(key1 = icon.res) {
-                delay(350)
-                atEnd = true
+        }
+    }
+
+    @Composable
+    private fun EditTile(
+        tileViewModel: EditTileViewModel,
+        iconOnly: Boolean,
+        modifier: Modifier = Modifier,
+    ) {
+        val label = tileViewModel.label.load() ?: tileViewModel.tileSpec.spec
+        val colors = ActiveTileColorAttributes
+
+        Row(
+            modifier = modifier.tileModifier(colors).semantics { this.contentDescription = label },
+            verticalAlignment = Alignment.CenterVertically,
+            horizontalArrangement = tileHorizontalArrangement(iconOnly)
+        ) {
+            TileContent(
+                label = label,
+                secondaryLabel = tileViewModel.appName?.load(),
+                colors = colors,
+                icon = tileViewModel.icon,
+                iconOnly = iconOnly,
+                animateIconToEnd = true,
+            )
+        }
+    }
+
+    private enum class ClickAction {
+        ADD,
+        REMOVE,
+    }
+}
+
+@OptIn(ExperimentalAnimationGraphicsApi::class)
+@Composable
+private fun TileIcon(
+    icon: Icon,
+    color: Color,
+    animateToEnd: Boolean = false,
+) {
+    val modifier = Modifier.size(dimensionResource(id = R.dimen.qs_icon_size))
+    val context = LocalContext.current
+    val loadedDrawable =
+        remember(icon, context) {
+            when (icon) {
+                is Icon.Loaded -> icon.drawable
+                is Icon.Resource -> AppCompatResources.getDrawable(context, icon.res)
             }
-            val painter = rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
-            Image(
-                painter = painter,
-                contentDescription = null,
-                colorFilter = ColorFilter.tint(color = color),
-                modifier = modifier
+        }
+    if (loadedDrawable !is Animatable) {
+        Icon(
+            icon = icon,
+            tint = color,
+            modifier = modifier,
+        )
+    } else if (icon is Icon.Resource) {
+        val image = AnimatedImageVector.animatedVectorResource(id = icon.res)
+        val painter =
+            if (animateToEnd) {
+                rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = true)
+            } else {
+                var atEnd by remember(icon.res) { mutableStateOf(false) }
+                LaunchedEffect(key1 = icon.res) {
+                    delay(350)
+                    atEnd = true
+                }
+                rememberAnimatedVectorPainter(animatedImageVector = image, atEnd = atEnd)
+            }
+        Image(
+            painter = painter,
+            contentDescription = null,
+            colorFilter = ColorFilter.tint(color = color),
+            modifier = modifier
+        )
+    }
+}
+
+@Composable
+private fun TileLazyGrid(
+    modifier: Modifier = Modifier,
+    content: LazyGridScope.() -> Unit,
+) {
+    LazyVerticalGrid(
+        columns =
+            GridCells.Fixed(integerResource(R.integer.quick_settings_infinite_grid_num_columns)),
+        verticalArrangement = spacedBy(dimensionResource(R.dimen.qs_tile_margin_vertical)),
+        horizontalArrangement = spacedBy(dimensionResource(R.dimen.qs_tile_margin_horizontal)),
+        modifier = modifier,
+        content = content,
+    )
+}
+
+@Composable
+private fun Modifier.tileModifier(colors: TileColorAttributes): Modifier {
+    return fillMaxWidth()
+        .clip(RoundedCornerShape(dimensionResource(R.dimen.qs_corner_radius)))
+        .background(colorAttr(colors.background))
+        .padding(horizontal = dimensionResource(id = R.dimen.qs_label_container_margin))
+}
+
+@Composable
+private fun tileHorizontalArrangement(iconOnly: Boolean): Arrangement.Horizontal {
+    val horizontalAlignment =
+        if (iconOnly) {
+            Alignment.CenterHorizontally
+        } else {
+            Alignment.Start
+        }
+    return spacedBy(
+        space = dimensionResource(id = R.dimen.qs_label_container_margin),
+        alignment = horizontalAlignment
+    )
+}
+
+@Composable
+private fun TileContent(
+    label: String,
+    secondaryLabel: String?,
+    icon: Icon,
+    colors: TileColorAttributes,
+    iconOnly: Boolean,
+    animateIconToEnd: Boolean = false,
+) {
+    TileIcon(icon, colorAttr(colors.icon), animateIconToEnd)
+
+    if (!iconOnly) {
+        Column(verticalArrangement = Arrangement.Center, modifier = Modifier.fillMaxHeight()) {
+            Text(
+                label,
+                color = colorAttr(colors.label),
+                modifier = Modifier.basicMarquee(),
             )
+            if (!TextUtils.isEmpty(secondaryLabel)) {
+                Text(
+                    secondaryLabel ?: "",
+                    color = colorAttr(colors.secondaryLabel),
+                    modifier = Modifier.basicMarquee(),
+                )
+            }
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt
new file mode 100644
index 0000000..69f50a7
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModel.kt
@@ -0,0 +1,154 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.ui.viewmodel
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.qs.panels.domain.interactor.EditTilesListInteractor
+import com.android.systemui.qs.panels.domain.interactor.GridLayoutTypeInteractor
+import com.android.systemui.qs.panels.shared.model.GridLayoutType
+import com.android.systemui.qs.panels.ui.compose.GridLayout
+import com.android.systemui.qs.panels.ui.compose.InfiniteGridLayout
+import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor
+import com.android.systemui.qs.pipeline.domain.interactor.CurrentTilesInteractor.Companion.POSITION_AT_END
+import com.android.systemui.qs.pipeline.domain.interactor.MinimumTilesInteractor
+import com.android.systemui.qs.pipeline.shared.TileSpec
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+@SysUISingleton
+@OptIn(ExperimentalCoroutinesApi::class)
+class EditModeViewModel
+@Inject
+constructor(
+    private val editTilesListInteractor: EditTilesListInteractor,
+    private val currentTilesInteractor: CurrentTilesInteractor,
+    private val minTilesInteractor: MinimumTilesInteractor,
+    private val defaultGridLayout: InfiniteGridLayout,
+    @Application private val applicationScope: CoroutineScope,
+    gridLayoutTypeInteractor: GridLayoutTypeInteractor,
+    gridLayoutMap: Map<GridLayoutType, @JvmSuppressWildcards GridLayout>,
+) {
+    private val _isEditing = MutableStateFlow(false)
+
+    /**
+     * Whether we should be editing right now. Use [startEditing] and [stopEditing] to change this
+     */
+    val isEditing = _isEditing.asStateFlow()
+    private val minimumTiles: Int
+        get() = minTilesInteractor.minNumberOfTiles
+
+    val gridLayout: StateFlow<GridLayout> =
+        gridLayoutTypeInteractor.layout
+            .map { gridLayoutMap[it] ?: defaultGridLayout }
+            .stateIn(
+                applicationScope,
+                SharingStarted.WhileSubscribed(),
+                defaultGridLayout,
+            )
+
+    /**
+     * Flow of view models for each tile that should be visible in edit mode (or empty flow when not
+     * editing).
+     *
+     * Guarantees of the data:
+     * * The data for the tiles is fetched once whenever [isEditing] goes from `false` to `true`.
+     *   This prevents icons/labels changing while in edit mode.
+     * * It tracks the current tiles as they are added/removed/moved by the user.
+     * * The tiles that are current will be in the same relative order as the user sees them in
+     *   Quick Settings.
+     * * The tiles that are not current will preserve their relative order even when the current
+     *   tiles change.
+     */
+    val tiles =
+        isEditing.flatMapLatest {
+            if (it) {
+                val editTilesData = editTilesListInteractor.getTilesToEdit()
+                currentTilesInteractor.currentTiles.map { tiles ->
+                    val currentSpecs = tiles.map { it.spec }
+                    val canRemoveTiles = currentSpecs.size > minimumTiles
+                    val allTiles = editTilesData.stockTiles + editTilesData.customTiles
+                    val allTilesMap = allTiles.associate { it.tileSpec to it }
+                    val currentTiles = currentSpecs.map { allTilesMap.get(it) }.filterNotNull()
+                    val nonCurrentTiles = allTiles.filter { it.tileSpec !in currentSpecs }
+
+                    (currentTiles + nonCurrentTiles).map {
+                        val current = it.tileSpec in currentSpecs
+                        val availableActions = buildSet {
+                            if (current) {
+                                add(AvailableEditActions.MOVE)
+                                if (canRemoveTiles) {
+                                    add(AvailableEditActions.REMOVE)
+                                }
+                            } else {
+                                add(AvailableEditActions.ADD)
+                            }
+                        }
+                        EditTileViewModel(
+                            it.tileSpec,
+                            it.icon,
+                            it.label,
+                            it.appName,
+                            current,
+                            availableActions
+                        )
+                    }
+                }
+            } else {
+                emptyFlow()
+            }
+        }
+
+    /** @see isEditing */
+    fun startEditing() {
+        _isEditing.value = true
+    }
+
+    /** @see isEditing */
+    fun stopEditing() {
+        _isEditing.value = false
+    }
+
+    /** Immediately moves [tileSpec] to [position]. */
+    fun moveTile(tileSpec: TileSpec, position: Int) {
+        throw NotImplementedError("This is not supported yet")
+    }
+
+    /** Immediately adds [tileSpec] to the current tiles at [position]. */
+    fun addTile(tileSpec: TileSpec, position: Int = POSITION_AT_END) {
+        currentTilesInteractor.addTile(tileSpec, position)
+    }
+
+    /** Immediately removes [tileSpec] from the current tiles. */
+    fun removeTile(tileSpec: TileSpec) {
+        currentTilesInteractor.removeTiles(listOf(tileSpec))
+    }
+
+    /** Immediately resets the current tiles to the default list. */
+    fun resetCurrentTilesToDefault() {
+        throw NotImplementedError("This is not supported yet")
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditTileViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditTileViewModel.kt
new file mode 100644
index 0000000..ba9a044
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/panels/ui/viewmodel/EditTileViewModel.kt
@@ -0,0 +1,42 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.ui.viewmodel
+
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.common.shared.model.Text
+import com.android.systemui.qs.pipeline.shared.TileSpec
+
+/**
+ * View model for each tile that is available to be added/removed/moved in Edit mode.
+ *
+ * [isCurrent] indicates whether this tile is part of the current set of tiles that the user sees in
+ * Quick Settings.
+ */
+class EditTileViewModel(
+    val tileSpec: TileSpec,
+    val icon: Icon,
+    val label: Text,
+    val appName: Text?,
+    val isCurrent: Boolean,
+    val availableEditActions: Set<AvailableEditActions>,
+)
+
+enum class AvailableEditActions {
+    ADD,
+    REMOVE,
+    MOVE,
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepository.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepository.kt
index cfcea98..c5b2737 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/data/repository/InstalledTilesComponentRepository.kt
@@ -23,6 +23,7 @@
 import android.content.Intent
 import android.content.pm.PackageManager
 import android.content.pm.PackageManager.ResolveInfoFlags
+import android.content.pm.ServiceInfo
 import android.os.UserHandle
 import android.service.quicksettings.TileService
 import androidx.annotation.GuardedBy
@@ -36,14 +37,17 @@
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
-import kotlinx.coroutines.flow.shareIn
+import kotlinx.coroutines.flow.stateIn
 
 interface InstalledTilesComponentRepository {
 
     fun getInstalledTilesComponents(userId: Int): Flow<Set<ComponentName>>
+
+    fun getInstalledTilesServiceInfos(userId: Int): List<ServiceInfo>
 }
 
 @SysUISingleton
@@ -55,38 +59,45 @@
     private val packageChangeRepository: PackageChangeRepository
 ) : InstalledTilesComponentRepository {
 
-    @GuardedBy("userMap") private val userMap = mutableMapOf<Int, Flow<Set<ComponentName>>>()
+    @GuardedBy("userMap") private val userMap = mutableMapOf<Int, StateFlow<List<ServiceInfo>>>()
 
     override fun getInstalledTilesComponents(userId: Int): Flow<Set<ComponentName>> =
-        synchronized(userMap) {
-            userMap.getOrPut(userId) {
-                /*
-                 * In order to query [PackageManager] for different users, this implementation will
-                 * call [Context.createContextAsUser] and retrieve the [PackageManager] from that
-                 * context.
-                 */
-                val packageManager =
-                    if (applicationContext.userId == userId) {
-                        applicationContext.packageManager
-                    } else {
-                        applicationContext
-                            .createContextAsUser(
-                                UserHandle.of(userId),
-                                /* flags */ 0,
-                            )
-                            .packageManager
-                    }
-                packageChangeRepository
-                    .packageChanged(UserHandle.of(userId))
-                    .onStart { emit(PackageChangeModel.Empty) }
-                    .map { reloadComponents(userId, packageManager) }
-                    .distinctUntilChanged()
-                    .shareIn(backgroundScope, SharingStarted.WhileSubscribed(), replay = 1)
-            }
+        synchronized(userMap) { getForUserLocked(userId) }
+            .map { it.mapTo(mutableSetOf()) { it.componentName } }
+
+    override fun getInstalledTilesServiceInfos(userId: Int): List<ServiceInfo> {
+        return synchronized(userMap) { getForUserLocked(userId).value }
+    }
+
+    private fun getForUserLocked(userId: Int): StateFlow<List<ServiceInfo>> {
+        return userMap.getOrPut(userId) {
+            /*
+             * In order to query [PackageManager] for different users, this implementation will
+             * call [Context.createContextAsUser] and retrieve the [PackageManager] from that
+             * context.
+             */
+            val packageManager =
+                if (applicationContext.userId == userId) {
+                    applicationContext.packageManager
+                } else {
+                    applicationContext
+                        .createContextAsUser(
+                            UserHandle.of(userId),
+                            /* flags */ 0,
+                        )
+                        .packageManager
+                }
+            packageChangeRepository
+                .packageChanged(UserHandle.of(userId))
+                .onStart { emit(PackageChangeModel.Empty) }
+                .map { reloadComponents(userId, packageManager) }
+                .distinctUntilChanged()
+                .stateIn(backgroundScope, SharingStarted.WhileSubscribed(), emptyList())
         }
+    }
 
     @WorkerThread
-    private fun reloadComponents(userId: Int, packageManager: PackageManager): Set<ComponentName> {
+    private fun reloadComponents(userId: Int, packageManager: PackageManager): List<ServiceInfo> {
         return packageManager
             .queryIntentServicesAsUser(INTENT, FLAGS, userId)
             .mapNotNull { it.serviceInfo }
@@ -100,7 +111,6 @@
                     false
                 }
             }
-            .mapTo(mutableSetOf()) { it.componentName }
     }
 
     companion object {
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
index 61896f0..b7fcef4 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/CurrentTilesInteractor.kt
@@ -115,6 +115,10 @@
      * @see TileSpecRepository.setTiles
      */
     fun setTiles(specs: List<TileSpec>)
+
+    companion object {
+        val POSITION_AT_END: Int = TileSpecRepository.POSITION_AT_END
+    }
 }
 
 /**
diff --git a/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractor.kt
new file mode 100644
index 0000000..2ae3f07
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractor.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.pipeline.domain.interactor
+
+import com.android.systemui.qs.pipeline.data.repository.MinimumTilesRepository
+import javax.inject.Inject
+
+class MinimumTilesInteractor
+@Inject
+constructor(
+    private val minimumTilesRepository: MinimumTilesRepository,
+) {
+    val minNumberOfTiles: Int
+        get() = minimumTilesRepository.minNumberOfTiles
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsContainerViewModel.kt b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsContainerViewModel.kt
index d6325c0..a04fa38 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsContainerViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/ui/viewmodel/QuickSettingsContainerViewModel.kt
@@ -18,6 +18,7 @@
 
 import com.android.systemui.brightness.ui.viewmodel.BrightnessSliderViewModel
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.qs.panels.ui.viewmodel.EditModeViewModel
 import com.android.systemui.qs.panels.ui.viewmodel.TileGridViewModel
 import javax.inject.Inject
 
@@ -27,4 +28,5 @@
 constructor(
     val brightnessSliderViewModel: BrightnessSliderViewModel,
     val tileGridViewModel: TileGridViewModel,
+    val editModeViewModel: EditModeViewModel,
 )
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationActivityStarter.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationActivityStarter.java
index 0c341cc..ec3c7d0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationActivityStarter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationActivityStarter.java
@@ -27,6 +27,9 @@
  * (e.g. clicking on a notification, tapping on the settings icon in the notification guts)
  */
 public interface NotificationActivityStarter {
+    /** Called when the user clicks on the notification bubble icon. */
+    void onNotificationBubbleIconClicked(NotificationEntry entry);
+
     /** Called when the user clicks on the surface of a notification. */
     void onNotificationClicked(NotificationEntry entry, ExpandableNotificationRow row);
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationClicker.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationClicker.java
index d10fac6..6487d55 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationClicker.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationClicker.java
@@ -117,11 +117,14 @@
         Notification notification = sbn.getNotification();
         if (notification.contentIntent != null || notification.fullScreenIntent != null
                 || row.getEntry().isBubble()) {
+            row.setBubbleClickListener(v ->
+                    mNotificationActivityStarter.onNotificationBubbleIconClicked(row.getEntry()));
             row.setOnClickListener(this);
             row.setOnDragSuccessListener(mOnDragSuccessListener);
         } else {
             row.setOnClickListener(null);
             row.setOnDragSuccessListener(null);
+            row.setBubbleClickListener(null);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 5e3df7b..23c0a0d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -375,6 +375,8 @@
             };
 
     private OnClickListener mOnClickListener;
+    @Nullable
+    private OnClickListener mBubbleClickListener;
     private OnDragSuccessListener mOnDragSuccessListener;
     private boolean mHeadsupDisappearRunning;
     private View mChildAfterViewWhenDismissed;
@@ -1234,14 +1236,19 @@
     /**
      * The click listener for the bubble button.
      */
+    @Nullable
     public View.OnClickListener getBubbleClickListener() {
-        return v -> {
-            if (mBubblesManagerOptional.isPresent()) {
-                mBubblesManagerOptional.get()
-                        .onUserChangedBubble(mEntry, !mEntry.isBubble() /* createBubble */);
-            }
-            mHeadsUpManager.removeNotification(mEntry.getKey(), true /* releaseImmediately */);
-        };
+        return mBubbleClickListener;
+    }
+
+    /**
+     * Sets the click listener for the bubble button.
+     */
+    public void setBubbleClickListener(@Nullable OnClickListener l) {
+        mBubbleClickListener = l;
+        // ensure listener is passed to the content views
+        mPrivateLayout.updateBubbleButton(mEntry);
+        mPublicLayout.updateBubbleButton(mEntry);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
index e1a7f22..e92058b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarter.java
@@ -96,6 +96,20 @@
 @SysUISingleton
 public class StatusBarNotificationActivityStarter implements NotificationActivityStarter {
 
+    /**
+     * Helps to avoid recalculation of values provided to
+     * {@link #onDismiss(PendingIntent, boolean, boolean, boolean)}} method
+     */
+    private interface OnKeyguardDismissedAction {
+        /**
+         * Invoked when keyguard is dismissed
+         *
+         * @return is used as return value for {@link ActivityStarter.OnDismissAction#onDismiss()}
+         */
+        boolean onDismiss(PendingIntent intent, boolean isActivityIntent, boolean animate,
+                boolean showOverTheLockScreen);
+    }
+
     private final Context mContext;
     private final int mDisplayId;
 
@@ -207,6 +221,30 @@
     }
 
     /**
+     * Called when the user clicks on the notification bubble icon.
+     *
+     * @param entry notification that bubble icon was clicked
+     */
+    @Override
+    public void onNotificationBubbleIconClicked(NotificationEntry entry) {
+        Runnable action = () -> {
+            mBubblesManagerOptional.ifPresent(bubblesManager ->
+                    bubblesManager.onUserChangedBubble(entry, !entry.isBubble()));
+            mHeadsUpManager.removeNotification(entry.getKey(), /* releaseImmediately= */ true);
+        };
+        if (entry.isBubble()) {
+            // entry is being un-bubbled, no need to unlock
+            action.run();
+        } else {
+            performActionAfterKeyguardDismissed(entry,
+                    (intent, isActivityIntent, animate, showOverTheLockScreen) -> {
+                        action.run();
+                        return false;
+                    });
+        }
+    }
+
+    /**
      * Called when a notification is clicked.
      *
      * @param entry notification that was clicked
@@ -217,7 +255,15 @@
         mLogger.logStartingActivityFromClick(entry, row.isHeadsUpState(),
                 mKeyguardStateController.isVisible(),
                 mNotificationShadeWindowController.getPanelExpanded());
+        OnKeyguardDismissedAction action =
+                (intent, isActivityIntent, animate, showOverTheLockScreen) ->
+                        performActionOnKeyguardDismissed(entry, row, intent, isActivityIntent,
+                                animate, showOverTheLockScreen);
+        performActionAfterKeyguardDismissed(entry, action);
+    }
 
+    private void performActionAfterKeyguardDismissed(NotificationEntry entry,
+            OnKeyguardDismissedAction action) {
         if (mRemoteInputManager.isRemoteInputActive(entry)) {
             // We have an active remote input typed and the user clicked on the notification.
             // this was probably unintentional, so we're closing the edit text instead.
@@ -251,8 +297,7 @@
         ActivityStarter.OnDismissAction postKeyguardAction = new ActivityStarter.OnDismissAction() {
             @Override
             public boolean onDismiss() {
-                return handleNotificationClickAfterKeyguardDismissed(
-                        entry, row, intent, isActivityIntent, animate, showOverLockscreen);
+                return action.onDismiss(intent, isActivityIntent, animate, showOverLockscreen);
             }
 
             @Override
@@ -271,7 +316,7 @@
         }
     }
 
-    private boolean handleNotificationClickAfterKeyguardDismissed(
+    private boolean performActionOnKeyguardDismissed(
             NotificationEntry entry,
             ExpandableNotificationRow row,
             PendingIntent intent,
@@ -282,7 +327,6 @@
 
         final Runnable runnable = () -> handleNotificationClickAfterPanelCollapsed(
                 entry, row, intent, isActivityIntent, animate);
-
         if (showOverLockscreen) {
             mShadeController.addPostCollapseAction(runnable);
             mShadeController.collapseShade(true /* animate */);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryTest.kt
index bcaad01..f5b5261 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/data/repository/KeyguardBlueprintRepositoryTest.kt
@@ -19,24 +19,20 @@
 
 package com.android.systemui.keyguard.data.repository
 
-import android.os.fakeExecutorHandler
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.ui.data.repository.ConfigurationRepository
-import com.android.systemui.concurrency.fakeExecutor
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.ui.view.layout.blueprints.DefaultKeyguardBlueprint
-import com.android.systemui.keyguard.ui.view.layout.blueprints.DefaultKeyguardBlueprint.Companion.DEFAULT
+import com.android.systemui.keyguard.ui.view.layout.blueprints.SplitShadeKeyguardBlueprint
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testScope
 import com.android.systemui.testKosmos
 import com.android.systemui.util.ThreadAssert
-import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
-import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -50,31 +46,32 @@
 class KeyguardBlueprintRepositoryTest : SysuiTestCase() {
     private lateinit var underTest: KeyguardBlueprintRepository
     @Mock lateinit var configurationRepository: ConfigurationRepository
-    @Mock lateinit var defaultLockscreenBlueprint: DefaultKeyguardBlueprint
     @Mock lateinit var threadAssert: ThreadAssert
+
     private val testScope = TestScope(StandardTestDispatcher())
     private val kosmos: Kosmos = testKosmos()
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
-        with(kosmos) {
-            whenever(defaultLockscreenBlueprint.id).thenReturn(DEFAULT)
-            underTest =
-                KeyguardBlueprintRepository(
-                    setOf(defaultLockscreenBlueprint),
-                    fakeExecutorHandler,
-                    threadAssert,
-                )
-        }
+        underTest = kosmos.keyguardBlueprintRepository
     }
 
     @Test
     fun testApplyBlueprint_DefaultLayout() {
         testScope.runTest {
             val blueprint by collectLastValue(underTest.blueprint)
-            underTest.applyBlueprint(defaultLockscreenBlueprint)
-            assertThat(blueprint).isEqualTo(defaultLockscreenBlueprint)
+            underTest.applyBlueprint(DefaultKeyguardBlueprint.DEFAULT)
+            assertThat(blueprint).isEqualTo(kosmos.defaultKeyguardBlueprint)
+        }
+    }
+
+    @Test
+    fun testApplyBlueprint_SplitShadeLayout() {
+        testScope.runTest {
+            val blueprint by collectLastValue(underTest.blueprint)
+            underTest.applyBlueprint(SplitShadeKeyguardBlueprint.ID)
+            assertThat(blueprint).isEqualTo(kosmos.splitShadeBlueprint)
         }
     }
 
@@ -83,33 +80,22 @@
         testScope.runTest {
             val blueprint by collectLastValue(underTest.blueprint)
             underTest.refreshBlueprint()
-            assertThat(blueprint).isEqualTo(defaultLockscreenBlueprint)
+            assertThat(blueprint).isEqualTo(kosmos.defaultKeyguardBlueprint)
         }
     }
 
     @Test
-    fun testTransitionToLayout_validId() {
-        assertThat(underTest.applyBlueprint(DEFAULT)).isTrue()
+    fun testTransitionToDefaultLayout_validId() {
+        assertThat(underTest.applyBlueprint(DefaultKeyguardBlueprint.DEFAULT)).isTrue()
+    }
+
+    @Test
+    fun testTransitionToSplitShadeLayout_validId() {
+        assertThat(underTest.applyBlueprint(SplitShadeKeyguardBlueprint.ID)).isTrue()
     }
 
     @Test
     fun testTransitionToLayout_invalidId() {
         assertThat(underTest.applyBlueprint("abc")).isFalse()
     }
-
-    @Test
-    fun testTransitionToSameBlueprint_refreshesBlueprint() =
-        with(kosmos) {
-            testScope.runTest {
-                val transition by collectLastValue(underTest.refreshTransition)
-                fakeExecutor.runAllReady()
-                runCurrent()
-
-                underTest.applyBlueprint(defaultLockscreenBlueprint)
-                fakeExecutor.runAllReady()
-                runCurrent()
-
-                assertThat(transition).isNotNull()
-            }
-        }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractorTest.kt
index ac5823e..0bdf47a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/domain/interactor/KeyguardBlueprintInteractorTest.kt
@@ -29,6 +29,7 @@
 import com.android.systemui.common.ui.data.repository.fakeConfigurationRepository
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyguard.data.repository.fakeKeyguardClockRepository
+import com.android.systemui.keyguard.data.repository.keyguardBlueprintRepository
 import com.android.systemui.keyguard.ui.view.layout.blueprints.DefaultKeyguardBlueprint
 import com.android.systemui.keyguard.ui.view.layout.blueprints.SplitShadeKeyguardBlueprint
 import com.android.systemui.kosmos.testScope
@@ -40,6 +41,7 @@
 import com.android.systemui.util.mockito.whenever
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
@@ -54,7 +56,7 @@
 class KeyguardBlueprintInteractorTest : SysuiTestCase() {
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
-    private val underTest by lazy { kosmos.keyguardBlueprintInteractor }
+    private val underTest = kosmos.keyguardBlueprintInteractor
     private val clockRepository by lazy { kosmos.fakeKeyguardClockRepository }
     private val configurationRepository by lazy { kosmos.fakeConfigurationRepository }
     private val fingerprintPropertyRepository by lazy { kosmos.fakeFingerprintPropertyRepository }
@@ -79,8 +81,9 @@
             val blueprintId by collectLastValue(underTest.blueprintId)
             kosmos.shadeRepository.setShadeMode(ShadeMode.Single)
             configurationRepository.onConfigurationChange()
-            runCurrent()
 
+            runCurrent()
+            advanceUntilIdle()
             assertThat(blueprintId).isEqualTo(DefaultKeyguardBlueprint.Companion.DEFAULT)
         }
     }
@@ -92,8 +95,9 @@
             val blueprintId by collectLastValue(underTest.blueprintId)
             kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
             configurationRepository.onConfigurationChange()
-            runCurrent()
 
+            runCurrent()
+            advanceUntilIdle()
             assertThat(blueprintId).isEqualTo(SplitShadeKeyguardBlueprint.Companion.ID)
         }
     }
@@ -102,12 +106,13 @@
     @DisableFlags(Flags.FLAG_COMPOSE_LOCKSCREEN)
     fun fingerprintPropertyInitialized_updatesBlueprint() {
         testScope.runTest {
-            val blueprintId by collectLastValue(underTest.blueprintId)
-            kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
-            fingerprintPropertyRepository.supportsUdfps() // initialize properties
-            runCurrent()
+            assertThat(kosmos.keyguardBlueprintRepository.targetTransitionConfig).isNull()
 
-            assertThat(blueprintId).isEqualTo(SplitShadeKeyguardBlueprint.Companion.ID)
+            fingerprintPropertyRepository.supportsUdfps() // initialize properties
+
+            runCurrent()
+            advanceUntilIdle()
+            assertThat(kosmos.keyguardBlueprintRepository.targetTransitionConfig).isNotNull()
         }
     }
 
@@ -119,9 +124,23 @@
             kosmos.shadeRepository.setShadeMode(ShadeMode.Split)
             clockRepository.setCurrentClock(clockController)
             configurationRepository.onConfigurationChange()
-            runCurrent()
 
+            runCurrent()
+            advanceUntilIdle()
             assertThat(blueprintId).isEqualTo(DefaultKeyguardBlueprint.DEFAULT)
         }
     }
+
+    @Test
+    fun testRefreshFromConfigChange() {
+        testScope.runTest {
+            assertThat(kosmos.keyguardBlueprintRepository.targetTransitionConfig).isNull()
+
+            configurationRepository.onConfigurationChange()
+
+            runCurrent()
+            advanceUntilIdle()
+            assertThat(kosmos.keyguardBlueprintRepository.targetTransitionConfig).isNotNull()
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListenerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListenerTest.kt
index dbf6a29..8a0613f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListenerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/ui/view/layout/KeyguardBlueprintCommandListenerTest.kt
@@ -66,25 +66,19 @@
     fun testHelp() {
         command().execute(pw, listOf("help"))
         verify(pw, atLeastOnce()).println(anyString())
-        verify(keyguardBlueprintInteractor, never()).transitionToBlueprint(anyString())
+        verify(keyguardBlueprintInteractor, never()).transitionOrRefreshBlueprint(anyString())
     }
 
     @Test
     fun testBlank() {
         command().execute(pw, listOf())
         verify(pw, atLeastOnce()).println(anyString())
-        verify(keyguardBlueprintInteractor, never()).transitionToBlueprint(anyString())
+        verify(keyguardBlueprintInteractor, never()).transitionOrRefreshBlueprint(anyString())
     }
 
     @Test
     fun testValidArg() {
         command().execute(pw, listOf("fake"))
-        verify(keyguardBlueprintInteractor).transitionToBlueprint("fake")
-    }
-
-    @Test
-    fun testValidArg_Int() {
-        command().execute(pw, listOf("1"))
-        verify(keyguardBlueprintInteractor).transitionToBlueprint(1)
+        verify(keyguardBlueprintInteractor).transitionOrRefreshBlueprint("fake")
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/domain/repository/IconTilesRepositoryImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/qs/panels/domain/repository/IconTilesRepositoryImplTest.kt
deleted file mode 100644
index 8cc3a85..0000000
--- a/packages/SystemUI/tests/src/com/android/systemui/qs/panels/domain/repository/IconTilesRepositoryImplTest.kt
+++ /dev/null
@@ -1,61 +0,0 @@
-/*
- * Copyright (C) 2024 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.android.systemui.qs.panels.domain.repository
-
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.qs.panels.data.repository.IconTilesRepositoryImpl
-import com.android.systemui.qs.pipeline.shared.TileSpec
-import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.test.runTest
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidJUnit4::class)
-class IconTilesRepositoryImplTest : SysuiTestCase() {
-
-    private val underTest = IconTilesRepositoryImpl()
-
-    @Test
-    fun iconTilesSpecsIsValid() = runTest {
-        val tilesSpecs by collectLastValue(underTest.iconTilesSpecs)
-        assertThat(tilesSpecs).isEqualTo(ICON_ONLY_TILES_SPECS)
-    }
-
-    companion object {
-        private val ICON_ONLY_TILES_SPECS =
-            setOf(
-                TileSpec.create("airplane"),
-                TileSpec.create("battery"),
-                TileSpec.create("cameratoggle"),
-                TileSpec.create("cast"),
-                TileSpec.create("color_correction"),
-                TileSpec.create("inversion"),
-                TileSpec.create("saver"),
-                TileSpec.create("dnd"),
-                TileSpec.create("flashlight"),
-                TileSpec.create("location"),
-                TileSpec.create("mictoggle"),
-                TileSpec.create("nfc"),
-                TileSpec.create("night"),
-                TileSpec.create("rotation")
-            )
-    }
-}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
index 127a3d7..269510e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/StatusBarNotificationActivityStarterTest.java
@@ -168,6 +168,8 @@
     private FakePowerRepository mPowerRepository;
     @Mock
     private UserTracker mUserTracker;
+    @Mock
+    private HeadsUpManager mHeadsUpManager;
     private final FakeExecutor mUiBgExecutor = new FakeExecutor(new FakeSystemClock());
     private ExpandableNotificationRow mNotificationRow;
     private ExpandableNotificationRow mBubbleNotificationRow;
@@ -222,13 +224,12 @@
                 mScreenOffAnimationController,
                 mStatusBarStateController).getPowerInteractor();
 
-        HeadsUpManager headsUpManager = mock(HeadsUpManager.class);
         NotificationLaunchAnimatorControllerProvider notificationAnimationProvider =
                 new NotificationLaunchAnimatorControllerProvider(
                         new NotificationLaunchAnimationInteractor(
                                 new NotificationLaunchAnimationRepository()),
                         mock(NotificationListContainer.class),
-                        headsUpManager,
+                        mHeadsUpManager,
                         mJankMonitor);
         mNotificationActivityStarter =
                 new StatusBarNotificationActivityStarter(
@@ -237,7 +238,7 @@
                         mHandler,
                         mUiBgExecutor,
                         mVisibilityProvider,
-                        headsUpManager,
+                        mHeadsUpManager,
                         mActivityStarter,
                         mCommandQueue,
                         mClickNotifier,
@@ -417,6 +418,51 @@
     }
 
     @Test
+    public void testOnNotificationBubbleIconClicked_unbubble_keyGuardShowing()
+            throws RemoteException {
+        NotificationEntry entry = mBubbleNotificationRow.getEntry();
+        StatusBarNotification sbn = entry.getSbn();
+
+        // Given
+        sbn.getNotification().contentIntent = mContentIntent;
+        when(mKeyguardStateController.isShowing()).thenReturn(true);
+        when(mKeyguardStateController.isOccluded()).thenReturn(true);
+
+        // When
+        mNotificationActivityStarter.onNotificationBubbleIconClicked(entry);
+
+        // Then
+        verify(mBubblesManager).onUserChangedBubble(entry, false);
+
+        verify(mHeadsUpManager).removeNotification(entry.getKey(), true);
+
+        verifyNoMoreInteractions(mContentIntent);
+        verifyNoMoreInteractions(mShadeController);
+    }
+
+    @Test
+    public void testOnNotificationBubbleIconClicked_bubble_keyGuardShowing() {
+        NotificationEntry entry = mNotificationRow.getEntry();
+        StatusBarNotification sbn = entry.getSbn();
+
+        // Given
+        sbn.getNotification().contentIntent = mContentIntent;
+        when(mKeyguardStateController.isShowing()).thenReturn(true);
+        when(mKeyguardStateController.isOccluded()).thenReturn(true);
+
+        // When
+        mNotificationActivityStarter.onNotificationBubbleIconClicked(entry);
+
+        // Then
+        verify(mBubblesManager).onUserChangedBubble(entry, true);
+
+        verify(mHeadsUpManager).removeNotification(entry.getKey(), true);
+
+        verify(mContentIntent, atLeastOnce()).isActivity();
+        verifyNoMoreInteractions(mContentIntent);
+    }
+
+    @Test
     public void testOnFullScreenIntentWhenDozing_wakeUpDevice() {
         // GIVEN entry that can has a full screen intent that can show
         PendingIntent fullScreenIntent = PendingIntent.getActivity(mContext, 1,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/FakeQSTile.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/FakeQSTile.kt
similarity index 97%
rename from packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/FakeQSTile.kt
rename to packages/SystemUI/tests/utils/src/com/android/systemui/qs/FakeQSTile.kt
index 302ac35..093ebd6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/pipeline/domain/interactor/FakeQSTile.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/FakeQSTile.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.qs.pipeline.domain.interactor
+package com.android.systemui.qs
 
 import com.android.internal.logging.InstanceId
 import com.android.systemui.animation.Expandable
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryKosmos.kt
new file mode 100644
index 0000000..d686699
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/IconAndNameCustomRepositoryKosmos.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.qs.panels.data.repository
+
+import android.content.packageManager
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.backgroundCoroutineContext
+import com.android.systemui.qs.pipeline.data.repository.installedTilesRepository
+import com.android.systemui.settings.userTracker
+import com.android.systemui.util.mockito.whenever
+
+val Kosmos.iconAndNameCustomRepository by
+    Kosmos.Fixture {
+        whenever(userTracker.userContext.packageManager).thenReturn(packageManager)
+        IconAndNameCustomRepository(
+            installedTilesRepository,
+            userTracker,
+            backgroundCoroutineContext,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryKosmos.kt
new file mode 100644
index 0000000..ff33650
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/data/repository/StockTilesRepositoryKosmos.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.data.repository
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.testCase
+
+var Kosmos.stockTilesRepository by
+    Kosmos.Fixture {
+        testCase.context.orCreateTestableResources
+        StockTilesRepository(testCase.context.resources)
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorKosmos.kt
new file mode 100644
index 0000000..bd54fd4
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/domain/interactor/EditTilesListInteractorKosmos.kt
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.panels.data.repository.iconAndNameCustomRepository
+import com.android.systemui.qs.panels.data.repository.stockTilesRepository
+import com.android.systemui.qs.tiles.viewmodel.qSTileConfigProvider
+
+val Kosmos.editTilesListInteractor by
+    Kosmos.Fixture {
+        EditTilesListInteractor(
+            stockTilesRepository,
+            qSTileConfigProvider,
+            iconAndNameCustomRepository,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelKosmos.kt
new file mode 100644
index 0000000..612a5d9
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/panels/ui/viewmodel/EditModeViewModelKosmos.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.panels.ui.viewmodel
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.kosmos.applicationCoroutineScope
+import com.android.systemui.qs.panels.domain.interactor.editTilesListInteractor
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutMap
+import com.android.systemui.qs.panels.domain.interactor.gridLayoutTypeInteractor
+import com.android.systemui.qs.panels.domain.interactor.infiniteGridLayout
+import com.android.systemui.qs.pipeline.domain.interactor.currentTilesInteractor
+import com.android.systemui.qs.pipeline.domain.interactor.minimumTilesInteractor
+
+val Kosmos.editModeViewModel by
+    Kosmos.Fixture {
+        EditModeViewModel(
+            editTilesListInteractor,
+            currentTilesInteractor,
+            minimumTilesInteractor,
+            infiniteGridLayout,
+            applicationCoroutineScope,
+            gridLayoutTypeInteractor,
+            gridLayoutMap,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeInstalledTilesComponentRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeInstalledTilesComponentRepository.kt
index ff6b7d0..ed4c67e 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeInstalledTilesComponentRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/data/repository/FakeInstalledTilesComponentRepository.kt
@@ -17,23 +17,78 @@
 package com.android.systemui.qs.pipeline.data.repository
 
 import android.content.ComponentName
+import android.content.pm.ApplicationInfo
+import android.content.pm.PackageManager
+import android.content.pm.ServiceInfo
+import android.graphics.drawable.Drawable
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.map
 
 class FakeInstalledTilesComponentRepository : InstalledTilesComponentRepository {
 
-    private val installedComponentsPerUser =
-        mutableMapOf<Int, MutableStateFlow<Set<ComponentName>>>()
+    private val installedServicesPerUser = mutableMapOf<Int, MutableStateFlow<List<ServiceInfo>>>()
 
     override fun getInstalledTilesComponents(userId: Int): Flow<Set<ComponentName>> {
-        return getFlow(userId).asStateFlow()
+        return getFlow(userId).map { it.map { it.componentName }.toSet() }
+    }
+
+    override fun getInstalledTilesServiceInfos(userId: Int): List<ServiceInfo> {
+        return getFlow(userId).value
     }
 
     fun setInstalledPackagesForUser(userId: Int, components: Set<ComponentName>) {
-        getFlow(userId).value = components
+        getFlow(userId).value =
+            components.map {
+                ServiceInfo().apply {
+                    packageName = it.packageName
+                    name = it.className
+                    applicationInfo = ApplicationInfo()
+                }
+            }
     }
 
-    private fun getFlow(userId: Int): MutableStateFlow<Set<ComponentName>> =
-        installedComponentsPerUser.getOrPut(userId) { MutableStateFlow(emptySet()) }
+    fun setInstalledServicesForUser(userId: Int, services: List<ServiceInfo>) {
+        getFlow(userId).value = services.toList()
+    }
+
+    private fun getFlow(userId: Int): MutableStateFlow<List<ServiceInfo>> =
+        installedServicesPerUser.getOrPut(userId) { MutableStateFlow(emptyList()) }
+
+    companion object {
+        fun ServiceInfo(
+            componentName: ComponentName,
+            serviceName: String,
+            serviceIcon: Drawable? = null,
+            appName: String = "",
+            appIcon: Drawable? = null
+        ): ServiceInfo {
+            val appInfo =
+                object : ApplicationInfo() {
+                    override fun loadLabel(pm: PackageManager): CharSequence {
+                        return appName
+                    }
+
+                    override fun loadIcon(pm: PackageManager?): Drawable? {
+                        return appIcon
+                    }
+                }
+            val serviceInfo =
+                object : ServiceInfo() {
+                        override fun loadLabel(pm: PackageManager): CharSequence {
+                            return serviceName
+                        }
+
+                        override fun loadIcon(pm: PackageManager?): Drawable? {
+                            return serviceIcon ?: getApplicationInfo().loadIcon(pm)
+                        }
+                    }
+                    .apply {
+                        packageName = componentName.packageName
+                        name = componentName.className
+                        applicationInfo = appInfo
+                    }
+            return serviceInfo
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractorKosmos.kt
new file mode 100644
index 0000000..ef1189f
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/qs/pipeline/domain/interactor/MinimumTilesInteractorKosmos.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.qs.pipeline.domain.interactor
+
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.qs.pipeline.data.repository.minimumTilesRepository
+
+var Kosmos.minimumTilesInteractor by
+    Kosmos.Fixture { MinimumTilesInteractor(minimumTilesRepository) }
diff --git a/services/accessibility/java/com/android/server/accessibility/gestures/TouchExplorer.java b/services/accessibility/java/com/android/server/accessibility/gestures/TouchExplorer.java
index 2f54f8c..2a7458f 100644
--- a/services/accessibility/java/com/android/server/accessibility/gestures/TouchExplorer.java
+++ b/services/accessibility/java/com/android/server/accessibility/gestures/TouchExplorer.java
@@ -1469,10 +1469,7 @@
             int policyFlags = mState.getLastReceivedPolicyFlags();
             if (mState.isDragging()) {
                 // Send an event to the end of the drag gesture.
-                int pointerIdBits = ALL_POINTER_ID_BITS;
-                if (Flags.fixDragPointerWhenEndingDrag()) {
-                    pointerIdBits = 1 << mDraggingPointerId;
-                }
+                int pointerIdBits = 1 << mDraggingPointerId;
                 mDispatcher.sendMotionEvent(event, ACTION_UP, rawEvent, pointerIdBits, policyFlags);
             }
             mState.startDelegating();
diff --git a/services/companion/java/com/android/server/companion/virtual/InputController.java b/services/companion/java/com/android/server/companion/virtual/InputController.java
index 9b72288..d7c65c7 100644
--- a/services/companion/java/com/android/server/companion/virtual/InputController.java
+++ b/services/companion/java/com/android/server/companion/virtual/InputController.java
@@ -41,7 +41,6 @@
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.Slog;
-import android.view.Display;
 import android.view.InputDevice;
 import android.view.WindowManager;
 
@@ -169,7 +168,6 @@
         createDeviceInternal(InputDeviceDescriptor.TYPE_MOUSE, deviceName, vendorId, productId,
                 deviceToken, displayId, phys,
                 () -> mNativeWrapper.openUinputMouse(deviceName, vendorId, productId, phys));
-        setVirtualMousePointerDisplayId(displayId);
     }
 
     void createTouchscreen(@NonNull String deviceName, int vendorId, int productId,
@@ -236,15 +234,6 @@
         if (inputDeviceDescriptor.getType() == InputDeviceDescriptor.TYPE_KEYBOARD) {
             mInputManagerInternal.removeKeyboardLayoutAssociation(phys);
         }
-
-        // Reset values to the default if all virtual mice are unregistered, or set display
-        // id if there's another mouse (choose the most recent). The inputDeviceDescriptor must be
-        // removed from the mInputDeviceDescriptors instance variable prior to this point.
-        if (inputDeviceDescriptor.isMouse()) {
-            if (getVirtualMousePointerDisplayId() == inputDeviceDescriptor.getDisplayId()) {
-                updateActivePointerDisplayIdLocked();
-            }
-        }
     }
 
     /**
@@ -276,29 +265,6 @@
         mWindowManager.setDisplayImePolicy(displayId, policy);
     }
 
-    // TODO(b/293587049): Remove after pointer icon refactor is complete.
-    @GuardedBy("mLock")
-    private void updateActivePointerDisplayIdLocked() {
-        InputDeviceDescriptor mostRecentlyCreatedMouse = null;
-        for (int i = 0; i < mInputDeviceDescriptors.size(); ++i) {
-            InputDeviceDescriptor otherInputDeviceDescriptor = mInputDeviceDescriptors.valueAt(i);
-            if (otherInputDeviceDescriptor.isMouse()) {
-                if (mostRecentlyCreatedMouse == null
-                        || (otherInputDeviceDescriptor.getCreationOrderNumber()
-                        > mostRecentlyCreatedMouse.getCreationOrderNumber())) {
-                    mostRecentlyCreatedMouse = otherInputDeviceDescriptor;
-                }
-            }
-        }
-        if (mostRecentlyCreatedMouse != null) {
-            setVirtualMousePointerDisplayId(
-                    mostRecentlyCreatedMouse.getDisplayId());
-        } else {
-            // All mice have been unregistered
-            setVirtualMousePointerDisplayId(Display.INVALID_DISPLAY);
-        }
-    }
-
     /**
      * Validates a device name by checking whether a device with the same name already exists.
      * @param deviceName The name of the device to be validated
@@ -355,9 +321,6 @@
             if (inputDeviceDescriptor == null) {
                 return false;
             }
-            if (inputDeviceDescriptor.getDisplayId() != getVirtualMousePointerDisplayId()) {
-                setVirtualMousePointerDisplayId(inputDeviceDescriptor.getDisplayId());
-            }
             return mNativeWrapper.writeButtonEvent(inputDeviceDescriptor.getNativePointer(),
                     event.getButtonCode(), event.getAction(), event.getEventTimeNanos());
         }
@@ -384,9 +347,6 @@
             if (inputDeviceDescriptor == null) {
                 return false;
             }
-            if (inputDeviceDescriptor.getDisplayId() != getVirtualMousePointerDisplayId()) {
-                setVirtualMousePointerDisplayId(inputDeviceDescriptor.getDisplayId());
-            }
             return mNativeWrapper.writeRelativeEvent(inputDeviceDescriptor.getNativePointer(),
                     event.getRelativeX(), event.getRelativeY(), event.getEventTimeNanos());
         }
@@ -399,9 +359,6 @@
             if (inputDeviceDescriptor == null) {
                 return false;
             }
-            if (inputDeviceDescriptor.getDisplayId() != getVirtualMousePointerDisplayId()) {
-                setVirtualMousePointerDisplayId(inputDeviceDescriptor.getDisplayId());
-            }
             return mNativeWrapper.writeScrollEvent(inputDeviceDescriptor.getNativePointer(),
                     event.getXAxisMovement(), event.getYAxisMovement(), event.getEventTimeNanos());
         }
@@ -415,9 +372,6 @@
                 throw new IllegalArgumentException(
                         "Could not get cursor position for input device for given token");
             }
-            if (inputDeviceDescriptor.getDisplayId() != getVirtualMousePointerDisplayId()) {
-                setVirtualMousePointerDisplayId(inputDeviceDescriptor.getDisplayId());
-            }
             return LocalServices.getService(InputManagerInternal.class).getCursorPosition(
                     inputDeviceDescriptor.getDisplayId());
         }
@@ -878,22 +832,4 @@
         /** Returns true if the calling thread is a valid thread for device creation. */
         boolean isValidThread();
     }
-
-    // TODO(b/293587049): Remove after pointer icon refactor is complete.
-    private void setVirtualMousePointerDisplayId(int displayId) {
-        if (com.android.input.flags.Flags.enablePointerChoreographer()) {
-            // We no longer need to set the pointer display when pointer choreographer is enabled.
-            return;
-        }
-        mInputManagerInternal.setVirtualMousePointerDisplayId(displayId);
-    }
-
-    // TODO(b/293587049): Remove after pointer icon refactor is complete.
-    private int getVirtualMousePointerDisplayId() {
-        if (com.android.input.flags.Flags.enablePointerChoreographer()) {
-            // We no longer need to get the pointer display when pointer choreographer is enabled.
-            return Display.INVALID_DISPLAY;
-        }
-        return mInputManagerInternal.getVirtualMousePointerDisplayId();
-    }
 }
diff --git a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
index 3cea014..a182a10 100644
--- a/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
+++ b/services/core/java/com/android/server/am/ActivityManagerShellCommand.java
@@ -115,6 +115,7 @@
 import android.util.ArraySet;
 import android.util.DebugUtils;
 import android.util.DisplayMetrics;
+import android.util.SparseArray;
 import android.util.TeeWriter;
 import android.util.proto.ProtoOutputStream;
 import android.view.Choreographer;
@@ -275,9 +276,9 @@
                 case "compact":
                     return runCompact(pw);
                 case "freeze":
-                    return runFreeze(pw);
+                    return runFreeze(pw, true);
                 case "unfreeze":
-                    return runUnfreeze(pw);
+                    return runFreeze(pw, false);
                 case "instrument":
                     getOutPrintWriter().println("Error: must be invoked through 'am instrument'.");
                     return -1;
@@ -1203,45 +1204,27 @@
     }
 
     @NeverCompile
-    int runFreeze(PrintWriter pw) throws RemoteException {
+    int runFreeze(PrintWriter pw, boolean freeze) throws RemoteException {
         String freezerOpt = getNextOption();
         boolean isSticky = false;
-        if (freezerOpt != null) {
-            isSticky = freezerOpt.equals("--sticky");
-        }
-        ProcessRecord app = getProcessFromShell();
-        if (app == null) {
-            getErrPrintWriter().println("Error: could not find process");
-            return -1;
-        }
-        pw.println("Freezing pid: " + app.mPid + " sticky=" + isSticky);
-        synchronized (mInternal) {
-            synchronized (mInternal.mProcLock) {
-                app.mOptRecord.setFreezeSticky(isSticky);
-                mInternal.mOomAdjuster.mCachedAppOptimizer.forceFreezeAppAsyncLSP(app);
-            }
-        }
-        return 0;
-    }
 
-    @NeverCompile
-    int runUnfreeze(PrintWriter pw) throws RemoteException {
-        String freezerOpt = getNextOption();
-        boolean isSticky = false;
         if (freezerOpt != null) {
             isSticky = freezerOpt.equals("--sticky");
         }
-        ProcessRecord app = getProcessFromShell();
-        if (app == null) {
-            getErrPrintWriter().println("Error: could not find process");
+        ProcessRecord proc = getProcessFromShell();
+        if (proc == null) {
             return -1;
         }
-        pw.println("Unfreezing pid: " + app.mPid);
+        pw.print(freeze ? "Freezing" : "Unfreezing");
+        pw.print(" process " + proc.processName);
+        pw.println(" (" + proc.mPid + ") sticky=" + isSticky);
         synchronized (mInternal) {
             synchronized (mInternal.mProcLock) {
-                synchronized (mInternal.mOomAdjuster.mCachedAppOptimizer.mFreezerLock) {
-                    app.mOptRecord.setFreezeSticky(isSticky);
-                    mInternal.mOomAdjuster.mCachedAppOptimizer.unfreezeAppInternalLSP(app, 0,
+                proc.mOptRecord.setFreezeSticky(isSticky);
+                if (freeze) {
+                    mInternal.mOomAdjuster.mCachedAppOptimizer.forceFreezeAppAsyncLSP(proc);
+                } else {
+                    mInternal.mOomAdjuster.mCachedAppOptimizer.unfreezeAppInternalLSP(proc, 0,
                             true);
                 }
             }
@@ -1250,43 +1233,42 @@
     }
 
     /**
-     * Parses from the shell the process name and user id if provided and provides the corresponding
-     * {@link ProcessRecord)} If no user is provided, it will fallback to current user.
-     * Example usage: {@code <processname> --user current} or {@code <processname>}
-     * @return process record of process, null if none found.
+     * Parses from the shell the pid or process name and provides the corresponding
+     * {@link ProcessRecord}.
+     * Example usage: {@code <processname>} or {@code <pid>}
+     * @return process record of process, null if none or more than one found.
      * @throws RemoteException
      */
     @NeverCompile
     ProcessRecord getProcessFromShell() throws RemoteException {
-        ProcessRecord app;
-        String processName = getNextArgRequired();
-        synchronized (mInternal.mProcLock) {
-            // Default to current user
-            int userId = getUserIdFromShellOrFallback();
-            final int uid =
-                    mInternal.getPackageManagerInternal().getPackageUid(processName, 0, userId);
-            app = mInternal.getProcessRecordLocked(processName, uid);
+        ProcessRecord proc = null;
+        String process = getNextArgRequired();
+        try {
+            int pid = Integer.parseInt(process);
+            synchronized (mInternal.mPidsSelfLocked) {
+                proc = mInternal.mPidsSelfLocked.get(pid);
+            }
+        } catch (NumberFormatException e) {
+            // Fallback to process name if it's not a valid pid
         }
-        return app;
-    }
 
-    /**
-     * @return User id from command line provided in the form of
-     *  {@code --user <userid|current|all>} and if the argument is not found it will fallback
-     *  to current user.
-     * @throws RemoteException
-     */
-    @NeverCompile
-    int getUserIdFromShellOrFallback() throws RemoteException {
-        int userId = mInterface.getCurrentUserId();
-        String userOpt = getNextOption();
-        if (userOpt != null && "--user".equals(userOpt)) {
-            int inputUserId = UserHandle.parseUserArg(getNextArgRequired());
-            if (inputUserId != UserHandle.USER_CURRENT) {
-                userId = inputUserId;
+        if (proc == null) {
+            synchronized (mInternal.mProcLock) {
+                ArrayMap<String, SparseArray<ProcessRecord>> all =
+                        mInternal.mProcessList.getProcessNamesLOSP().getMap();
+                SparseArray<ProcessRecord> procs = all.get(process);
+                if (procs == null || procs.size() == 0) {
+                    getErrPrintWriter().println("Error: could not find process");
+                    return null;
+                } else if (procs.size() > 1) {
+                    getErrPrintWriter().println("Error: more than one processes found");
+                    return null;
+                }
+                proc = procs.valueAt(0);
             }
         }
-        return userId;
+
+        return proc;
     }
 
     int runDumpHeap(PrintWriter pw) throws RemoteException {
@@ -4306,24 +4288,26 @@
             pw.println("      --allow-background-activity-starts: The receiver may start activities");
             pw.println("          even if in the background.");
             pw.println("      --async: Send without waiting for the completion of the receiver.");
-            pw.println("  compact [some|full] <process_name> [--user <USER_ID>]");
-            pw.println("      Perform a single process compaction.");
+            pw.println("  compact {some|full} <PROCESS>");
+            pw.println("      Perform a single process compaction. The given <PROCESS> argument");
+            pw.println("          may be either a process name or pid.");
             pw.println("      some: execute file compaction.");
             pw.println("      full: execute anon + file compaction.");
-            pw.println("      system: system compaction.");
             pw.println("  compact system");
             pw.println("      Perform a full system compaction.");
-            pw.println("  compact native [some|full] <pid>");
+            pw.println("  compact native {some|full} <pid>");
             pw.println("      Perform a native compaction for process with <pid>.");
             pw.println("      some: execute file compaction.");
             pw.println("      full: execute anon + file compaction.");
-            pw.println("  freeze [--sticky] <processname> [--user <USER_ID>]");
-            pw.println("      Freeze a process.");
-            pw.println("        --sticky: persists the frozen state for the process lifetime or");
+            pw.println("  freeze [--sticky] <PROCESS>");
+            pw.println("      Freeze a process. The given <PROCESS> argument");
+            pw.println("          may be either a process name or pid.  Options are:");
+            pw.println("      --sticky: persists the frozen state for the process lifetime or");
             pw.println("                  until an unfreeze is triggered via shell");
-            pw.println("  unfreeze [--sticky] <processname> [--user <USER_ID>]");
-            pw.println("      Unfreeze a process.");
-            pw.println("        --sticky: persists the unfrozen state for the process lifetime or");
+            pw.println("  unfreeze [--sticky] <PROCESS>");
+            pw.println("      Unfreeze a process. The given <PROCESS> argument");
+            pw.println("          may be either a process name or pid.  Options are:");
+            pw.println("      --sticky: persists the unfrozen state for the process lifetime or");
             pw.println("                  until a freeze is triggered via shell");
             pw.println("  instrument [-r] [-e <NAME> <VALUE>] [-p <FILE>] [-w]");
             pw.println("          [--user <USER_ID> | current]");
diff --git a/services/core/java/com/android/server/input/InputManagerInternal.java b/services/core/java/com/android/server/input/InputManagerInternal.java
index 4e9cf51..b47631c3 100644
--- a/services/core/java/com/android/server/input/InputManagerInternal.java
+++ b/services/core/java/com/android/server/input/InputManagerInternal.java
@@ -84,29 +84,6 @@
             @NonNull IBinder toChannelToken);
 
     /**
-     * Sets the display id that the MouseCursorController will be forced to target. Pass
-     * {@link android.view.Display#INVALID_DISPLAY} to clear the override.
-     *
-     * Note: This method generally blocks until the pointer display override has propagated.
-     * When setting a new override, the caller should ensure that an input device that can control
-     * the mouse pointer is connected. If a new override is set when no such input device is
-     * connected, the caller may be blocked for an arbitrary period of time.
-     *
-     * @return true if the pointer displayId was set successfully, or false if it fails.
-     *
-     * @deprecated TODO(b/293587049): Not needed - remove after Pointer Icon Refactor is complete.
-     */
-    public abstract boolean setVirtualMousePointerDisplayId(int pointerDisplayId);
-
-    /**
-     * Gets the display id that the MouseCursorController is being forced to target. Returns
-     * {@link android.view.Display#INVALID_DISPLAY} if there is no override
-     *
-     * @deprecated TODO(b/293587049): Not needed - remove after Pointer Icon Refactor is complete.
-     */
-    public abstract int getVirtualMousePointerDisplayId();
-
-    /**
      * Gets the current position of the mouse cursor.
      *
      * Returns NaN-s as the coordinates if the cursor is not available.
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index eb71952..cbd309e 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -282,33 +282,9 @@
     // WARNING: Do not call other services outside of input while holding this lock.
     private final Object mAdditionalDisplayInputPropertiesLock = new Object();
 
-    // Forces the PointerController to target a specific display id.
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private int mOverriddenPointerDisplayId = Display.INVALID_DISPLAY;
-
-    // PointerController is the source of truth of the pointer display. This is the value of the
-    // latest pointer display id reported by PointerController.
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private int mAcknowledgedPointerDisplayId = Display.INVALID_DISPLAY;
-    // This is the latest display id that IMS has requested PointerController to use. If there are
-    // no devices that can control the pointer, PointerController may end up disregarding this
-    // value.
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private int mRequestedPointerDisplayId = Display.INVALID_DISPLAY;
     @GuardedBy("mAdditionalDisplayInputPropertiesLock")
     private final SparseArray<AdditionalDisplayInputProperties> mAdditionalDisplayInputProperties =
             new SparseArray<>();
-    // This contains the per-display properties that are currently applied by native code. It should
-    // be kept in sync with the properties for mRequestedPointerDisplayId.
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private final AdditionalDisplayInputProperties mCurrentDisplayProperties =
-            new AdditionalDisplayInputProperties();
-    // TODO(b/293587049): Pointer Icon Refactor: There can be more than one pointer icon
-    // visible at once. Update this to support multi-pointer use cases.
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private int mPointerIconType = PointerIcon.TYPE_NOT_SPECIFIED;
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private PointerIcon mPointerIcon;
 
     // Holds all the registered gesture monitors that are implemented as spy windows. The spy
     // windows are mapped by their InputChannel tokens.
@@ -617,14 +593,9 @@
         }
         mNative.setDisplayViewports(vArray);
 
-        // Attempt to update the pointer display when viewports change when there is no override.
+        // Attempt to update the default pointer display when the viewports change.
         // Take care to not make calls to window manager while holding internal locks.
-        final int pointerDisplayId = mWindowManagerCallbacks.getPointerDisplayId();
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            if (mOverriddenPointerDisplayId == Display.INVALID_DISPLAY) {
-                updatePointerDisplayIdLocked(pointerDisplayId);
-            }
-        }
+        mNative.setPointerDisplayId(mWindowManagerCallbacks.getPointerDisplayId());
     }
 
     /**
@@ -1353,84 +1324,11 @@
                 properties -> properties.pointerIconVisible = visible);
     }
 
-    /**
-     * Update the display on which the mouse pointer is shown.
-     *
-     * @return true if the pointer displayId changed, false otherwise.
-     */
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private boolean updatePointerDisplayIdLocked(int pointerDisplayId) {
-        if (mRequestedPointerDisplayId == pointerDisplayId) {
-            return false;
-        }
-        mRequestedPointerDisplayId = pointerDisplayId;
-        mNative.setPointerDisplayId(pointerDisplayId);
-        applyAdditionalDisplayInputProperties();
-        return true;
-    }
-
     private void handlePointerDisplayIdChanged(PointerDisplayIdChangedArgs args) {
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            mAcknowledgedPointerDisplayId = args.mPointerDisplayId;
-            // Notify waiting threads that the display of the mouse pointer has changed.
-            mAdditionalDisplayInputPropertiesLock.notifyAll();
-        }
         mWindowManagerCallbacks.notifyPointerDisplayIdChanged(
                 args.mPointerDisplayId, args.mXPosition, args.mYPosition);
     }
 
-    private boolean setVirtualMousePointerDisplayIdBlocking(int overrideDisplayId) {
-        if (com.android.input.flags.Flags.enablePointerChoreographer()) {
-            throw new IllegalStateException(
-                    "This must not be used when PointerChoreographer is enabled");
-        }
-        final boolean isRemovingOverride = overrideDisplayId == Display.INVALID_DISPLAY;
-
-        // Take care to not make calls to window manager while holding internal locks.
-        final int resolvedDisplayId = isRemovingOverride
-                ? mWindowManagerCallbacks.getPointerDisplayId()
-                : overrideDisplayId;
-
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            mOverriddenPointerDisplayId = overrideDisplayId;
-
-            if (!updatePointerDisplayIdLocked(resolvedDisplayId)
-                    && mAcknowledgedPointerDisplayId == resolvedDisplayId) {
-                // The requested pointer display is already set.
-                return true;
-            }
-            if (isRemovingOverride && mAcknowledgedPointerDisplayId == Display.INVALID_DISPLAY) {
-                // The pointer display override is being removed, but the current pointer display
-                // is already invalid. This can happen when the PointerController is destroyed as a
-                // result of the removal of all input devices that can control the pointer.
-                return true;
-            }
-            try {
-                // The pointer display changed, so wait until the change has propagated.
-                mAdditionalDisplayInputPropertiesLock.wait(5_000 /*mills*/);
-            } catch (InterruptedException ignored) {
-            }
-            // This request succeeds in two cases:
-            // - This request was to remove the override, in which case the new pointer display
-            //   could be anything that WM has set.
-            // - We are setting a new override, in which case the request only succeeds if the
-            //   reported new displayId is the one we requested. This check ensures that if two
-            //   competing overrides are requested in succession, the caller can be notified if one
-            //   of them fails.
-            return  isRemovingOverride || mAcknowledgedPointerDisplayId == overrideDisplayId;
-        }
-    }
-
-    private int getVirtualMousePointerDisplayId() {
-        if (com.android.input.flags.Flags.enablePointerChoreographer()) {
-            throw new IllegalStateException(
-                    "This must not be used when PointerChoreographer is enabled");
-        }
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            return mOverriddenPointerDisplayId;
-        }
-    }
-
     private void setDisplayEligibilityForPointerCapture(int displayId, boolean isEligible) {
         mNative.setDisplayEligibilityForPointerCapture(displayId, isEligible);
     }
@@ -1715,31 +1613,13 @@
     // Binder call
     @Override
     public void setPointerIconType(int iconType) {
-        if (iconType == PointerIcon.TYPE_CUSTOM) {
-            throw new IllegalArgumentException("Use setCustomPointerIcon to set custom pointers");
-        }
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            mPointerIcon = null;
-            mPointerIconType = iconType;
-
-            if (!mCurrentDisplayProperties.pointerIconVisible) return;
-
-            mNative.setPointerIconType(mPointerIconType);
-        }
+        // TODO(b/311416205): Remove.
     }
 
     // Binder call
     @Override
     public void setCustomPointerIcon(PointerIcon icon) {
-        Objects.requireNonNull(icon);
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            mPointerIconType = PointerIcon.TYPE_CUSTOM;
-            mPointerIcon = icon;
-
-            if (!mCurrentDisplayProperties.pointerIconVisible) return;
-
-            mNative.setCustomPointerIcon(mPointerIcon);
-        }
+        // TODO(b/311416205): Remove.
     }
 
     // Binder call
@@ -1747,12 +1627,7 @@
     public boolean setPointerIcon(PointerIcon icon, int displayId, int deviceId, int pointerId,
             IBinder inputToken) {
         Objects.requireNonNull(icon);
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            mPointerIconType = icon.getType();
-            mPointerIcon = mPointerIconType == PointerIcon.TYPE_CUSTOM ? icon : null;
-
-            return mNative.setPointerIcon(icon, displayId, deviceId, pointerId, inputToken);
-        }
+        return mNative.setPointerIcon(icon, displayId, deviceId, pointerId, inputToken);
     }
 
     /**
@@ -2281,28 +2156,24 @@
 
     private void dumpDisplayInputPropertiesValues(IndentingPrintWriter pw) {
         synchronized (mAdditionalDisplayInputPropertiesLock) {
-            if (mAdditionalDisplayInputProperties.size() != 0) {
-                pw.println("mAdditionalDisplayInputProperties:");
-                pw.increaseIndent();
+            pw.println("mAdditionalDisplayInputProperties:");
+            pw.increaseIndent();
+            try {
+                if (mAdditionalDisplayInputProperties.size() == 0) {
+                    pw.println("<none>");
+                    return;
+                }
                 for (int i = 0; i < mAdditionalDisplayInputProperties.size(); i++) {
-                    pw.println("displayId: "
-                            + mAdditionalDisplayInputProperties.keyAt(i));
+                    pw.println("displayId: " + mAdditionalDisplayInputProperties.keyAt(i));
                     final AdditionalDisplayInputProperties properties =
                             mAdditionalDisplayInputProperties.valueAt(i);
                     pw.println("mousePointerAccelerationEnabled: "
                             + properties.mousePointerAccelerationEnabled);
                     pw.println("pointerIconVisible: " + properties.pointerIconVisible);
                 }
+            } finally {
                 pw.decreaseIndent();
             }
-            if (mOverriddenPointerDisplayId != Display.INVALID_DISPLAY) {
-                pw.println("mOverriddenPointerDisplayId: " + mOverriddenPointerDisplayId);
-            }
-
-            pw.println("mAcknowledgedPointerDisplayId=" + mAcknowledgedPointerDisplayId);
-            pw.println("mRequestedPointerDisplayId=" + mRequestedPointerDisplayId);
-            pw.println("mPointerIconType=" + PointerIcon.typeToString(mPointerIconType));
-            pw.println("mPointerIcon=" + mPointerIcon);
         }
     }
     private boolean checkCallingPermission(String permission, String func) {
@@ -3267,17 +3138,6 @@
         }
 
         @Override
-        public boolean setVirtualMousePointerDisplayId(int pointerDisplayId) {
-            return InputManagerService.this
-                    .setVirtualMousePointerDisplayIdBlocking(pointerDisplayId);
-        }
-
-        @Override
-        public int getVirtualMousePointerDisplayId() {
-            return InputManagerService.this.getVirtualMousePointerDisplayId();
-        }
-
-        @Override
         public PointF getCursorPosition(int displayId) {
             final float[] p = mNative.getMouseCursorPosition(displayId);
             if (p == null || p.length != 2) {
@@ -3413,44 +3273,6 @@
         }
     }
 
-    private void applyAdditionalDisplayInputProperties() {
-        synchronized (mAdditionalDisplayInputPropertiesLock) {
-            AdditionalDisplayInputProperties properties =
-                    mAdditionalDisplayInputProperties.get(mRequestedPointerDisplayId);
-            if (properties == null) properties = DEFAULT_ADDITIONAL_DISPLAY_INPUT_PROPERTIES;
-            applyAdditionalDisplayInputPropertiesLocked(properties);
-        }
-    }
-
-    @GuardedBy("mAdditionalDisplayInputPropertiesLock")
-    private void applyAdditionalDisplayInputPropertiesLocked(
-            AdditionalDisplayInputProperties properties) {
-        // Handle changes to each of the individual properties.
-        // TODO(b/293587049): This approach for updating pointer display properties is only for when
-        //  PointerChoreographer is disabled. Remove this logic when PointerChoreographer is
-        //  permanently enabled.
-
-        if (properties.pointerIconVisible != mCurrentDisplayProperties.pointerIconVisible) {
-            mCurrentDisplayProperties.pointerIconVisible = properties.pointerIconVisible;
-            if (properties.pointerIconVisible) {
-                if (mPointerIconType == PointerIcon.TYPE_CUSTOM) {
-                    Objects.requireNonNull(mPointerIcon);
-                    mNative.setCustomPointerIcon(mPointerIcon);
-                } else {
-                    mNative.setPointerIconType(mPointerIconType);
-                }
-            } else {
-                mNative.setPointerIconType(PointerIcon.TYPE_NULL);
-            }
-        }
-
-        if (properties.mousePointerAccelerationEnabled
-                != mCurrentDisplayProperties.mousePointerAccelerationEnabled) {
-            mCurrentDisplayProperties.mousePointerAccelerationEnabled =
-                    properties.mousePointerAccelerationEnabled;
-        }
-    }
-
     private void updateAdditionalDisplayInputProperties(int displayId,
             Consumer<AdditionalDisplayInputProperties> updater) {
         synchronized (mAdditionalDisplayInputPropertiesLock) {
@@ -3473,13 +3295,6 @@
             if (properties.allDefaults()) {
                 mAdditionalDisplayInputProperties.remove(displayId);
             }
-            if (displayId != mRequestedPointerDisplayId) {
-                Log.i(TAG, "Not applying additional properties for display " + displayId
-                        + " because the pointer is currently targeting display "
-                        + mRequestedPointerDisplayId + ".");
-                return;
-            }
-            applyAdditionalDisplayInputPropertiesLocked(properties);
         }
     }
 
diff --git a/services/core/java/com/android/server/input/NativeInputManagerService.java b/services/core/java/com/android/server/input/NativeInputManagerService.java
index 32d5044..f742360 100644
--- a/services/core/java/com/android/server/input/NativeInputManagerService.java
+++ b/services/core/java/com/android/server/input/NativeInputManagerService.java
@@ -189,12 +189,8 @@
 
     void disableInputDevice(int deviceId);
 
-    void setPointerIconType(int iconId);
-
     void reloadPointerIcons();
 
-    void setCustomPointerIcon(@NonNull PointerIcon icon);
-
     boolean setPointerIcon(@NonNull PointerIcon icon, int displayId, int deviceId, int pointerId,
             @NonNull IBinder inputToken);
 
@@ -467,15 +463,9 @@
         public native void disableInputDevice(int deviceId);
 
         @Override
-        public native void setPointerIconType(int iconId);
-
-        @Override
         public native void reloadPointerIcons();
 
         @Override
-        public native void setCustomPointerIcon(PointerIcon icon);
-
-        @Override
         public native boolean setPointerIcon(PointerIcon icon, int displayId, int deviceId,
                 int pointerId, IBinder inputToken);
 
diff --git a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
index 869b89a..73647db 100644
--- a/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
+++ b/services/core/java/com/android/server/media/MediaRouter2ServiceImpl.java
@@ -1305,22 +1305,20 @@
                         route.getId(),
                         requestId));
 
+        UserHandler userHandler = routerRecord.mUserRecord.mHandler;
         if (managerRequestId != MediaRoute2ProviderService.REQUEST_ID_NONE) {
-            ManagerRecord manager = routerRecord.mUserRecord.mHandler.findManagerWithId(
-                    toRequesterId(managerRequestId));
+            ManagerRecord manager = userHandler.findManagerWithId(toRequesterId(managerRequestId));
             if (manager == null || manager.mLastSessionCreationRequest == null) {
                 Slog.w(TAG, "requestCreateSessionWithRouter2Locked: "
                         + "Ignoring unknown request.");
-                routerRecord.mUserRecord.mHandler.notifySessionCreationFailedToRouter(
-                        routerRecord, requestId);
+                userHandler.notifySessionCreationFailedToRouter(routerRecord, requestId);
                 return;
             }
             if (!TextUtils.equals(manager.mLastSessionCreationRequest.mOldSession.getId(),
                     oldSession.getId())) {
                 Slog.w(TAG, "requestCreateSessionWithRouter2Locked: "
                         + "Ignoring unmatched routing session.");
-                routerRecord.mUserRecord.mHandler.notifySessionCreationFailedToRouter(
-                        routerRecord, requestId);
+                userHandler.notifySessionCreationFailedToRouter(routerRecord, requestId);
                 return;
             }
             if (!TextUtils.equals(manager.mLastSessionCreationRequest.mRoute.getId(),
@@ -1333,29 +1331,28 @@
                 } else {
                     Slog.w(TAG, "requestCreateSessionWithRouter2Locked: "
                             + "Ignoring unmatched route.");
-                    routerRecord.mUserRecord.mHandler.notifySessionCreationFailedToRouter(
-                            routerRecord, requestId);
+                    userHandler.notifySessionCreationFailedToRouter(routerRecord, requestId);
                     return;
                 }
             }
             manager.mLastSessionCreationRequest = null;
         } else {
+            String defaultRouteId = userHandler.mSystemProvider.getDefaultRoute().getId();
             if (route.isSystemRoute()
                     && !routerRecord.hasSystemRoutingPermission()
-                    && !TextUtils.equals(route.getId(), MediaRoute2Info.ROUTE_ID_DEFAULT)) {
+                    && !TextUtils.equals(route.getId(), defaultRouteId)) {
                 Slog.w(TAG, "MODIFY_AUDIO_ROUTING permission is required to transfer to"
                         + route);
-                routerRecord.mUserRecord.mHandler.notifySessionCreationFailedToRouter(
-                        routerRecord, requestId);
+                userHandler.notifySessionCreationFailedToRouter(routerRecord, requestId);
                 return;
             }
         }
 
         long uniqueRequestId = toUniqueRequestId(routerRecord.mRouterId, requestId);
-        routerRecord.mUserRecord.mHandler.sendMessage(
+        userHandler.sendMessage(
                 obtainMessage(
                         UserHandler::requestCreateSessionWithRouter2OnHandler,
-                        routerRecord.mUserRecord.mHandler,
+                        userHandler,
                         uniqueRequestId,
                         managerRequestId,
                         transferInitiatorUserHandle,
@@ -1429,18 +1426,22 @@
                         "transferToRouteWithRouter2 | router: %s(id: %d), route: %s",
                         routerRecord.mPackageName, routerRecord.mRouterId, route.getId()));
 
+        UserHandler userHandler = routerRecord.mUserRecord.mHandler;
+        String defaultRouteId = userHandler.mSystemProvider.getDefaultRoute().getId();
         if (route.isSystemRoute()
                 && !routerRecord.hasSystemRoutingPermission()
-                && !TextUtils.equals(route.getId(), MediaRoute2Info.ROUTE_ID_DEFAULT)) {
-            routerRecord.mUserRecord.mHandler.sendMessage(
-                    obtainMessage(UserHandler::notifySessionCreationFailedToRouter,
-                            routerRecord.mUserRecord.mHandler,
-                            routerRecord, toOriginalRequestId(DUMMY_REQUEST_ID)));
+                && !TextUtils.equals(route.getId(), defaultRouteId)) {
+            userHandler.sendMessage(
+                    obtainMessage(
+                            UserHandler::notifySessionCreationFailedToRouter,
+                            userHandler,
+                            routerRecord,
+                            toOriginalRequestId(DUMMY_REQUEST_ID)));
         } else {
-            routerRecord.mUserRecord.mHandler.sendMessage(
+            userHandler.sendMessage(
                     obtainMessage(
                             UserHandler::transferToRouteOnHandler,
-                            routerRecord.mUserRecord.mHandler,
+                            userHandler,
                             DUMMY_REQUEST_ID,
                             transferInitiatorUserHandle,
                             routerRecord.mPackageName,
diff --git a/services/core/java/com/android/server/media/SystemMediaRoute2Provider.java b/services/core/java/com/android/server/media/SystemMediaRoute2Provider.java
index c105b9c..6ce3ab4 100644
--- a/services/core/java/com/android/server/media/SystemMediaRoute2Provider.java
+++ b/services/core/java/com/android/server/media/SystemMediaRoute2Provider.java
@@ -232,10 +232,16 @@
             String sessionId,
             String routeId,
             @RoutingSessionInfo.TransferReason int transferReason) {
+        String selectedDeviceRouteId = mDeviceRouteController.getSelectedRoute().getId();
         if (TextUtils.equals(routeId, MediaRoute2Info.ROUTE_ID_DEFAULT)) {
-            // The currently selected route is the default route.
-            Log.w(TAG, "Ignoring transfer to " + MediaRoute2Info.ROUTE_ID_DEFAULT);
-            return;
+            if (Flags.enableBuiltInSpeakerRouteSuitabilityStatuses()) {
+                // Transfer to the default route (which is the selected route). We replace the id to
+                // be the selected route id so that the transfer reason gets updated.
+                routeId = selectedDeviceRouteId;
+            } else {
+                Log.w(TAG, "Ignoring transfer to " + MediaRoute2Info.ROUTE_ID_DEFAULT);
+                return;
+            }
         }
 
         if (Flags.enableBuiltInSpeakerRouteSuitabilityStatuses()) {
@@ -250,11 +256,11 @@
             }
         }
 
-        MediaRoute2Info selectedDeviceRoute = mDeviceRouteController.getSelectedRoute();
+        String finalRouteId = routeId; // Make a final copy to use it in the lambda.
         boolean isAvailableDeviceRoute =
                 mDeviceRouteController.getAvailableRoutes().stream()
-                        .anyMatch(it -> it.getId().equals(routeId));
-        boolean isSelectedDeviceRoute = TextUtils.equals(routeId, selectedDeviceRoute.getId());
+                        .anyMatch(it -> it.getId().equals(finalRouteId));
+        boolean isSelectedDeviceRoute = TextUtils.equals(routeId, selectedDeviceRouteId);
 
         if (isSelectedDeviceRoute || isAvailableDeviceRoute) {
             // The requested route is managed by the device route controller. Note that the selected
diff --git a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
index 1793794..7a36f6d 100644
--- a/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
+++ b/services/core/java/com/android/server/pm/PackageManagerShellCommand.java
@@ -35,6 +35,8 @@
 import android.app.ActivityManager;
 import android.app.ActivityManagerInternal;
 import android.app.role.RoleManager;
+import android.app.usage.StorageStats;
+import android.app.usage.StorageStatsManager;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.IIntentReceiver;
@@ -136,6 +138,7 @@
 import java.io.PrintWriter;
 import java.net.URISyntaxException;
 import java.security.SecureRandom;
+import java.text.DecimalFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Base64;
@@ -275,6 +278,8 @@
                     return runClear();
                 case "get-archived-package-metadata":
                     return runGetArchivedPackageMetadata();
+                case "get-package-storage-stats":
+                    return runGetPackageStorageStats();
                 case "install-archived":
                     return runArchivedInstall();
                 case "enable":
@@ -1861,6 +1866,103 @@
         return 0;
     }
 
+    /**
+     * Returns a string that shows the number of bytes in b, Kb, Mb or Gb.
+     */
+    protected static String getFormattedBytes(long size) {
+        double k = size/1024.0;
+        double m = size/1048576.0;
+        double g = size/1073741824.0;
+
+        DecimalFormat dec = new DecimalFormat("0.00");
+        if (g > 1) {
+            return dec.format(g).concat(" Gb");
+        } else if (m > 1) {
+            return dec.format(m).concat(" Mb");
+        } else if (k > 1) {
+            return dec.format(k).concat(" Kb");
+        }
+        return "";
+    }
+
+    /**
+     * Return the string that displays the data size.
+     */
+    private String getDataSizeDisplay(long size) {
+        String formattedOutput = getFormattedBytes(size);
+        if (!formattedOutput.isEmpty()) {
+           formattedOutput = " (" + formattedOutput + ")";
+        }
+        return Long.toString(size) + " bytes" + formattedOutput;
+    }
+
+    /**
+     * Display storage stats of the specified package.
+     *
+     * Usage: get-package-storage-stats [--usr USER_ID] PACKAGE
+     */
+    private int runGetPackageStorageStats() throws RemoteException {
+        final PrintWriter pw = getOutPrintWriter();
+        if (!android.content.pm.Flags.getPackageStorageStats()) {
+            pw.println("Error: get_package_storage_stats flag is not enabled");
+            return 1;
+        }
+        if (!android.app.usage.Flags.getAppBytesByDataTypeApi()) {
+            pw.println("Error: get_app_bytes_by_data_type_api flag is not enabled");
+            return 1;
+        }
+        int userId = UserHandle.USER_CURRENT;
+
+        String opt;
+        while ((opt = getNextOption()) != null) {
+            switch (opt) {
+                case "--user":
+                    userId = UserHandle.parseUserArg(getNextArgRequired());
+                    break;
+                default:
+                    pw.println("Error: Unknown option: " + opt);
+                    return 1;
+            }
+        }
+
+        final String packageName = getNextArg();
+        if (packageName == null) {
+            pw.println("Error: package name not specified");
+            return 1;
+        }
+        try {
+            StorageStatsManager storageStatsManager =
+                mContext.getSystemService(StorageStatsManager.class);
+            final int translatedUserId = translateUserId(userId, UserHandle.USER_NULL,
+                "runGetPackageStorageStats");
+            StorageStats stats =
+                storageStatsManager.queryStatsForPackage(StorageManager.UUID_DEFAULT,
+                    packageName, UserHandle.of(translatedUserId));
+
+            pw.println("code: " + getDataSizeDisplay(stats.getAppBytes()));
+            pw.println("data: " + getDataSizeDisplay(stats.getDataBytes()));
+            pw.println("cache: " + getDataSizeDisplay(stats.getCacheBytes()));
+            pw.println("apk: " + getDataSizeDisplay(stats.getAppBytesByDataType(
+                StorageStats.APP_DATA_TYPE_FILE_TYPE_APK)));
+            pw.println("lib: " + getDataSizeDisplay(
+                stats.getAppBytesByDataType(StorageStats.APP_DATA_TYPE_LIB)));
+            pw.println("dm: " + getDataSizeDisplay(stats.getAppBytesByDataType(
+                StorageStats.APP_DATA_TYPE_FILE_TYPE_DM)));
+            pw.println("dexopt artifacts: " + getDataSizeDisplay(stats.getAppBytesByDataType(
+                StorageStats.APP_DATA_TYPE_FILE_TYPE_DEXOPT_ARTIFACT)));
+            pw.println("current profile : " + getDataSizeDisplay(stats.getAppBytesByDataType(
+                StorageStats.APP_DATA_TYPE_FILE_TYPE_CURRENT_PROFILE)));
+            pw.println("reference profile: " + getDataSizeDisplay(stats.getAppBytesByDataType(
+                StorageStats.APP_DATA_TYPE_FILE_TYPE_REFERENCE_PROFILE)));
+            pw.println("external cache: " + getDataSizeDisplay(stats.getExternalCacheBytes()));
+        } catch (Exception e) {
+            getErrPrintWriter().println("Failed to get storage stats, reason: " + e);
+            pw.println("Failure [failed to get storage stats], reason: " + e);
+            return -1;
+        }
+        return 0;
+    }
+
     private int runInstallExisting() throws RemoteException {
         final PrintWriter pw = getOutPrintWriter();
         int userId = UserHandle.USER_CURRENT;
@@ -4869,6 +4971,8 @@
         pw.println("    Displays the component name of the domain verification agent on device.");
         pw.println("    If the component isn't enabled, an error message will be displayed.");
         pw.println("      --user: return the agent of the given user (SYSTEM_USER if unspecified)");
+        pw.println("  get-package-storage-stats [--user <USER_ID>] <PACKAGE>");
+        pw.println("    Return the storage stats for the given app, if present");
         pw.println("");
         printArtServiceHelp();
         pw.println("");
diff --git a/services/core/jni/com_android_server_input_InputManagerService.cpp b/services/core/jni/com_android_server_input_InputManagerService.cpp
index fac08eb..62f5b89 100644
--- a/services/core/jni/com_android_server_input_InputManagerService.cpp
+++ b/services/core/jni/com_android_server_input_InputManagerService.cpp
@@ -2463,20 +2463,12 @@
     im->setInputDeviceEnabled(deviceId, false);
 }
 
-static void nativeSetPointerIconType(JNIEnv* env, jobject nativeImplObj, jint iconId) {
-    // TODO(b/311416205): Remove
-}
-
 static void nativeReloadPointerIcons(JNIEnv* env, jobject nativeImplObj) {
     NativeInputManager* im = getNativeInputManager(env, nativeImplObj);
 
     im->reloadPointerIcons();
 }
 
-static void nativeSetCustomPointerIcon(JNIEnv* env, jobject nativeImplObj, jobject iconObj) {
-    // TODO(b/311416205): Remove
-}
-
 static bool nativeSetPointerIcon(JNIEnv* env, jobject nativeImplObj, jobject iconObj,
                                  jint displayId, jint deviceId, jint pointerId,
                                  jobject inputTokenObj) {
@@ -2796,10 +2788,7 @@
         {"isInputDeviceEnabled", "(I)Z", (void*)nativeIsInputDeviceEnabled},
         {"enableInputDevice", "(I)V", (void*)nativeEnableInputDevice},
         {"disableInputDevice", "(I)V", (void*)nativeDisableInputDevice},
-        {"setPointerIconType", "(I)V", (void*)nativeSetPointerIconType},
         {"reloadPointerIcons", "()V", (void*)nativeReloadPointerIcons},
-        {"setCustomPointerIcon", "(Landroid/view/PointerIcon;)V",
-         (void*)nativeSetCustomPointerIcon},
         {"setPointerIcon", "(Landroid/view/PointerIcon;IIILandroid/os/IBinder;)Z",
          (void*)nativeSetPointerIcon},
         {"setPointerIconVisibility", "(IZ)V", (void*)nativeSetPointerIconVisibility},
diff --git a/services/tests/mockingservicestests/src/com/android/server/RescuePartyTest.java b/services/tests/mockingservicestests/src/com/android/server/RescuePartyTest.java
index 4a21645..42814e7 100644
--- a/services/tests/mockingservicestests/src/com/android/server/RescuePartyTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/RescuePartyTest.java
@@ -239,6 +239,9 @@
 
     @Test
     public void testBootLoopDetectionWithExecutionForAllRescueLevels() {
+        // this is old test where the flag needs to be disabled
+        mSetFlagsRule.disableFlags(Flags.FLAG_RECOVERABILITY_DETECTION);
+
         RescueParty.onSettingsProviderPublished(mMockContext);
         verify(() -> DeviceConfig.setMonitorCallback(eq(mMockContentResolver),
                 any(Executor.class),
@@ -449,6 +452,9 @@
 
     @Test
     public void testNonPersistentAppCrashDetectionWithScopedResets() {
+        // this is old test where the flag needs to be disabled
+        mSetFlagsRule.disableFlags(Flags.FLAG_RECOVERABILITY_DETECTION);
+
         RescueParty.onSettingsProviderPublished(mMockContext);
         verify(() -> DeviceConfig.setMonitorCallback(eq(mMockContentResolver),
                 any(Executor.class),
@@ -506,6 +512,9 @@
 
     @Test
     public void testNonDeviceConfigSettingsOnlyResetOncePerLevel() {
+        // this is old test where the flag needs to be disabled
+        mSetFlagsRule.disableFlags(Flags.FLAG_RECOVERABILITY_DETECTION);
+
         RescueParty.onSettingsProviderPublished(mMockContext);
         verify(() -> DeviceConfig.setMonitorCallback(eq(mMockContentResolver),
                 any(Executor.class),
@@ -879,6 +888,9 @@
 
     @Test
     public void testBootLoopLevels() {
+        // this is old test where the flag needs to be disabled
+        mSetFlagsRule.disableFlags(Flags.FLAG_RECOVERABILITY_DETECTION);
+
         RescuePartyObserver observer = RescuePartyObserver.getInstance(mMockContext);
 
         assertEquals(observer.onBootLoop(0), PackageHealthObserverImpact.USER_IMPACT_LEVEL_0);
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/JobSchedulerServiceTest.java b/services/tests/mockingservicestests/src/com/android/server/job/JobSchedulerServiceTest.java
index 11f20e3..d15c24b 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/JobSchedulerServiceTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/JobSchedulerServiceTest.java
@@ -31,6 +31,7 @@
 import static com.android.server.job.JobSchedulerService.sUptimeMillisClock;
 import static com.android.server.job.Flags.FLAG_BATCH_ACTIVE_BUCKET_JOBS;
 import static com.android.server.job.Flags.FLAG_BATCH_CONNECTIVITY_JOBS_PER_NETWORK;
+import static com.android.server.job.Flags.FLAG_THERMAL_RESTRICTIONS_TO_FGS_JOBS;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -74,6 +75,9 @@
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.SystemClock;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.platform.test.flag.junit.SetFlagsRule;
 
 import com.android.server.AppStateTracker;
@@ -85,6 +89,8 @@
 import com.android.server.job.controllers.ConnectivityController;
 import com.android.server.job.controllers.JobStatus;
 import com.android.server.job.controllers.QuotaController;
+import com.android.server.job.restrictions.JobRestriction;
+import com.android.server.job.restrictions.ThermalStatusRestriction;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.usage.AppStandbyInternal;
 
@@ -121,6 +127,9 @@
     @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
     private ChargingPolicyChangeListener mChargingPolicyChangeListener;
 
     private class TestJobSchedulerService extends JobSchedulerService {
@@ -2385,6 +2394,108 @@
         assertEquals(JobScheduler.PENDING_JOB_REASON_USER, mService.getPendingJobReason(job2b));
     }
 
+    /**
+     * Unit tests {@link JobSchedulerService#checkIfRestricted(JobStatus)} with single {@link
+     * JobRestriction} registered.
+     */
+    @Test
+    public void testCheckIfRestrictedSingleRestriction() {
+        int bias = JobInfo.BIAS_BOUND_FOREGROUND_SERVICE;
+        JobStatus fgsJob =
+                createJobStatus(
+                        "testCheckIfRestrictedSingleRestriction", createJobInfo(1).setBias(bias));
+        ThermalStatusRestriction mockThermalStatusRestriction =
+                mock(ThermalStatusRestriction.class);
+        mService.mJobRestrictions.clear();
+        mService.mJobRestrictions.add(mockThermalStatusRestriction);
+        when(mockThermalStatusRestriction.isJobRestricted(fgsJob, bias)).thenReturn(true);
+
+        synchronized (mService.mLock) {
+            assertEquals(mService.checkIfRestricted(fgsJob), mockThermalStatusRestriction);
+        }
+
+        when(mockThermalStatusRestriction.isJobRestricted(fgsJob, bias)).thenReturn(false);
+        synchronized (mService.mLock) {
+            assertNull(mService.checkIfRestricted(fgsJob));
+        }
+    }
+
+    /**
+     * Unit tests {@link JobSchedulerService#checkIfRestricted(JobStatus)} with multiple {@link
+     * JobRestriction} registered.
+     */
+    @Test
+    public void testCheckIfRestrictedMultipleRestrictions() {
+        int bias = JobInfo.BIAS_BOUND_FOREGROUND_SERVICE;
+        JobStatus fgsJob =
+                createJobStatus(
+                        "testGetMinJobExecutionGuaranteeMs", createJobInfo(1).setBias(bias));
+        JobRestriction mock1JobRestriction = mock(JobRestriction.class);
+        JobRestriction mock2JobRestriction = mock(JobRestriction.class);
+        mService.mJobRestrictions.clear();
+        mService.mJobRestrictions.add(mock1JobRestriction);
+        mService.mJobRestrictions.add(mock2JobRestriction);
+
+        // Jobs will be restricted if any one of the registered {@link JobRestriction}
+        // reports true.
+        when(mock1JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(true);
+        when(mock2JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(false);
+        synchronized (mService.mLock) {
+            assertEquals(mService.checkIfRestricted(fgsJob), mock1JobRestriction);
+        }
+
+        when(mock1JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(false);
+        when(mock2JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(true);
+        synchronized (mService.mLock) {
+            assertEquals(mService.checkIfRestricted(fgsJob), mock2JobRestriction);
+        }
+
+        when(mock1JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(false);
+        when(mock2JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(false);
+        synchronized (mService.mLock) {
+            assertNull(mService.checkIfRestricted(fgsJob));
+        }
+
+        when(mock1JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(true);
+        when(mock2JobRestriction.isJobRestricted(fgsJob, bias)).thenReturn(true);
+        synchronized (mService.mLock) {
+            assertNotEquals(mService.checkIfRestricted(fgsJob), mock1JobRestriction);
+        }
+    }
+
+    /**
+     * Jobs with foreground service and top app biases must not be restricted when the flag is
+     * disabled.
+     */
+    @Test
+    @RequiresFlagsDisabled(FLAG_THERMAL_RESTRICTIONS_TO_FGS_JOBS)
+    public void testCheckIfRestricted_highJobBias_flagThermalRestrictionsToFgsJobsDisabled() {
+        JobStatus fgsJob =
+                createJobStatus(
+                        "testCheckIfRestrictedJobBiasFgs",
+                        createJobInfo(1).setBias(JobInfo.BIAS_FOREGROUND_SERVICE));
+        JobStatus topAppJob =
+                createJobStatus(
+                        "testCheckIfRestrictedJobBiasTopApp",
+                        createJobInfo(2).setBias(JobInfo.BIAS_TOP_APP));
+
+        synchronized (mService.mLock) {
+            assertNull(mService.checkIfRestricted(fgsJob));
+            assertNull(mService.checkIfRestricted(topAppJob));
+        }
+    }
+
+    /** Jobs with top app biases must not be restricted. */
+    @Test
+    public void testCheckIfRestricted_highJobBias() {
+        JobStatus topAppJob = createJobStatus(
+                "testCheckIfRestrictedJobBiasTopApp",
+                createJobInfo(1).setBias(JobInfo.BIAS_TOP_APP));
+        synchronized (mService.mLock) {
+            assertNull(mService.checkIfRestricted(topAppJob));
+        }
+    }
+
     private void setBatteryLevel(int level) {
         doReturn(level).when(mBatteryManagerInternal).getBatteryLevel();
         mService.mBatteryStateTracker
diff --git a/services/tests/mockingservicestests/src/com/android/server/job/restrictions/ThermalStatusRestrictionTest.java b/services/tests/mockingservicestests/src/com/android/server/job/restrictions/ThermalStatusRestrictionTest.java
index 754f409..c2c67e6 100644
--- a/services/tests/mockingservicestests/src/com/android/server/job/restrictions/ThermalStatusRestrictionTest.java
+++ b/services/tests/mockingservicestests/src/com/android/server/job/restrictions/ThermalStatusRestrictionTest.java
@@ -28,6 +28,7 @@
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mock;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.when;
+import static com.android.server.job.Flags.FLAG_THERMAL_RESTRICTIONS_TO_FGS_JOBS;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -43,7 +44,12 @@
 import android.content.ComponentName;
 import android.content.Context;
 import android.os.PowerManager;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 import android.provider.DeviceConfig;
+import android.util.DebugUtils;
 
 import androidx.test.runner.AndroidJUnit4;
 
@@ -53,6 +59,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
@@ -76,6 +83,157 @@
     @Mock
     private JobSchedulerService mJobSchedulerService;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
+
+    class JobStatusContainer {
+        public final JobStatus jobMinPriority;
+        public final JobStatus jobLowPriority;
+        public final JobStatus jobLowPriorityRunning;
+        public final JobStatus jobLowPriorityRunningLong;
+        public final JobStatus jobDefaultPriority;
+        public final JobStatus jobHighPriority;
+        public final JobStatus jobHighPriorityRunning;
+        public final JobStatus jobHighPriorityRunningLong;
+        public final JobStatus ejDowngraded;
+        public final JobStatus ej;
+        public final JobStatus ejRetried;
+        public final JobStatus ejRunning;
+        public final JobStatus ejRunningLong;
+        public final JobStatus ui;
+        public final JobStatus uiRetried;
+        public final JobStatus uiRunning;
+        public final JobStatus uiRunningLong;
+        public final JobStatus importantWhileForeground;
+        public final JobStatus importantWhileForegroundRunning;
+        public final JobStatus importantWhileForegroundRunningLong;
+        public final int[] allJobBiases = {
+            JobInfo.BIAS_ADJ_ALWAYS_RUNNING,
+            JobInfo.BIAS_ADJ_OFTEN_RUNNING,
+            JobInfo.BIAS_DEFAULT,
+            JobInfo.BIAS_SYNC_EXPEDITED,
+            JobInfo.BIAS_SYNC_INITIALIZATION,
+            JobInfo.BIAS_BOUND_FOREGROUND_SERVICE,
+            JobInfo.BIAS_FOREGROUND_SERVICE,
+            JobInfo.BIAS_TOP_APP
+        };
+        public final int[] biasesBelowFgs = {
+            JobInfo.BIAS_ADJ_ALWAYS_RUNNING,
+            JobInfo.BIAS_ADJ_OFTEN_RUNNING,
+            JobInfo.BIAS_DEFAULT,
+            JobInfo.BIAS_SYNC_EXPEDITED,
+            JobInfo.BIAS_SYNC_INITIALIZATION,
+            JobInfo.BIAS_BOUND_FOREGROUND_SERVICE
+        };
+        public final int[] thermalStatuses = {
+            THERMAL_STATUS_NONE,
+            THERMAL_STATUS_LIGHT,
+            THERMAL_STATUS_MODERATE,
+            THERMAL_STATUS_SEVERE,
+            THERMAL_STATUS_CRITICAL,
+            THERMAL_STATUS_EMERGENCY,
+            THERMAL_STATUS_SHUTDOWN
+        };
+
+        JobStatusContainer(String jobName, JobSchedulerService mJobSchedulerService) {
+            jobMinPriority =
+                    createJobStatus(
+                            jobName, createJobBuilder(1).setPriority(JobInfo.PRIORITY_MIN).build());
+            jobLowPriority =
+                    createJobStatus(
+                            jobName, createJobBuilder(2).setPriority(JobInfo.PRIORITY_LOW).build());
+            jobLowPriorityRunning =
+                    createJobStatus(
+                            jobName, createJobBuilder(3).setPriority(JobInfo.PRIORITY_LOW).build());
+            jobLowPriorityRunningLong =
+                    createJobStatus(
+                            jobName, createJobBuilder(9).setPriority(JobInfo.PRIORITY_LOW).build());
+            jobDefaultPriority =
+                    createJobStatus(
+                            jobName,
+                            createJobBuilder(4).setPriority(JobInfo.PRIORITY_DEFAULT).build());
+            jobHighPriority =
+                    createJobStatus(
+                            jobName,
+                            createJobBuilder(5).setPriority(JobInfo.PRIORITY_HIGH).build());
+            jobHighPriorityRunning =
+                    createJobStatus(
+                            jobName,
+                            createJobBuilder(6).setPriority(JobInfo.PRIORITY_HIGH).build());
+            jobHighPriorityRunningLong =
+                    createJobStatus(
+                            jobName,
+                            createJobBuilder(10).setPriority(JobInfo.PRIORITY_HIGH).build());
+            ejDowngraded = createJobStatus(jobName, createJobBuilder(7).setExpedited(true).build());
+            ej = spy(createJobStatus(jobName, createJobBuilder(8).setExpedited(true).build()));
+            ejRetried =
+                    spy(createJobStatus(jobName, createJobBuilder(11).setExpedited(true).build()));
+            ejRunning =
+                    spy(createJobStatus(jobName, createJobBuilder(12).setExpedited(true).build()));
+            ejRunningLong =
+                    spy(createJobStatus(jobName, createJobBuilder(13).setExpedited(true).build()));
+            ui = spy(createJobStatus(jobName, createJobBuilder(14).build()));
+            uiRetried = spy(createJobStatus(jobName, createJobBuilder(15).build()));
+            uiRunning = spy(createJobStatus(jobName, createJobBuilder(16).build()));
+            uiRunningLong = spy(createJobStatus(jobName, createJobBuilder(17).build()));
+            importantWhileForeground = spy(createJobStatus(jobName, createJobBuilder(18)
+                    .setImportantWhileForeground(true)
+                    .build()));
+            importantWhileForegroundRunning = spy(createJobStatus(jobName, createJobBuilder(20)
+                    .setImportantWhileForeground(true)
+                     .build()));
+            importantWhileForegroundRunningLong = spy(createJobStatus(jobName, createJobBuilder(19)
+                     .setImportantWhileForeground(true)
+                     .build()));
+
+            when(ej.shouldTreatAsExpeditedJob()).thenReturn(true);
+            when(ejRetried.shouldTreatAsExpeditedJob()).thenReturn(true);
+            when(ejRunning.shouldTreatAsExpeditedJob()).thenReturn(true);
+            when(ejRunningLong.shouldTreatAsExpeditedJob()).thenReturn(true);
+            when(ui.shouldTreatAsUserInitiatedJob()).thenReturn(true);
+            when(uiRetried.shouldTreatAsUserInitiatedJob()).thenReturn(true);
+            when(uiRunning.shouldTreatAsUserInitiatedJob()).thenReturn(true);
+            when(uiRunningLong.shouldTreatAsUserInitiatedJob()).thenReturn(true);
+            when(ejRetried.getNumPreviousAttempts()).thenReturn(1);
+            when(uiRetried.getNumPreviousAttempts()).thenReturn(2);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(jobLowPriorityRunning))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(jobHighPriorityRunning))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(jobLowPriorityRunningLong))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(jobHighPriorityRunningLong))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(ejRunning)).thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(ejRunningLong)).thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(uiRunning)).thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(uiRunningLong)).thenReturn(true);
+            when(mJobSchedulerService.isJobInOvertimeLocked(jobLowPriorityRunningLong))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(importantWhileForegroundRunning))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isJobInOvertimeLocked(jobHighPriorityRunningLong))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isJobInOvertimeLocked(ejRunningLong)).thenReturn(true);
+            when(mJobSchedulerService.isJobInOvertimeLocked(uiRunningLong)).thenReturn(true);
+            when(mJobSchedulerService.isCurrentlyRunningLocked(importantWhileForegroundRunningLong))
+                    .thenReturn(true);
+            when(mJobSchedulerService.isJobInOvertimeLocked(importantWhileForegroundRunningLong))
+                    .thenReturn(true);
+        }
+    }
+
+    private boolean isJobRestricted(JobStatus status, int bias) {
+        return mThermalStatusRestriction.isJobRestricted(status, bias);
+    }
+
+    private static String debugTag(int bias, @PowerManager.ThermalStatus int status) {
+        return "Bias = "
+                + JobInfo.getBiasString(bias)
+                + " Thermal Status = "
+                + DebugUtils.valueToString(PowerManager.class, "THERMAL_STATUS_", status);
+    }
+
     @Before
     public void setUp() {
         mMockingSession = mockitoSession()
@@ -156,169 +314,302 @@
         assertEquals(THERMAL_STATUS_EMERGENCY, mThermalStatusRestriction.getThermalStatus());
     }
 
+    /**
+     * Test {@link JobSchedulerService#isJobRestricted(JobStatus)} when Thermal is in default state
+     */
     @Test
-    public void testIsJobRestricted() {
+    public void testIsJobRestrictedDefaultStates() {
         mStatusChangedListener.onThermalStatusChanged(THERMAL_STATUS_NONE);
+        JobStatusContainer jc = new JobStatusContainer("testIsJobRestricted", mJobSchedulerService);
 
-        final JobStatus jobMinPriority = createJobStatus("testIsJobRestricted",
-                createJobBuilder(1).setPriority(JobInfo.PRIORITY_MIN).build());
-        final JobStatus jobLowPriority = createJobStatus("testIsJobRestricted",
-                createJobBuilder(2).setPriority(JobInfo.PRIORITY_LOW).build());
-        final JobStatus jobLowPriorityRunning = createJobStatus("testIsJobRestricted",
-                createJobBuilder(3).setPriority(JobInfo.PRIORITY_LOW).build());
-        final JobStatus jobLowPriorityRunningLong = createJobStatus("testIsJobRestricted",
-                createJobBuilder(9).setPriority(JobInfo.PRIORITY_LOW).build());
-        final JobStatus jobDefaultPriority = createJobStatus("testIsJobRestricted",
-                createJobBuilder(4).setPriority(JobInfo.PRIORITY_DEFAULT).build());
-        final JobStatus jobHighPriority = createJobStatus("testIsJobRestricted",
-                createJobBuilder(5).setPriority(JobInfo.PRIORITY_HIGH).build());
-        final JobStatus jobHighPriorityRunning = createJobStatus("testIsJobRestricted",
-                createJobBuilder(6).setPriority(JobInfo.PRIORITY_HIGH).build());
-        final JobStatus jobHighPriorityRunningLong = createJobStatus("testIsJobRestricted",
-                createJobBuilder(10).setPriority(JobInfo.PRIORITY_HIGH).build());
-        final JobStatus ejDowngraded = createJobStatus("testIsJobRestricted",
-                createJobBuilder(7).setExpedited(true).build());
-        final JobStatus ej = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(8).setExpedited(true).build()));
-        final JobStatus ejRetried = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(11).setExpedited(true).build()));
-        final JobStatus ejRunning = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(12).setExpedited(true).build()));
-        final JobStatus ejRunningLong = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(13).setExpedited(true).build()));
-        final JobStatus ui = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(14).build()));
-        final JobStatus uiRetried = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(15).build()));
-        final JobStatus uiRunning = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(16).build()));
-        final JobStatus uiRunningLong = spy(createJobStatus("testIsJobRestricted",
-                createJobBuilder(17).build()));
-        when(ej.shouldTreatAsExpeditedJob()).thenReturn(true);
-        when(ejRetried.shouldTreatAsExpeditedJob()).thenReturn(true);
-        when(ejRunning.shouldTreatAsExpeditedJob()).thenReturn(true);
-        when(ejRunningLong.shouldTreatAsExpeditedJob()).thenReturn(true);
-        when(ui.shouldTreatAsUserInitiatedJob()).thenReturn(true);
-        when(uiRetried.shouldTreatAsUserInitiatedJob()).thenReturn(true);
-        when(uiRunning.shouldTreatAsUserInitiatedJob()).thenReturn(true);
-        when(uiRunningLong.shouldTreatAsUserInitiatedJob()).thenReturn(true);
-        when(ejRetried.getNumPreviousAttempts()).thenReturn(1);
-        when(uiRetried.getNumPreviousAttempts()).thenReturn(2);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(jobLowPriorityRunning)).thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(jobHighPriorityRunning))
-                .thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(jobLowPriorityRunningLong))
-                .thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(jobHighPriorityRunningLong))
-                .thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(ejRunning)).thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(ejRunningLong)).thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(uiRunning)).thenReturn(true);
-        when(mJobSchedulerService.isCurrentlyRunningLocked(uiRunningLong)).thenReturn(true);
-        when(mJobSchedulerService.isJobInOvertimeLocked(jobLowPriorityRunningLong))
-                .thenReturn(true);
-        when(mJobSchedulerService.isJobInOvertimeLocked(jobHighPriorityRunningLong))
-                .thenReturn(true);
-        when(mJobSchedulerService.isJobInOvertimeLocked(ejRunningLong)).thenReturn(true);
-        when(mJobSchedulerService.isJobInOvertimeLocked(uiRunningLong)).thenReturn(true);
+        for (int jobBias : jc.allJobBiases) {
+            assertFalse(isJobRestricted(jc.jobMinPriority, jobBias));
+            assertFalse(isJobRestricted(jc.jobLowPriority, jobBias));
+            assertFalse(isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+            assertFalse(isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+            assertFalse(isJobRestricted(jc.jobDefaultPriority, jobBias));
+            assertFalse(isJobRestricted(jc.jobHighPriority, jobBias));
+            assertFalse(isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+            assertFalse(isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+            assertFalse(isJobRestricted(jc.importantWhileForeground, jobBias));
+            assertFalse(isJobRestricted(jc.importantWhileForegroundRunning, jobBias));
+            assertFalse(isJobRestricted(jc.importantWhileForegroundRunningLong, jobBias));
+            assertFalse(isJobRestricted(jc.ej, jobBias));
+            assertFalse(isJobRestricted(jc.ejDowngraded, jobBias));
+            assertFalse(isJobRestricted(jc.ejRetried, jobBias));
+            assertFalse(isJobRestricted(jc.ejRunning, jobBias));
+            assertFalse(isJobRestricted(jc.ejRunningLong, jobBias));
+            assertFalse(isJobRestricted(jc.ui, jobBias));
+            assertFalse(isJobRestricted(jc.uiRetried, jobBias));
+            assertFalse(isJobRestricted(jc.uiRunning, jobBias));
+            assertFalse(isJobRestricted(jc.uiRunningLong, jobBias));
+        }
+    }
 
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobMinPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobLowPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobDefaultPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ej));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejDowngraded));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ui));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunningLong));
+    /**
+     * Test {@link JobSchedulerService#isJobRestricted(JobStatus)} when Job Bias is Top App and all
+     * Thermal states.
+     */
+    @Test
+    public void testIsJobRestrictedBiasTopApp() {
+        JobStatusContainer jc =
+                new JobStatusContainer("testIsJobRestrictedBiasTopApp", mJobSchedulerService);
 
-        mStatusChangedListener.onThermalStatusChanged(THERMAL_STATUS_LIGHT);
+        int jobBias = JobInfo.BIAS_TOP_APP;
+        for (int thermalStatus : jc.thermalStatuses) {
+            String msg = "Thermal Status = " + DebugUtils.valueToString(
+                    PowerManager.class, "THERMAL_STATUS_", thermalStatus);
+            mStatusChangedListener.onThermalStatusChanged(thermalStatus);
 
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobMinPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobDefaultPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejDowngraded));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ej));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ui));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunningLong));
+            // No restrictions on any jobs
+            assertFalse(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.importantWhileForeground, jobBias));
+            assertFalse(msg, isJobRestricted(jc.importantWhileForegroundRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.importantWhileForegroundRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRetried, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+        }
+    }
 
-        mStatusChangedListener.onThermalStatusChanged(THERMAL_STATUS_MODERATE);
+    /**
+     * Test {@link JobSchedulerService#isJobRestricted(JobStatus)} when Job Bias is Foreground
+     * Service and all Thermal states.
+     */
+    @Test
+    @RequiresFlagsDisabled(FLAG_THERMAL_RESTRICTIONS_TO_FGS_JOBS)
+    public void testIsJobRestrictedBiasFgs_flagThermalRestrictionsToFgsJobsDisabled() {
+        JobStatusContainer jc =
+                new JobStatusContainer("testIsJobRestrictedBiasFgs", mJobSchedulerService);
 
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobMinPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobDefaultPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriority));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejDowngraded));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ej));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ejRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRunningLong));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(ui));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRetried));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunning));
-        assertFalse(mThermalStatusRestriction.isJobRestricted(uiRunningLong));
+        int jobBias = JobInfo.BIAS_FOREGROUND_SERVICE;
+        for (int thermalStatus : jc.thermalStatuses) {
+            String msg = "Thermal Status = " + DebugUtils.valueToString(
+                    PowerManager.class, "THERMAL_STATUS_", thermalStatus);
+            mStatusChangedListener.onThermalStatusChanged(thermalStatus);
+            // No restrictions on any jobs
+            assertFalse(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRetried, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+            assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+            assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+        }
+    }
 
-        mStatusChangedListener.onThermalStatusChanged(THERMAL_STATUS_SEVERE);
+    /**
+     * Test {@link JobSchedulerService#isJobRestricted(JobStatus)} when Job Bias is Foreground
+     * Service and all Thermal states.
+     */
+    @Test
+    @RequiresFlagsEnabled(FLAG_THERMAL_RESTRICTIONS_TO_FGS_JOBS)
+    public void testIsJobRestrictedBiasFgs_flagThermalRestrictionsToFgsJobsEnabled() {
+        JobStatusContainer jc =
+                new JobStatusContainer("testIsJobRestrictedBiasFgs", mJobSchedulerService);
+        int jobBias = JobInfo.BIAS_FOREGROUND_SERVICE;
+        for (int thermalStatus : jc.thermalStatuses) {
+            String msg = debugTag(jobBias, thermalStatus);
+            mStatusChangedListener.onThermalStatusChanged(thermalStatus);
+            if (thermalStatus >= THERMAL_STATUS_SEVERE) {
+                // Full restrictions on all jobs
+                assertTrue(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ej, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejRetried, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ui, jobBias));
+                assertTrue(msg, isJobRestricted(jc.uiRetried, jobBias));
+                assertTrue(msg, isJobRestricted(jc.uiRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+            } else if (thermalStatus >= THERMAL_STATUS_MODERATE) {
+                // No restrictions on user related jobs
+                assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                // Some restrictions on expedited jobs
+                assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejRetried, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                // Some restrictions on high priority jobs
+                assertTrue(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                // Some restructions on important while foreground jobs
+                assertFalse(isJobRestricted(jc.importantWhileForeground, jobBias));
+                assertFalse(isJobRestricted(jc.importantWhileForegroundRunning, jobBias));
+                assertTrue(isJobRestricted(jc.importantWhileForegroundRunningLong, jobBias));
+                // Full restriction on default priority jobs
+                assertTrue(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                // Full restriction on low priority jobs
+                assertTrue(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                // Full restriction on min priority jobs
+                assertTrue(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+            } else {
+                // thermalStatus < THERMAL_STATUS_MODERATE
+                // No restrictions on any job type
+                assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ejRetried, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                assertFalse(msg, isJobRestricted(jc.importantWhileForeground, jobBias));
+                assertFalse(msg, isJobRestricted(jc.importantWhileForegroundRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.importantWhileForegroundRunningLong, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                assertFalse(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+            }
+        }
+    }
 
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobMinPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobDefaultPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejDowngraded));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ej));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRetried));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ui));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRetried));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRunningLong));
+    /**
+     * Test {@link JobSchedulerService#isJobRestricted(JobStatus)} when Job Bias is less than
+     * Foreground Service and all Thermal states.
+     */
+    @Test
+    public void testIsJobRestrictedBiasLessThanFgs() {
+        JobStatusContainer jc =
+                new JobStatusContainer("testIsJobRestrictedBiasLessThanFgs", mJobSchedulerService);
 
-        mStatusChangedListener.onThermalStatusChanged(THERMAL_STATUS_CRITICAL);
-
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobMinPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobLowPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobDefaultPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriority));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(jobHighPriorityRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejDowngraded));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ej));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRetried));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ejRunningLong));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(ui));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRetried));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRunning));
-        assertTrue(mThermalStatusRestriction.isJobRestricted(uiRunningLong));
+        for (int jobBias : jc.biasesBelowFgs) {
+            for (int thermalStatus : jc.thermalStatuses) {
+                String msg = debugTag(jobBias, thermalStatus);
+                mStatusChangedListener.onThermalStatusChanged(thermalStatus);
+                if (thermalStatus >= THERMAL_STATUS_SEVERE) {
+                    // Full restrictions on all jobs
+                    assertTrue(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ej, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejRetried, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ui, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.uiRetried, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.uiRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                } else if (thermalStatus >= THERMAL_STATUS_MODERATE) {
+                    // No restrictions on user related jobs
+                    assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                    // Some restrictions on expedited jobs
+                    assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                    // Some restrictions on high priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                    // Full restriction on default priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                    // Full restriction on low priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                    // Full restriction on min priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+                } else if (thermalStatus >= THERMAL_STATUS_LIGHT) {
+                    // No restrictions on any user related jobs
+                    assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                    // No restrictions on any expedited jobs
+                    assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                    // No restrictions on any high priority jobs
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                    // No restrictions on default priority jobs
+                    assertFalse(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                    // Some restrictions on low priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                    assertTrue(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                    // Full restriction on min priority jobs
+                    assertTrue(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+                } else { // THERMAL_STATUS_NONE
+                    // No restrictions on any jobs
+                    assertFalse(msg, isJobRestricted(jc.jobMinPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobLowPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobLowPriorityRunningLong, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobDefaultPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriority, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.jobHighPriorityRunningLong, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ej, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejDowngraded, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ejRunningLong, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.ui, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRetried, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunning, jobBias));
+                    assertFalse(msg, isJobRestricted(jc.uiRunningLong, jobBias));
+                }
+            }
+        }
     }
 
     private JobInfo.Builder createJobBuilder(int jobId) {
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
index fd880dd..178e7ec 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/InputControllerTest.java
@@ -33,7 +33,6 @@
 import android.os.Handler;
 import android.os.IBinder;
 import android.platform.test.annotations.Presubmit;
-import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
 import android.view.DisplayInfo;
@@ -41,13 +40,11 @@
 
 import androidx.test.InstrumentationRegistry;
 
-import com.android.input.flags.Flags;
 import com.android.server.LocalServices;
 import com.android.server.input.InputManagerInternal;
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
@@ -60,9 +57,6 @@
     private static final String LANGUAGE_TAG = "en-US";
     private static final String LAYOUT_TYPE = "qwerty";
 
-    @Rule
-    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
-
     @Mock
     private InputManagerInternal mInputManagerInternalMock;
     @Mock
@@ -77,8 +71,6 @@
 
     @Before
     public void setUp() throws Exception {
-        mSetFlagsRule.enableFlags(Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER);
-
         MockitoAnnotations.initMocks(this);
         mInputManagerMockHelper = new InputManagerMockHelper(
                 TestableLooper.get(this), mNativeWrapperMock, mIInputManagerMock);
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
index 2b81d78..da8961d 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceManagerServiceTest.java
@@ -339,8 +339,6 @@
         LocalServices.removeServiceForTest(DisplayManagerInternal.class);
         LocalServices.addService(DisplayManagerInternal.class, mDisplayManagerInternalMock);
 
-        mSetFlagsRule.enableFlags(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER);
-
         doNothing().when(mInputManagerInternalMock)
                 .setMousePointerAccelerationEnabled(anyBoolean(), anyInt());
         doNothing().when(mInputManagerInternalMock).setPointerIconVisible(anyBoolean(), anyInt());
diff --git a/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt b/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
index 0fc9d6f..3b9ee80 100644
--- a/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
+++ b/tests/Input/src/com/android/server/input/InputManagerServiceTests.kt
@@ -28,14 +28,11 @@
 import android.os.SystemClock
 import android.os.test.TestLooper
 import android.platform.test.annotations.Presubmit
-import android.platform.test.annotations.RequiresFlagsDisabled
 import android.platform.test.flag.junit.DeviceFlagsValueProvider
 import android.provider.Settings
 import android.view.View.OnKeyListener
-import android.view.Display
 import android.view.InputDevice
 import android.view.KeyEvent
-import android.view.PointerIcon
 import android.view.SurfaceHolder
 import android.view.SurfaceView
 import android.test.mock.MockContentResolver
@@ -44,7 +41,6 @@
 import com.google.common.truth.Truth.assertThat
 import com.android.compatibility.common.util.SystemUtil.runWithShellPermissionIdentity
 import org.junit.After
-import org.junit.Assert.assertFalse
 import org.junit.Assert.assertTrue
 import org.junit.Before
 import org.junit.Rule
@@ -53,22 +49,16 @@
 import org.mockito.ArgumentMatchers.anyBoolean
 import org.mockito.ArgumentMatchers.anyFloat
 import org.mockito.ArgumentMatchers.anyInt
-import org.mockito.ArgumentMatchers.eq
 import org.mockito.Mock
 import org.mockito.Mockito.`when`
-import org.mockito.Mockito.clearInvocations
-import org.mockito.Mockito.doAnswer
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.never
 import org.mockito.Mockito.spy
 import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
-import org.mockito.Mockito.verifyNoMoreInteractions
 import org.mockito.Mockito.verifyZeroInteractions
 import org.mockito.junit.MockitoJUnit
 import org.mockito.stubbing.OngoingStubbing
-import java.util.concurrent.CountDownLatch
-import java.util.concurrent.TimeUnit
 
 /**
  * Tests for {@link InputManagerService}.
@@ -187,197 +177,6 @@
         verify(wmCallbacks).notifyPointerDisplayIdChanged(displayId, x, y)
     }
 
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun testSetVirtualMousePointerDisplayId() {
-        // Set the virtual mouse pointer displayId, and ensure that the calling thread is blocked
-        // until the native callback happens.
-        var countDownLatch = CountDownLatch(1)
-        val overrideDisplayId = 123
-        Thread {
-            assertTrue("Setting virtual pointer display should succeed",
-                localService.setVirtualMousePointerDisplayId(overrideDisplayId))
-            countDownLatch.countDown()
-        }.start()
-        assertFalse("Setting virtual pointer display should block",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-
-        val x = 42f
-        val y = 314f
-        service.onPointerDisplayIdChanged(overrideDisplayId, x, y)
-        testLooper.dispatchNext()
-        verify(wmCallbacks).notifyPointerDisplayIdChanged(overrideDisplayId, x, y)
-        assertTrue("Native callback unblocks calling thread",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-        verify(native).setPointerDisplayId(overrideDisplayId)
-
-        // Ensure that setting the same override again succeeds immediately.
-        assertTrue("Setting the same virtual mouse pointer displayId again should succeed",
-            localService.setVirtualMousePointerDisplayId(overrideDisplayId))
-
-        // Ensure that we did not query WM for the pointerDisplayId when setting the override
-        verify(wmCallbacks, never()).pointerDisplayId
-
-        // Unset the virtual mouse pointer displayId, and ensure that we query WM for the new
-        // pointer displayId and the calling thread is blocked until the native callback happens.
-        countDownLatch = CountDownLatch(1)
-        val pointerDisplayId = 42
-        `when`(wmCallbacks.pointerDisplayId).thenReturn(pointerDisplayId)
-        Thread {
-            assertTrue("Unsetting virtual mouse pointer displayId should succeed",
-                localService.setVirtualMousePointerDisplayId(Display.INVALID_DISPLAY))
-            countDownLatch.countDown()
-        }.start()
-        assertFalse("Unsetting virtual mouse pointer displayId should block",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-
-        service.onPointerDisplayIdChanged(pointerDisplayId, x, y)
-        testLooper.dispatchNext()
-        verify(wmCallbacks).notifyPointerDisplayIdChanged(pointerDisplayId, x, y)
-        assertTrue("Native callback unblocks calling thread",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-        verify(native).setPointerDisplayId(pointerDisplayId)
-    }
-
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun testSetVirtualMousePointerDisplayId_unsuccessfulUpdate() {
-        // Set the virtual mouse pointer displayId, and ensure that the calling thread is blocked
-        // until the native callback happens.
-        val countDownLatch = CountDownLatch(1)
-        val overrideDisplayId = 123
-        Thread {
-            assertFalse("Setting virtual pointer display should be unsuccessful",
-                localService.setVirtualMousePointerDisplayId(overrideDisplayId))
-            countDownLatch.countDown()
-        }.start()
-        assertFalse("Setting virtual pointer display should block",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-
-        val x = 42f
-        val y = 314f
-        // Assume the native callback updates the pointerDisplayId to the incorrect value.
-        service.onPointerDisplayIdChanged(Display.INVALID_DISPLAY, x, y)
-        testLooper.dispatchNext()
-        verify(wmCallbacks).notifyPointerDisplayIdChanged(Display.INVALID_DISPLAY, x, y)
-        assertTrue("Native callback unblocks calling thread",
-            countDownLatch.await(100, TimeUnit.MILLISECONDS))
-        verify(native).setPointerDisplayId(overrideDisplayId)
-    }
-
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun testSetVirtualMousePointerDisplayId_competingRequests() {
-        val firstRequestSyncLatch = CountDownLatch(1)
-        doAnswer {
-            firstRequestSyncLatch.countDown()
-        }.`when`(native).setPointerDisplayId(anyInt())
-
-        val firstRequestLatch = CountDownLatch(1)
-        val firstOverride = 123
-        Thread {
-            assertFalse("Setting virtual pointer display from thread 1 should be unsuccessful",
-                localService.setVirtualMousePointerDisplayId(firstOverride))
-            firstRequestLatch.countDown()
-        }.start()
-        assertFalse("Setting virtual pointer display should block",
-            firstRequestLatch.await(100, TimeUnit.MILLISECONDS))
-
-        assertTrue("Wait for first thread's request should succeed",
-            firstRequestSyncLatch.await(100, TimeUnit.MILLISECONDS))
-
-        val secondRequestLatch = CountDownLatch(1)
-        val secondOverride = 42
-        Thread {
-            assertTrue("Setting virtual mouse pointer from thread 2 should be successful",
-                localService.setVirtualMousePointerDisplayId(secondOverride))
-            secondRequestLatch.countDown()
-        }.start()
-        assertFalse("Setting virtual mouse pointer should block",
-            secondRequestLatch.await(100, TimeUnit.MILLISECONDS))
-
-        val x = 42f
-        val y = 314f
-        // Assume the native callback updates directly to the second request.
-        service.onPointerDisplayIdChanged(secondOverride, x, y)
-        testLooper.dispatchNext()
-        verify(wmCallbacks).notifyPointerDisplayIdChanged(secondOverride, x, y)
-        assertTrue("Native callback unblocks first thread",
-            firstRequestLatch.await(100, TimeUnit.MILLISECONDS))
-        assertTrue("Native callback unblocks second thread",
-            secondRequestLatch.await(100, TimeUnit.MILLISECONDS))
-        verify(native, times(2)).setPointerDisplayId(anyInt())
-    }
-
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun onDisplayRemoved_resetAllAdditionalInputProperties() {
-        setVirtualMousePointerDisplayIdAndVerify(10)
-
-        localService.setPointerIconVisible(false, 10)
-        verify(native).setPointerIconVisibility(10, false)
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NULL))
-        localService.setMousePointerAccelerationEnabled(false, 10)
-        verify(native).setMousePointerAccelerationEnabled(10, false)
-
-        service.onDisplayRemoved(10)
-        verify(native).setPointerIconVisibility(10, true)
-        verify(native).displayRemoved(eq(10))
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NOT_SPECIFIED))
-        verify(native).setMousePointerAccelerationEnabled(10, true)
-        verifyNoMoreInteractions(native)
-
-        // This call should not block because the virtual mouse pointer override was never removed.
-        localService.setVirtualMousePointerDisplayId(10)
-
-        verify(native).setPointerDisplayId(eq(10))
-        verifyNoMoreInteractions(native)
-    }
-
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun updateAdditionalInputPropertiesForOverrideDisplay() {
-        setVirtualMousePointerDisplayIdAndVerify(10)
-
-        localService.setPointerIconVisible(false, 10)
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NULL))
-        verify(native).setPointerIconVisibility(10, false)
-        localService.setMousePointerAccelerationEnabled(false, 10)
-        verify(native).setMousePointerAccelerationEnabled(10, false)
-
-        localService.setPointerIconVisible(true, 10)
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NOT_SPECIFIED))
-        verify(native).setPointerIconVisibility(10, true)
-        localService.setMousePointerAccelerationEnabled(true, 10)
-        verify(native).setMousePointerAccelerationEnabled(10, true)
-
-        localService.setPointerIconVisible(false, 20)
-        verify(native).setPointerIconVisibility(20, false)
-        localService.setMousePointerAccelerationEnabled(false, 20)
-        verify(native).setMousePointerAccelerationEnabled(20, false)
-        verifyNoMoreInteractions(native)
-
-        clearInvocations(native)
-        setVirtualMousePointerDisplayIdAndVerify(20)
-
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NULL))
-    }
-
-    @RequiresFlagsDisabled(com.android.input.flags.Flags.FLAG_ENABLE_POINTER_CHOREOGRAPHER)
-    @Test
-    fun setAdditionalInputPropertiesBeforeOverride() {
-        localService.setPointerIconVisible(false, 10)
-        localService.setMousePointerAccelerationEnabled(false, 10)
-
-        verify(native).setPointerIconVisibility(10, false)
-        verify(native).setMousePointerAccelerationEnabled(10, false)
-        verifyNoMoreInteractions(native)
-
-        setVirtualMousePointerDisplayIdAndVerify(10)
-
-        verify(native).setPointerIconType(eq(PointerIcon.TYPE_NULL))
-    }
-
     @Test
     fun setDeviceTypeAssociation_setsDeviceTypeAssociation() {
         val inputPort = "inputPort"
@@ -412,20 +211,6 @@
         verify(native, times(2)).changeKeyboardLayoutAssociation()
     }
 
-    private fun setVirtualMousePointerDisplayIdAndVerify(overrideDisplayId: Int) {
-        val thread = Thread { localService.setVirtualMousePointerDisplayId(overrideDisplayId) }
-        thread.start()
-
-        // Allow some time for the set override call to park while waiting for the native callback.
-        Thread.sleep(100 /*millis*/)
-        verify(native).setPointerDisplayId(overrideDisplayId)
-
-        service.onPointerDisplayIdChanged(overrideDisplayId, 0f, 0f)
-        testLooper.dispatchNext()
-        verify(wmCallbacks).notifyPointerDisplayIdChanged(overrideDisplayId, 0f, 0f)
-        thread.join(100 /*millis*/)
-    }
-
     private fun createVirtualDisplays(count: Int): List<VirtualDisplay> {
         val displayManager: DisplayManager = context.getSystemService(
                 DisplayManager::class.java
diff --git a/tools/hoststubgen/TEST_MAPPING b/tools/hoststubgen/TEST_MAPPING
index f6885e1..856e6ee 100644
--- a/tools/hoststubgen/TEST_MAPPING
+++ b/tools/hoststubgen/TEST_MAPPING
@@ -1,63 +1,7 @@
-// Keep the following two TEST_MAPPINGs in sync:
-// frameworks/base/ravenwood/TEST_MAPPING
-// frameworks/base/tools/hoststubgen/TEST_MAPPING
 {
-  "presubmit": [
-    { "name": "tiny-framework-dump-test" },
-    { "name": "hoststubgentest" },
-    { "name": "hoststubgen-invoke-test" },
+  "imports": [
     {
-      "name": "RavenwoodMockitoTest_device"
-    },
-    {
-      "name": "RavenwoodBivalentTest_device"
-    },
-    // The sysui tests should match vendor/unbundled_google/packages/SystemUIGoogle/TEST_MAPPING
-    {
-      "name": "SystemUIGoogleTests",
-      "options": [
-        {
-          "exclude-annotation": "org.junit.Ignore"
-        },
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        }
-      ]
-    }
-  ],
-  "presubmit-large": [
-    {
-      "name": "SystemUITests",
-      "options": [
-        {
-          "exclude-annotation": "androidx.test.filters.FlakyTest"
-        },
-        {
-          "exclude-annotation": "org.junit.Ignore"
-        }
-      ]
-    }
-  ],
-  "ravenwood-presubmit": [
-    {
-      "name": "RavenwoodMinimumTest",
-      "host": true
-    },
-    {
-      "name": "RavenwoodMockitoTest",
-      "host": true
-    },
-    {
-      "name": "CtsUtilTestCasesRavenwood",
-      "host": true
-    },
-    {
-      "name": "RavenwoodCoreTest",
-      "host": true
-    },
-    {
-      "name": "RavenwoodBivalentTest",
-      "host": true
+      "path": "frameworks/base/ravenwood"
     }
   ]
 }