Merge "Fix error when backing up channels" into main
diff --git a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
index c1894f0..a37779e 100644
--- a/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
+++ b/apex/jobscheduler/service/java/com/android/server/DeviceIdleController.java
@@ -3568,7 +3568,7 @@
             Slog.i(TAG, "becomeActiveLocked, reason=" + activeReason
                     + ", changeLightIdle=" + changeLightIdle);
         }
-        if (mState != STATE_ACTIVE || mLightState != STATE_ACTIVE) {
+        if (mState != STATE_ACTIVE || mLightState != LIGHT_STATE_ACTIVE) {
             moveToStateLocked(STATE_ACTIVE, activeReason);
             mInactiveTimeout = newInactiveTimeout;
             resetIdleManagementLocked();
diff --git a/core/java/android/hardware/display/DisplayManager.java b/core/java/android/hardware/display/DisplayManager.java
index 97f6899..b0ea92d 100644
--- a/core/java/android/hardware/display/DisplayManager.java
+++ b/core/java/android/hardware/display/DisplayManager.java
@@ -18,6 +18,7 @@
 
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.Display.HdrCapabilities.HdrType;
+import static android.view.Display.INVALID_DISPLAY;
 
 import android.Manifest;
 import android.annotation.FlaggedApi;
@@ -47,6 +48,7 @@
 import android.os.Process;
 import android.os.RemoteException;
 import android.os.ServiceManager;
+import android.os.UserManager;
 import android.util.Log;
 import android.util.Pair;
 import android.util.Slog;
@@ -96,6 +98,8 @@
     @GuardedBy("mLock")
     private final WeakDisplayCache mDisplayCache = new WeakDisplayCache();
 
+    private int mDisplayIdToMirror = INVALID_DISPLAY;
+
     /**
      * Broadcast receiver that indicates when the Wifi display status changes.
      * <p>
@@ -1086,6 +1090,7 @@
         if (surface != null) {
             builder.setSurface(surface);
         }
+        builder.setDisplayIdToMirror(getDisplayIdToMirror());
         return createVirtualDisplay(builder.build(), handler, callback);
     }
 
@@ -1163,6 +1168,7 @@
         if (surface != null) {
             builder.setSurface(surface);
         }
+        builder.setDisplayIdToMirror(getDisplayIdToMirror());
         return createVirtualDisplay(projection, builder.build(), callback, handler);
     }
 
@@ -1708,6 +1714,16 @@
         return mGlobal.getDefaultDozeBrightness(displayId);
     }
 
+    private int getDisplayIdToMirror() {
+        if (mDisplayIdToMirror == INVALID_DISPLAY) {
+            final UserManager userManager = mContext.getSystemService(UserManager.class);
+            mDisplayIdToMirror = userManager.isVisibleBackgroundUsersSupported()
+                    ? userManager.getMainDisplayIdAssignedToUser()
+                    : DEFAULT_DISPLAY;
+        }
+        return mDisplayIdToMirror;
+    }
+
     /**
      * Listens for changes in available display devices.
      */
diff --git a/core/java/android/os/StatsBootstrapAtomValue.aidl b/core/java/android/os/StatsBootstrapAtomValue.aidl
index b59bc06..b31eb6f 100644
--- a/core/java/android/os/StatsBootstrapAtomValue.aidl
+++ b/core/java/android/os/StatsBootstrapAtomValue.aidl
@@ -19,12 +19,36 @@
  *
  * @hide
  */
-union StatsBootstrapAtomValue {
-    boolean boolValue;
-    int intValue;
-    long longValue;
-    float floatValue;
-    String stringValue;
-    byte[] bytesValue;
-    String[] stringArrayValue;
-}
\ No newline at end of file
+parcelable StatsBootstrapAtomValue {
+    union Primitive {
+        boolean boolValue;
+        int intValue;
+        long longValue;
+        float floatValue;
+        String stringValue;
+        byte[] bytesValue;
+	String[] stringArrayValue;
+    }
+
+    Primitive value;
+
+    parcelable Annotation {
+        // Match the definitions in
+        // packages/modules/StatsD/framework/java/android/util/StatsLog.java
+        // Only supports UIDs for now.
+        @Backing(type="byte")
+        enum Id {
+            NONE,
+            IS_UID,
+        }
+        Id id;
+
+        union Primitive {
+            boolean boolValue;
+            int intValue;
+        }
+        Primitive value;
+    }
+
+    Annotation[] annotations;
+}
diff --git a/core/jni/android_view_InputDevice.cpp b/core/jni/android_view_InputDevice.cpp
index f5992d9..ce40e51 100644
--- a/core/jni/android_view_InputDevice.cpp
+++ b/core/jni/android_view_InputDevice.cpp
@@ -60,7 +60,7 @@
                                                                   ? layoutInfo->layoutType.c_str()
                                                                   : NULL));
 
