diff --git a/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java b/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
index c6537c0..b6ea9dd 100644
--- a/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/app/ChooserActivityWorkProfileTest.java
@@ -16,11 +16,14 @@
 
 package com.android.internal.app;
 
+import static android.util.PollingCheck.waitFor;
+
 import static androidx.test.espresso.Espresso.onView;
 import static androidx.test.espresso.action.ViewActions.click;
 import static androidx.test.espresso.action.ViewActions.swipeUp;
 import static androidx.test.espresso.assertion.ViewAssertions.matches;
 import static androidx.test.espresso.matcher.ViewMatchers.isDisplayed;
+import static androidx.test.espresso.matcher.ViewMatchers.isSelected;
 import static androidx.test.espresso.matcher.ViewMatchers.withId;
 import static androidx.test.espresso.matcher.ViewMatchers.withText;
 
@@ -49,6 +52,8 @@
 import com.android.internal.app.ResolverActivity.ResolvedComponentInfo;
 import com.android.internal.app.ChooserActivityWorkProfileTest.TestCase.Tab;
 
+import junit.framework.AssertionFailedError;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -331,8 +336,17 @@
         final int stringId = tab == Tab.WORK ? R.string.resolver_work_tab
                 : R.string.resolver_personal_tab;
 
-        onView(withText(stringId)).perform(click());
-        waitForIdle();
+        waitFor(() -> {
+            onView(withText(stringId)).perform(click());
+            waitForIdle();
+
+            try {
+                onView(withText(stringId)).check(matches(isSelected()));
+                return true;
+            } catch (AssertionFailedError e) {
+                return false;
+            }
+        });
 
         onView(withId(R.id.contentPanel))
                 .perform(swipeUp());
diff --git a/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java b/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
index ce68906..4a79f5e 100644
--- a/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
+++ b/core/tests/coretests/src/com/android/internal/app/ResolverActivityWorkProfileTest.java
@@ -16,11 +16,14 @@
 
 package com.android.internal.app;
 
+import static android.util.PollingCheck.waitFor;
+
 import static androidx.test.espresso.Espresso.onView;
 import static androidx.test.espresso.action.ViewActions.click;
 import static androidx.test.espresso.action.ViewActions.swipeUp;
 import static androidx.test.espresso.assertion.ViewAssertions.matches;
 import static androidx.test.espresso.matcher.ViewMatchers.isDisplayed;
+import static androidx.test.espresso.matcher.ViewMatchers.isSelected;
 import static androidx.test.espresso.matcher.ViewMatchers.withId;
 import static androidx.test.espresso.matcher.ViewMatchers.withText;
 
@@ -48,6 +51,8 @@
 import com.android.internal.app.ResolverActivity.ResolvedComponentInfo;
 import com.android.internal.app.ResolverActivityWorkProfileTest.TestCase.Tab;
 
+import junit.framework.AssertionFailedError;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -309,8 +314,17 @@
         final int stringId = tab == Tab.WORK ? R.string.resolver_work_tab
                 : R.string.resolver_personal_tab;
 
-        onView(withText(stringId)).perform(click());
-        waitForIdle();
+        waitFor(() -> {
+            onView(withText(stringId)).perform(click());
+            waitForIdle();
+
+            try {
+                onView(withText(stringId)).check(matches(isSelected()));
+                return true;
+            } catch (AssertionFailedError e) {
+                return false;
+            }
+        });
 
         onView(withId(R.id.contentPanel))
                 .perform(swipeUp());
diff --git a/libs/WindowManager/Jetpack/Android.bp b/libs/WindowManager/Jetpack/Android.bp
index dc4b563..a5b192c 100644
--- a/libs/WindowManager/Jetpack/Android.bp
+++ b/libs/WindowManager/Jetpack/Android.bp
@@ -63,6 +63,12 @@
     sdk_version: "current",
 }
 
