Merge "[waitForAllWindowsDrawn] Wait for transitions only on default displays" into main
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index 76012bb..03eab7c 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3494,7 +3494,7 @@
     method @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void registerIntentInterceptor(@NonNull android.content.IntentFilter, @NonNull java.util.concurrent.Executor, @NonNull android.companion.virtual.VirtualDeviceManager.IntentInterceptorCallback);
     method public void removeActivityListener(@NonNull android.companion.virtual.VirtualDeviceManager.ActivityListener);
     method @FlaggedApi("android.companion.virtual.flags.dynamic_policy") @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void removeActivityPolicyExemption(@NonNull android.content.ComponentName);
-    method @FlaggedApi("android.companion.virtual.flags.dynamic_policy") @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void removeActivityPolicyExemption(@NonNull String);
+    method @FlaggedApi("android.companion.virtualdevice.flags.activity_control_api") @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void removeActivityPolicyExemption(@NonNull String);
     method @FlaggedApi("android.companion.virtualdevice.flags.activity_control_api") @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void removeActivityPolicyExemption(@NonNull android.content.ComponentName, int);
     method @FlaggedApi("android.companion.virtualdevice.flags.activity_control_api") @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE) public void removeActivityPolicyExemption(@NonNull String, int);
     method public void removeSoundEffectListener(@NonNull android.companion.virtual.VirtualDeviceManager.SoundEffectListener);
diff --git a/core/java/android/app/admin/Provisioning_OWNERS b/core/java/android/app/admin/Provisioning_OWNERS
index 91b9761..09ebb26 100644
--- a/core/java/android/app/admin/Provisioning_OWNERS
+++ b/core/java/android/app/admin/Provisioning_OWNERS
@@ -1,4 +1,7 @@
 # Assign bugs to android-enterprise-triage@google.com
 ae-provisioning-reviews@google.com
 acjohnston@google.com #{LAST_RESORT_SUGGESTION}
+sinduran@google.com #{LAST_RESORT_SUGGESTION}
+nupursn@google.com #{LAST_RESORT_SUGGESTION}
+shreyacsingh@google.com #{LAST_RESORT_SUGGESTION}
 file:EnterprisePlatform_OWNERS
diff --git a/core/java/android/companion/virtual/VirtualDeviceManager.java b/core/java/android/companion/virtual/VirtualDeviceManager.java
index 68a864d..9e32cba 100644
--- a/core/java/android/companion/virtual/VirtualDeviceManager.java
+++ b/core/java/android/companion/virtual/VirtualDeviceManager.java
@@ -847,7 +847,7 @@
          * @see #addActivityPolicyExemption(String)
          * @see #setDevicePolicy
          */