-    std::shared_ptr<KeyCharacterMap> map = deviceInfo.getKeyCharacterMap();
+    const KeyCharacterMap* map = deviceInfo.getKeyCharacterMap();
     std::unique_ptr<KeyCharacterMap> mapCopy;
     if (map != nullptr) {
         mapCopy = std::make_unique<KeyCharacterMap>(*map);
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
index 8201bbe..e7d2ef1 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorTest.kt
@@ -30,8 +30,11 @@
 import com.android.systemui.education.data.repository.contextualEducationRepository
 import com.android.systemui.education.data.repository.fakeEduClock
 import com.android.systemui.education.shared.model.EducationUiType
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
+import com.android.systemui.inputdevice.tutorial.tutorialSchedulerRepository
 import com.android.systemui.keyboard.data.repository.keyboardRepository
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.recents.OverviewProxyService.OverviewProxyListener
 import com.android.systemui.testKosmos
 import com.android.systemui.touchpad.data.repository.touchpadRepository
 import com.android.systemui.user.data.repository.fakeUserRepository
@@ -42,10 +45,13 @@
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runCurrent
 import kotlinx.coroutines.test.runTest
+import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.verify
 import platform.test.runner.parameterized.ParameterizedAndroidJunit4
 import platform.test.runner.parameterized.Parameters
 
@@ -56,14 +62,19 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val contextualEduInteractor = kosmos.contextualEducationInteractor
+    private val repository = kosmos.contextualEducationRepository
     private val touchpadRepository = kosmos.touchpadRepository
     private val keyboardRepository = kosmos.keyboardRepository
+    private val tutorialSchedulerRepository = kosmos.tutorialSchedulerRepository
     private val userRepository = kosmos.fakeUserRepository
+    private val overviewProxyService = kosmos.mockOverviewProxyService
 
     private val underTest: KeyboardTouchpadEduInteractor = kosmos.keyboardTouchpadEduInteractor
     private val eduClock = kosmos.fakeEduClock
     private val minDurationForNextEdu =
         KeyboardTouchpadEduInteractor.minIntervalBetweenEdu + 1.seconds
+    private val initialDelayElapsedDuration =
+        KeyboardTouchpadEduInteractor.initialDelayDuration + 1.seconds
 
     @Before
     fun setup() {
@@ -271,6 +282,131 @@
             assertThat(model?.lastShortcutTriggeredTime).isEqualTo(eduClock.instant())
         }
 
+    @Test
+    fun dataUpdatedOnIncrementSignalCountWhenTouchpadConnected() =
+        testScope.runTest {
+            assumeTrue(gestureType != ALL_APPS)
+            setUpForInitialDelayElapse()
+            touchpadRepository.setIsAnyTouchpadConnected(true)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenTouchpadDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            touchpadRepository.setIsAnyTouchpadConnected(false)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountWhenKeyboardConnected() =
+        testScope.runTest {
+            assumeTrue(gestureType == ALL_APPS)
+            setUpForInitialDelayElapse()
+            keyboardRepository.setIsAnyKeyboardConnected(true)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWhenKeyboardDisconnected() =
+        testScope.runTest {
+            setUpForInitialDelayElapse()
+            keyboardRepository.setIsAnyKeyboardConnected(false)
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataAddedOnUpdateShortcutTriggerTime() =
+        testScope.runTest {
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            assertThat(model?.lastShortcutTriggeredTime).isNull()
+
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ true, gestureType)
+
+            assertThat(model?.lastShortcutTriggeredTime).isEqualTo(kosmos.fakeEduClock.instant())
+        }
+
+    @Test
+    fun dataUpdatedOnIncrementSignalCountAfterInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            eduClock.offset(initialDelayElapsedDuration)
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountBeforeInitialDelay() =
+        testScope.runTest {
+            setUpForDeviceConnection()
+            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            // No offset to the clock to simulate update before initial delay
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    @Test
+    fun dataUnchangedOnIncrementSignalCountWithoutOobeLaunchTime() =
+        testScope.runTest {
+            // No update to OOBE launch time to simulate no OOBE is launched yet
+            setUpForDeviceConnection()
+
+            val model by collectLastValue(repository.readGestureEduModelFlow(gestureType))
+            val originalValue = model!!.signalCount
+            val listener = getOverviewProxyListener()
+            listener.updateContextualEduStats(/* isTrackpadGesture= */ false, gestureType)
+
+            assertThat(model?.signalCount).isEqualTo(originalValue)
+        }
+
+    private suspend fun setUpForInitialDelayElapse() {
+        tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, eduClock.instant())
+        tutorialSchedulerRepository.updateLaunchTime(DeviceType.KEYBOARD, eduClock.instant())
+        eduClock.offset(initialDelayElapsedDuration)
+    }
+
+    @After
+    fun clear() {
+        testScope.launch { tutorialSchedulerRepository.clearDataStore() }
+    }
+
     private suspend fun triggerMaxEducationSignals(gestureType: GestureType) {
         // Increment max number of signal to try triggering education
         for (i in 1..KeyboardTouchpadEduInteractor.MAX_SIGNAL_COUNT) {
@@ -288,9 +424,15 @@
         runCurrent()
     }
 
-    private suspend fun setUpForDeviceConnection() {
-        contextualEduInteractor.updateKeyboardFirstConnectionTime()
-        contextualEduInteractor.updateTouchpadFirstConnectionTime()
+    private fun setUpForDeviceConnection() {
+        touchpadRepository.setIsAnyTouchpadConnected(true)
+        keyboardRepository.setIsAnyKeyboardConnected(true)
+    }
+
+    private fun getOverviewProxyListener(): OverviewProxyListener {
+        val listenerCaptor = argumentCaptor<OverviewProxyListener>()
+        verify(overviewProxyService).addCallback(listenerCaptor.capture())
+        return listenerCaptor.firstValue
     }
 
     companion object {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
deleted file mode 100644
index 98e0947..0000000
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadStatsInteractorTest.kt
+++ /dev/null
@@ -1,172 +0,0 @@
-/*
- * Copyright 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.education.domain.interactor
-
-import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.filters.SmallTest
-import com.android.systemui.SysuiTestCase
-import com.android.systemui.contextualeducation.GestureType.ALL_APPS
-import com.android.systemui.contextualeducation.GestureType.BACK
-import com.android.systemui.coroutines.collectLastValue
-import com.android.systemui.education.data.repository.contextualEducationRepository
-import com.android.systemui.education.data.repository.fakeEduClock
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
-import com.android.systemui.inputdevice.tutorial.tutorialSchedulerRepository
-import com.android.systemui.keyboard.data.repository.keyboardRepository
-import com.android.systemui.kosmos.testScope
-import com.android.systemui.testKosmos
-import com.android.systemui.touchpad.data.repository.touchpadRepository
-import com.google.common.truth.Truth.assertThat
-import kotlin.time.Duration.Companion.seconds
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.test.runTest
-import org.junit.After
-import org.junit.Test
-import org.junit.runner.RunWith
-
-@SmallTest
-@RunWith(AndroidJUnit4::class)
-class KeyboardTouchpadStatsInteractorTest : SysuiTestCase() {
-    private val kosmos = testKosmos()
-    private val testScope = kosmos.testScope
-    private val underTest = kosmos.keyboardTouchpadEduStatsInteractor
-    private val keyboardRepository = kosmos.keyboardRepository
-    private val touchpadRepository = kosmos.touchpadRepository
-    private val repository = kosmos.contextualEducationRepository
-    private val fakeClock = kosmos.fakeEduClock
-    private val tutorialSchedulerRepository = kosmos.tutorialSchedulerRepository
-    private val initialDelayElapsedDuration =
-        KeyboardTouchpadEduStatsInteractorImpl.initialDelayDuration + 1.seconds
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountWhenTouchpadConnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            touchpadRepository.setIsAnyTouchpadConnected(true)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWhenTouchpadDisconnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            touchpadRepository.setIsAnyTouchpadConnected(false)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountWhenKeyboardConnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            keyboardRepository.setIsAnyKeyboardConnected(true)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(ALL_APPS)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWhenKeyboardDisconnected() =
-        testScope.runTest {
-            setUpForInitialDelayElapse()
-            keyboardRepository.setIsAnyKeyboardConnected(false)
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(ALL_APPS))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(ALL_APPS)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataAddedOnUpdateShortcutTriggerTime() =
-        testScope.runTest {
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            assertThat(model?.lastShortcutTriggeredTime).isNull()
-            underTest.updateShortcutTriggerTime(BACK)
-            assertThat(model?.lastShortcutTriggeredTime).isEqualTo(kosmos.fakeEduClock.instant())
-        }
-
-    @Test
-    fun dataUpdatedOnIncrementSignalCountAfterInitialDelay() =
-        testScope.runTest {
-            setUpForDeviceConnection()
-            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-
-            fakeClock.offset(initialDelayElapsedDuration)
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue + 1)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountBeforeInitialDelay() =
-        testScope.runTest {
-            setUpForDeviceConnection()
-            tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-
-            // No offset to the clock to simulate update before initial delay
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    @Test
-    fun dataUnchangedOnIncrementSignalCountWithoutOobeLaunchTime() =
-        testScope.runTest {
-            // No update to OOBE launch time to simulate no OOBE is launched yet
-            setUpForDeviceConnection()
-
-            val model by collectLastValue(repository.readGestureEduModelFlow(BACK))
-            val originalValue = model!!.signalCount
-            underTest.incrementSignalCount(BACK)
-
-            assertThat(model?.signalCount).isEqualTo(originalValue)
-        }
-
-    private suspend fun setUpForInitialDelayElapse() {
-        tutorialSchedulerRepository.updateLaunchTime(DeviceType.TOUCHPAD, fakeClock.instant())
-        tutorialSchedulerRepository.updateLaunchTime(DeviceType.KEYBOARD, fakeClock.instant())
-        fakeClock.offset(initialDelayElapsedDuration)
-    }
-
-    private fun setUpForDeviceConnection() {
-        touchpadRepository.setIsAnyTouchpadConnected(true)
-        keyboardRepository.setIsAnyKeyboardConnected(true)
-    }
-
-    @After
-    fun clear() {
-        testScope.launch { tutorialSchedulerRepository.clearDataStore() }
-    }
-}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
index 1d80826..de3dc57 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractorTest.kt
@@ -28,7 +28,6 @@
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.SysuiTestableContext
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.shared.model.asIcon
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.coroutines.collectValues
@@ -145,13 +144,13 @@
 
             // Tile starts with the generic Modes icon.
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Add an inactive mode -> Still modes icon
             zenModeRepository.addMode(id = "Mode", active = false)
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Add an active mode with a default icon: icon should be the mode icon, and the
@@ -159,7 +158,7 @@
             zenModeRepository.addMode(
                 id = "Bedtime with default icon",
                 type = AutomaticZenRule.TYPE_BEDTIME,
-                active = true,
+                active = true
             )
             runCurrent()
             assertThat(tileData?.icon).isEqualTo(BEDTIME_ICON)
@@ -190,7 +189,7 @@
             // Deactivate remaining mode: back to the default modes icon
             zenModeRepository.deactivateMode("Driving with custom icon")
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
         }
 
@@ -205,18 +204,18 @@
                 )
 
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             // Activate a Mode -> Icon doesn't change.
             zenModeRepository.addMode(id = "Mode", active = true)
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
 
             zenModeRepository.deactivateMode(id = "Mode")
             runCurrent()
-            assertThat(tileData?.icon).isEqualTo(MODES_RESOURCE_ICON)
+            assertThat(tileData?.icon).isEqualTo(MODES_ICON)
             assertThat(tileData?.iconResId).isEqualTo(MODES_DRAWABLE_ID)
         }
 
@@ -264,7 +263,7 @@
         val BEDTIME_DRAWABLE = TestStubDrawable("bedtime")
         val CUSTOM_DRAWABLE = TestStubDrawable("custom")
 
-        val MODES_RESOURCE_ICON = Icon.Resource(MODES_DRAWABLE_ID, null)
+        val MODES_ICON = MODES_DRAWABLE.asIcon()
         val BEDTIME_ICON = BEDTIME_DRAWABLE.asIcon()
         val CUSTOM_ICON = CUSTOM_DRAWABLE.asIcon()
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
index a58cb9c..c3d45db 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapperTest.kt
@@ -22,9 +22,7 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.common.shared.model.asIcon
-import com.android.systemui.qs.tiles.ModesTile
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
 import com.android.systemui.qs.tiles.viewmodel.QSTileConfigTestBuilder
 import com.android.systemui.qs.tiles.viewmodel.QSTileState
@@ -53,11 +51,6 @@
                 .apply {
                     addOverride(R.drawable.qs_dnd_icon_on, TestStubDrawable())
                     addOverride(R.drawable.qs_dnd_icon_off, TestStubDrawable())
-                    addOverride(
-                        ModesTile.ICON_RES_ID,
-                        TestStubDrawable(ModesTile.ICON_RES_ID.toString()),
-                    )
-                    addOverride(123, TestStubDrawable("123"))
                 }
                 .resources,
             context.theme,
@@ -66,7 +59,12 @@
     @Test
     fun inactiveState() {
         val icon = TestStubDrawable("res123").asIcon()
-        val model = ModesTileModel(isActivated = false, activeModes = emptyList(), icon = icon)
+        val model =
+            ModesTileModel(
+                isActivated = false,
+                activeModes = emptyList(),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
@@ -78,7 +76,12 @@
     @Test
     fun activeState_oneMode() {
         val icon = TestStubDrawable("res123").asIcon()
-        val model = ModesTileModel(isActivated = true, activeModes = listOf("DND"), icon = icon)
+        val model =
+            ModesTileModel(
+                isActivated = true,
+                activeModes = listOf("DND"),
+                icon = icon,
+            )
 
         val state = underTest.map(config, model)
 
@@ -105,36 +108,19 @@
     }
 
     @Test
-    fun resourceIconModel_whenResIdsIdentical_mapsToLoadedIconWithInputResId() {
-        val icon = Icon.Resource(123, null)
+    fun state_modelHasIconResId_includesIconResId() {
+        val icon = TestStubDrawable("res123").asIcon()
         val model =
             ModesTileModel(
                 isActivated = false,
                 activeModes = emptyList(),
                 icon = icon,
-                iconResId = 123,
+                iconResId = 123
             )
 
         val state = underTest.map(config, model)
 
-        assertThat(state.icon()).isEqualTo(TestStubDrawable("123").asIcon())
-        assertThat(state.iconRes).isEqualTo(123)
-    }
-
-    @Test
-    fun resourceIconModel_whenResIdsNonIdentical_mapsToLoadedIconWithIconResourceId() {
-        val icon = Icon.Resource(123, null)
-        val model =
-            ModesTileModel(
-                isActivated = false,
-                activeModes = emptyList(),
-                icon = icon,
-                iconResId = 321, // Note: NOT 123. This will be ignored.
-            )
-
-        val state = underTest.map(config, model)
-
-        assertThat(state.icon()).isEqualTo(TestStubDrawable("123").asIcon())
+        assertThat(state.icon()).isEqualTo(icon)
         assertThat(state.iconRes).isEqualTo(123)
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
index 254f1e1..4d71dc4 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/data/model/DisplayContentScenarios.kt
@@ -21,8 +21,8 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.PIP
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_BOTTOM
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.HORIZONTAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.emptyRootSplit
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.freeForm
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks.fullScreen
@@ -39,16 +39,14 @@
 
     data class TaskSpec(val taskId: Int, val userId: Int, val name: String)
 
+    val emptyDisplayContent = DisplayContentModel(0, SystemUiState(shadeExpanded = false), listOf())
+
     /** Home screen, with only the launcher visible */
     fun launcherOnly(shadeExpanded: Boolean = false) =
         DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks =
-                listOf(
-                    launcher(visible = true),
-                    emptyRootSplit,
-                )
+            rootTasks = listOf(launcher(visible = true), emptyRootSplit),
         )
 
     /** A Full screen activity for the personal (primary) user, with launcher behind it */
@@ -57,48 +55,72 @@
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
-                listOf(
-                    fullScreen(spec, visible = true),
-                    launcher(visible = false),
-                    emptyRootSplit,
-                )
+                listOf(fullScreen(spec, visible = true), launcher(visible = false), emptyRootSplit),
         )
 
+    enum class Orientation {
+        HORIZONTAL,
+        VERTICAL,
+    }
+
+    internal fun Rect.splitLeft(margin: Int = 0) = Rect(left, top, centerX() - margin, bottom)
+
+    internal fun Rect.splitRight(margin: Int = 0) = Rect(centerX() + margin, top, right, bottom)
+
+    internal fun Rect.splitTop(margin: Int = 0) = Rect(left, top, right, centerY() - margin)
+
+    internal fun Rect.splitBottom(margin: Int = 0) = Rect(left, centerY() + margin, right, bottom)
+
     fun splitScreenApps(
-        top: TaskSpec,
-        bottom: TaskSpec,
+        displayId: Int = 0,
+        parentBounds: Rect = FULL_SCREEN,
+        taskMargin: Int = 0,
+        orientation: Orientation = VERTICAL,
+        first: TaskSpec,
+        second: TaskSpec,
         focusedTaskId: Int,
+        parentTaskId: Int = 2,
         shadeExpanded: Boolean = false,
     ): DisplayContentModel {
-        val topBounds = SPLIT_TOP
-        val bottomBounds = SPLIT_BOTTOM
+
+        val firstBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitTop(taskMargin)
+                HORIZONTAL -> parentBounds.splitLeft(taskMargin)
+            }
+        val secondBounds =
+            when (orientation) {
+                VERTICAL -> parentBounds.splitBottom(taskMargin)
+                HORIZONTAL -> parentBounds.splitRight(taskMargin)
+            }
+
         return DisplayContentModel(
-            displayId = 0,
+            displayId = displayId,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
             rootTasks =
                 listOf(
                     newRootTaskInfo(
-                        taskId = 2,
+                        taskId = parentTaskId,
                         userId = TestUserIds.PERSONAL,
-                        bounds = FULL_SCREEN,
+                        bounds = parentBounds,
                         topActivity =
                             ComponentName.unflattenFromString(
-                                if (top.taskId == focusedTaskId) top.name else bottom.name
+                                if (first.taskId == focusedTaskId) first.name else second.name
                             ),
                     ) {
                         listOf(
                                 newChildTask(
-                                    taskId = top.taskId,
-                                    bounds = topBounds,
-                                    userId = top.userId,
-                                    name = top.name
+                                    taskId = first.taskId,
+                                    bounds = firstBounds,
+                                    userId = first.userId,
+                                    name = first.name,
                                 ),
                                 newChildTask(
-                                    taskId = bottom.taskId,
-                                    bounds = bottomBounds,
-                                    userId = bottom.userId,
-                                    name = bottom.name
-                                )
+                                    taskId = second.taskId,
+                                    bounds = secondBounds,
+                                    userId = second.userId,
+                                    name = second.name,
+                                ),
                             )
                             // Child tasks are ordered bottom-up in RootTaskInfo.
                             // Sort 'focusedTaskId' last.
@@ -106,7 +128,7 @@
                             .sortedBy { it.id == focusedTaskId }
                     },
                     launcher(visible = false),
-                )
+                ),
         )
     }
 
@@ -124,7 +146,7 @@
                     fullScreen?.also { add(fullScreen(it, visible = true)) }
                     add(launcher(visible = (fullScreen == null)))
                     add(emptyRootSplit)
-                }
+                },
         )
     }
 
@@ -142,7 +164,7 @@
         return DisplayContentModel(
             displayId = 0,
             systemUiState = SystemUiState(shadeExpanded = shadeExpanded),
-            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit
+            rootTasks = freeFormTasks + launcher(visible = true) + emptyRootSplit,
         )
     }
 
@@ -153,11 +175,18 @@
      * somewhat sensible in terms of logical position (Re: PIP, SPLIT, etc).
      */
     object Bounds {
+        // "Phone" size
         val FULL_SCREEN = Rect(0, 0, 1080, 2400)
         val PIP = Rect(440, 1458, 1038, 1794)
         val SPLIT_TOP = Rect(0, 0, 1080, 1187)
         val SPLIT_BOTTOM = Rect(0, 1213, 1080, 2400)
         val FREE_FORM = Rect(119, 332, 1000, 1367)
+
+        // "Tablet" size
+        val FREEFORM_FULL_SCREEN = Rect(0, 0, 2560, 1600)
+        val FREEFORM_MAXIMIZED = Rect(0, 48, 2560, 1480)
+        val FREEFORM_SPLIT_LEFT = Rect(0, 0, 1270, 1600)
+        val FREEFORM_SPLIT_RIGHT = Rect(1290, 0, 2560, 1600)
     }
 
     /** A collection of task names used in test scenarios */
@@ -177,6 +206,8 @@
             "com.google.android.youtube/" +
                 "com.google.android.apps.youtube.app.watchwhile.WatchWhileActivity"
 
+        const val MESSAGES = "com.google.android.apps.messaging/.ui.ConversationListActivity"
+
         /** The NexusLauncher activity */
         const val LAUNCHER =
             "com.google.android.apps.nexuslauncher/" +
@@ -220,7 +251,7 @@
             }
 
         /** NexusLauncher on the default display. Usually below all other visible tasks */
-        fun launcher(visible: Boolean) =
+        fun launcher(visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = 1,
                 activityType = ActivityType.Home,
@@ -229,43 +260,63 @@
                 topActivity = ComponentName.unflattenFromString(ActivityNames.LAUNCHER),
                 topActivityType = ActivityType.Home,
             ) {
-                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER))
+                listOf(newChildTask(taskId = 1002, name = ActivityNames.LAUNCHER, bounds = bounds))
             }
 
         /** A full screen Activity */