+android_library_import {
+    name: "window-extensions-core",
+    aars: ["window-extensions-core-release.aar"],
+    sdk_version: "current",
+}
+
 java_library {
     name: "androidx.window.extensions",
     srcs: [
@@ -70,7 +76,10 @@
         "src/androidx/window/util/**/*.java",
         "src/androidx/window/common/**/*.java",
     ],
-    static_libs: ["window-extensions"],
+    static_libs: [
+        "window-extensions",
+        "window-extensions-core",
+    ],
     installable: true,
     sdk_version: "core_platform",
     system_ext_specific: true,
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
index 1af1313..85812ea 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitController.java
@@ -75,7 +75,10 @@
 import androidx.annotation.Nullable;
 import androidx.window.common.CommonFoldingFeature;
 import androidx.window.common.EmptyLifecycleCallbacksAdapter;
+import androidx.window.extensions.WindowExtensionsImpl;
 import androidx.window.extensions.WindowExtensionsProvider;
+import androidx.window.extensions.core.util.function.Consumer;
+import androidx.window.extensions.core.util.function.Function;
 import androidx.window.extensions.embedding.TransactionManager.TransactionRecord;
 import androidx.window.extensions.layout.WindowLayoutComponentImpl;
 
@@ -86,7 +89,6 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.Executor;
-import java.util.function.Consumer;
 
 /**
  * Main controller class that manages split states and presentation.
@@ -112,7 +114,7 @@
     /**
      * A developer-defined {@link SplitAttributes} calculator to compute the current
      * {@link SplitAttributes} with the current device and window states.
-     * It is registered via {@link #setSplitAttributesCalculator(SplitAttributesCalculator)}
+     * It is registered via {@link #setSplitAttributesCalculator(Function)}
      * and unregistered via {@link #clearSplitAttributesCalculator()}.
      * This is called when:
      * <ul>
@@ -125,7 +127,7 @@
      */
     @GuardedBy("mLock")
     @Nullable
-    private SplitAttributesCalculator mSplitAttributesCalculator;
+    private Function<SplitAttributesCalculatorParams, SplitAttributes> mSplitAttributesCalculator;
 
     /**
      * Map from Task id to {@link TaskContainer} which contains all TaskFragment and split pair info
@@ -138,6 +140,7 @@
     final SparseArray<TaskContainer> mTaskContainers = new SparseArray<>();
 
     /** Callback to Jetpack to notify about changes to split states. */
+    @GuardedBy("mLock")
     @Nullable
     private Consumer<List<SplitInfo>> mEmbeddingCallback;
     private final List<SplitInfo> mLastReportedSplitStates = new ArrayList<>();
@@ -171,7 +174,8 @@
         mWindowLayoutComponent.addFoldingStateChangedCallback(new FoldingFeatureListener());
     }
 