-        @FlaggedApi(Flags.FLAG_DYNAMIC_POLICY)
+        @FlaggedApi(android.companion.virtualdevice.flags.Flags.FLAG_ACTIVITY_CONTROL_API)
         @RequiresPermission(android.Manifest.permission.CREATE_VIRTUAL_DEVICE)
         public void removeActivityPolicyExemption(@NonNull String packageName) {
             mVirtualDeviceInternal.removeActivityPolicyPackageExemption(
diff --git a/core/res/res/layout/app_language_picker_current_locale_item.xml b/core/res/res/layout/app_language_picker_current_locale_item.xml
index 01b9cc5..edd6d64 100644
--- a/core/res/res/layout/app_language_picker_current_locale_item.xml
+++ b/core/res/res/layout/app_language_picker_current_locale_item.xml
@@ -18,26 +18,31 @@
 <LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:app="http://schemas.android.com/apk/res-auto"
     android:layout_width="match_parent"
-    android:layout_height="match_parent">
-    <FrameLayout
+    android:layout_height="wrap_content"
+    android:gravity="center_vertical">
+    <RelativeLayout
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
-        android:layout_weight=".8">
+        android:layout_marginEnd="6dip"
+        android:layout_marginTop="6dip"
+        android:layout_marginBottom="6dip"
+        android:layout_weight="1">
         <include
             android:id="@+id/language_picker_item"
             layout="@layout/language_picker_item" />
-    </FrameLayout>
+    </RelativeLayout>
 
     <LinearLayout
         android:layout_width="wrap_content"
-        android:layout_height="wrap_content"
-        android:layout_weight=".2"
+        android:layout_height="match_parent"
         android:gravity="center"
         android:minHeight="?android:attr/listPreferredItemHeight">
         <ImageView
             android:id="@+id/imageView"
-            android:layout_width="24dp"
-            android:layout_height="24dp"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_centerVertical="true"
+            android:layout_marginHorizontal="16dp"
             android:src="@drawable/ic_check_24dp"
             app:tint="?attr/colorAccentPrimaryVariant"
             android:contentDescription="@*android:string/checked"/>
diff --git a/core/res/res/layout/language_picker_item.xml b/core/res/res/layout/language_picker_item.xml
index 88012a9..3e55f12 100644
--- a/core/res/res/layout/language_picker_item.xml
+++ b/core/res/res/layout/language_picker_item.xml
@@ -21,7 +21,6 @@
           android:gravity="center_vertical"
           android:minHeight="?android:attr/listPreferredItemHeight"
           android:paddingStart="?android:attr/listPreferredItemPaddingStart"
-          android:paddingEnd="?android:attr/listPreferredItemPaddingEnd"
           android:textAppearance="?android:attr/textAppearanceListItem"
           android:layoutDirection="locale"
           android:textDirection="locale"
diff --git a/libs/WindowManager/Shell/res/values/config.xml b/libs/WindowManager/Shell/res/values/config.xml
index 39f6d8c..fe8b818 100644
--- a/libs/WindowManager/Shell/res/values/config.xml
+++ b/libs/WindowManager/Shell/res/values/config.xml
@@ -183,4 +183,7 @@
     <!-- This is to be overridden to define a list of packages mapped to web links which will be
          parsed and utilized for desktop windowing's app-to-web feature. -->
     <string name="generic_links_list" translatable="false"/>
+
+    <!-- Apps that can trigger Desktop Windowing App handle Education -->
+    <string-array name="desktop_windowing_app_handle_education_allowlist_apps"></string-array>
 </resources>
diff --git a/libs/WindowManager/Shell/res/values/integers.xml b/libs/WindowManager/Shell/res/values/integers.xml
index 583bf33..300baea 100644
--- a/libs/WindowManager/Shell/res/values/integers.xml
+++ b/libs/WindowManager/Shell/res/values/integers.xml
@@ -22,4 +22,16 @@
     <integer name="bubbles_overflow_columns">4</integer>
     <!-- Maximum number of bubbles we allow in overflow before we dismiss the oldest one. -->
     <integer name="bubbles_max_overflow">16</integer>
+    <!-- App Handle Education - Minimum number of times an app should have been launched, in order
+         to be eligible to show education in it -->
+    <integer name="desktop_windowing_education_min_app_launch_count">3</integer>
+    <!-- App Handle Education - Interval at which app usage stats should be queried and updated in
+         cache periodically -->
+    <integer name="desktop_windowing_education_app_usage_cache_interval_seconds">86400</integer>
+    <!-- App Handle Education - Time interval in seconds for which we'll analyze app usage
+         stats to determine if minimum usage requirements are met.  -->
+    <integer name="desktop_windowing_education_app_launch_interval_seconds">2592000</integer>
+    <!-- App Handle Education - Required time passed in seconds since device has been setup
+         in order to be eligible to show education -->
+    <integer name="desktop_windowing_education_required_time_since_setup_seconds">604800</integer>
 </resources>
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index ce054a8..d947326 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -72,6 +72,7 @@
 import com.android.wm.shell.desktopmode.ReturnToDragStartAnimator;
 import com.android.wm.shell.desktopmode.SpringDragToDesktopTransitionHandler;
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler;
+import com.android.wm.shell.desktopmode.education.AppHandleEducationFilter;
 import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.draganddrop.GlobalDragListener;
@@ -711,6 +712,14 @@
         return new AppHandleEducationDatastoreRepository(context);
     }
 
+    @WMSingleton
+    @Provides
+    static AppHandleEducationFilter provideAppHandleEducationFilter(
+            Context context,
+            AppHandleEducationDatastoreRepository appHandleEducationDatastoreRepository) {
+        return new AppHandleEducationFilter(context, appHandleEducationDatastoreRepository);
+    }
+
     //
     // Drag and drop
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilter.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilter.kt
new file mode 100644
index 0000000..51bdb40
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilter.kt
@@ -0,0 +1,123 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode.education
+
+import android.annotation.IntegerRes
+import android.app.usage.UsageStatsManager
+import android.content.Context
+import android.os.SystemClock
+import android.provider.Settings.Secure
+import com.android.wm.shell.R
+import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository
+import com.android.wm.shell.desktopmode.education.data.WindowingEducationProto
+import java.time.Duration
+
+/** Filters incoming app handle education triggers based on set conditions. */
+class AppHandleEducationFilter(
+    private val context: Context,
+    private val appHandleEducationDatastoreRepository: AppHandleEducationDatastoreRepository
+) {
+  private val usageStatsManager =
+      context.getSystemService(Context.USAGE_STATS_SERVICE) as UsageStatsManager
+
+  /** Returns true if conditions to show app handle education are met, returns false otherwise. */
+  suspend fun shouldShowAppHandleEducation(focusAppPackageName: String): Boolean {
+    val windowingEducationProto = appHandleEducationDatastoreRepository.windowingEducationProto()
+    return isFocusAppInAllowlist(focusAppPackageName) &&
+        !isOtherEducationShowing() &&
+        hasSufficientTimeSinceSetup() &&
+        !isEducationViewedBefore(windowingEducationProto) &&
+        !isFeatureUsedBefore(windowingEducationProto) &&
+        hasMinAppUsage(windowingEducationProto, focusAppPackageName)
+  }
+
+  private fun isFocusAppInAllowlist(focusAppPackageName: String): Boolean =
+      focusAppPackageName in
+          context.resources.getStringArray(
+              R.array.desktop_windowing_app_handle_education_allowlist_apps)
+
+  // TODO: b/350953004 - Add checks based on App compat
+  // TODO: b/350951797 - Add checks based on PKT tips education
+  private fun isOtherEducationShowing(): Boolean = isTaskbarEducationShowing()
+
+  private fun isTaskbarEducationShowing(): Boolean =
+      Secure.getInt(context.contentResolver, Secure.LAUNCHER_TASKBAR_EDUCATION_SHOWING, 0) == 1
+
+  private fun hasSufficientTimeSinceSetup(): Boolean =
+      Duration.ofMillis(SystemClock.elapsedRealtime()) >
+          convertIntegerResourceToDuration(
+              R.integer.desktop_windowing_education_required_time_since_setup_seconds)
+
+  private fun isEducationViewedBefore(windowingEducationProto: WindowingEducationProto): Boolean =
+      windowingEducationProto.hasEducationViewedTimestampMillis()
+
+  private fun isFeatureUsedBefore(windowingEducationProto: WindowingEducationProto): Boolean =
+      windowingEducationProto.hasFeatureUsedTimestampMillis()
+
+  private suspend fun hasMinAppUsage(
+      windowingEducationProto: WindowingEducationProto,
+      focusAppPackageName: String
+  ): Boolean =
+      (launchCountByPackageName(windowingEducationProto)[focusAppPackageName] ?: 0) >=
+          context.resources.getInteger(R.integer.desktop_windowing_education_min_app_launch_count)
+
+  private suspend fun launchCountByPackageName(
+      windowingEducationProto: WindowingEducationProto
+  ): Map<String, Int> =
+      if (isAppUsageCacheStale(windowingEducationProto)) {
+        // Query and return user stats, update cache in datastore
+        getAndCacheAppUsageStats()
+      } else {
+        // Return cached usage stats
+        windowingEducationProto.appHandleEducation.appUsageStatsMap
+      }
+
+  private fun isAppUsageCacheStale(windowingEducationProto: WindowingEducationProto): Boolean {
+    val currentTime = currentTimeInDuration()
+    val lastUpdateTime =
+        Duration.ofMillis(
+            windowingEducationProto.appHandleEducation.appUsageStatsLastUpdateTimestampMillis)
+    val appUsageStatsCachingInterval =
+        convertIntegerResourceToDuration(
+            R.integer.desktop_windowing_education_app_usage_cache_interval_seconds)
+    return (currentTime - lastUpdateTime) > appUsageStatsCachingInterval
+  }
+
+  private suspend fun getAndCacheAppUsageStats(): Map<String, Int> {
+    val currentTime = currentTimeInDuration()
+    val appUsageStats = queryAppUsageStats()
+    appHandleEducationDatastoreRepository.updateAppUsageStats(appUsageStats, currentTime)
+    return appUsageStats
+  }
+
+  private fun queryAppUsageStats(): Map<String, Int> {
+    val endTime = currentTimeInDuration()
+    val appLaunchInterval =
+        convertIntegerResourceToDuration(
+            R.integer.desktop_windowing_education_app_launch_interval_seconds)
+    val startTime = endTime - appLaunchInterval
+
+    return usageStatsManager
+        .queryAndAggregateUsageStats(startTime.toMillis(), endTime.toMillis())
+        .mapValues { it.value.appLaunchCount }
+  }
+
+  private fun convertIntegerResourceToDuration(@IntegerRes resourceId: Int): Duration =
+      Duration.ofSeconds(context.resources.getInteger(resourceId).toLong())
+
+  private fun currentTimeInDuration(): Duration = Duration.ofMillis(System.currentTimeMillis())
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
index bf4a2ab..a7fff8a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
@@ -22,12 +22,12 @@
 import androidx.datastore.core.DataStore
 import androidx.datastore.core.DataStoreFactory
 import androidx.datastore.core.Serializer
-import androidx.datastore.dataStore
 import androidx.datastore.dataStoreFile
 import com.android.framework.protobuf.InvalidProtocolBufferException
 import com.android.internal.annotations.VisibleForTesting
 import java.io.InputStream
 import java.io.OutputStream
+import java.time.Duration
 import kotlinx.coroutines.flow.first
 
 /**
@@ -58,6 +58,24 @@
         WindowingEducationProto.getDefaultInstance()
       }
 
+  /**
+   * Updates [AppHandleEducation.appUsageStats] and
+   * [AppHandleEducation.appUsageStatsLastUpdateTimestampMillis] fields in datastore with
+   * [appUsageStats] and [appUsageStatsLastUpdateTimestamp].
+   */
+  suspend fun updateAppUsageStats(
+      appUsageStats: Map<String, Int>,
+      appUsageStatsLastUpdateTimestamp: Duration
+  ) {
+    val currentAppHandleProto = windowingEducationProto().appHandleEducation.toBuilder()
+    currentAppHandleProto
+        .putAllAppUsageStats(appUsageStats)
+        .setAppUsageStatsLastUpdateTimestampMillis(appUsageStatsLastUpdateTimestamp.toMillis())
+    dataStore.updateData { preferences: WindowingEducationProto ->
+      preferences.toBuilder().setAppHandleEducation(currentAppHandleProto).build()
+    }
+  }
+
   companion object {
     private const val TAG = "AppHandleEducationDatastoreRepository"
     private const val APP_HANDLE_EDUCATION_DATASTORE_FILEPATH = "app_handle_education.pb"
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt
index 4d40738..765021f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt
@@ -26,13 +26,16 @@
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository
 import com.android.wm.shell.desktopmode.education.data.WindowingEducationProto
+import com.android.wm.shell.util.createWindowingEducationProto
 import com.google.common.truth.Truth.assertThat
 import java.io.File
+import java.time.Duration
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
@@ -88,35 +91,22 @@
         assertThat(resultProto).isEqualTo(windowingEducationProto)
       }
 
-  private fun createWindowingEducationProto(
-      educationViewedTimestampMillis: Long? = null,
-      featureUsedTimestampMillis: Long? = null,
-      appUsageStats: Map<String, Int>? = null,
-      appUsageStatsLastUpdateTimestampMillis: Long? = null
-  ): WindowingEducationProto =
-      WindowingEducationProto.newBuilder()
-          .apply {
-            if (educationViewedTimestampMillis != null)
-                setEducationViewedTimestampMillis(educationViewedTimestampMillis)
-            if (featureUsedTimestampMillis != null)
-                setFeatureUsedTimestampMillis(featureUsedTimestampMillis)
-            setAppHandleEducation(
-                createAppHandleEducationProto(
-                    appUsageStats, appUsageStatsLastUpdateTimestampMillis))
-          }
-          .build()
+  @Test
+  fun updateAppUsageStats_updatesDatastoreProto() =
+      runTest(StandardTestDispatcher()) {
+        val appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 3)
+        val appUsageStatsLastUpdateTimestamp = Duration.ofMillis(123L)
+        val windowingEducationProto =
+            createWindowingEducationProto(
+                appUsageStats = appUsageStats,
+                appUsageStatsLastUpdateTimestampMillis =
+                    appUsageStatsLastUpdateTimestamp.toMillis())
 
-  private fun createAppHandleEducationProto(
-      appUsageStats: Map<String, Int>? = null,
-      appUsageStatsLastUpdateTimestampMillis: Long? = null
-  ): WindowingEducationProto.AppHandleEducation =
-      WindowingEducationProto.AppHandleEducation.newBuilder()
-          .apply {
-            if (appUsageStats != null) putAllAppUsageStats(appUsageStats)
-            if (appUsageStatsLastUpdateTimestampMillis != null)
-                setAppUsageStatsLastUpdateTimestampMillis(appUsageStatsLastUpdateTimestampMillis)
-          }
-          .build()
+        datastoreRepository.updateAppUsageStats(appUsageStats, appUsageStatsLastUpdateTimestamp)
+
+        val result = testDatastore.data.first()
+        assertThat(result).isEqualTo(windowingEducationProto)
+      }
 
   companion object {
     private const val GMAIL_PACKAGE_NAME = "com.google.android.gm"
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilterTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilterTest.kt
new file mode 100644
index 0000000..c0d71c0
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationFilterTest.kt
@@ -0,0 +1,193 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.desktopmode.education
+
+import android.app.usage.UsageStats
+import android.app.usage.UsageStatsManager
+import android.content.Context
+import android.testing.AndroidTestingRunner
+import android.testing.TestableContext
+import android.testing.TestableResources
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.R
+import com.android.wm.shell.ShellTestCase
+import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository
+import com.android.wm.shell.util.createWindowingEducationProto
+import com.google.common.truth.Truth.assertThat
+import kotlin.Int.Companion.MAX_VALUE
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyLong
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+class AppHandleEducationFilterTest : ShellTestCase() {
+  @Mock private lateinit var datastoreRepository: AppHandleEducationDatastoreRepository
+  @Mock private lateinit var mockUsageStatsManager: UsageStatsManager
+  private lateinit var educationFilter: AppHandleEducationFilter
+  private lateinit var testableResources: TestableResources
+  private lateinit var testableContext: TestableContext
+
+  @Before
+  fun setup() {
+    MockitoAnnotations.initMocks(this)
+    testableContext = TestableContext(mContext)
+    testableResources =
+        testableContext.orCreateTestableResources.apply {
+          addOverride(
+              R.array.desktop_windowing_app_handle_education_allowlist_apps,
+              arrayOf(GMAIL_PACKAGE_NAME))
+          addOverride(R.integer.desktop_windowing_education_required_time_since_setup_seconds, 0)
+          addOverride(R.integer.desktop_windowing_education_min_app_launch_count, 3)
+          addOverride(
+              R.integer.desktop_windowing_education_app_usage_cache_interval_seconds, MAX_VALUE)
+          addOverride(R.integer.desktop_windowing_education_app_launch_interval_seconds, 100)
+        }
+    testableContext.addMockSystemService(Context.USAGE_STATS_SERVICE, mockUsageStatsManager)
+    educationFilter = AppHandleEducationFilter(testableContext, datastoreRepository)
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_isTriggerValid_returnsTrue() = runTest {
+    // setup() makes sure that all of the conditions satisfy and #shouldShowAppHandleEducation
+    // should return true
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 4),
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    assertThat(result).isTrue()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_focusAppNotInAllowlist_returnsFalse() = runTest {
+    // Pass Youtube as current focus app, it is not in allowlist hence #shouldShowAppHandleEducation
+    // should return false
+    testableResources.addOverride(
+        R.array.desktop_windowing_app_handle_education_allowlist_apps, arrayOf(GMAIL_PACKAGE_NAME))
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(YOUTUBE_PACKAGE_NAME to 4),
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(YOUTUBE_PACKAGE_NAME)
+
+    assertThat(result).isFalse()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_timeSinceSetupIsNotSufficient_returnsFalse() = runTest {
+    // Time required to have passed setup is > 100 years, hence #shouldShowAppHandleEducation should
+    // return false
+    testableResources.addOverride(
+        R.integer.desktop_windowing_education_required_time_since_setup_seconds, MAX_VALUE)
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 4),
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    assertThat(result).isFalse()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_educationViewedBefore_returnsFalse() = runTest {
+    // Education has been viewed before, hence #shouldShowAppHandleEducation should return false
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 4),
+            educationViewedTimestampMillis = 123L,
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    assertThat(result).isFalse()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_featureUsedBefore_returnsFalse() = runTest {
+    // Feature has been used before, hence #shouldShowAppHandleEducation should return false
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 4),
+            featureUsedTimestampMillis = 123L,
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    assertThat(result).isFalse()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_doesNotHaveMinAppUsage_returnsFalse() = runTest {
+    // Simulate that gmail app has been launched twice before, minimum app launch count is 3, hence
+    // #shouldShowAppHandleEducation should return false
+    testableResources.addOverride(R.integer.desktop_windowing_education_min_app_launch_count, 3)
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 2),
+            appUsageStatsLastUpdateTimestampMillis = Long.MAX_VALUE)
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    assertThat(result).isFalse()
+  }
+
+  @Test
+  fun shouldShowAppHandleEducation_appUsageStatsStale_queryAppUsageStats() = runTest {
+    // UsageStats caching interval is set to 0ms, that means caching should happen very frequently
+    testableResources.addOverride(
+        R.integer.desktop_windowing_education_app_usage_cache_interval_seconds, 0)
+    // The DataStore currently holds a proto object where Gmail's app launch count is recorded as 4.
+    // This value exceeds the minimum required count of 3.
+    testableResources.addOverride(R.integer.desktop_windowing_education_min_app_launch_count, 3)
+    val windowingEducationProto =
+        createWindowingEducationProto(
+            appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 4),
+            appUsageStatsLastUpdateTimestampMillis = 0)
+    // The mocked UsageStatsManager is configured to return a launch count of 2 for Gmail.
+    // This value is below the minimum required count of 3.
+    `when`(mockUsageStatsManager.queryAndAggregateUsageStats(anyLong(), anyLong()))
+        .thenReturn(mapOf(GMAIL_PACKAGE_NAME to UsageStats().apply { mAppLaunchCount = 2 }))
+    `when`(datastoreRepository.windowingEducationProto()).thenReturn(windowingEducationProto)
+
+    val result = educationFilter.shouldShowAppHandleEducation(GMAIL_PACKAGE_NAME)
+
+    // Result should be false as queried usage stats should be considered to determine the result
+    // instead of cached stats
+    assertThat(result).isFalse()
+  }
+
+  companion object {
+    private const val GMAIL_PACKAGE_NAME = "com.google.android.gm"
+    private const val YOUTUBE_PACKAGE_NAME = "com.google.android.youtube"
+  }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/util/WindowingEducationProtoUtils.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/util/WindowingEducationProtoUtils.kt
new file mode 100644
index 0000000..def4b91
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/util/WindowingEducationProtoUtils.kt
@@ -0,0 +1,63 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.util
+
+import com.android.wm.shell.desktopmode.education.data.WindowingEducationProto
+
+/**
+ * Constructs a [WindowingEducationProto] object, populating its fields with the provided
+ * parameters.
+ *
+ * Any fields without corresponding parameters will retain their default values.
+ */
+fun createWindowingEducationProto(
+    educationViewedTimestampMillis: Long? = null,
+    featureUsedTimestampMillis: Long? = null,
+    appUsageStats: Map<String, Int>? = null,
+    appUsageStatsLastUpdateTimestampMillis: Long? = null
+): WindowingEducationProto =
+    WindowingEducationProto.newBuilder()
+        .apply {
+          if (educationViewedTimestampMillis != null) {
+            setEducationViewedTimestampMillis(educationViewedTimestampMillis)
+          }
+          if (featureUsedTimestampMillis != null) {
+            setFeatureUsedTimestampMillis(featureUsedTimestampMillis)
+          }
+          setAppHandleEducation(
+              createAppHandleEducationProto(appUsageStats, appUsageStatsLastUpdateTimestampMillis))
+        }
+        .build()
+
+/**
+ * Constructs a [WindowingEducationProto.AppHandleEducation] object, populating its fields with the
+ * provided parameters.
+ *
+ * Any fields without corresponding parameters will retain their default values.
+ */
+fun createAppHandleEducationProto(
+    appUsageStats: Map<String, Int>? = null,
+    appUsageStatsLastUpdateTimestampMillis: Long? = null
+): WindowingEducationProto.AppHandleEducation =
+    WindowingEducationProto.AppHandleEducation.newBuilder()
+        .apply {
+          if (appUsageStats != null) putAllAppUsageStats(appUsageStats)
+          if (appUsageStatsLastUpdateTimestampMillis != null) {
+            setAppUsageStatsLastUpdateTimestampMillis(appUsageStatsLastUpdateTimestampMillis)
+          }
+        }
+        .build()
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index c5c3a62..e619814 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -1106,7 +1106,10 @@
             val directionSign = if (transition.isUpOrLeft) -1 else 1
             val isToContent = overscroll.scene == transition.toContent
             val linearProgress = transition.progress.let { if (isToContent) it - 1f else it }
-            val progress = directionSign * overscroll.progressConverter.convert(linearProgress)
+            val progressConverter =
+                overscroll.progressConverter
+                    ?: layoutImpl.state.transitions.defaultProgressConverter
+            val progress = directionSign * progressConverter.convert(linearProgress)
             val rangeProgress = propertySpec.range?.progress(progress) ?: progress
 
             // Interpolate between the value at rest and the over scrolled value.
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
index ae5344f..8f1a4141 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/ObservableTransitionState.kt
@@ -74,6 +74,12 @@
          * the transition completes/settles.
          */
         val isUserInputOngoing: Flow<Boolean>,
+
+        /** Current progress of the preview part of the transition */
+        val previewProgress: Flow<Float> = flowOf(0f),
+
+        /** Whether the transition is currently in the preview stage or not */
+        val isInPreviewStage: Flow<Boolean> = flowOf(false),
     ) : ObservableTransitionState {
         override fun toString(): String =
             """Transition
@@ -113,6 +119,8 @@
                         progress = snapshotFlow { state.progress },
                         isInitiatedByUserInput = state.isInitiatedByUserInput,
                         isUserInputOngoing = snapshotFlow { state.isUserInputOngoing },
+                        previewProgress = snapshotFlow { state.previewProgress },
+                        isInPreviewStage = snapshotFlow { state.isInPreviewStage }
                     )
                 }
             }
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
index 2fbdf7c..cc53a28 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/PredictiveBackHandler.kt
@@ -85,7 +85,7 @@
         get() = 0f // Currently, velocity is not exposed by predictive back API
 
     override val isInPreviewStage: Boolean
-        get() = progressAnimatable == null && previewTransformationSpec != null
+        get() = previewTransformationSpec != null && currentScene == fromScene
 
     override val progress: Float
         get() = progressAnimatable?.value ?: previewTransformationSpec?.let { 0f } ?: dragProgress
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
index a063438..d35d956 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/SceneTransitions.kt
@@ -45,6 +45,7 @@
     internal val transitionSpecs: List<TransitionSpecImpl>,
     internal val overscrollSpecs: List<OverscrollSpecImpl>,
     internal val interruptionHandler: InterruptionHandler,
+    internal val defaultProgressConverter: ProgressConverter,
 ) {
     private val transitionCache =
         mutableMapOf<
@@ -147,6 +148,7 @@
                 transitionSpecs = emptyList(),
                 overscrollSpecs = emptyList(),
                 interruptionHandler = DefaultInterruptionHandler,
+                defaultProgressConverter = ProgressConverter.Default,
             )
     }
 }
@@ -282,14 +284,14 @@
      * - 1, the user overscrolled by exactly the [OverscrollBuilder.distance].
      * - Greater than 1, the user overscrolled more than the [OverscrollBuilder.distance].
      */
-    val progressConverter: ProgressConverter
+    val progressConverter: ProgressConverter?
 }
 
 internal class OverscrollSpecImpl(
     override val scene: SceneKey,
     override val orientation: Orientation,
     override val transformationSpec: TransformationSpecImpl,
-    override val progressConverter: ProgressConverter,
+    override val progressConverter: ProgressConverter?,
 ) : OverscrollSpec
 
 /**
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
index ad1fd96..e38c849 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDsl.kt
@@ -50,6 +50,12 @@
     var interruptionHandler: InterruptionHandler
 
     /**
+     * Default [ProgressConverter] used during overscroll. It lets you change a linear progress into
+     * a function of your choice. Defaults to [ProgressConverter.Default].
+     */
+    var defaultOverscrollProgressConverter: ProgressConverter
+
+    /**
      * Define the default animation to be played when transitioning [to] the specified content, from
      * any content. For the animation specification to apply only when transitioning between two
      * specific contents, use [from] instead.
@@ -217,7 +223,7 @@
      * - 1, the user overscrolled by exactly the [distance].
      * - Greater than 1, the user overscrolled more than the [distance].
      */
-    var progressConverter: ProgressConverter
+    var progressConverter: ProgressConverter?
 
     /** Translate the element(s) matching [matcher] by ([x], [y]) pixels. */
     fun translate(
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
index 771d1dd..523e5bdd7 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/TransitionDslImpl.kt
@@ -50,12 +50,14 @@
         impl.transitionSpecs,
         impl.transitionOverscrollSpecs,
         impl.interruptionHandler,
+        impl.defaultOverscrollProgressConverter,
     )
 }
 
 private class SceneTransitionsBuilderImpl : SceneTransitionsBuilder {
     override var defaultSwipeSpec: SpringSpec<Float> = SceneTransitions.DefaultSwipeSpec
     override var interruptionHandler: InterruptionHandler = DefaultInterruptionHandler
+    override var defaultOverscrollProgressConverter: ProgressConverter = ProgressConverter.Default
 
     val transitionSpecs = mutableListOf<TransitionSpecImpl>()
     val transitionOverscrollSpecs = mutableListOf<OverscrollSpecImpl>()
@@ -271,7 +273,7 @@
 }
 
 internal open class OverscrollBuilderImpl : BaseTransitionBuilderImpl(), OverscrollBuilder {
-    override var progressConverter: ProgressConverter = ProgressConverter.Default
+    override var progressConverter: ProgressConverter? = null
 
     override fun translate(
         matcher: ElementMatcher,
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
index 34e6090..20b9b49 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ElementTest.kt
@@ -961,6 +961,97 @@
     }
 
     @Test
+    fun elementTransitionWithDistanceDuringOverscrollWithDefaultProgressConverter() {
+        val layoutWidth = 200.dp
+        val layoutHeight = 400.dp
+        var animatedFloat = 0f
+        val state =
+            setupOverscrollScenario(
+                layoutWidth = layoutWidth,
+                layoutHeight = layoutHeight,
+                sceneTransitions = {
+                    // Overscroll progress will be halved
+                    defaultOverscrollProgressConverter = ProgressConverter { it / 2f }
+
+                    overscroll(SceneB, Orientation.Vertical) {
+                        // On overscroll 100% -> Foo should translate by layoutHeight
+                        translate(TestElements.Foo, y = { absoluteDistance })
+                    }
+                },
+                firstScroll = 1f, // 100% scroll
+                animatedFloatRange = 0f..100f,
+                onAnimatedFloat = { animatedFloat = it },
+            )
+
+        val fooElement = rule.onNodeWithTag(TestElements.Foo.testTag)
+        fooElement.assertTopPositionInRootIsEqualTo(0.dp)
+        assertThat(animatedFloat).isEqualTo(100f)
+
+        rule.onRoot().performTouchInput {
+            // Scroll another 100%
+            moveBy(Offset(0f, layoutHeight.toPx()), delayMillis = 1_000)
+        }
+
+        val transition = assertThat(state.transitionState).isTransition()
+        assertThat(animatedFloat).isEqualTo(100f)
+
+        // Scroll 200% (100% scroll + 100% overscroll)
+        assertThat(transition).hasProgress(2f)
+        assertThat(transition).hasOverscrollSpec()
+
+        // Overscroll progress is halved, we are at 50% of the overscroll progress.
+        fooElement.assertTopPositionInRootIsEqualTo(layoutHeight * 0.5f)
+        assertThat(animatedFloat).isEqualTo(100f)
+    }
+
+    @Test
+    fun elementTransitionWithDistanceDuringOverscrollWithOverrideDefaultProgressConverter() {
+        val layoutWidth = 200.dp
+        val layoutHeight = 400.dp
+        var animatedFloat = 0f
+        val state =
+            setupOverscrollScenario(
+                layoutWidth = layoutWidth,
+                layoutHeight = layoutHeight,
+                sceneTransitions = {
+                    // Overscroll progress will be linear (by default)
+                    defaultOverscrollProgressConverter = ProgressConverter { it }
+
+                    overscroll(SceneB, Orientation.Vertical) {
+                        // This override the defaultOverscrollProgressConverter
+                        // Overscroll progress will be halved
+                        progressConverter = ProgressConverter { it / 2f }
+                        // On overscroll 100% -> Foo should translate by layoutHeight
+                        translate(TestElements.Foo, y = { absoluteDistance })
+                    }
+                },
+                firstScroll = 1f, // 100% scroll
+                animatedFloatRange = 0f..100f,
+                onAnimatedFloat = { animatedFloat = it },
+            )
+
+        val fooElement = rule.onNodeWithTag(TestElements.Foo.testTag)
+        fooElement.assertTopPositionInRootIsEqualTo(0.dp)
+        assertThat(animatedFloat).isEqualTo(100f)
+
+        rule.onRoot().performTouchInput {
+            // Scroll another 100%
+            moveBy(Offset(0f, layoutHeight.toPx()), delayMillis = 1_000)
+        }
+
+        val transition = assertThat(state.transitionState).isTransition()
+        assertThat(animatedFloat).isEqualTo(100f)
+
+        // Scroll 200% (100% scroll + 100% overscroll)
+        assertThat(transition).hasProgress(2f)
+        assertThat(transition).hasOverscrollSpec()
+
+        // Overscroll progress is halved, we are at 50% of the overscroll progress.
+        fooElement.assertTopPositionInRootIsEqualTo(layoutHeight * 0.5f)
+        assertThat(animatedFloat).isEqualTo(100f)
+    }
+
+    @Test
     fun elementTransitionWithDistanceDuringOverscrollWithProgressConverter() {
         val layoutWidth = 200.dp
         val layoutHeight = 400.dp
diff --git a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
index f717301..0543e7f 100644
--- a/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
+++ b/packages/SystemUI/compose/scene/tests/src/com/android/compose/animation/scene/ObservableTransitionStateTest.kt
@@ -16,11 +16,16 @@
 
 package com.android.compose.animation.scene
 
+import androidx.activity.BackEventCompat
+import androidx.activity.ComponentActivity
+import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.setValue
 import androidx.compose.runtime.snapshots.Snapshot
-import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.test.junit4.createAndroidComposeRule
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.compose.animation.scene.TestScenes.SceneA
 import com.android.compose.animation.scene.TestScenes.SceneB
@@ -36,7 +41,7 @@
 
 @RunWith(AndroidJUnit4::class)
 class ObservableTransitionStateTest {
-    @get:Rule val rule = createComposeRule()
+    @get:Rule val rule = createAndroidComposeRule<ComponentActivity>()
 
     @Test
     fun testObservableTransitionState() = runTest {
@@ -145,6 +150,82 @@
         assertThat(currentScene.value).isEqualTo(SceneA)
     }
 
+    @Test
+    fun testObservablePreviewTransitionState() = runTest {
+        val layoutState =
+            rule.runOnUiThread {
+                MutableSceneTransitionLayoutState(
+                    SceneA,
+                    transitions = transitions { from(SceneA, to = SceneB, preview = {}) }
+                )
+            }
+        rule.setContent {
+            SceneTransitionLayout(layoutState) {
+                scene(SceneA, mapOf(Back to SceneB)) { Box(Modifier.fillMaxSize()) }
+                scene(SceneB) { Box(Modifier.fillMaxSize()) }
+            }
+        }
+
+        var observableState: ObservableTransitionState? = null
+        backgroundScope.launch {
+            layoutState.observableTransitionState().collect { observableState = it }
+        }
+
+        fun observableState(): ObservableTransitionState {
+            runCurrent()
+            return observableState!!
+        }
+
+        fun ObservableTransitionState.Transition.previewProgress(): Float {
+            var lastProgress = -1f
+            backgroundScope.launch { previewProgress.collect { lastProgress = it } }
+            runCurrent()
+            return lastProgress
+        }
+
+        fun ObservableTransitionState.Transition.isInPreviewStage(): Boolean {
+            var lastIsInPreviewStage = false
+            backgroundScope.launch { isInPreviewStage.collect { lastIsInPreviewStage = it } }
+            runCurrent()
+            return lastIsInPreviewStage
+        }
+
+        // Start back.
+        val dispatcher = rule.activity.onBackPressedDispatcher
+        rule.runOnUiThread {
+            dispatcher.dispatchOnBackStarted(backEvent())
+            dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
+        }
+
+        var state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
+        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat(state.previewProgress()).isEqualTo(0.4f)
+        assertThat(state.isInPreviewStage()).isEqualTo(true)
+
+        // Cancel it.
+        rule.runOnUiThread { dispatcher.dispatchOnBackCancelled() }
+        rule.waitForIdle()
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Idle::class.java)
+
+        // Start again and commit it.
+        rule.runOnUiThread {
+            dispatcher.dispatchOnBackStarted(backEvent())
+            dispatcher.dispatchOnBackProgressed(backEvent(progress = 0.4f))
+            dispatcher.onBackPressed()
+        }
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Transition::class.java)
+        assertThat((state as ObservableTransitionState.Transition).fromScene).isEqualTo(SceneA)
+        assertThat(state.previewProgress()).isEqualTo(0.4f)
+        assertThat(state.isInPreviewStage()).isEqualTo(false)
+
+        rule.waitForIdle()
+        state = observableState()
+        assertThat(state).isInstanceOf(ObservableTransitionState.Idle::class.java)
+    }
+
     // See http://shortn/_hj4Mhikmos for inspiration.
     private fun runTestWithSnapshots(testBody: suspend TestScope.() -> Unit) {
         val globalWriteObserverHandle =
@@ -159,4 +240,13 @@
             globalWriteObserverHandle.dispose()
         }
     }
+
+    private fun backEvent(progress: Float = 0f): BackEventCompat {
+        return BackEventCompat(
+            touchX = 0f,
+            touchY = 0f,
+            progress = progress,
+            swipeEdge = BackEventCompat.EDGE_LEFT,
+        )
+    }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
index 2f17ca8..53d3c01 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/scene/data/repository/SceneContainerRepositoryUtil.kt
@@ -69,6 +69,8 @@
     progress: Flow<Float> = flowOf(0f),
     isInitiatedByUserInput: Boolean = false,
     isUserInputOngoing: Flow<Boolean> = flowOf(false),
+    previewProgress: Flow<Float> = flowOf(0f),
+    isInPreviewStage: Flow<Boolean> = flowOf(false)
 ): ObservableTransitionState.Transition {
     return ObservableTransitionState.Transition(
         fromScene = from,
@@ -76,7 +78,9 @@
         currentScene = currentScene,
         progress = progress,
         isInitiatedByUserInput = isInitiatedByUserInput,
-        isUserInputOngoing = isUserInputOngoing
+        isUserInputOngoing = isUserInputOngoing,
+        previewProgress = previewProgress,
+        isInPreviewStage = isInPreviewStage
     )
 }
 
diff --git a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
index e91097c..1c786e6 100644
--- a/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
+++ b/services/core/java/com/android/server/rollback/RollbackPackageHealthObserver.java
@@ -492,23 +492,19 @@
         PackageManager pm = mContext.getPackageManager();
 
         if (Flags.refactorCrashrecovery() && provideInfoOfApkInApex()) {
-            // Check if the package is listed among the system modules.
-            boolean isApex = false;
-            try {
-                isApex = (pm.getModuleInfo(packageName, 0 /* flags */) != null);
-            } catch (PackageManager.NameNotFoundException e) {
-                //pass
-            }
-
-            // Check if the package is an APK inside an APEX.
-            boolean isApkInApex = false;
+            // Check if the package is listed among the system modules or is an
+            // APK inside an updatable APEX.
             try {
                 final PackageInfo pkg = pm.getPackageInfo(packageName, 0 /* flags */);
-                isApkInApex = (pkg.getApexPackageName() != null);
+                String apexPackageName = pkg.getApexPackageName();
+                if (apexPackageName != null) {
+                    packageName = apexPackageName;
+                }
+
+                return pm.getModuleInfo(packageName, 0 /* flags */) != null;
             } catch (PackageManager.NameNotFoundException e) {
-                // pass
+                return false;
             }
-            return isApex || isApkInApex;
         } else {
             // Check if the package is an APK inside an APEX. If it is, use the parent APEX package
             // when querying PackageManager.
diff --git a/services/core/java/com/android/server/wm/DimmerAnimationHelper.java b/services/core/java/com/android/server/wm/DimmerAnimationHelper.java
index 3dba57f..4abf806 100644
--- a/services/core/java/com/android/server/wm/DimmerAnimationHelper.java
+++ b/services/core/java/com/android/server/wm/DimmerAnimationHelper.java
@@ -56,9 +56,14 @@
         Change() {}
 
         Change(@NonNull Change other) {
+            copyFrom(other);
+        }
+
+        void copyFrom(@NonNull Change other) {
             mAlpha = other.mAlpha;
             mBlurRadius = other.mBlurRadius;
             mDimmingContainer = other.mDimmingContainer;
+            mGeometryParent = other.mGeometryParent;
             mRelativeLayer = other.mRelativeLayer;
         }
 
@@ -83,8 +88,8 @@
         }
     }
 
-    private Change mCurrentProperties = new Change();
-    private Change mRequestedProperties = new Change();
+    private final Change mCurrentProperties = new Change();
+    private final Change mRequestedProperties = new Change();
     private AnimationSpec mAlphaAnimationSpec;
 
     private final AnimationAdapterFactory mAnimationAdapterFactory;
@@ -123,12 +128,15 @@
      * {@link Change#setRequestedAppearance(float, int)}
      */
     void applyChanges(@NonNull SurfaceControl.Transaction t, @NonNull Dimmer.DimState dim) {
+        final Change startProperties = new Change(mCurrentProperties);
+        mCurrentProperties.copyFrom(mRequestedProperties);
+
         if (mRequestedProperties.mDimmingContainer == null) {
             Log.e(TAG, this + " does not have a dimming container. Have you forgotten to "
                     + "call adjustRelativeLayer?");
             return;
         }
-        if (mRequestedProperties.mDimmingContainer.mSurfaceControl == null) {
+        if (mRequestedProperties.mDimmingContainer.getSurfaceControl() == null) {
             Log.w(TAG, "container " + mRequestedProperties.mDimmingContainer
                     + "does not have a surface");
             dim.remove(t);
@@ -137,52 +145,49 @@
 
         dim.ensureVisible(t);
         reparent(dim.mDimSurface,
-                mRequestedProperties.mGeometryParent != mCurrentProperties.mGeometryParent
+                startProperties.mGeometryParent != mRequestedProperties.mGeometryParent
                         ? mRequestedProperties.mGeometryParent.getSurfaceControl() : null,
                 mRequestedProperties.mDimmingContainer.getSurfaceControl(),
                 mRequestedProperties.mRelativeLayer, t);
 
-        if (!mCurrentProperties.hasSameVisualProperties(mRequestedProperties)) {
+        if (!startProperties.hasSameVisualProperties(mRequestedProperties)) {
             stopCurrentAnimation(dim.mDimSurface);
 
             if (dim.mSkipAnimation
                     // If the container doesn't change but requests a dim change, then it is
                     // directly providing us the animated values
-                    || (mRequestedProperties.hasSameDimmingContainer(mCurrentProperties)
+                    || (startProperties.hasSameDimmingContainer(mRequestedProperties)
                     && dim.isDimming())) {
                 ProtoLog.d(WM_DEBUG_DIMMER,
                         "%s skipping animation and directly setting alpha=%f, blur=%d",
-                        dim, mRequestedProperties.mAlpha,
+                        dim, startProperties.mAlpha,
                         mRequestedProperties.mBlurRadius);
-                setAlphaBlur(dim.mDimSurface, mRequestedProperties.mAlpha,
-                        mRequestedProperties.mBlurRadius, t);
+                setCurrentAlphaBlur(dim.mDimSurface, t);
                 dim.mSkipAnimation = false;
             } else {
-                startAnimation(t, dim);
+                startAnimation(t, dim, startProperties, mRequestedProperties);
             }
-
         } else if (!dim.isDimming()) {
             // We are not dimming, so we tried the exit animation but the alpha is already 0,
             // therefore, let's just remove this surface
             dim.remove(t);
         }
-        mCurrentProperties = new Change(mRequestedProperties);
     }
 
     private void startAnimation(
-            @NonNull SurfaceControl.Transaction t, @NonNull Dimmer.DimState dim) {
+            @NonNull SurfaceControl.Transaction t, @NonNull Dimmer.DimState dim,
+            @NonNull Change from, @NonNull Change to) {
         ProtoLog.v(WM_DEBUG_DIMMER, "Starting animation on %s", dim);
-        mAlphaAnimationSpec = getRequestedAnimationSpec();
+        mAlphaAnimationSpec = getRequestedAnimationSpec(from, to);
         mLocalAnimationAdapter = mAnimationAdapterFactory.get(mAlphaAnimationSpec,
                 dim.mHostContainer.mWmService.mSurfaceAnimationRunner);
 
-        float targetAlpha = mRequestedProperties.mAlpha;
-        int targetBlur = mRequestedProperties.mBlurRadius;
+        float targetAlpha = to.mAlpha;
 
         mLocalAnimationAdapter.startAnimation(dim.mDimSurface, t,
                 ANIMATION_TYPE_DIMMER, /* finishCallback */ (type, animator) -> {
                     synchronized (dim.mHostContainer.mWmService.mGlobalLock) {
-                        setAlphaBlur(dim.mDimSurface, targetAlpha, targetBlur, t);
+                        setCurrentAlphaBlur(dim.mDimSurface, t);
                         if (targetAlpha == 0f && !dim.isDimming()) {
                             dim.remove(t);
                         }
@@ -207,15 +212,15 @@
     }
 
     @NonNull
-    private AnimationSpec getRequestedAnimationSpec() {
-        final float startAlpha = Math.max(mCurrentProperties.mAlpha, 0f);
-        final int startBlur = Math.max(mCurrentProperties.mBlurRadius, 0);
-        long duration = (long) (getDimDuration(mRequestedProperties.mDimmingContainer)
-                * Math.abs(mRequestedProperties.mAlpha - startAlpha));
+    private static AnimationSpec getRequestedAnimationSpec(Change from, Change to) {
+        final float startAlpha = Math.max(from.mAlpha, 0f);
+        final int startBlur = Math.max(from.mBlurRadius, 0);
+        long duration = (long) (getDimDuration(to.mDimmingContainer)
+                * Math.abs(to.mAlpha - startAlpha));
 
         final AnimationSpec spec =  new AnimationSpec(
-                new AnimationSpec.AnimationExtremes<>(startAlpha, mRequestedProperties.mAlpha),
-                new AnimationSpec.AnimationExtremes<>(startBlur, mRequestedProperties.mBlurRadius),
+                new AnimationSpec.AnimationExtremes<>(startAlpha, to.mAlpha),
+                new AnimationSpec.AnimationExtremes<>(startBlur, to.mBlurRadius),
                 duration
         );
         ProtoLog.v(WM_DEBUG_DIMMER, "Dim animation requested: %s", spec);
@@ -225,7 +230,7 @@
     /**
      * Change the geometry and relative parent of this dim layer
      */
-    void reparent(@NonNull SurfaceControl dimLayer,
+    static void reparent(@NonNull SurfaceControl dimLayer,
                   @Nullable SurfaceControl newGeometryParent,
                   @NonNull SurfaceControl relativeParent,
                   int relativePosition,
@@ -240,17 +245,16 @@
         }
     }
 
-    void setAlphaBlur(@NonNull SurfaceControl sc, float alpha, int blur,
-                      @NonNull SurfaceControl.Transaction t) {
+    void setCurrentAlphaBlur(@NonNull SurfaceControl sc, @NonNull SurfaceControl.Transaction t) {
         try {
-            t.setAlpha(sc, alpha);
-            t.setBackgroundBlurRadius(sc, blur);
+            t.setAlpha(sc, mCurrentProperties.mAlpha);
+            t.setBackgroundBlurRadius(sc, mCurrentProperties.mBlurRadius);
         } catch (NullPointerException e) {
             Log.w(TAG , "Tried to change look of dim " + sc + " after remove",  e);
         }
     }
 
-    private long getDimDuration(@NonNull WindowContainer<?> container) {
+    private static long getDimDuration(@NonNull WindowContainer<?> container) {
         // Use the same duration as the animation on the WindowContainer
         AnimationAdapter animationAdapter = container.mSurfaceAnimator.getAnimation();
         final float durationScale = container.mWmService.getTransitionAnimationScaleLocked();