-        fun fullScreen(task: TaskSpec, visible: Boolean) =
+        fun fullScreen(task: TaskSpec, visible: Boolean, bounds: Rect = FULL_SCREEN) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
                 visible = visible,
-                bounds = FULL_SCREEN,
+                bounds = bounds,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = task.userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = task.userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in Picture-in-Picture mode */
-        fun pictureInPicture(task: TaskSpec) =
+        fun pictureInPicture(task: TaskSpec, bounds: Rect = PIP) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = PIP,
                 windowingMode = WindowingMode.PictureInPicture,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
 
         /** An activity in FreeForm mode */
-        fun freeForm(task: TaskSpec) =
+        fun freeForm(task: TaskSpec, bounds: Rect = FREE_FORM) =
             newRootTaskInfo(
                 taskId = task.taskId,
                 userId = task.userId,
-                bounds = FREE_FORM,
+                bounds = bounds,
                 windowingMode = WindowingMode.Freeform,
                 topActivity = ComponentName.unflattenFromString(task.name),
             ) {
-                listOf(newChildTask(taskId = task.taskId, userId = userId, name = task.name))
+                listOf(
+                    newChildTask(
+                        taskId = task.taskId,
+                        userId = userId,
+                        name = task.name,
+                        bounds = bounds,
+                    )
+                )
             }
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
index 6c35b23..cedf0c8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/NewRootTaskInfo.kt
@@ -69,7 +69,7 @@
     taskId: Int,
     name: String,
     bounds: Rect? = null,
-    userId: Int? = null
+    userId: Int? = null,
 ): ChildTaskModel {
     return ChildTaskModel(taskId, name, bounds ?: this.bounds, userId ?: this.userId)
 }
@@ -83,7 +83,7 @@
     running: Boolean = true,
     activityType: ActivityType = Standard,
     windowingMode: WindowingMode = FullScreen,
-    bounds: Rect? = null,
+    bounds: Rect = Rect(),
     topActivity: ComponentName? = null,
     topActivityType: ActivityType = Standard,
     numActivities: Int? = null,
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
index 6e57761..b7f565d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/PrivateProfilePolicyTest.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.screenshot.policy
 
 import android.content.ComponentName
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.data.model.DisplayContentModel
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
@@ -59,7 +59,7 @@
             policy.check(
                 singleFullScreen(
                     spec = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
@@ -93,8 +93,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -110,25 +110,20 @@
                         listOf(
                             fullScreen(
                                 TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                                visible = true
+                                visible = true,
                             ),
                             fullScreen(
                                 TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                                visible = false
+                                visible = false,
                             ),
                             launcher(visible = false),
                             emptyRootSplit,
-                        )
+                        ),
                 )
             )
 
         assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    PrivateProfilePolicy.NAME,
-                    PrivateProfilePolicy.NO_VISIBLE_TASKS,
-                )
-            )
+            .isEqualTo(NotMatched(PrivateProfilePolicy.NAME, PrivateProfilePolicy.NO_VISIBLE_TASKS))
     }
 
     @Test
@@ -136,9 +131,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -150,8 +145,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -161,9 +156,9 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
-                    focusedTaskId = 1002
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -175,8 +170,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(FILES),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -196,8 +191,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
@@ -220,8 +215,8 @@
                     CaptureParameters(
                         type = FullScreen(displayId = 0),
                         component = ComponentName.unflattenFromString(YOUTUBE_PIP),
-                        owner = UserHandle.of(PRIVATE)
-                    )
+                        owner = UserHandle.of(PRIVATE),
+                    ),
                 )
             )
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
new file mode 100644
index 0000000..28eb9fc
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/ScreenshotPolicyTest.kt
@@ -0,0 +1,351 @@
+/*
+ * 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.screenshot.policy
+
+import android.content.ComponentName
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.kosmos.Kosmos
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.LAUNCHER
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.MESSAGES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREEFORM_FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Orientation.VERTICAL
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
+import com.android.systemui.screenshot.policy.TestUserIds.PRIVATE
+import com.android.systemui.screenshot.policy.TestUserIds.WORK
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class ScreenshotPolicyTest {
+    private val kosmos = Kosmos()
+
+    private val defaultComponent = ComponentName("default", "default")
+    private val defaultOwner = UserHandle.SYSTEM
+
+    @Test
+    fun fullScreen_work() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = FILES, userId = WORK)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_private() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PRIVATE)),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPersonal() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_personalAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = PERSONAL),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_workAndPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun splitScreen_twoWorkTasks() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                splitScreenApps(
+                    parentTaskId = 1,
+                    parentBounds = FREEFORM_FULL_SCREEN,
+                    orientation = VERTICAL,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type =
+                        RootTask(
+                            parentTaskId = 1,
+                            taskBounds = FREEFORM_FULL_SCREEN,
+                            childTaskIds = listOf(1002, 1003),
+                        ),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_maximized() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PERSONAL),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floatingWindows_withPrivate() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    TaskSpec(taskId = 1003, name = YOUTUBE, userId = PRIVATE),
+                    TaskSpec(taskId = 1004, name = MESSAGES, userId = PERSONAL),
+                    focusedTaskId = 1004,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(YOUTUBE),
+                    owner = UserHandle.of(PRIVATE),
+                )
+            )
+    }
+
+    @Test
+    fun freeform_floating_workOnly() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                freeFormApps(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    focusedTaskId = 1002,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = ComponentName.unflattenFromString(LAUNCHER),
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_shadeExpanded() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                singleFullScreen(
+                    TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    shadeExpanded = true,
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = FullScreen(displayId = 0),
+                    component = defaultComponent,
+                    owner = defaultOwner,
+                )
+            )
+    }
+
+    @Test
+    fun fullScreen_with_PictureInPicture() = runTest {
+        val policy = ScreenshotPolicy(kosmos.profileTypeRepository)
+
+        val result =
+            policy.apply(
+                pictureInPictureApp(
+                    pip = TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
+                    fullScreen = TaskSpec(taskId = 1003, name = FILES, userId = WORK),
+                ),
+                defaultComponent,
+                defaultOwner,
+            )
+
+        assertThat(result)
+            .isEqualTo(
+                CaptureParameters(
+                    type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
+                    component = ComponentName.unflattenFromString(FILES),
+                    owner = UserHandle.of(WORK),
+                )
+            )
+    }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
index be9fcc2..30a786c 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/screenshot/policy/WorkProfilePolicyTest.kt
@@ -31,13 +31,13 @@
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.YOUTUBE
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FREE_FORM
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
-import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.SPLIT_TOP
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.RootTasks
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.freeFormApps
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.pictureInPictureApp
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitScreenApps
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.splitTop
 import com.android.systemui.screenshot.data.model.SystemUiState
 import com.android.systemui.screenshot.data.repository.profileTypeRepository
 import com.android.systemui.screenshot.policy.CapturePolicy.PolicyResult
@@ -69,6 +69,7 @@
     @JvmField @Rule(order = 2) val mockitoRule: MockitoRule = MockitoJUnit.rule()
 
     @Mock lateinit var mContext: Context
+
     @Mock lateinit var mResources: Resources
 
     private val kosmos = Kosmos()
@@ -94,17 +95,11 @@
                 DisplayContentModel(
                     displayId = 0,
                     systemUiState = SystemUiState(shadeExpanded = false),
-                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks)
+                    rootTasks = listOf(RootTasks.emptyWithNoChildTasks),
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -114,13 +109,7 @@
                 singleFullScreen(TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL))
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -129,17 +118,11 @@
             policy.check(
                 singleFullScreen(
                     TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    shadeExpanded = true
+                    shadeExpanded = true,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    SHADE_EXPANDED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, SHADE_EXPANDED))
     }
 
     @Test
@@ -156,7 +139,7 @@
                         type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -166,9 +149,11 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1002
+                    parentBounds = FULL_SCREEN,
+                    taskMargin = 20,
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1002,
                 )
             )
 
@@ -178,10 +163,10 @@
                     policy = WorkProfilePolicy.NAME,
                     reason = WORK_TASK_IS_TOP,
                     CaptureParameters(
-                        type = IsolatedTask(taskId = 1002, taskBounds = SPLIT_TOP),
+                        type = IsolatedTask(taskId = 1002, taskBounds = FULL_SCREEN.splitTop(20)),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -191,19 +176,13 @@
         val result =
             policy.check(
                 splitScreenApps(
-                    top = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
-                    bottom = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
-                    focusedTaskId = 1003
+                    first = TaskSpec(taskId = 1002, name = FILES, userId = WORK),
+                    second = TaskSpec(taskId = 1003, name = YOUTUBE, userId = PERSONAL),
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    WORK_TASK_NOT_TOP,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, WORK_TASK_NOT_TOP))
     }
 
     @Test
@@ -225,7 +204,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FULL_SCREEN),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -238,7 +217,7 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
@@ -251,7 +230,7 @@
                         type = IsolatedTask(taskId = 1003, taskBounds = FREE_FORM),
                         component = ComponentName.unflattenFromString(FILES),
                         owner = UserHandle.of(WORK),
-                    )
+                    ),
                 )
             )
     }
@@ -264,16 +243,10 @@
                 freeFormApps(
                     TaskSpec(taskId = 1002, name = YOUTUBE, userId = PERSONAL),
                     TaskSpec(taskId = 1003, name = FILES, userId = WORK),
-                    focusedTaskId = 1003
+                    focusedTaskId = 1003,
                 )
             )
 
-        assertThat(result)
-            .isEqualTo(
-                NotMatched(
-                    WorkProfilePolicy.NAME,
-                    DESKTOP_MODE_ENABLED,
-                )
-            )
+        assertThat(result).isEqualTo(NotMatched(WorkProfilePolicy.NAME, DESKTOP_MODE_ENABLED))
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
index a5b2277..c6be0dd 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/ReferenceSystemUIModule.java
@@ -34,6 +34,7 @@
 import com.android.systemui.dock.DockManager;
 import com.android.systemui.dock.DockManagerImpl;
 import com.android.systemui.doze.DozeHost;
+import com.android.systemui.education.dagger.ContextualEducationModule;
 import com.android.systemui.inputdevice.tutorial.KeyboardTouchpadTutorialModule;
 import com.android.systemui.keyboard.shortcut.ShortcutHelperModule;
 import com.android.systemui.keyguard.ui.composable.blueprint.DefaultBlueprintModule;
@@ -153,6 +154,7 @@
         VolumeModule.class,
         WallpaperModule.class,
         ShortcutHelperModule.class,
+        ContextualEducationModule.class,
 })
 public abstract class ReferenceSystemUIModule {
 
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
index b55108d..450863f 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUIModule.java
@@ -63,7 +63,6 @@
 import com.android.systemui.doze.dagger.DozeComponent;
 import com.android.systemui.dreams.dagger.DreamModule;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.education.dagger.ContextualEducationModule;
 import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.flags.FlagDependenciesModule;
 import com.android.systemui.flags.FlagsModule;
@@ -272,8 +271,7 @@
         UserModule.class,
         UtilModule.class,
         NoteTaskModule.class,
-        WalletModule.class,
-        ContextualEducationModule.class
+        WalletModule.class
 },
         subcomponents = {
                 ComplicationComponent.class,
diff --git a/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
index 7fa7da1..abe0289 100644
--- a/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/dagger/ContextualEducationModule.kt
@@ -18,15 +18,12 @@
 
 import com.android.systemui.CoreStartable
 import com.android.systemui.Flags
-import com.android.systemui.contextualeducation.GestureType
 import com.android.systemui.coroutines.newTracingContext
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.education.data.repository.ContextualEducationRepository
 import com.android.systemui.education.data.repository.UserContextualEducationRepository
 import com.android.systemui.education.domain.interactor.ContextualEducationInteractor
 import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduInteractor
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractorImpl
 import com.android.systemui.education.ui.view.ContextualEduUiCoordinator
 import dagger.Binds
 import dagger.Lazy
@@ -83,18 +80,6 @@
         }
 
         @Provides
-        fun provideKeyboardTouchpadEduStatsInteractor(
-            implLazy: Lazy<KeyboardTouchpadEduStatsInteractorImpl>
-        ): KeyboardTouchpadEduStatsInteractor {
-            return if (Flags.keyboardTouchpadContextualEducation()) {
-                implLazy.get()
-            } else {
-                // No-op implementation when the flag is disabled.
-                return NoOpKeyboardTouchpadEduStatsInteractor
-            }
-        }
-
-        @Provides
         @IntoMap
         @ClassKey(KeyboardTouchpadEduInteractor::class)
         fun provideKeyboardTouchpadEduInteractor(
@@ -124,12 +109,6 @@
     }
 }
 
-private object NoOpKeyboardTouchpadEduStatsInteractor : KeyboardTouchpadEduStatsInteractor {
-    override fun incrementSignalCount(gestureType: GestureType) {}
-
-    override fun updateShortcutTriggerTime(gestureType: GestureType) {}
-}
-
 private object NoOpCoreStartable : CoreStartable {
     override fun start() {}
 }
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
index faee326..c17f3fb 100644
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractor.kt
@@ -18,6 +18,9 @@
 
 import android.os.SystemProperties
 import com.android.systemui.CoreStartable
+import com.android.systemui.common.coroutine.ChannelExt.trySendWithFailureLogging
+import com.android.systemui.contextualeducation.GestureType
+import com.android.systemui.contextualeducation.GestureType.ALL_APPS
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.education.dagger.ContextualEducationModule.EduClock
@@ -25,6 +28,13 @@
 import com.android.systemui.education.shared.model.EducationInfo
 import com.android.systemui.education.shared.model.EducationUiType
 import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.KEYBOARD
+import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.TOUCHPAD
+import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
+import com.android.systemui.recents.OverviewProxyService
+import com.android.systemui.recents.OverviewProxyService.OverviewProxyListener
+import com.android.systemui.utils.coroutines.flow.conflatedCallbackFlow
 import java.time.Clock
 import javax.inject.Inject
 import kotlin.time.Duration
@@ -33,9 +43,11 @@
 import kotlin.time.toDuration
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.merge
 import kotlinx.coroutines.launch
@@ -48,6 +60,8 @@
     @Background private val backgroundScope: CoroutineScope,
     private val contextualEducationInteractor: ContextualEducationInteractor,
     private val userInputDeviceRepository: UserInputDeviceRepository,
+    private val tutorialRepository: TutorialSchedulerRepository,
+    private val overviewProxyService: OverviewProxyService,
     @EduClock private val clock: Clock,
 ) : CoreStartable {
 
@@ -59,14 +73,16 @@
             getDurationForConfig("persist.contextual_edu.usage_session_sec", 3.days)
         val minIntervalBetweenEdu =
             getDurationForConfig("persist.contextual_edu.edu_interval_sec", 7.days)
+        val initialDelayDuration =
+            getDurationForConfig("persist.contextual_edu.initial_delay_sec", 7.days)
 
         private fun getDurationForConfig(
             systemPropertyKey: String,
-            defaultDuration: Duration
+            defaultDuration: Duration,
         ): Duration =
             SystemProperties.getLong(
                     systemPropertyKey,
-                    /* defaultValue= */ defaultDuration.inWholeSeconds
+                    /* defaultValue= */ defaultDuration.inWholeSeconds,
                 )
                 .toDuration(DurationUnit.SECONDS)
     }