-    private class FoldingFeatureListener implements Consumer<List<CommonFoldingFeature>> {
+    private class FoldingFeatureListener
+            implements java.util.function.Consumer<List<CommonFoldingFeature>> {
         @Override
         public void accept(List<CommonFoldingFeature> foldingFeatures) {
             synchronized (mLock) {
@@ -212,7 +216,8 @@
     }
 
     @Override
-    public void setSplitAttributesCalculator(@NonNull SplitAttributesCalculator calculator) {
+    public void setSplitAttributesCalculator(
+            @NonNull Function<SplitAttributesCalculatorParams, SplitAttributes> calculator) {
         synchronized (mLock) {
             mSplitAttributesCalculator = calculator;
         }
@@ -227,7 +232,7 @@
 
     @GuardedBy("mLock")
     @Nullable
-    SplitAttributesCalculator getSplitAttributesCalculator() {
+    Function<SplitAttributesCalculatorParams, SplitAttributes> getSplitAttributesCalculator() {
         return mSplitAttributesCalculator;
     }
 
@@ -240,9 +245,22 @@
 
     /**
      * Registers the split organizer callback to notify about changes to active splits.
+     * @deprecated Use {@link #setSplitInfoCallback(Consumer)} starting with
+     * {@link WindowExtensionsImpl#getVendorApiLevel()} 2.
      */
+    @Deprecated
     @Override
-    public void setSplitInfoCallback(@NonNull Consumer<List<SplitInfo>> callback) {
+    public void setSplitInfoCallback(
+            @NonNull java.util.function.Consumer<List<SplitInfo>> callback) {
+        Consumer<List<SplitInfo>> oemConsumer = callback::accept;
+        setSplitInfoCallback(oemConsumer);
+    }
+
+    /**
+     * Registers the split organizer callback to notify about changes to active splits.
+     * @since {@link WindowExtensionsImpl#getVendorApiLevel()} 2
+     */
+    public void setSplitInfoCallback(Consumer<List<SplitInfo>> callback) {
         synchronized (mLock) {
             mEmbeddingCallback = callback;
             updateCallbackIfNecessary();
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
index 9db9f87..41580a0 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/embedding/SplitPresenter.java
@@ -42,11 +42,11 @@
 import androidx.annotation.IntDef;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.window.extensions.core.util.function.Function;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType.ExpandContainersSplitType;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType.HingeSplitType;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType.RatioSplitType;
-import androidx.window.extensions.embedding.SplitAttributesCalculator.SplitAttributesCalculatorParams;
 import androidx.window.extensions.embedding.TaskContainer.TaskProperties;
 import androidx.window.extensions.layout.DisplayFeature;
 import androidx.window.extensions.layout.FoldingFeature;
@@ -522,7 +522,8 @@
             @NonNull SplitRule rule, @Nullable Pair<Size, Size> minDimensionsPair) {
         final Configuration taskConfiguration = taskProperties.getConfiguration();
         final WindowMetrics taskWindowMetrics = getTaskWindowMetrics(taskConfiguration);
-        final SplitAttributesCalculator calculator = mController.getSplitAttributesCalculator();
+        final Function<SplitAttributesCalculatorParams, SplitAttributes> calculator =
+                mController.getSplitAttributesCalculator();
         final SplitAttributes defaultSplitAttributes = rule.getDefaultSplitAttributes();
         final boolean isDefaultMinSizeSatisfied = rule.checkParentMetrics(taskWindowMetrics);
         if (calculator == null) {
@@ -538,7 +539,7 @@
         final SplitAttributesCalculatorParams params = new SplitAttributesCalculatorParams(
                 taskWindowMetrics, taskConfiguration, defaultSplitAttributes,
                 isDefaultMinSizeSatisfied, windowLayoutInfo, rule.getTag());
-        final SplitAttributes splitAttributes = calculator.computeSplitAttributesForParams(params);
+        final SplitAttributes splitAttributes = calculator.apply(params);
         return sanitizeSplitAttributes(taskProperties, splitAttributes, minDimensionsPair);
     }
 
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
index 84b2bfc..32b915c 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/extensions/layout/WindowLayoutComponentImpl.java
@@ -44,6 +44,7 @@
 import androidx.window.common.DeviceStateManagerFoldingFeatureProducer;
 import androidx.window.common.EmptyLifecycleCallbacksAdapter;
 import androidx.window.common.RawFoldingFeatureProducer;
+import androidx.window.extensions.core.util.function.Consumer;
 import androidx.window.util.DataProducer;
 
 import java.util.ArrayList;
@@ -51,7 +52,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Consumer;
 
 /**
  * Reference implementation of androidx.window.extensions.layout OEM interface for use with
@@ -80,6 +80,10 @@
     private final Map<IBinder, ConfigurationChangeListener> mConfigurationChangeListeners =
             new ArrayMap<>();
 
+    @GuardedBy("mLock")
+    private final Map<java.util.function.Consumer<WindowLayoutInfo>, Consumer<WindowLayoutInfo>>
+            mJavaToExtConsumers = new ArrayMap<>();
+
     public WindowLayoutComponentImpl(@NonNull Context context) {
         ((Application) context.getApplicationContext())
                 .registerActivityLifecycleCallbacks(new NotifyOnConfigurationChanged());
@@ -90,7 +94,8 @@
     }
 
     /** Registers to listen to {@link CommonFoldingFeature} changes */
-    public void addFoldingStateChangedCallback(Consumer<List<CommonFoldingFeature>> consumer) {
+    public void addFoldingStateChangedCallback(
+            java.util.function.Consumer<List<CommonFoldingFeature>> consumer) {
         synchronized (mLock) {
             mFoldingFeatureProducer.addDataChangedCallback(consumer);
         }
@@ -104,13 +109,17 @@
      */
     @Override
     public void addWindowLayoutInfoListener(@NonNull Activity activity,
-            @NonNull Consumer<WindowLayoutInfo> consumer) {
-        addWindowLayoutInfoListener((Context) activity, consumer);
+            @NonNull java.util.function.Consumer<WindowLayoutInfo> consumer) {
+        final Consumer<WindowLayoutInfo> extConsumer = consumer::accept;
+        synchronized (mLock) {
+            mJavaToExtConsumers.put(consumer, extConsumer);
+        }
+        addWindowLayoutInfoListener(activity, extConsumer);
     }
 
     /**
-     * Similar to {@link #addWindowLayoutInfoListener(Activity, Consumer)}, but takes a UI Context
-     * as a parameter.
+     * Similar to {@link #addWindowLayoutInfoListener(Activity, java.util.function.Consumer)}, but
+     * takes a UI Context as a parameter.
      *
      * Jetpack {@link androidx.window.layout.ExtensionWindowLayoutInfoBackend} makes sure all
      * consumers related to the same {@link Context} gets updated {@link WindowLayoutInfo}
@@ -151,6 +160,18 @@
         }
     }
 
+    @Override
+    public void removeWindowLayoutInfoListener(
+            @NonNull java.util.function.Consumer<WindowLayoutInfo> consumer) {
+        final Consumer<WindowLayoutInfo> extConsumer;
+        synchronized (mLock) {
+            extConsumer = mJavaToExtConsumers.remove(consumer);
+        }
+        if (extConsumer != null) {
+            removeWindowLayoutInfoListener(extConsumer);
+        }
+    }
+
     /**
      * Removes a listener no longer interested in receiving updates.
      *
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
index 2f92a57..459ec9f 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/EmbeddingTestUtils.java
@@ -34,9 +34,11 @@
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.util.Pair;
+import android.view.WindowMetrics;
 import android.window.TaskFragmentInfo;
 import android.window.WindowContainerToken;
 
+import androidx.window.extensions.core.util.function.Predicate;
 import androidx.window.extensions.embedding.SplitAttributes.SplitType;
 import androidx.window.extensions.layout.DisplayFeature;
 import androidx.window.extensions.layout.FoldingFeature;
@@ -107,7 +109,7 @@
     static SplitRule createSplitRule(@NonNull Activity primaryActivity,
             @NonNull Intent secondaryIntent, boolean clearTop) {
         final Pair<Activity, Intent> targetPair = new Pair<>(primaryActivity, secondaryIntent);
-        return new SplitPairRule.Builder(
+        return createSplitPairRuleBuilder(
                 activityPair -> false,
                 targetPair::equals,
                 w -> true)
@@ -144,7 +146,7 @@
             @NonNull Activity secondaryActivity, int finishPrimaryWithSecondary,
             int finishSecondaryWithPrimary, boolean clearTop) {
         final Pair<Activity, Activity> targetPair = new Pair<>(primaryActivity, secondaryActivity);
-        return new SplitPairRule.Builder(
+        return createSplitPairRuleBuilder(
                 targetPair::equals,
                 activityIntentPair -> false,
                 w -> true)
@@ -223,4 +225,26 @@
         displayFeatures.add(foldingFeature);
         return new WindowLayoutInfo(displayFeatures);
     }
+
+    static ActivityRule.Builder createActivityBuilder(
+            @NonNull Predicate<Activity> activityPredicate,
+            @NonNull Predicate<Intent> intentPredicate) {
+        return new ActivityRule.Builder(activityPredicate, intentPredicate);
+    }
+
+    static SplitPairRule.Builder createSplitPairRuleBuilder(
+            @NonNull Predicate<Pair<Activity, Activity>> activitiesPairPredicate,
+            @NonNull Predicate<Pair<Activity, Intent>> activityIntentPairPredicate,
+            @NonNull Predicate<WindowMetrics> windowMetricsPredicate) {
+        return new SplitPairRule.Builder(activitiesPairPredicate, activityIntentPairPredicate,
+                windowMetricsPredicate);
+    }
+
+    static SplitPlaceholderRule.Builder createSplitPlaceholderRuleBuilder(
+            @NonNull Intent placeholderIntent, @NonNull Predicate<Activity> activityPredicate,
+            @NonNull Predicate<Intent> intentPredicate,
+            @NonNull Predicate<WindowMetrics> windowMetricsPredicate) {
+        return new SplitPlaceholderRule.Builder(placeholderIntent, activityPredicate,
+                intentPredicate, windowMetricsPredicate);
+    }
 }
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
index 3cc31f9..f9cbb5c 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitControllerTest.java
@@ -34,8 +34,11 @@
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.TASK_BOUNDS;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.TASK_ID;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.TEST_TAG;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createActivityBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createActivityInfoWithMinDimensions;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createMockTaskFragmentInfo;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPairRuleBuilder;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPlaceholderRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitRule;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createTestTaskContainer;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.getSplitBounds;
@@ -429,7 +432,7 @@
     @Test
     public void testResolveStartActivityIntent_withoutLaunchingActivity() {
         final Intent intent = new Intent();
-        final ActivityRule expandRule = new ActivityRule.Builder(r -> false, i -> i == intent)
+        final ActivityRule expandRule = createActivityBuilder(r -> false, i -> i == intent)
                 .setShouldAlwaysExpand(true)
                 .build();
         mSplitController.setEmbeddingRules(Collections.singleton(expandRule));
@@ -1167,7 +1170,7 @@
 
     @Test
     public void testHasSamePresentation() {
-        SplitPairRule splitRule1 = new SplitPairRule.Builder(
+        SplitPairRule splitRule1 = createSplitPairRuleBuilder(
                 activityPair -> true,
                 activityIntentPair -> true,
                 windowMetrics -> true)
@@ -1175,7 +1178,7 @@
                 .setFinishPrimaryWithSecondary(DEFAULT_FINISH_PRIMARY_WITH_SECONDARY)
                 .setDefaultSplitAttributes(SPLIT_ATTRIBUTES)
                 .build();
-        SplitPairRule splitRule2 = new SplitPairRule.Builder(
+        SplitPairRule splitRule2 = createSplitPairRuleBuilder(
                 activityPair -> true,
                 activityIntentPair -> true,
                 windowMetrics -> true)
@@ -1188,7 +1191,7 @@
                 SplitController.haveSamePresentation(splitRule1, splitRule2,
                         new WindowMetrics(TASK_BOUNDS, WindowInsets.CONSUMED)));
 
-        splitRule2 = new SplitPairRule.Builder(
+        splitRule2 = createSplitPairRuleBuilder(
                 activityPair -> true,
                 activityIntentPair -> true,
                 windowMetrics -> true)
@@ -1352,7 +1355,7 @@
 
     /** Setups a rule to always expand the given intent. */
     private void setupExpandRule(@NonNull Intent expandIntent) {
-        final ActivityRule expandRule = new ActivityRule.Builder(r -> false, expandIntent::equals)
+        final ActivityRule expandRule = createActivityBuilder(r -> false, expandIntent::equals)
                 .setShouldAlwaysExpand(true)
                 .build();
         mSplitController.setEmbeddingRules(Collections.singleton(expandRule));
@@ -1360,7 +1363,7 @@
 
     /** Setups a rule to always expand the given activity. */
     private void setupExpandRule(@NonNull Activity expandActivity) {
-        final ActivityRule expandRule = new ActivityRule.Builder(expandActivity::equals, i -> false)
+        final ActivityRule expandRule = createActivityBuilder(expandActivity::equals, i -> false)
                 .setShouldAlwaysExpand(true)
                 .build();
         mSplitController.setEmbeddingRules(Collections.singleton(expandRule));
@@ -1368,7 +1371,7 @@
 
     /** Setups a rule to launch placeholder for the given activity. */
     private void setupPlaceholderRule(@NonNull Activity primaryActivity) {
-        final SplitRule placeholderRule = new SplitPlaceholderRule.Builder(PLACEHOLDER_INTENT,
+        final SplitRule placeholderRule = createSplitPlaceholderRuleBuilder(PLACEHOLDER_INTENT,
                 primaryActivity::equals, i -> false, w -> true)
                 .setDefaultSplitAttributes(SPLIT_ATTRIBUTES)
                 .build();
diff --git a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
index 6dae0a1..cbb6e31 100644
--- a/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
+++ b/libs/WindowManager/Jetpack/tests/unittest/src/androidx/window/extensions/embedding/SplitPresenterTest.java
@@ -27,6 +27,7 @@
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.TASK_ID;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createActivityInfoWithMinDimensions;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createMockTaskFragmentInfo;
+import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitPairRuleBuilder;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createSplitRule;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.createWindowLayoutInfo;
 import static androidx.window.extensions.embedding.EmbeddingTestUtils.getSplitBounds;
@@ -473,7 +474,7 @@
         final Activity secondaryActivity = createMockActivity();
         final TaskFragmentContainer bottomTf = mController.newContainer(secondaryActivity, TASK_ID);
         final TaskFragmentContainer primaryTf = mController.newContainer(mActivity, TASK_ID);
-        final SplitPairRule rule = new SplitPairRule.Builder(pair ->
+        final SplitPairRule rule = createSplitPairRuleBuilder(pair ->
                 pair.first == mActivity && pair.second == secondaryActivity, pair -> false,
                 metrics -> true)
                 .setDefaultSplitAttributes(SPLIT_ATTRIBUTES)
@@ -491,7 +492,7 @@
 
     @Test
     public void testComputeSplitAttributes() {
-        final SplitPairRule splitPairRule = new SplitPairRule.Builder(
+        final SplitPairRule splitPairRule = createSplitPairRuleBuilder(
                 activityPair -> true,
                 activityIntentPair -> true,
                 windowMetrics -> windowMetrics.getBounds().equals(TASK_BOUNDS))
diff --git a/libs/WindowManager/Jetpack/window-extensions-core-release.aar b/libs/WindowManager/Jetpack/window-extensions-core-release.aar
new file mode 100644
index 0000000..4c08074
--- /dev/null
+++ b/libs/WindowManager/Jetpack/window-extensions-core-release.aar
Binary files differ
diff --git a/libs/WindowManager/Jetpack/window-extensions-release.aar b/libs/WindowManager/Jetpack/window-extensions-release.aar
index 4978e04..6fd8d29 100644
--- a/libs/WindowManager/Jetpack/window-extensions-release.aar
+++ b/libs/WindowManager/Jetpack/window-extensions-release.aar
Binary files differ
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt b/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
index 209d93f..1482cfc 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/SysUIUnfoldModule.kt
@@ -20,6 +20,7 @@
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.shade.NotificationPanelUnfoldAnimationController
 import com.android.systemui.statusbar.phone.StatusBarMoveFromCenterAnimationController
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityManager
 import com.android.systemui.unfold.util.NaturalRotationUnfoldProgressProvider
 import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider
 import com.android.systemui.util.kotlin.getOrNull
@@ -95,4 +96,6 @@
     fun getUnfoldHapticsPlayer(): UnfoldHapticsPlayer
 
     fun getUnfoldLightRevealOverlayAnimation(): UnfoldLightRevealOverlayAnimation
+
+    fun getUnfoldKeyguardVisibilityManager(): UnfoldKeyguardVisibilityManager
 }
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldKeyguardVisibilityListener.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldKeyguardVisibilityListener.kt
new file mode 100644
index 0000000..59558ac
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldKeyguardVisibilityListener.kt
@@ -0,0 +1,39 @@
+package com.android.systemui.unfold
+
+import android.util.Log
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.policy.KeyguardStateController
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityManager
+import com.android.systemui.util.kotlin.getOrNull
+import java.util.Optional
+import javax.inject.Inject
+
+/**
+ * Used to set the keyguard visibility state to [UnfoldKeyguardVisibilityManager].
+ *
+ * It is not possible to directly inject a sysui class (e.g. [KeyguardStateController]) into
+ * [DeviceStateProvider], as it can't depend on google sysui directly. So,
+ * [UnfoldKeyguardVisibilityManager] is provided to clients, that can set the keyguard visibility
+ * accordingly.
+ */
+@SysUISingleton
+class UnfoldKeyguardVisibilityListener
+@Inject
+constructor(
+    keyguardStateController: KeyguardStateController,
+    unfoldComponent: Optional<SysUIUnfoldComponent>,
+) {
+
+    private val unfoldKeyguardVisibilityManager =
+        unfoldComponent.getOrNull()?.getUnfoldKeyguardVisibilityManager()
+
+    private val delegate = { keyguardStateController.isVisible }
+
+    fun init() {
+        unfoldKeyguardVisibilityManager?.setKeyguardVisibleDelegate(delegate).also {
+            Log.d(TAG, "setKeyguardVisibleDelegate set")
+        }
+    }
+}
+
+private const val TAG = "UnfoldKeyguardVisibilityListener"
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
index 7d5f06c..6086e16 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceFoldStateProviderTest.kt
@@ -29,6 +29,7 @@
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider.ScreenListener
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityProvider
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.capture
 import com.google.common.truth.Truth.assertThat
@@ -56,6 +57,9 @@
     @Mock
     private lateinit var rotationChangeProvider: RotationChangeProvider
 
+    @Mock
+    private lateinit var unfoldKeyguardVisibilityProvider: UnfoldKeyguardVisibilityProvider
+
     @Captor
     private lateinit var rotationListener: ArgumentCaptor<RotationListener>
 
@@ -87,6 +91,7 @@
                 screenOnStatusProvider,
                 foldProvider,
                 activityTypeProvider,
+                unfoldKeyguardVisibilityProvider,
                 rotationChangeProvider,
                 context.mainExecutor,
                 handler
@@ -380,6 +385,47 @@
     }
 
     @Test
+    fun startClosingEvent_whileNotOnKeyguardAndNotOnLauncher_doesNotTriggerBeforeThreshold() {
+        setKeyguardVisibility(visible = false)
+        setupForegroundActivityType(isHomeActivity = false)
+        sendHingeAngleEvent(180)
+
+        sendHingeAngleEvent(START_CLOSING_ON_APPS_THRESHOLD_DEGREES + 1)
+
+        assertThat(foldUpdates).isEmpty()
+    }
+
+    @Test
+    fun startClosingEvent_whileKeyguardStateNotAvailable_triggerBeforeThreshold() {
+        setKeyguardVisibility(visible = null)
+        sendHingeAngleEvent(180)
+
+        sendHingeAngleEvent(START_CLOSING_ON_APPS_THRESHOLD_DEGREES + 1)
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_START_CLOSING)
+    }
+
+    @Test
+    fun startClosingEvent_whileonKeyguard_doesTriggerBeforeThreshold() {
+        setKeyguardVisibility(visible = true)
+        sendHingeAngleEvent(180)
+
+        sendHingeAngleEvent(START_CLOSING_ON_APPS_THRESHOLD_DEGREES + 1)
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_START_CLOSING)
+    }
+
+    @Test
+    fun startClosingEvent_whileNotOnKeyguard_triggersAfterThreshold() {
+        setKeyguardVisibility(visible = false)
+        sendHingeAngleEvent(START_CLOSING_ON_APPS_THRESHOLD_DEGREES)
+
+        sendHingeAngleEvent(START_CLOSING_ON_APPS_THRESHOLD_DEGREES - 1)
+
+        assertThat(foldUpdates).containsExactly(FOLD_UPDATE_START_CLOSING)
+    }
+
+    @Test
     fun screenOff_whileFolded_hingeAngleProviderRemainsOff() {
         setFoldState(folded = true)
         assertThat(testHingeAngleProvider.isStarted).isFalse()
@@ -445,6 +491,10 @@
         whenever(activityTypeProvider.isHomeActivity).thenReturn(isHomeActivity)
     }
 
+    private fun setKeyguardVisibility(visible: Boolean?) {
+        whenever(unfoldKeyguardVisibilityProvider.isKeyguardVisible).thenReturn(visible)
+    }
+
     private fun simulateTimeout(waitTime: Long = HALF_OPENED_TIMEOUT_MILLIS) {
         val runnableDelay = scheduledRunnableDelay ?: throw Exception("No runnable scheduled.")
         if (waitTime >= runnableDelay) {
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
index 8f4ee4d..3fa5469 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
@@ -28,6 +28,9 @@
 import com.android.systemui.unfold.updates.hinge.HingeSensorAngleProvider
 import com.android.systemui.unfold.util.ATraceLoggerTransitionProgressListener
 import com.android.systemui.unfold.util.ScaleAwareTransitionProgressProvider
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityManager
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityManagerImpl
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityProvider
 import dagger.Module
 import dagger.Provides
 import java.util.Optional
@@ -57,7 +60,8 @@
                 scaleAwareProviderFactory.wrap(baseProgressProvider).apply {
                     // Always present callback that logs animation beginning and end.
                     addCallback(tracingListener)
-                })
+                }
+            )
         }
 
     @Provides
@@ -77,4 +81,16 @@
         } else {
             EmptyHingeAngleProvider
         }
+
+    @Provides
+    @Singleton
+    fun unfoldKeyguardVisibilityProvider(
+        impl: UnfoldKeyguardVisibilityManagerImpl
+    ): UnfoldKeyguardVisibilityProvider = impl
+
+    @Provides
+    @Singleton
+    fun unfoldKeyguardVisibilityManager(
+        impl: UnfoldKeyguardVisibilityManagerImpl
+    ): UnfoldKeyguardVisibilityManager = impl
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
index 808128d..5b45897 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
@@ -31,6 +31,7 @@
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.screen.ScreenStatusProvider
 import com.android.systemui.unfold.util.CurrentActivityTypeProvider
+import com.android.systemui.unfold.util.UnfoldKeyguardVisibilityProvider
 import java.util.concurrent.Executor
 import javax.inject.Inject
 
@@ -42,6 +43,7 @@
     private val screenStatusProvider: ScreenStatusProvider,
     private val foldProvider: FoldProvider,
     private val activityTypeProvider: CurrentActivityTypeProvider,
+    private val unfoldKeyguardVisibilityProvider: UnfoldKeyguardVisibilityProvider,
     private val rotationChangeProvider: RotationChangeProvider,
     @UnfoldMain private val mainExecutor: Executor,
     @UnfoldMain private val handler: Handler
@@ -152,12 +154,13 @@
      */
     private fun getClosingThreshold(): Int? {
         val isHomeActivity = activityTypeProvider.isHomeActivity ?: return null
+        val isKeyguardVisible = unfoldKeyguardVisibilityProvider.isKeyguardVisible == true
 
         if (DEBUG) {
-            Log.d(TAG, "isHomeActivity=$isHomeActivity")
+            Log.d(TAG, "isHomeActivity=$isHomeActivity, isOnKeyguard=$isKeyguardVisible")
         }
 
-        return if (isHomeActivity) {
+        return if (isHomeActivity || isKeyguardVisible) {
             null
         } else {
             START_CLOSING_ON_APPS_THRESHOLD_DEGREES
@@ -257,7 +260,7 @@
     }
 
 private const val TAG = "DeviceFoldProvider"
-private const val DEBUG = false
+private val DEBUG = Log.isLoggable(TAG, Log.DEBUG)
 
 /** Threshold after which we consider the device fully unfolded. */
 @VisibleForTesting const val FULLY_OPEN_THRESHOLD_DEGREES = 15f
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/UnfoldKeyguardVisibilityProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/UnfoldKeyguardVisibilityProvider.kt
new file mode 100644
index 0000000..9f0efa0
--- /dev/null
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/util/UnfoldKeyguardVisibilityProvider.kt
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2022 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.unfold.util
+
+import java.lang.ref.WeakReference
+import javax.inject.Inject
+import javax.inject.Singleton
+
+interface UnfoldKeyguardVisibilityProvider {
+    /**
+     * True when the keyguard is visible.
+     *
+     * Might be [null] when it is not known.
+     */
+    val isKeyguardVisible: Boolean?
+}
+
+/** Used to notify keyguard visibility. */
+interface UnfoldKeyguardVisibilityManager {
+    /** Sets the delegate. [delegate] should return true when the keyguard is visible. */
+    fun setKeyguardVisibleDelegate(delegate: () -> Boolean)
+}
+
+/**
+ * Keeps a [WeakReference] for the keyguard visibility provider.
+ *
+ * It is a weak reference because this is in the global scope, while the delegate might be set from
+ * another subcomponent (that might have shorter lifespan).
+ */
+@Singleton
+class UnfoldKeyguardVisibilityManagerImpl @Inject constructor() :
+    UnfoldKeyguardVisibilityProvider, UnfoldKeyguardVisibilityManager {
+
+    private var delegatedProvider: WeakReference<() -> Boolean?>? = null
+
+    override fun setKeyguardVisibleDelegate(delegate: () -> Boolean) {
+        delegatedProvider = WeakReference(delegate)
+    }
+
+    override val isKeyguardVisible: Boolean?
+        get() = delegatedProvider?.get()?.invoke()
+}
diff --git a/services/core/java/com/android/server/attention/AttentionManagerService.java b/services/core/java/com/android/server/attention/AttentionManagerService.java
index d4ef638..658e38b 100644
--- a/services/core/java/com/android/server/attention/AttentionManagerService.java
+++ b/services/core/java/com/android/server/attention/AttentionManagerService.java
@@ -668,8 +668,8 @@
             mIProximityUpdateCallback = new IProximityUpdateCallback.Stub() {
                 @Override
                 public void onProximityUpdate(double distance) {
+                    mCallbackInternal.onProximityUpdate(distance);
                     synchronized (mLock) {
-                        mCallbackInternal.onProximityUpdate(distance);
                         freeIfInactiveLocked();
                     }
                 }