@@ -74,6 +90,24 @@
     private val _educationTriggered = MutableStateFlow<EducationInfo?>(null)
     val educationTriggered = _educationTriggered.asStateFlow()
 
+    private val statsUpdateRequests: Flow<StatsUpdateRequest> = conflatedCallbackFlow {
+        val listener: OverviewProxyListener =
+            object : OverviewProxyListener {
+                override fun updateContextualEduStats(
+                    isTrackpadGesture: Boolean,
+                    gestureType: GestureType,
+                ) {
+                    trySendWithFailureLogging(
+                        StatsUpdateRequest(isTrackpadGesture, gestureType),
+                        TAG,
+                    )
+                }
+            }
+
+        overviewProxyService.addCallback(listener)
+        awaitClose { overviewProxyService.removeCallback(listener) }
+    }
+
     @OptIn(ExperimentalCoroutinesApi::class)
     override fun start() {
         backgroundScope.launch {
@@ -133,6 +167,16 @@
                 contextualEducationInteractor.updateShortcutTriggerTime(it)
             }
         }
+
+        backgroundScope.launch {
+            statsUpdateRequests.collect {
+                if (it.isTrackpadGesture) {
+                    contextualEducationInteractor.updateShortcutTriggerTime(it.gestureType)
+                } else {
+                    incrementSignalCount(it.gestureType)
+                }
+            }
+        }
     }
 
     private fun isEducationNeeded(model: GestureEduModel): Boolean {
@@ -160,4 +204,41 @@
 
     private fun getEduType(model: GestureEduModel) =
         if (model.educationShownCount > 0) EducationUiType.Notification else EducationUiType.Toast
+
+    private suspend fun incrementSignalCount(gestureType: GestureType) {
+        val targetDevice = getTargetDevice(gestureType)
+        if (isTargetDeviceConnected(targetDevice) && hasInitialDelayElapsed(targetDevice)) {
+            contextualEducationInteractor.incrementSignalCount(gestureType)
+        }
+    }
+
+    private suspend fun isTargetDeviceConnected(deviceType: DeviceType): Boolean {
+        return when (deviceType) {
+            KEYBOARD -> userInputDeviceRepository.isAnyKeyboardConnectedForUser.first().isConnected
+            TOUCHPAD -> userInputDeviceRepository.isAnyTouchpadConnectedForUser.first().isConnected
+        }
+    }
+
+    /**
+     * Keyboard shortcut education would be provided for All Apps. Touchpad gesture education would
+     * be provided for the rest of the gesture types (i.e. Home, Overview, Back). This method maps
+     * gesture to its target education device.
+     */
+    private fun getTargetDevice(gestureType: GestureType) =
+        when (gestureType) {
+            ALL_APPS -> KEYBOARD
+            else -> TOUCHPAD
+        }
+
+    private suspend fun hasInitialDelayElapsed(deviceType: DeviceType): Boolean {
+        val oobeLaunchTime = tutorialRepository.launchTime(deviceType) ?: return false
+        return clock
+            .instant()
+            .isAfter(oobeLaunchTime.plusSeconds(initialDelayDuration.inWholeSeconds))
+    }
+
+    private data class StatsUpdateRequest(
+        val isTrackpadGesture: Boolean,
+        val gestureType: GestureType,
+    )
 }
diff --git a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt b/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
deleted file mode 100644
index 43e39cf..0000000
--- a/packages/SystemUI/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduStatsInteractor.kt
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Copyright 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.education.domain.interactor
-
-import android.os.SystemProperties
-import com.android.systemui.contextualeducation.GestureType
-import com.android.systemui.contextualeducation.GestureType.ALL_APPS
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.education.dagger.ContextualEducationModule.EduClock
-import com.android.systemui.inputdevice.data.repository.UserInputDeviceRepository
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.KEYBOARD
-import com.android.systemui.inputdevice.tutorial.data.repository.DeviceType.TOUCHPAD
-import com.android.systemui.inputdevice.tutorial.data.repository.TutorialSchedulerRepository
-import java.time.Clock
-import javax.inject.Inject
-import kotlin.time.Duration
-import kotlin.time.Duration.Companion.days
-import kotlin.time.DurationUnit
-import kotlin.time.toDuration
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.launch
-
-/**
- * Encapsulates the update functions of KeyboardTouchpadEduStatsInteractor. This encapsulation is
- * for having a different implementation of interactor when the feature flag is off.
- */
-interface KeyboardTouchpadEduStatsInteractor {
-    fun incrementSignalCount(gestureType: GestureType)
-
-    fun updateShortcutTriggerTime(gestureType: GestureType)
-}
-
-/** Allow update to education data related to keyboard/touchpad. */
-@SysUISingleton
-class KeyboardTouchpadEduStatsInteractorImpl
-@Inject
-constructor(
-    @Background private val backgroundScope: CoroutineScope,
-    private val contextualEducationInteractor: ContextualEducationInteractor,
-    private val inputDeviceRepository: UserInputDeviceRepository,
-    private val tutorialRepository: TutorialSchedulerRepository,
-    @EduClock private val clock: Clock,
-) : KeyboardTouchpadEduStatsInteractor {
-
-    companion object {
-        val initialDelayDuration: Duration
-            get() =
-                SystemProperties.getLong(
-                        "persist.contextual_edu.initial_delay_sec",
-                        /* defaultValue= */ 7.days.inWholeSeconds,
-                    )
-                    .toDuration(DurationUnit.SECONDS)
-    }
-
-    override fun incrementSignalCount(gestureType: GestureType) {
-        backgroundScope.launch {
-            val targetDevice = getTargetDevice(gestureType)
-            if (isTargetDeviceConnected(targetDevice) && hasInitialDelayElapsed(targetDevice)) {
-                contextualEducationInteractor.incrementSignalCount(gestureType)
-            }
-        }
-    }
-
-    override fun updateShortcutTriggerTime(gestureType: GestureType) {
-        backgroundScope.launch {
-            contextualEducationInteractor.updateShortcutTriggerTime(gestureType)
-        }
-    }
-
-    private suspend fun isTargetDeviceConnected(deviceType: DeviceType): Boolean {
-        return when (deviceType) {
-            KEYBOARD -> inputDeviceRepository.isAnyKeyboardConnectedForUser.first().isConnected
-            TOUCHPAD -> inputDeviceRepository.isAnyTouchpadConnectedForUser.first().isConnected
-        }
-    }
-
-    /**
-     * Keyboard shortcut education would be provided for All Apps. Touchpad gesture education would
-     * be provided for the rest of the gesture types (i.e. Home, Overview, Back). This method maps
-     * gesture to its target education device.
-     */
-    private fun getTargetDevice(gestureType: GestureType) =
-        when (gestureType) {
-            ALL_APPS -> KEYBOARD
-            else -> TOUCHPAD
-        }
-
-    private suspend fun hasInitialDelayElapsed(deviceType: DeviceType): Boolean {
-        val oobeLaunchTime = tutorialRepository.launchTime(deviceType) ?: return false
-        return clock
-            .instant()
-            .isAfter(oobeLaunchTime.plusSeconds(initialDelayDuration.inWholeSeconds))
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt b/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
index 62694ce..ef7e7eb 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSTileIcon.kt
@@ -22,18 +22,13 @@
 
 /**
  * Creates a [QSTile.Icon] from an [Icon].
- * * [Icon.Loaded] && [resId] null -> [QSTileImpl.DrawableIcon]
- * * [Icon.Loaded] && [resId] available -> [QSTileImpl.DrawableIconWithRes]
+ * * [Icon.Loaded] -> [QSTileImpl.DrawableIcon]
  * * [Icon.Resource] -> [QSTileImpl.ResourceIcon]
  */
-fun Icon.asQSTileIcon(resId: Int?): QSTile.Icon {
+fun Icon.asQSTileIcon(): QSTile.Icon {
     return when (this) {
         is Icon.Loaded -> {
-            if (resId != null) {
-                QSTileImpl.DrawableIconWithRes(this.drawable, resId)
-            } else {
-                QSTileImpl.DrawableIcon(this.drawable)
-            }
+            QSTileImpl.DrawableIcon(this.drawable)
         }
         is Icon.Resource -> {
             QSTileImpl.ResourceIcon.get(this.res)
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
index 3bbe624..cf2db6c 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/ModesTile.kt
@@ -121,7 +121,7 @@
         state?.apply {
             this.state = tileState.activationState.legacyState
             val tileStateIcon = tileState.icon()
-            icon = tileStateIcon?.asQSTileIcon(tileState.iconRes) ?: ResourceIcon.get(ICON_RES_ID)
+            icon = tileStateIcon?.asQSTileIcon() ?: ResourceIcon.get(ICON_RES_ID)
             label = tileLabel
             secondaryLabel = tileState.secondaryLabel
             contentDescription = tileState.contentDescription
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
index cc14e71..3e442582 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/interactor/ModesTileDataInteractor.kt
@@ -78,14 +78,14 @@
         } else {
             return ModesTileModel(
                 isActivated = activeModes.isAnyActive(),
-                icon = Icon.Resource(ModesTile.ICON_RES_ID, null),
+                icon = context.getDrawable(ModesTile.ICON_RES_ID)!!.asIcon(),
                 iconResId = ModesTile.ICON_RES_ID,
                 activeModes = activeModes.modeNames,
             )
         }
     }
 
-    private data class TileIcon(val icon: Icon, val resId: Int?)
+    private data class TileIcon(val icon: Icon.Loaded, val resId: Int?)
 
     private fun getTileIcon(activeMode: ZenModeInfo?): TileIcon {
         return if (activeMode != null) {
@@ -96,7 +96,7 @@
                 TileIcon(activeMode.icon.drawable.asIcon(), null)
             }
         } else {
-            TileIcon(Icon.Resource(ModesTile.ICON_RES_ID, null), ModesTile.ICON_RES_ID)
+            TileIcon(context.getDrawable(ModesTile.ICON_RES_ID)!!.asIcon(), ModesTile.ICON_RES_ID)
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
index 9c31e32..db48123 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/domain/model/ModesTileModel.kt
@@ -21,12 +21,12 @@
 data class ModesTileModel(
     val isActivated: Boolean,
     val activeModes: List<String>,
-    val icon: Icon,
+    val icon: Icon.Loaded,
 
     /**
      * Resource id corresponding to [icon]. Will only be present if it's know to correspond to a
      * resource with a known id in SystemUI (such as resources from `android.R`,
      * `com.android.internal.R`, or `com.android.systemui.res` itself).
      */
-    val iconResId: Int? = null,
+    val iconResId: Int? = null
 )
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
index 537b56b..69da313 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/impl/modes/ui/ModesTileMapper.kt
@@ -18,9 +18,7 @@
 
 import android.content.res.Resources
 import android.icu.text.MessageFormat
-import android.util.Log
 import android.widget.Button
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.qs.tiles.base.interactor.QSTileDataToStateMapper
 import com.android.systemui.qs.tiles.impl.modes.domain.model.ModesTileModel
@@ -32,30 +30,14 @@
 
 class ModesTileMapper
 @Inject
-constructor(@Main private val resources: Resources, val theme: Resources.Theme) :
-    QSTileDataToStateMapper<ModesTileModel> {
+constructor(
+    @Main private val resources: Resources,
+    val theme: Resources.Theme,
+) : QSTileDataToStateMapper<ModesTileModel> {
     override fun map(config: QSTileConfig, data: ModesTileModel): QSTileState =
         QSTileState.build(resources, theme, config.uiConfig) {
-            val loadedIcon: Icon.Loaded =
-                when (val dataIcon = data.icon) {
-                    is Icon.Resource -> {
-                        if (data.iconResId != dataIcon.res) {
-                            Log.wtf(
-                                "ModesTileMapper",
-                                "Icon.Resource.res & iconResId are not identical",
-                            )
-                        }
-                        iconRes = dataIcon.res
-                        Icon.Loaded(resources.getDrawable(dataIcon.res, theme), null)
-                    }
-                    is Icon.Loaded -> {
-                        iconRes = data.iconResId
-                        dataIcon
-                    }
-                }
-
-            icon = { loadedIcon }
-
+            iconRes = data.iconResId
+            icon = { data.icon }
             activationState =
                 if (data.isActivated) {
                     QSTileState.ActivationState.ACTIVE
@@ -65,7 +47,10 @@
             secondaryLabel = getModesStatus(data, resources)
             contentDescription = "$label. $secondaryLabel"
             supportedActions =
-                setOf(QSTileState.UserAction.CLICK, QSTileState.UserAction.LONG_CLICK)
+                setOf(
+                    QSTileState.UserAction.CLICK,
+                    QSTileState.UserAction.LONG_CLICK,
+                )
             sideViewIcon = QSTileState.SideViewIcon.Chevron
             expandedAccessibilityClass = Button::class
         }
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index 559c263..ce9c441 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -87,7 +87,6 @@
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.dump.DumpManager;
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor;
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController;
 import com.android.systemui.keyguard.KeyguardWmStateRefactor;
 import com.android.systemui.keyguard.WakefulnessLifecycle;
@@ -160,8 +159,6 @@
     private final Provider<SceneInteractor> mSceneInteractor;
     private final Provider<ShadeInteractor> mShadeInteractor;
 
-    private final KeyboardTouchpadEduStatsInteractor mKeyboardTouchpadEduStatsInteractor;
-
     private final Runnable mConnectionRunnable = () ->
             internalConnectToCurrentUser("runnable: startConnectionToCurrentUser");
     private final ComponentName mRecentsComponentName;
@@ -660,8 +657,7 @@
             AssistUtils assistUtils,
             DumpManager dumpManager,
             Optional<UnfoldTransitionProgressForwarder> unfoldTransitionProgressForwarder,
-            BroadcastDispatcher broadcastDispatcher,
-            KeyboardTouchpadEduStatsInteractor keyboardTouchpadEduStatsInteractor
+            BroadcastDispatcher broadcastDispatcher
     ) {
         // b/241601880: This component should only be running for primary users or
         // secondaryUsers when visibleBackgroundUsers are supported.
@@ -699,7 +695,6 @@
         mDisplayTracker = displayTracker;
         mUnfoldTransitionProgressForwarder = unfoldTransitionProgressForwarder;
         mBroadcastDispatcher = broadcastDispatcher;
-        mKeyboardTouchpadEduStatsInteractor = keyboardTouchpadEduStatsInteractor;
 
         if (!KeyguardWmStateRefactor.isEnabled()) {
             mSysuiUnlockAnimationController = sysuiUnlockAnimationController;
@@ -940,19 +935,6 @@
         return isEnabled() && !QuickStepContract.isLegacyMode(mNavBarMode);
     }
 
-    /**
-     * Updates contextual education stats when a gesture is triggered
-     * @param isTrackpadGesture indicates if the gesture is triggered by trackpad
-     * @param gestureType type of gesture triggered
-     */
-    public void updateContextualEduStats(boolean isTrackpadGesture, GestureType gestureType) {
-        if (isTrackpadGesture) {
-            mKeyboardTouchpadEduStatsInteractor.updateShortcutTriggerTime(gestureType);
-        } else {
-            mKeyboardTouchpadEduStatsInteractor.incrementSignalCount(gestureType);
-        }
-    }
-
     public boolean isEnabled() {
         return mIsEnabled;
     }
@@ -978,6 +960,17 @@
         }
     }
 
+    /**
+     * Updates contextual education stats when a gesture is triggered
+     * @param isTrackpadGesture indicates if the gesture is triggered by trackpad
+     * @param gestureType type of gesture triggered
+     */
+    public void updateContextualEduStats(boolean isTrackpadGesture, GestureType gestureType) {
+        for (int i = mConnectionCallbacks.size() - 1; i >= 0; --i) {
+            mConnectionCallbacks.get(i).updateContextualEduStats(isTrackpadGesture, gestureType);
+        }
+    }
+
     private void notifyHomeRotationEnabled(boolean enabled) {
         for (int i = mConnectionCallbacks.size() - 1; i >= 0; --i) {
             mConnectionCallbacks.get(i).onHomeRotationEnabled(enabled);
@@ -1207,6 +1200,9 @@
         /** Set override of home button long press duration, touch slop multiplier, and haptic. */
         default void setOverrideHomeButtonLongPress(
                 long override, float slopMultiplier, boolean haptic) {}
+        /** Updates contextual education stats when target gesture type is triggered. */
+        default void updateContextualEduStats(
+                boolean isTrackpadGesture, GestureType gestureType) {}
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
index 0ef5207..9455201 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/CaptureType.kt
@@ -24,8 +24,8 @@
     data class FullScreen(val displayId: Int) : CaptureType
 
     /** Capture the contents of the task only. */
-    data class IsolatedTask(
-        val taskId: Int,
-        val taskBounds: Rect?,
-    ) : CaptureType
+    data class IsolatedTask(val taskId: Int, val taskBounds: Rect?) : CaptureType
+
+    data class RootTask(val parentTaskId: Int, val taskBounds: Rect?, val childTaskIds: List<Int>) :
+        CaptureType
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
index 039143a..e840668 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/PolicyRequestProcessor.kt
@@ -26,6 +26,7 @@
 import android.util.Log
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import android.view.WindowManager.TAKE_SCREENSHOT_PROVIDED_IMAGE
+import com.android.systemui.Flags.screenshotPolicySplitAndDesktopMode
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
@@ -47,14 +48,17 @@
     private val capture: ImageCapture,
     /** Provides information about the tasks on a given display */
     private val displayTasks: DisplayContentRepository,
-    /** The list of policies to apply, in order of priority */
+    /** The legacy list of policy implementations to apply, in order of priority */
     private val policies: List<CapturePolicy>,
+    /** Implements the combined policy rules for all profile types. */
+    private val policy: ScreenshotPolicy,
     /** The owner to assign for screenshot when a focused task isn't visible */
     private val defaultOwner: UserHandle = myUserHandle(),
     /** The assigned component when no application has focus, or not visible */
     private val defaultComponent: ComponentName,
 ) : ScreenshotRequestProcessor {
     override suspend fun process(original: ScreenshotData): ScreenshotData {
+
         if (original.type == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             // The request contains an already captured screenshot, accept it as is.
             Log.i(TAG, "Screenshot bitmap provided. No modifications applied.")
@@ -62,6 +66,12 @@
         }
         val displayContent = displayTasks.getDisplayContent(original.displayId)
 
+        if (screenshotPolicySplitAndDesktopMode()) {
+            Log.i(TAG, "Applying screenshot policy....")
+            val type = policy.apply(displayContent, defaultComponent, defaultOwner)
+            return modify(original, type)
+        }
+
         // If policies yield explicit modifications, apply them and return the result
         Log.i(TAG, "Applying policy checks....")
         policies.map { policy ->
@@ -79,10 +89,8 @@
     }
 
     /** Produce a new [ScreenshotData] using [CaptureParameters] */
-    private suspend fun modify(
-        original: ScreenshotData,
-        updates: CaptureParameters,
-    ): ScreenshotData {
+    suspend fun modify(original: ScreenshotData, updates: CaptureParameters): ScreenshotData {
+        Log.d(TAG, "[modify] CaptureParameters = $updates")
         // Update and apply bitmap capture depending on the parameters.
         val updated =
             when (val type = updates.type) {
@@ -94,6 +102,14 @@
                         type.taskId,
                         type.taskBounds,
                     )
+                is CaptureType.RootTask ->
+                    replaceWithTaskSnapshot(
+                        original,
+                        updates.component,
+                        updates.owner,
+                        type.parentTaskId,
+                        type.taskBounds,
+                    )
                 is FullScreen ->
                     replaceWithScreenshot(
                         original,
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
index f768cfb..dd39f92 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/RootTaskInfoExt.kt
@@ -26,9 +26,11 @@
             childTaskIds[index],
             childTaskNames[index],
             childTaskBounds[index],
-            childTaskUserIds[index]
+            childTaskUserIds[index],
         )
     }
 }
 
 internal fun RootTaskInfo.hasChildTasks() = childTaskUserIds.isNotEmpty()
+
+internal fun RootTaskInfo.childTaskCount() = childTaskIds.size
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
new file mode 100644
index 0000000..9967aff
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicy.kt
@@ -0,0 +1,155 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.screenshot.policy
+
+import android.app.ActivityTaskManager.RootTaskInfo
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.app.WindowConfiguration.WINDOWING_MODE_PINNED
+import android.content.ComponentName
+import android.os.UserHandle
+import android.util.Log
+import com.android.systemui.screenshot.data.model.DisplayContentModel
+import com.android.systemui.screenshot.data.model.ProfileType
+import com.android.systemui.screenshot.data.model.ProfileType.PRIVATE
+import com.android.systemui.screenshot.data.model.ProfileType.WORK
+import com.android.systemui.screenshot.data.repository.ProfileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
+import com.android.systemui.screenshot.policy.CaptureType.RootTask
+import javax.inject.Inject
+
+private const val TAG = "ScreenshotPolicy"
+
+/** Determines what to capture and which user owns the output. */
+class ScreenshotPolicy @Inject constructor(private val profileTypes: ProfileTypeRepository) {
+    /**
+     * Apply the policy to the content, resulting in [CaptureParameters].
+     *
+     * @param content the content of the display
+     * @param defaultComponent the component associated with the screenshot by default
+     * @param defaultOwner the user to own the screenshot by default
+     */
+    suspend fun apply(
+        content: DisplayContentModel,
+        defaultComponent: ComponentName,
+        defaultOwner: UserHandle,
+    ): CaptureParameters {
+        val defaultFullScreen by lazy {
+            CaptureParameters(
+                type = FullScreen(displayId = content.displayId),
+                component = defaultComponent,
+                owner = defaultOwner,
+            )
+        }
+
+        // When the systemUI notification shade is open, disregard tasks.
+        if (content.systemUiState.shadeExpanded) {
+            return defaultFullScreen
+        }
+
+        // find the first (top) RootTask which is visible and not Picture-in-Picture
+        val topRootTask =
+            content.rootTasks.firstOrNull {
+                it.isVisible && it.windowingMode != WindowConfiguration.WINDOWING_MODE_PINNED
+            } ?: return defaultFullScreen
+
+        Log.d(TAG, "topRootTask: $topRootTask")
+        val rootTaskOwners = topRootTask.childTaskUserIds.distinct()
+
+        // Special case: Only WORK in top root task which is full-screen or maximized freeform
+        if (
+            rootTaskOwners.size == 1 &&
+                profileTypes.getProfileType(rootTaskOwners.single()) == WORK &&
+                (topRootTask.isFullScreen() || topRootTask.isMaximizedFreeform())
+        ) {
+            val type =
+                if (topRootTask.childTaskCount() > 1) {
+                    RootTask(
+                        parentTaskId = topRootTask.taskId,
+                        taskBounds = topRootTask.bounds,
+                        childTaskIds = topRootTask.childTasksTopDown().map { it.id }.toList(),
+                    )
+                } else {
+                    IsolatedTask(
+                        taskId = topRootTask.childTasksTopDown().first().id,
+                        taskBounds = topRootTask.bounds,
+                    )
+                }
+            // Capture the RootTask (and all children)
+            return CaptureParameters(
+                type = type,
+                component = topRootTask.topActivity,
+                owner = UserHandle.of(rootTaskOwners.single()),
+            )
+        }
+
+        // In every other case the output will be a full screen capture regardless of content.
+        // For this reason, consider all owners of all visible content on the display (in all
+        // root tasks). This includes all root tasks in free-form mode.
+        val visibleChildTasks =
+            content.rootTasks.filter { it.isVisible }.flatMap { it.childTasksTopDown() }
+
+        val allVisibleProfileTypes =
+            visibleChildTasks
+                .map { it.userId }
+                .distinct()
+                .associate { profileTypes.getProfileType(it) to UserHandle.of(it) }
+
+        // If any visible content belongs to the private profile user -> private profile
+        // otherwise the personal user (including partial screen work content).
+        val ownerHandle =
+            allVisibleProfileTypes[PRIVATE]
+                ?: allVisibleProfileTypes[ProfileType.NONE]
+                ?: defaultOwner
+
+        // Attribute to the component of top-most task owned by this user (or fallback to default)
+        val topComponent =
+            visibleChildTasks.firstOrNull { it.userId == ownerHandle.identifier }?.componentName
+
+        return CaptureParameters(
+            type = FullScreen(content.displayId),
+            component = topComponent ?: topRootTask.topActivity ?: defaultComponent,
+            owner = ownerHandle,
+        )
+    }
+
+    private fun RootTaskInfo.isFullScreen(): Boolean =
+        configuration.windowConfiguration.windowingMode == WINDOWING_MODE_FULLSCREEN
+
+    private fun RootTaskInfo.isMaximizedFreeform(): Boolean {
+        val bounds = configuration.windowConfiguration.bounds
+        val maxBounds = configuration.windowConfiguration.maxBounds
+
+        if (
+            windowingMode != WINDOWING_MODE_FREEFORM ||
+                childTaskCount() != 1 ||
+                childTaskBounds[0] != bounds
+        ) {
+            return false
+        }
+
+        // Maximized floating windows fill maxBounds width
+        if (bounds.width() != maxBounds.width()) {
+            return false
+        }
+
+        // Maximized floating windows fill nearly all the height
+        return (bounds.height().toFloat() / maxBounds.height()) >= 0.89f
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
index 2cb9fe7..a9c6370 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/ScreenshotPolicyModule.kt
@@ -37,7 +37,6 @@
 
 @Module
 interface ScreenshotPolicyModule {
-
     @Binds
     @SysUISingleton
     fun bindProfileTypeRepository(impl: ProfileTypeRepositoryImpl): ProfileTypeRepository
@@ -67,6 +66,7 @@
             imageCapture: ImageCapture,
             displayContentRepo: DisplayContentRepository,
             policyListProvider: Provider<List<CapturePolicy>>,
+            standardPolicy: ScreenshotPolicy,
         ): ScreenshotRequestProcessor {
             return PolicyRequestProcessor(
                 background = background,
@@ -75,7 +75,8 @@
                 policies = policyListProvider.get(),
                 defaultOwner = Process.myUserHandle(),
                 defaultComponent =
-                    ComponentName(context.packageName, SystemUIService::class.java.toString())
+                    ComponentName(context.packageName, SystemUIService::class.java.toString()),
+                policy = standardPolicy,
             )
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
index 29450a2..cf90c0a 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/policy/WorkProfilePolicy.kt
@@ -28,7 +28,6 @@
 import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import javax.inject.Inject
-import kotlinx.coroutines.flow.first
 
 /**
  * Condition: When the top visible task (excluding PIP mode) belongs to a work user.
@@ -37,10 +36,8 @@
  */
 class WorkProfilePolicy
 @Inject
-constructor(
-    private val profileTypes: ProfileTypeRepository,
-    private val context: Context,
-) : CapturePolicy {
+constructor(private val profileTypes: ProfileTypeRepository, private val context: Context) :
+    CapturePolicy {
 
     override suspend fun check(content: DisplayContentModel): PolicyResult {
         // The systemUI notification shade isn't a work app, skip.
@@ -65,11 +62,7 @@
                 .map { it to it.childTasksTopDown().first() }
                 .firstOrNull { (_, child) ->
                     profileTypes.getProfileType(child.userId) == ProfileType.WORK
-                }
-                ?: return NotMatched(
-                    policy = NAME,
-                    reason = WORK_TASK_NOT_TOP,
-                )
+                } ?: return NotMatched(policy = NAME, reason = WORK_TASK_NOT_TOP)
 
         // If matched, return parameters needed to modify the request.
         return PolicyResult.Matched(
@@ -79,7 +72,7 @@
                 type = IsolatedTask(taskId = childTask.id, taskBounds = childTask.bounds),
                 component = childTask.componentName ?: rootTask.topActivity,
                 owner = UserHandle.of(childTask.userId),
-            )
+            ),
         )
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
index 4959224..3bfde68 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/recents/OverviewProxyServiceTest.kt
@@ -35,7 +35,6 @@
 import com.android.systemui.broadcast.BroadcastDispatcher
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.dump.DumpManager
-import com.android.systemui.education.domain.interactor.KeyboardTouchpadEduStatsInteractor
 import com.android.systemui.keyguard.KeyguardUnlockAnimationController
 import com.android.systemui.keyguard.WakefulnessLifecycle
 import com.android.systemui.keyguard.ui.view.InWindowLauncherUnlockAnimationManager
@@ -122,9 +121,6 @@
         Optional<UnfoldTransitionProgressForwarder>
     @Mock private lateinit var broadcastDispatcher: BroadcastDispatcher
 
-    @Mock
-    private lateinit var keyboardTouchpadEduStatsInteractor: KeyboardTouchpadEduStatsInteractor
-
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
@@ -293,7 +289,6 @@
             dumpManager,
             unfoldTransitionProgressForwarder,
             broadcastDispatcher,
-            keyboardTouchpadEduStatsInteractor,
         )
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
index 0d4cb4c..7709a65 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/policy/PolicyRequestProcessorTest.kt
@@ -20,30 +20,141 @@
 import android.graphics.Insets
 import android.graphics.Rect
 import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import android.view.Display.DEFAULT_DISPLAY
 import android.view.WindowManager.ScreenshotSource.SCREENSHOT_KEY_CHORD
 import android.view.WindowManager.TAKE_SCREENSHOT_FULLSCREEN
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.systemui.Flags
+import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.screenshot.ImageCapture
 import com.android.systemui.screenshot.ScreenshotData
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.ActivityNames.FILES
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.Bounds.FULL_SCREEN
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.TaskSpec
+import com.android.systemui.screenshot.data.model.DisplayContentScenarios.emptyDisplayContent
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.launcherOnly
 import com.android.systemui.screenshot.data.model.DisplayContentScenarios.singleFullScreen
 import com.android.systemui.screenshot.data.repository.DisplayContentRepository
+import com.android.systemui.screenshot.data.repository.profileTypeRepository
+import com.android.systemui.screenshot.policy.CaptureType.FullScreen
+import com.android.systemui.screenshot.policy.CaptureType.IsolatedTask
 import com.android.systemui.screenshot.policy.TestUserIds.PERSONAL
 import com.android.systemui.screenshot.policy.TestUserIds.WORK
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.Truth.assertWithMessage
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.runTest
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 
 @RunWith(AndroidJUnit4::class)
 class PolicyRequestProcessorTest {
+    private val kosmos = Kosmos()
+
+    private val screenshotRequest =
+        ScreenshotData(
+            TAKE_SCREENSHOT_FULLSCREEN,
+            SCREENSHOT_KEY_CHORD,
+            UserHandle.CURRENT,
+            topComponent = null,
+            originalScreenBounds = FULL_SCREEN,
+            taskId = -1,
+            originalInsets = Insets.NONE,
+            bitmap = null,
+            displayId = DEFAULT_DISPLAY,
+        )
+
+    val defaultComponent = ComponentName("default", "Component")
+    val defaultOwner = UserHandle.of(PERSONAL)
+
+    @get:Rule val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    /** Tests applying CaptureParameters with 'IsolatedTask' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_isolatedTask() = runTest {
+        val taskImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(task = taskImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(
+                    IsolatedTask(taskId = TASK_ID, taskBounds = null),
+                    ComponentName.unflattenFromString(FILES),
+                    UserHandle.of(WORK),
+                ),
+            )
+
+        assertWithMessage("The screenshot bitmap").that(result.bitmap).isSameInstanceAs(taskImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(UserHandle.of(WORK))
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(ComponentName.unflattenFromString(FILES))
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(TASK_ID)
+    }
+
+    /** Tests applying CaptureParameters with 'FullScreen' CaptureType */
+    @Test
+    @EnableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
+    fun testProcess_newPolicy_fullScreen() = runTest {
+        val screenImage = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+
+        /* Create a policy request processor with no capture policies */
+        val requestProcessor =
+            PolicyRequestProcessor(
+                Dispatchers.Unconfined,
+                createImageCapture(display = screenImage),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
+                policies = emptyList(),
+                defaultOwner = defaultOwner,
+                defaultComponent = defaultComponent,
+                displayTasks = { emptyDisplayContent },
+            )
+
+        val result =
+            requestProcessor.modify(
+                screenshotRequest,
+                CaptureParameters(FullScreen(displayId = 0), defaultComponent, defaultOwner),
+            )
+
+        assertWithMessage("The result bitmap").that(result.bitmap).isSameInstanceAs(screenImage)
+
+        assertWithMessage("The assigned owner of the screenshot")
+            .that(result.userHandle)
+            .isEqualTo(defaultOwner)
+
+        assertWithMessage("The topComponent of the screenshot")
+            .that(result.topComponent)
+            .isEqualTo(defaultComponent)
+
+        assertWithMessage("Task ID").that(result.taskId).isEqualTo(-1)
+    }
+
     /** Tests behavior when no policies are applied */
     @Test
+    @DisableFlags(Flags.FLAG_SCREENSHOT_POLICY_SPLIT_AND_DESKTOP_MODE)
     fun testProcess_defaultOwner_whenNoPolicyApplied() {
         val fullScreenWork = DisplayContentRepository {
             singleFullScreen(TaskSpec(taskId = TASK_ID, name = FILES, userId = WORK))
@@ -67,6 +178,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultOwner = UserHandle.of(PERSONAL),
                 defaultComponent = ComponentName("default", "Component"),
@@ -95,6 +207,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(display = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = emptyList(),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = DisplayContentRepository { launcherOnly() },
@@ -118,7 +231,7 @@
                 reason = "",
                 parameters =
                     CaptureParameters(
-                        CaptureType.IsolatedTask(taskId = 0, taskBounds = null),
+                        IsolatedTask(taskId = 0, taskBounds = null),
                         null,
                         UserHandle.CURRENT,
                     ),
@@ -130,6 +243,7 @@
             PolicyRequestProcessor(
                 Dispatchers.Unconfined,
                 createImageCapture(task = null),
+                policy = ScreenshotPolicy(kosmos.profileTypeRepository),
                 policies = listOf(captureTaskPolicy),
                 defaultComponent = ComponentName("default", "Component"),
                 displayTasks = fullScreenWork,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
index 2d275f9..3fd2503 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/education/domain/interactor/KeyboardTouchpadEduInteractorKosmos.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.kosmos.Kosmos
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.recents.OverviewProxyService
 import com.android.systemui.touchpad.data.repository.touchpadRepository
 import com.android.systemui.user.data.repository.userRepository
 import org.mockito.kotlin.mock
@@ -38,27 +39,13 @@
                     testDispatcher,
                     keyboardRepository,
                     touchpadRepository,
-                    userRepository
-                ),
-            clock = fakeEduClock
-        )
-    }
-
-var Kosmos.mockEduInputManager by Kosmos.Fixture { mock<InputManager>() }
-
-var Kosmos.keyboardTouchpadEduStatsInteractor by
-    Kosmos.Fixture {
-        KeyboardTouchpadEduStatsInteractorImpl(
-            backgroundScope = testScope.backgroundScope,
-            contextualEducationInteractor = contextualEducationInteractor,
-            inputDeviceRepository =
-                UserInputDeviceRepository(
-                    testDispatcher,
-                    keyboardRepository,
-                    touchpadRepository,
-                    userRepository
+                    userRepository,
                 ),
             tutorialSchedulerRepository,
-            fakeEduClock
+            mockOverviewProxyService,
+            clock = fakeEduClock,
         )
     }
+
+var Kosmos.mockOverviewProxyService by Kosmos.Fixture { mock<OverviewProxyService>() }
+var Kosmos.mockEduInputManager by Kosmos.Fixture { mock<InputManager>() }
diff --git a/services/core/java/com/android/server/net/OWNERS b/services/core/java/com/android/server/net/OWNERS
index bbc7c01..4596a44 100644
--- a/services/core/java/com/android/server/net/OWNERS
+++ b/services/core/java/com/android/server/net/OWNERS
@@ -2,7 +2,5 @@
 file:platform/packages/modules/Connectivity:main:/OWNERS_core_networking
 per-file NetworkPolicyManagerService.java=jackyu@google.com, sarahchin@google.com
 
-jsharkey@android.com
 sudheersai@google.com
-yamasani@google.com
 suprabh@google.com
diff --git a/services/core/java/com/android/server/notification/ManagedServices.java b/services/core/java/com/android/server/notification/ManagedServices.java
index 93482e7..122836e 100644
--- a/services/core/java/com/android/server/notification/ManagedServices.java
+++ b/services/core/java/com/android/server/notification/ManagedServices.java
@@ -21,6 +21,9 @@
 import static android.content.Context.BIND_AUTO_CREATE;
 import static android.content.Context.BIND_FOREGROUND_SERVICE;
 import static android.content.Context.DEVICE_POLICY_SERVICE;
+import static android.content.pm.PackageManager.MATCH_DIRECT_BOOT_AWARE;
+import static android.content.pm.PackageManager.MATCH_DIRECT_BOOT_UNAWARE;
+import static android.content.pm.PackageManager.MATCH_INSTANT;
 import static android.os.UserHandle.USER_ALL;
 import static android.os.UserHandle.USER_SYSTEM;
 import static android.service.notification.NotificationListenerService.META_DATA_DEFAULT_AUTOBIND;
@@ -106,7 +109,8 @@
     protected final String TAG = getClass().getSimpleName().replace('$', '.');
     protected final boolean DEBUG = Log.isLoggable(TAG, Log.DEBUG);
 
-    private static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
+    protected static final int ON_BINDING_DIED_REBIND_DELAY_MS = 10000;
+    protected static final int ON_BINDING_DIED_REBIND_MSG = 1234;
     protected static final String ENABLED_SERVICES_SEPARATOR = ":";
     private static final String DB_VERSION_1 = "1";
     private static final String DB_VERSION_2 = "2";
@@ -875,7 +879,21 @@
             String approvedItem = getApprovedValue(pkgOrComponent);
 
             if (approvedItem != null) {
+                final ComponentName component = ComponentName.unflattenFromString(approvedItem);
                 if (enabled) {
+                    if (Flags.notificationNlsRebind()) {
+                        if (component != null && !isValidService(component, userId)) {
+                            // Only fail if package is available
+                            // If not, it will be validated again in onPackagesChanged
+                            final PackageManager pm = mContext.getPackageManager();
+                            if (pm.isPackageAvailable(component.getPackageName())) {
+                                Slog.w(TAG, "Skip allowing " + mConfig.caption
+                                        + " " + pkgOrComponent + " (userSet: " + userSet
+                                        + ") for invalid service");
+                                return;
+                            }
+                        }
+                    }
                     approved.add(approvedItem);
                 } else {
                     approved.remove(approvedItem);
@@ -973,7 +991,7 @@
                 || isPackageOrComponentAllowed(component.getPackageName(), userId))) {
             return false;
         }
-        return componentHasBindPermission(component, userId);
+        return isValidService(component, userId);
     }
 
     private boolean componentHasBindPermission(ComponentName component, int userId) {
@@ -1220,12 +1238,21 @@
         if (!TextUtils.isEmpty(packageName)) {
             queryIntent.setPackage(packageName);
         }
+
+        if (Flags.notificationNlsRebind()) {
+            // Expand the package query
+            extraFlags |= MATCH_DIRECT_BOOT_AWARE | MATCH_DIRECT_BOOT_UNAWARE;
+            extraFlags |= MATCH_INSTANT;
+        }
+
         List<ResolveInfo> installedServices = pm.queryIntentServicesAsUser(
                 queryIntent,
                 PackageManager.GET_SERVICES | PackageManager.GET_META_DATA | extraFlags,
                 userId);
-        if (DEBUG)
-            Slog.v(TAG, mConfig.serviceInterface + " services: " + installedServices);
+        if (DEBUG) {
+            Slog.v(TAG, mConfig.serviceInterface + " pkg: " + packageName + " services: "
+                    + installedServices);
+        }
         if (installedServices != null) {
             for (int i = 0, count = installedServices.size(); i < count; i++) {
                 ResolveInfo resolveInfo = installedServices.get(i);
@@ -1325,11 +1352,12 @@
                     if (TextUtils.equals(getPackageName(approvedPackageOrComponent), packageName)) {
                         final ComponentName component = ComponentName.unflattenFromString(
                                 approvedPackageOrComponent);
-                        if (component != null && !componentHasBindPermission(component, userId)) {
+                        if (component != null && !isValidService(component, userId)) {
                             approved.removeAt(j);
                             if (DEBUG) {
                                 Slog.v(TAG, "Removing " + approvedPackageOrComponent
-                                        + " from approved list; no bind permission found "
+                                        + " from approved list; no bind permission or "
+                                        + "service interface filter found "
                                         + mConfig.bindPermission);
                             }
                         }
@@ -1348,6 +1376,15 @@
         }
     }
 
+    protected boolean isValidService(ComponentName component, int userId) {
+        if (Flags.notificationNlsRebind()) {
+            return componentHasBindPermission(component, userId) && queryPackageForServices(
+                    component.getPackageName(), userId).contains(component);
+        } else {
+            return componentHasBindPermission(component, userId);
+        }
+    }
+
     protected boolean isValidEntry(String packageOrComponent, int userId) {
         return hasMatchingServices(packageOrComponent, userId);
     }
@@ -1505,23 +1542,27 @@
      * Called when user switched to unbind all services from other users.
      */
     @VisibleForTesting
-    void unbindOtherUserServices(int currentUser) {
+    void unbindOtherUserServices(int switchedToUser) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + currentUser);
-        unbindServicesImpl(currentUser, true /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindOtherUserServices_current" + switchedToUser);
+        unbindServicesImpl(switchedToUser, true /* allExceptUser */);
         t.traceEnd();
     }
 
-    void unbindUserServices(int user) {
+    void unbindUserServices(int removedUser) {
         TimingsTraceAndSlog t = new TimingsTraceAndSlog();
-        t.traceBegin("ManagedServices.unbindUserServices" + user);
-        unbindServicesImpl(user, false /* allExceptUser */);
+        t.traceBegin("ManagedServices.unbindUserServices" + removedUser);
+        unbindServicesImpl(removedUser, false /* allExceptUser */);
         t.traceEnd();
     }
 
     void unbindServicesImpl(int user, boolean allExceptUser) {
         final SparseArray<Set<ComponentName>> componentsToUnbind = new SparseArray<>();
         synchronized (mMutex) {
+            if (Flags.notificationNlsRebind()) {
+                // Remove enqueued rebinds to avoid rebinding services for a switched user
+                mHandler.removeMessages(ON_BINDING_DIED_REBIND_MSG);
+            }
             final Set<ManagedServiceInfo> removableBoundServices = getRemovableConnectedServices();
             for (ManagedServiceInfo info : removableBoundServices) {
                 if ((allExceptUser && (info.userid != user))
@@ -1716,6 +1757,7 @@
                             mServicesRebinding.add(servicesBindingTag);
                             mHandler.postDelayed(() ->
                                     reregisterService(name, userid),
+                                    ON_BINDING_DIED_REBIND_MSG,
                                     ON_BINDING_DIED_REBIND_DELAY_MS);
                         } else {
                             Slog.v(TAG, getCaption() + " not rebinding in user " + userid
diff --git a/services/core/java/com/android/server/notification/flags.aconfig b/services/core/java/com/android/server/notification/flags.aconfig
index f79d9ef..c479acf 100644
--- a/services/core/java/com/android/server/notification/flags.aconfig
+++ b/services/core/java/com/android/server/notification/flags.aconfig
@@ -194,3 +194,13 @@
   description: "Enables sound uri with vibration source in notification channel"
   bug: "351975435"
 }
+
+flag {
+  name: "notification_nls_rebind"
+  namespace: "systemui"
+  description: "Check for NLS service intent filter when rebinding services"
+  bug: "347674739"
+  metadata {
+    purpose: PURPOSE_BUGFIX
+  }
+}
diff --git a/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java b/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
index dcb47a7..4c9cbc4 100644
--- a/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
+++ b/services/core/java/com/android/server/stats/bootstrap/StatsBootstrapAtomService.java
@@ -42,27 +42,28 @@
             return;
         }
         StatsEvent.Builder builder = StatsEvent.newBuilder().setAtomId(atom.atomId);
-        for (StatsBootstrapAtomValue value : atom.values) {
+        for (StatsBootstrapAtomValue atomValue : atom.values) {
+            StatsBootstrapAtomValue.Primitive value = atomValue.value;
             switch (value.getTag()) {
-                case StatsBootstrapAtomValue.boolValue:
+                case StatsBootstrapAtomValue.Primitive.boolValue:
                     builder.writeBoolean(value.getBoolValue());
                     break;
-                case StatsBootstrapAtomValue.intValue:
+                case StatsBootstrapAtomValue.Primitive.intValue:
                     builder.writeInt(value.getIntValue());
                     break;
-                case StatsBootstrapAtomValue.longValue:
+                case StatsBootstrapAtomValue.Primitive.longValue:
                     builder.writeLong(value.getLongValue());
                     break;
-                case StatsBootstrapAtomValue.floatValue:
+                case StatsBootstrapAtomValue.Primitive.floatValue:
                     builder.writeFloat(value.getFloatValue());
                     break;
-                case StatsBootstrapAtomValue.stringValue:
+                case StatsBootstrapAtomValue.Primitive.stringValue:
                     builder.writeString(value.getStringValue());
                     break;
-                case StatsBootstrapAtomValue.bytesValue:
+                case StatsBootstrapAtomValue.Primitive.bytesValue:
                     builder.writeByteArray(value.getBytesValue());
                     break;
-                case StatsBootstrapAtomValue.stringArrayValue:
+                case StatsBootstrapAtomValue.Primitive.stringArrayValue:
                     builder.writeStringArray(value.getStringArrayValue());
                     break;
                 default:
@@ -71,6 +72,25 @@
                     return;
 
             }
+            StatsBootstrapAtomValue.Annotation[] annotations = atomValue.annotations;
+            for (StatsBootstrapAtomValue.Annotation annotation : atomValue.annotations) {
+                if (annotation.id != StatsBootstrapAtomValue.Annotation.Id.IS_UID) {
+                    Slog.e(TAG, "Unexpected annotation ID: " + annotation.id
+                            + ", for atom " + atom.atomId + ": only UIDs are supported!");
+                    return;
+                }
+
+                switch (annotation.value.getTag()) {
+                    case StatsBootstrapAtomValue.Annotation.Primitive.boolValue:
+                        builder.addBooleanAnnotation(
+                                annotation.id, annotation.value.getBoolValue());
+                        break;
+                    default:
+                        Slog.e(TAG, "Unexpected value type " + annotation.value.getTag()
+                                + " when logging UID for atom " + atom.atomId);
+                        return;
+                }
+            }
         }
         StatsLog.write(builder.usePooledBuffer().build());
     }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
index 48bc9d7..b5724b5c 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/ManagedServicesTest.java
@@ -21,8 +21,10 @@
 import static android.os.UserManager.USER_TYPE_PROFILE_CLONE;
 import static android.os.UserManager.USER_TYPE_PROFILE_MANAGED;
 import static android.os.UserManager.USER_TYPE_PROFILE_PRIVATE;
+import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
 import static android.service.notification.NotificationListenerService.META_DATA_DEFAULT_AUTOBIND;
 
+import static com.android.server.notification.Flags.FLAG_NOTIFICATION_NLS_REBIND;
 import static com.android.server.notification.ManagedServices.APPROVAL_BY_COMPONENT;
 import static com.android.server.notification.ManagedServices.APPROVAL_BY_PACKAGE;
 import static com.android.server.notification.NotificationManagerService.privateSpaceFlagsEnabled;
@@ -63,11 +65,14 @@
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.IInterface;
+import android.os.Looper;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.provider.Settings;
+import android.testing.TestableLooper;
 import android.text.TextUtils;
 import android.util.ArrayMap;
 import android.util.ArraySet;
@@ -82,7 +87,9 @@
 
 import com.google.android.collect.Lists;
 
+import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
@@ -103,7 +110,10 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 
+
 public class ManagedServicesTest extends UiServiceTestCase {
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule(DEVICE_DEFAULT);
 
     @Mock
     private IPackageManager mIpm;
@@ -115,6 +125,7 @@
     private ManagedServices.UserProfiles mUserProfiles;
     @Mock private DevicePolicyManager mDpm;
     Object mLock = new Object();
+    private TestableLooper mTestableLooper;
 
     UserInfo mZero = new UserInfo(0, "zero", 0);
     UserInfo mTen = new UserInfo(10, "ten", 0);
@@ -142,6 +153,7 @@
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
+        mTestableLooper = new TestableLooper(Looper.getMainLooper());
 
         mContext.setMockPackageManager(mPm);
         mContext.addMockSystemService(Context.USER_SERVICE, mUm);
@@ -199,6 +211,11 @@
                 mIpm, APPROVAL_BY_COMPONENT);
     }
 
+    @After
+    public void tearDown() throws Exception {
+        mTestableLooper.destroy();
+    }
+
     @Test
     public void testBackupAndRestore_migration() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
@@ -888,7 +905,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -919,7 +936,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -950,7 +967,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a/a", 0, true);
 
         service.reregisterService(cn, 0);
@@ -981,7 +998,7 @@
             return true;
         });
 
-        mockServiceInfoWithMetaData(List.of(cn), service, new ArrayMap<>());
+        mockServiceInfoWithMetaData(List.of(cn), service, pm, new ArrayMap<>());
         service.addApprovedList("a/a", 0, false);
 
         service.reregisterService(cn, 0);
@@ -1053,6 +1070,78 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void registerService_bindingDied_rebindIsClearedOnUserSwitch() throws Exception {
+        Context context = mock(Context.class);
+        PackageManager pm = mock(PackageManager.class);
+        ApplicationInfo ai = new ApplicationInfo();
+        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
+
+        when(context.getPackageName()).thenReturn(mPkg);
+        when(context.getUserId()).thenReturn(mUser.getIdentifier());
+        when(context.getPackageManager()).thenReturn(pm);
+        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
+                APPROVAL_BY_PACKAGE);
+        service = spy(service);
+        ComponentName cn = ComponentName.unflattenFromString("a/a");
+
+        // Trigger onBindingDied for component when registering
+        //  => will schedule a rebind in 10 seconds
+        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
+            Object[] args = invocation.getArguments();
+            ServiceConnection sc = (ServiceConnection) args[1];
+            sc.onBindingDied(cn);
+            return true;
+        });
+        service.registerService(cn, 0);
+        assertThat(service.isBound(cn, 0)).isFalse();
+
+        // Switch to user 10
+        service.onUserSwitched(10);
+
+        // Check that the scheduled rebind for user 0 was cleared
+        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
+        mTestableLooper.processAllMessages();
+        verify(service, never()).reregisterService(any(), anyInt());
+    }
+
+    @Test
+    public void registerService_bindingDied_rebindIsExecutedAfterTimeout() throws Exception {
+        Context context = mock(Context.class);
+        PackageManager pm = mock(PackageManager.class);
+        ApplicationInfo ai = new ApplicationInfo();
+        ai.targetSdkVersion = Build.VERSION_CODES.CUR_DEVELOPMENT;
+
+        when(context.getPackageName()).thenReturn(mPkg);
+        when(context.getUserId()).thenReturn(mUser.getIdentifier());
+        when(context.getPackageManager()).thenReturn(pm);
+        when(pm.getApplicationInfo(anyString(), anyInt())).thenReturn(ai);
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles, mIpm,
+                APPROVAL_BY_PACKAGE);
+        service = spy(service);
+        ComponentName cn = ComponentName.unflattenFromString("a/a");
+
+        // Trigger onBindingDied for component when registering
+        //  => will schedule a rebind in 10 seconds
+        when(context.bindServiceAsUser(any(), any(), anyInt(), any())).thenAnswer(invocation -> {
+            Object[] args = invocation.getArguments();
+            ServiceConnection sc = (ServiceConnection) args[1];
+            sc.onBindingDied(cn);
+            return true;
+        });
+        service.registerService(cn, 0);
+        assertThat(service.isBound(cn, 0)).isFalse();
+
+        // Check that the scheduled rebind is run
+        mTestableLooper.moveTimeForward(ManagedServices.ON_BINDING_DIED_REBIND_DELAY_MS);
+        mTestableLooper.processAllMessages();
+        verify(service, times(1)).reregisterService(eq(cn), eq(0));
+    }
+
+    @Test
     public void testPackageUninstall_packageNoLongerInApprovedList() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1211,6 +1300,65 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testUpgradeAppNoIntentFilterNoRebind() throws Exception {
+        Context context = spy(getContext());
+        doReturn(true).when(context).bindServiceAsUser(any(), any(), anyInt(), any());
+
+        ManagedServices service = new TestManagedServices(context, mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        List<String> packages = new ArrayList<>();
+        packages.add("package");
+        addExpectedServices(service, packages, 0);
+
+        final ComponentName unapprovedComponent = ComponentName.unflattenFromString("package/C1");
+        final ComponentName approvedComponent = ComponentName.unflattenFromString("package/C2");
+
+        // Both components are approved initially
+        mExpectedPrimaryComponentNames.clear();
+        mExpectedPrimaryPackages.clear();
+        mExpectedPrimaryComponentNames.put(0, "package/C1:package/C2");
+        mExpectedSecondaryComponentNames.clear();
+        mExpectedSecondaryPackages.clear();
+
+        loadXml(service);
+
+        //Component package/C1 loses serviceInterface intent filter
+        ManagedServices.Config config = service.getConfig();
+        when(mPm.queryIntentServicesAsUser(any(), anyInt(), anyInt()))
+                .thenAnswer(new Answer<List<ResolveInfo>>() {
+                    @Override
+                    public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
+                            throws Throwable {
+                        Object[] args = invocationOnMock.getArguments();
+                        Intent invocationIntent = (Intent) args[0];
+                        if (invocationIntent != null) {
+                            if (invocationIntent.getAction().equals(config.serviceInterface)
+                                    && packages.contains(invocationIntent.getPackage())) {
+                                List<ResolveInfo> dummyServices = new ArrayList<>();
+                                ResolveInfo resolveInfo = new ResolveInfo();
+                                ServiceInfo serviceInfo = new ServiceInfo();
+                                serviceInfo.packageName = invocationIntent.getPackage();
+                                serviceInfo.name = approvedComponent.getClassName();
+                                serviceInfo.permission = service.getConfig().bindPermission;
+                                resolveInfo.serviceInfo = serviceInfo;
+                                dummyServices.add(resolveInfo);
+                                return dummyServices;
+                            }
+                        }
+                        return new ArrayList<>();
+                    }
+                });
+
+        // Trigger package update
+        service.onPackagesChanged(false, new String[]{"package"}, new int[]{0});
+
+        assertFalse(service.isComponentEnabledForCurrentProfiles(unapprovedComponent));
+        assertTrue(service.isComponentEnabledForCurrentProfiles(approvedComponent));
+    }
+
+    @Test
     public void testSetPackageOrComponentEnabled() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1223,6 +1371,21 @@
                             "user10package1/K", "user10.3/Component", "user10package2/L",
                             "user10.4/Component"}));
 
+            // mock permissions for services
+            PackageManager pm = mock(PackageManager.class);
+            when(getContext().getPackageManager()).thenReturn(pm);
+            List<ComponentName> enabledComponents = List.of(
+                    ComponentName.unflattenFromString("package/Comp"),
+                    ComponentName.unflattenFromString("package/C2"),
+                    ComponentName.unflattenFromString("again/M4"),
+                    ComponentName.unflattenFromString("user10package/B"),
+                    ComponentName.unflattenFromString("user10/Component"),
+                    ComponentName.unflattenFromString("user10package1/K"),
+                    ComponentName.unflattenFromString("user10.3/Component"),
+                    ComponentName.unflattenFromString("user10package2/L"),
+                    ComponentName.unflattenFromString("user10.4/Component"));
+            mockServiceInfoWithMetaData(enabledComponents, service, pm, new ArrayMap<>());
+
             for (int userId : expectedEnabled.keySet()) {
                 ArrayList<String> expectedForUser = expectedEnabled.get(userId);
                 for (int i = 0; i < expectedForUser.size(); i++) {
@@ -1284,6 +1447,90 @@
     }
 
     @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testSetPackageOrComponentEnabled_pkgInstalledAfterEnabling() throws Exception {
+        ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        final int userId = 0;
+        final String validComponent = "again/M4";
+        ArrayList<String> expectedEnabled = Lists.newArrayList("package/Comp", "package/C2",
+                validComponent);
+
+        PackageManager pm = mock(PackageManager.class);
+        when(getContext().getPackageManager()).thenReturn(pm);
+        service = spy(service);
+
+        // Component again/M4 is a valid service and the package is available
+        doReturn(true).when(service)
+                .isValidService(ComponentName.unflattenFromString(validComponent), userId);
+        when(pm.isPackageAvailable("again")).thenReturn(true);
+
+        // "package" is not available and its services are not valid
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString("package/Comp"), userId);
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString("package/C2"), userId);
+        when(pm.isPackageAvailable("package")).thenReturn(false);
+
+        // Enable all components
+        for (String component: expectedEnabled) {
+            service.setPackageOrComponentEnabled(component, userId, true, true);
+        }
+
+        // Verify everything added is approved
+        for (String component: expectedEnabled) {
+            assertTrue("Not allowed: user: " + userId + " entry: " + component
+                    + " for approval level " + APPROVAL_BY_COMPONENT,
+                    service.isPackageOrComponentAllowed(component, userId));
+        }
+
+        // Add missing package "package"
+        service.onPackagesChanged(false, new String[]{"package"}, new int[]{0});
+
+        // Check that component of "package" are not enabled
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString("package/Comp")));
+        assertFalse(service.isPackageOrComponentAllowed("package/Comp", userId));
+
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString("package/C2")));
+        assertFalse(service.isPackageOrComponentAllowed("package/C2", userId));
+
+        // Check that the valid components are still enabled
+        assertTrue(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString(validComponent)));
+        assertTrue(service.isPackageOrComponentAllowed(validComponent, userId));
+    }
+
+    @Test
+    @EnableFlags(FLAG_NOTIFICATION_NLS_REBIND)
+    public void testSetPackageOrComponentEnabled_invalidComponent() throws Exception {
+        ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
+                mIpm, APPROVAL_BY_COMPONENT);
+
+        final int userId = 0;
+        final String invalidComponent = "package/Comp";
+
+        PackageManager pm = mock(PackageManager.class);
+        when(getContext().getPackageManager()).thenReturn(pm);
+        service = spy(service);
+
+        // Component is an invalid service and the package is available
+        doReturn(false).when(service)
+                .isValidService(ComponentName.unflattenFromString(invalidComponent), userId);
+        when(pm.isPackageAvailable("package")).thenReturn(true);
+        service.setPackageOrComponentEnabled(invalidComponent, userId, true, true);
+
+        // Verify that the component was not enabled
+        assertFalse("Not allowed: user: " + userId + " entry: " + invalidComponent
+                    + " for approval level " + APPROVAL_BY_COMPONENT,
+                service.isPackageOrComponentAllowed(invalidComponent, userId));
+        assertFalse(service.isComponentEnabledForCurrentProfiles(
+                ComponentName.unflattenFromString(invalidComponent)));
+    }
+
+    @Test
     public void testGetAllowedPackages_byUser() throws Exception {
         for (int approvalLevel : new int[] {APPROVAL_BY_COMPONENT, APPROVAL_BY_PACKAGE}) {
             ManagedServices service = new TestManagedServices(getContext(), mLock, mUserProfiles,
@@ -1944,7 +2191,7 @@
         metaDataAutobindAllow.putBoolean(META_DATA_DEFAULT_AUTOBIND, true);
         metaDatas.put(cn_allowed, metaDataAutobindAllow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_allowed.flattenToString(), 0, true);
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
@@ -1989,7 +2236,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2028,7 +2275,7 @@
         metaDataAutobindDisallow.putBoolean(META_DATA_DEFAULT_AUTOBIND, false);
         metaDatas.put(cn_disallowed, metaDataAutobindDisallow);
 
-        mockServiceInfoWithMetaData(componentNames, service, metaDatas);
+        mockServiceInfoWithMetaData(componentNames, service, pm, metaDatas);
 
         service.addApprovedList(cn_disallowed.flattenToString(), 0, true);
 
@@ -2099,8 +2346,8 @@
     }
 
     private void mockServiceInfoWithMetaData(List<ComponentName> componentNames,
-            ManagedServices service, ArrayMap<ComponentName, Bundle> metaDatas)
-            throws RemoteException {
+            ManagedServices service, PackageManager packageManager,
+            ArrayMap<ComponentName, Bundle> metaDatas) throws RemoteException {
         when(mIpm.getServiceInfo(any(), anyLong(), anyInt())).thenAnswer(
                 (Answer<ServiceInfo>) invocation -> {
                     ComponentName invocationCn = invocation.getArgument(0);
@@ -2115,6 +2362,39 @@
                     return null;
                 }
         );
+
+        // add components to queryIntentServicesAsUser response
+        final List<String> packages = new ArrayList<>();
+        for (ComponentName cn: componentNames) {
+            packages.add(cn.getPackageName());
+        }
+        ManagedServices.Config config = service.getConfig();
+        when(packageManager.queryIntentServicesAsUser(any(), anyInt(), anyInt())).
+                thenAnswer(new Answer<List<ResolveInfo>>() {
+                @Override
+                public List<ResolveInfo> answer(InvocationOnMock invocationOnMock)
+                    throws Throwable {
+                    Object[] args = invocationOnMock.getArguments();
+                    Intent invocationIntent = (Intent) args[0];
+                    if (invocationIntent != null) {
+                        if (invocationIntent.getAction().equals(config.serviceInterface)
+                            && packages.contains(invocationIntent.getPackage())) {
+                            List<ResolveInfo> dummyServices = new ArrayList<>();
+                            for (ComponentName cn: componentNames) {
+                                ResolveInfo resolveInfo = new ResolveInfo();
+                                ServiceInfo serviceInfo = new ServiceInfo();
+                                serviceInfo.packageName = invocationIntent.getPackage();
+                                serviceInfo.name = cn.getClassName();
+                                serviceInfo.permission = service.getConfig().bindPermission;
+                                resolveInfo.serviceInfo = serviceInfo;
+                                dummyServices.add(resolveInfo);
+                            }
+                            return dummyServices;
+                        }
+                    }
+                    return new ArrayList<>();
+                }
+            });
     }
 
     private void resetComponentsAndPackages() {
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
index 0f7de7d..2c645e0 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationAssistantsTest.java
@@ -28,6 +28,7 @@
 import static junit.framework.Assert.assertTrue;
 
 import static org.junit.Assert.assertNull;
+
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.any;
@@ -197,6 +198,8 @@
     public void testWriteXml_userTurnedOffNAS() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
+
         mAssistants.loadDefaultsFromConfig(true);
 
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
@@ -432,6 +435,10 @@
     public void testSetPackageOrComponentEnabled_onlyOnePackage() throws Exception {
         ComponentName component1 = ComponentName.unflattenFromString("package/Component1");
         ComponentName component2 = ComponentName.unflattenFromString("package/Component2");
+
+        doReturn(true).when(mAssistants).isValidService(eq(component1), eq(mZero.id));
+        doReturn(true).when(mAssistants).isValidService(eq(component2), eq(mZero.id));
+
         mAssistants.setPackageOrComponentEnabled(component1.flattenToString(), mZero.id, true,
                 true, true);
         verify(mNm, never()).setNotificationAssistantAccessGrantedForUserInternal(
@@ -577,6 +584,7 @@
     public void testSetAdjustmentTypeSupportedState() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -600,6 +608,7 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_entries() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);
@@ -623,6 +632,7 @@
     public void testSetAdjustmentTypeSupportedState_readWriteXml_empty() throws Exception {
         int userId = ActivityManager.getCurrentUser();
 
+        doReturn(true).when(mAssistants).isValidService(eq(mCn), eq(userId));
         mAssistants.loadDefaultsFromConfig(true);
         mAssistants.setPackageOrComponentEnabled(mCn.flattenToString(), userId, true,
                 true, true);