diff --git a/core/java/com/android/internal/policy/FoldLockSettingsObserver.java b/core/java/com/android/internal/policy/FoldLockSettingsObserver.java
new file mode 100644
index 0000000..c6fba8a
--- /dev/null
+++ b/core/java/com/android/internal/policy/FoldLockSettingsObserver.java
@@ -0,0 +1,128 @@
+/*
+ * 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.internal.policy;
+
+import android.content.ContentResolver;
+import android.content.Context;
+import android.database.ContentObserver;
+import android.os.Handler;
+import android.os.UserHandle;
+import android.provider.Settings;
+
+import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
+
+import java.util.Set;
+
+/**
+ * A ContentObserver that listens for changes in the "Continue using apps on fold" setting. This
+ * setting determines a device's behavior when the user folds the device.
+ * @hide
+ *
+ * Keep the setting values in this class in sync with the values in
+ * {@link com.android.server.utils.FoldSettingProvider} and
+ * {@link com.android.settings.display.FoldLockBehaviorSettings}
+ */
+public class FoldLockSettingsObserver extends ContentObserver {
+    /** The setting for "stay awake on fold". */
+    public static final String SETTING_VALUE_STAY_AWAKE_ON_FOLD = "stay_awake_on_fold_key";
+    /** The setting for "swipe up to continue". */
+    public static final String SETTING_VALUE_SELECTIVE_STAY_AWAKE = "selective_stay_awake_key";
+    /** The setting for "always sleep on fold". */
+    public static final String SETTING_VALUE_SLEEP_ON_FOLD = "sleep_on_fold_key";
+    public static final String SETTING_VALUE_DEFAULT = SETTING_VALUE_SELECTIVE_STAY_AWAKE;
+    private static final Set<String> SETTING_VALUES = Set.of(SETTING_VALUE_STAY_AWAKE_ON_FOLD,
+            SETTING_VALUE_SELECTIVE_STAY_AWAKE, SETTING_VALUE_SLEEP_ON_FOLD);
+
+    private final Context mContext;
+
+    /** The cached value of the setting. */
+    @VisibleForTesting
+    String mFoldLockSetting;
+
+    public FoldLockSettingsObserver(Handler handler, Context context) {
+        super(handler);
+        mContext = context;
+    }
+
+    /** Registers the observer and updates the cache for the first time. */
+    public void register() {
+        final ContentResolver r = mContext.getContentResolver();
+        r.registerContentObserver(
+                Settings.System.getUriFor(Settings.System.FOLD_LOCK_BEHAVIOR),
+                false, this, UserHandle.USER_ALL);
+        requestAndCacheFoldLockSetting();
+    }
+
+    /** Unregisters the observer. */
+    public void unregister() {
+        mContext.getContentResolver().unregisterContentObserver(this);
+    }
+
+    /** Runs when settings changes. */
+    @Override
+    public void onChange(boolean selfChange) {
+        requestAndCacheFoldLockSetting();
+    }
+
+    /**
+     * Requests and caches the current FOLD_LOCK_BEHAVIOR setting, which should be one of three
+     * values: SETTING_VALUE_STAY_AWAKE_ON_FOLD, SETTING_VALUE_SELECTIVE_STAY_AWAKE,
+     * SETTING_VALUE_SLEEP_ON_FOLD. If null (not set), returns the system default.
+     */
+    @VisibleForTesting
+    void requestAndCacheFoldLockSetting() {
+        String currentSetting = request();
+
+        if (currentSetting == null || !SETTING_VALUES.contains(currentSetting)) {
+            currentSetting = SETTING_VALUE_DEFAULT;
+        }
+
+        setCurrentFoldSetting(currentSetting);
+    }
+
+    /**
+     * Makes a binder call to request the current FOLD_LOCK_BEHAVIOR setting.
+     */
+    @VisibleForTesting
+    @Nullable
+    String request() {
+        return Settings.System.getStringForUser(mContext.getContentResolver(),
+                Settings.System.FOLD_LOCK_BEHAVIOR, UserHandle.USER_CURRENT);
+    }
+
+    /** Caches the fold-lock behavior received from Settings. */
+    @VisibleForTesting
+    void setCurrentFoldSetting(String newSetting) {
+        mFoldLockSetting = newSetting;
+    }
+
+    /** Used by external requesters: checks if the current setting is "stay awake on fold". */
+    public boolean isStayAwakeOnFold() {
+        return mFoldLockSetting.equals(SETTING_VALUE_STAY_AWAKE_ON_FOLD);
+    }
+
+    /** Used by external requesters: checks if the current setting is "swipe up to continue". */
+    public boolean isSelectiveStayAwake() {
+        return mFoldLockSetting.equals(SETTING_VALUE_SELECTIVE_STAY_AWAKE);
+    }
+
+    /** Used by external requesters: checks if the current setting is "sleep on fold". */
+    public boolean isSleepOnFold() {
+        return mFoldLockSetting.equals(SETTING_VALUE_SLEEP_ON_FOLD);
+    }
+}
diff --git a/core/tests/coretests/src/com/android/internal/policy/FoldLockSettingsObserverTest.java b/core/tests/coretests/src/com/android/internal/policy/FoldLockSettingsObserverTest.java
new file mode 100644
index 0000000..537dd69
--- /dev/null
+++ b/core/tests/coretests/src/com/android/internal/policy/FoldLockSettingsObserverTest.java
@@ -0,0 +1,110 @@
+/*
+ * 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.internal.policy;
+
+import static com.android.internal.policy.FoldLockSettingsObserver.SETTING_VALUE_DEFAULT;
+import static com.android.internal.policy.FoldLockSettingsObserver.SETTING_VALUE_SLEEP_ON_FOLD;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+
+import android.content.ContentResolver;
+import android.content.Context;
+import android.os.Handler;
+import android.os.UserHandle;
+import android.provider.Settings;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for {@link FoldLockSettingsObserver}.
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public final class FoldLockSettingsObserverTest {
+    @Mock
+    private Context mContext;
+    @Mock
+    private Handler mHandler;
+    @Mock
+    private ContentResolver mContentResolver;
+
+    private FoldLockSettingsObserver mFoldLockSettingsObserver;
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+
+        mFoldLockSettingsObserver =
+                spy(new FoldLockSettingsObserver(mHandler, mContext));
+
+        doReturn(mContentResolver).when(mContext).getContentResolver();
+        doReturn(SETTING_VALUE_DEFAULT).when(mFoldLockSettingsObserver).request();
+
+        mFoldLockSettingsObserver.register();
+    }
+
+    @Test
+    public void shouldRegister() {
+        doReturn(mContentResolver).when(mContext).getContentResolver();
+
+        mFoldLockSettingsObserver.register();
+
+        verify(mContentResolver).registerContentObserver(
+                Settings.System.getUriFor(Settings.System.FOLD_LOCK_BEHAVIOR),
+                false /*notifyForDescendants */,
+                mFoldLockSettingsObserver,
+                UserHandle.USER_ALL
+        );
+    }
+
+    @Test
+    public void shouldUnregister() {
+        mFoldLockSettingsObserver.unregister();
+
+        verify(mContentResolver).unregisterContentObserver(mFoldLockSettingsObserver);
+    }
+
+    @Test
+    public void shouldCacheNewValue() {
+        // Reset the mock's behavior and call count to zero.
+        reset(mFoldLockSettingsObserver);
+        doReturn(SETTING_VALUE_SLEEP_ON_FOLD).when(mFoldLockSettingsObserver).request();
+
+        // Setting is DEFAULT at first.
+        assertEquals(SETTING_VALUE_DEFAULT, mFoldLockSettingsObserver.mFoldLockSetting);
+
+        // Cache new setting.
+        mFoldLockSettingsObserver.requestAndCacheFoldLockSetting();
+
+        // Check that setter was called once and change went through properly.
+        verify(mFoldLockSettingsObserver).setCurrentFoldSetting(anyString());
+        assertTrue(mFoldLockSettingsObserver.isSleepOnFold());
+    }
+}
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 45feff5..48e1376 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
@@ -366,13 +366,14 @@
             Optional<WindowDecorViewModel> windowDecorViewModel,
             Optional<DesktopTasksController> desktopTasksController,
             MultiInstanceHelper multiInstanceHelper,
-            @ShellMainThread ShellExecutor mainExecutor) {
+            @ShellMainThread ShellExecutor mainExecutor,
+            @ShellMainThread Handler mainHandler) {
         return new SplitScreenController(context, shellInit, shellCommandHandler, shellController,
                 shellTaskOrganizer, syncQueue, rootTaskDisplayAreaOrganizer, displayController,
                 displayImeController, displayInsetsController, dragAndDropController, transitions,
                 transactionPool, iconProvider, recentTasks, launchAdjacentController,
                 windowDecorViewModel, desktopTasksController, null /* stageCoordinator */,
-                multiInstanceHelper, mainExecutor);
+                multiInstanceHelper, mainExecutor, mainHandler);
     }
 
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
index 8df287d..06c57bd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java
@@ -113,6 +113,9 @@
     /** Called when device waking up finished. */
     void onFinishedWakingUp();
 
+    /** Called when device starts going to sleep (screen off). */
+    void onStartedGoingToSleep();
+
     /** Called when requested to go to fullscreen from the current active split app. */
     void goToFullscreenFromSplit();
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index e659151..b857556 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -50,6 +50,7 @@
 import android.content.pm.ShortcutInfo;
 import android.graphics.Rect;
 import android.os.Bundle;
+import android.os.Handler;
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.util.ArrayMap;
@@ -180,6 +181,7 @@
     private final LauncherApps mLauncherApps;
     private final RootTaskDisplayAreaOrganizer mRootTDAOrganizer;
     private final ShellExecutor mMainExecutor;
+    private final Handler mMainHandler;
     private final SplitScreenImpl mImpl = new SplitScreenImpl();
     private final DisplayController mDisplayController;
     private final DisplayImeController mDisplayImeController;
@@ -227,7 +229,8 @@
             Optional<DesktopTasksController> desktopTasksController,
             @Nullable StageCoordinator stageCoordinator,
             MultiInstanceHelper multiInstanceHelper,
-            ShellExecutor mainExecutor) {
+            ShellExecutor mainExecutor,
+            Handler mainHandler) {
         mShellCommandHandler = shellCommandHandler;
         mShellController = shellController;
         mTaskOrganizer = shellTaskOrganizer;
@@ -236,6 +239,7 @@
         mLauncherApps = context.getSystemService(LauncherApps.class);
         mRootTDAOrganizer = rootTDAOrganizer;
         mMainExecutor = mainExecutor;
+        mMainHandler = mainHandler;
         mDisplayController = displayController;
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
@@ -292,7 +296,7 @@
         return new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
                 mTaskOrganizer, mDisplayController, mDisplayImeController,
                 mDisplayInsetsController, mTransitions, mTransactionPool, mIconProvider,
-                mMainExecutor, mRecentTasksOptional, mLaunchAdjacentController,
+                mMainExecutor, mMainHandler, mRecentTasksOptional, mLaunchAdjacentController,
                 mWindowDecorViewModel);
     }
 
@@ -448,13 +452,17 @@
     @Override
     public void onKeyguardVisibilityChanged(boolean visible, boolean occluded,
             boolean animatingDismiss) {
-        mStageCoordinator.onKeyguardVisibilityChanged(visible);
+        mStageCoordinator.onKeyguardStateChanged(visible, occluded);
     }
 
     public void onFinishedWakingUp() {
         mStageCoordinator.onFinishedWakingUp();
     }
 
+    public void onStartedGoingToSleep() {
+        mStageCoordinator.onStartedGoingToSleep();
+    }
+
     public void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
         mStageCoordinator.exitSplitScreenOnHide(exitSplitScreenOnHide);
     }
@@ -1201,6 +1209,11 @@
         }
 
         @Override
+        public void onStartedGoingToSleep() {
+            mMainExecutor.execute(SplitScreenController.this::onStartedGoingToSleep);
+        }
+
+        @Override
         public void goToFullscreenFromSplit() {
             mMainExecutor.execute(SplitScreenController.this::goToFullscreenFromSplit);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index 9bcd9b0..dd7228a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -92,6 +92,7 @@
 import android.hardware.devicestate.DeviceStateManager;
 import android.os.Bundle;
 import android.os.Debug;
+import android.os.Handler;
 import android.os.IBinder;
 import android.os.RemoteException;
 import android.os.ServiceManager;
@@ -119,6 +120,7 @@
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.InstanceId;
+import com.android.internal.policy.FoldLockSettingsObserver;
 import com.android.internal.protolog.ProtoLog;
 import com.android.internal.util.ArrayUtils;
 import com.android.launcher3.icons.IconProvider;
@@ -191,7 +193,7 @@
     private SplitLayout mSplitLayout;
     private ValueAnimator mDividerFadeInAnimator;
     private boolean mDividerVisible;
-    private boolean mKeyguardShowing;
+    private boolean mKeyguardActive;
     private boolean mShowDecorImmediately;
     private final SyncTransactionQueue mSyncQueue;
     private final ShellTaskOrganizer mTaskOrganizer;
@@ -205,6 +207,7 @@
     private SplitScreenTransitions mSplitTransitions;
     private final SplitscreenEventLogger mLogger;
     private final ShellExecutor mMainExecutor;
+    private final Handler mMainHandler;
     // Cache live tile tasks while entering recents, evict them from stages in finish transaction
     // if user is opening another task(s).
     private final ArrayList<Integer> mPausingTasks = new ArrayList<>();
@@ -233,7 +236,10 @@
     private boolean mIsExiting;
     private boolean mIsRootTranslucent;
     @VisibleForTesting
-    int mTopStageAfterFoldDismiss;
+    @StageType int mLastActiveStage;
+    private boolean mBreakOnNextWake;
+    /** Used to get the Settings value for "Continue using apps on fold". */
+    private FoldLockSettingsObserver mFoldLockSettingsObserver;
 
     private DefaultMixedHandler mMixedHandler;
     private final Toast mSplitUnsupportedToast;
@@ -313,9 +319,8 @@
             ShellTaskOrganizer taskOrganizer, DisplayController displayController,
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, Transitions transitions,
-            TransactionPool transactionPool,
-            IconProvider iconProvider, ShellExecutor mainExecutor,
-            Optional<RecentTasksController> recentTasks,
+            TransactionPool transactionPool, IconProvider iconProvider, ShellExecutor mainExecutor,
+            Handler mainHandler, Optional<RecentTasksController> recentTasks,
             LaunchAdjacentController launchAdjacentController,
             Optional<WindowDecorViewModel> windowDecorViewModel) {
         mContext = context;
@@ -324,6 +329,7 @@
         mTaskOrganizer = taskOrganizer;
         mLogger = new SplitscreenEventLogger();
         mMainExecutor = mainExecutor;
+        mMainHandler = mainHandler;
         mRecentTasks = recentTasks;
         mLaunchAdjacentController = launchAdjacentController;
         mWindowDecorViewModel = windowDecorViewModel;
@@ -366,6 +372,9 @@
         // With shell transition, we should update recents tile each callback so set this to true by
         // default.
         mShouldUpdateRecents = ENABLE_SHELL_TRANSITIONS;
+        mFoldLockSettingsObserver =
+                new FoldLockSettingsObserver(mainHandler, context);
+        mFoldLockSettingsObserver.register();
     }
 
     @VisibleForTesting
@@ -373,9 +382,8 @@
             ShellTaskOrganizer taskOrganizer, MainStage mainStage, SideStage sideStage,
             DisplayController displayController, DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController, SplitLayout splitLayout,
-            Transitions transitions, TransactionPool transactionPool,
-            ShellExecutor mainExecutor,
-            Optional<RecentTasksController> recentTasks,
+            Transitions transitions, TransactionPool transactionPool, ShellExecutor mainExecutor,
+            Handler mainHandler, Optional<RecentTasksController> recentTasks,
             LaunchAdjacentController launchAdjacentController,
             Optional<WindowDecorViewModel> windowDecorViewModel) {
         mContext = context;
@@ -393,6 +401,7 @@
                 this::onTransitionAnimationComplete, this);
         mLogger = new SplitscreenEventLogger();
         mMainExecutor = mainExecutor;
+        mMainHandler = mainHandler;
         mRecentTasks = recentTasks;
         mLaunchAdjacentController = launchAdjacentController;
         mWindowDecorViewModel = windowDecorViewModel;
@@ -400,6 +409,9 @@
         transitions.addHandler(this);
         mSplitUnsupportedToast = Toast.makeText(mContext,
                 R.string.dock_non_resizeble_failed_to_dock_text, Toast.LENGTH_SHORT);
+        mFoldLockSettingsObserver =
+                new FoldLockSettingsObserver(context.getMainThreadHandler(), context);
+        mFoldLockSettingsObserver.register();
     }
 
     public void setMixedHandler(DefaultMixedHandler mixedHandler) {
@@ -1504,51 +1516,80 @@
         }
     }
 
-    void onKeyguardVisibilityChanged(boolean showing) {
-        mKeyguardShowing = showing;
+    /**
+     * Runs when keyguard state changes. The booleans here are a bit complicated, so for reference:
+     * @param active {@code true} if we are in a state where the keyguard *should* be shown
+     *                           -- still true when keyguard is "there" but is behind an app, or
+     *                           screen is off.
+     * @param occludingTaskRunning {@code true} when there is a running task that has
+     *                                         FLAG_SHOW_WHEN_LOCKED -- also true when the task is
+     *                                         just running on its own and keyguard is not active
+     *                                         at all.
+     */
+    void onKeyguardStateChanged(boolean active, boolean occludingTaskRunning) {
+        mKeyguardActive = active;
         if (!mMainStage.isActive()) {
             return;
         }
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onKeyguardVisibilityChanged: showing=%b", showing);
-        setDividerVisibility(!mKeyguardShowing, null);
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
+                "onKeyguardVisibilityChanged: active=%b occludingTaskRunning=%b",
+                active, occludingTaskRunning);
+        setDividerVisibility(!mKeyguardActive, null);
+
+        if (active && occludingTaskRunning) {
+            dismissSplitKeepingLastActiveStage(EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
+        }
     }
 
     void onFinishedWakingUp() {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onFinishedWakingUp");
-        if (!mMainStage.isActive()) {
+        if (mBreakOnNextWake) {
+            dismissSplitKeepingLastActiveStage(EXIT_REASON_DEVICE_FOLDED);
+        }
+    }
+
+    void onStartedGoingToSleep() {
+        recordLastActiveStage();
+    }
+
+    /**
+     * Records the user's last focused stage -- main stage or side stage. Used to determine which
+     * stage of a split pair should be kept, in cases where system focus has moved elsewhere.
+     */
+    void recordLastActiveStage() {
+        if (!isSplitActive() || !isSplitScreenVisible()) {
+            mLastActiveStage = STAGE_TYPE_UNDEFINED;
+        } else if (mMainStage.isFocused()) {
+            mLastActiveStage = STAGE_TYPE_MAIN;
+        } else if (mSideStage.isFocused()) {
+            mLastActiveStage = STAGE_TYPE_SIDE;
+        }
+    }
+
+    /**
+     * Dismisses split, keeping the app that the user focused last in split screen. If the user was
+     * not in split screen, {@link #mLastActiveStage} should be set to STAGE_TYPE_UNDEFINED, and we
+     * will do a no-op.
+     */
+    void dismissSplitKeepingLastActiveStage(@ExitReason int reason) {
+        if (!mMainStage.isActive() || mLastActiveStage == STAGE_TYPE_UNDEFINED) {
+            // no-op
             return;
         }
 
-        // Check if there's only one stage visible while keyguard occluded.
-        final boolean mainStageVisible = mMainStage.mRootTaskInfo.isVisible;
-        final boolean oneStageVisible =
-                mMainStage.mRootTaskInfo.isVisible != mSideStage.mRootTaskInfo.isVisible;
-        if (oneStageVisible && !ENABLE_SHELL_TRANSITIONS) {
-            // Dismiss split because there's show-when-locked activity showing on top of keyguard.
-            // Also make sure the task contains show-when-locked activity remains on top after split
-            // dismissed.
-            final StageTaskListener toTop = mainStageVisible ? mMainStage : mSideStage;
-            exitSplitScreen(toTop, EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
+        if (ENABLE_SHELL_TRANSITIONS) {
+            // Need manually clear here due to this transition might be aborted due to keyguard
+            // on top and lead to no visible change.
+            clearSplitPairedInRecents(reason);
+            final WindowContainerTransaction wct = new WindowContainerTransaction();
+            prepareExitSplitScreen(mLastActiveStage, wct);
+            mSplitTransitions.startDismissTransition(wct, this, mLastActiveStage, reason);
+            setSplitsVisible(false);
+        } else {
+            exitSplitScreen(mLastActiveStage == STAGE_TYPE_MAIN ? mMainStage : mSideStage, reason);
         }
 
-        // Dismiss split if the flag record any side of stages.
-        if (mTopStageAfterFoldDismiss != STAGE_TYPE_UNDEFINED) {
-            if (ENABLE_SHELL_TRANSITIONS) {
-                // Need manually clear here due to this transition might be aborted due to keyguard
-                // on top and lead to no visible change.
-                clearSplitPairedInRecents(EXIT_REASON_DEVICE_FOLDED);
-                final WindowContainerTransaction wct = new WindowContainerTransaction();
-                prepareExitSplitScreen(mTopStageAfterFoldDismiss, wct);
-                mSplitTransitions.startDismissTransition(wct, this,
-                        mTopStageAfterFoldDismiss, EXIT_REASON_DEVICE_FOLDED);
-                setSplitsVisible(false);
-            } else {
-                exitSplitScreen(
-                        mTopStageAfterFoldDismiss == STAGE_TYPE_MAIN ? mMainStage : mSideStage,
-                        EXIT_REASON_DEVICE_FOLDED);
-            }
-            mTopStageAfterFoldDismiss = STAGE_TYPE_UNDEFINED;
-        }
+        mBreakOnNextWake = false;
     }
 
     void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
@@ -2223,11 +2264,11 @@
 
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
                 "setDividerVisibility: visible=%b keyguardShowing=%b dividerAnimating=%b caller=%s",
-                visible, mKeyguardShowing, mIsDividerRemoteAnimating, Debug.getCaller());
+                visible, mKeyguardActive, mIsDividerRemoteAnimating, Debug.getCaller());
 
         // Defer showing divider bar after keyguard dismissed, so it won't interfere with keyguard
         // dismissing animation.
-        if (visible && mKeyguardShowing) {
+        if (visible && mKeyguardActive) {
             ProtoLog.d(WM_SHELL_SPLIT_SCREEN,
                     "   Defer showing divider bar due to keyguard showing.");
             return;
@@ -2597,21 +2638,24 @@
     @VisibleForTesting
     void onFoldedStateChanged(boolean folded) {
         ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onFoldedStateChanged: folded=%b", folded);
-        mTopStageAfterFoldDismiss = STAGE_TYPE_UNDEFINED;
-        if (!folded) return;
 
-        if (!isSplitActive() || !isSplitScreenVisible()) return;
-
-        // To avoid split dismiss when user fold the device and unfold to use later, we only
-        // record the flag here and try to dismiss on wakeUp callback to ensure split dismiss
-        // when user interact on phone folded.
-        if (mMainStage.isFocused()) {
-            mTopStageAfterFoldDismiss = STAGE_TYPE_MAIN;
-        } else if (mSideStage.isFocused()) {
-            mTopStageAfterFoldDismiss = STAGE_TYPE_SIDE;
+        if (folded) {
+            recordLastActiveStage();
+            // If user folds and has the setting "Continue using apps on fold = NEVER", we assume
+            // they don't want to continue using split on the outer screen (i.e. we break split if
+            // they wake the device in its folded state).
+            mBreakOnNextWake = willSleepOnFold();
+        } else {
+            mBreakOnNextWake = false;
         }
     }
 
+    /** Returns true if the phone will sleep when it folds. */
+    @VisibleForTesting
+    boolean willSleepOnFold() {
+        return mFoldLockSettingsObserver != null && mFoldLockSettingsObserver.isSleepOnFold();
+    }
+
     private Rect getSideStageBounds() {
         return mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
                 ? mSplitLayout.getBounds1() : mSplitLayout.getBounds2();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvSplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvSplitScreenController.java
index e330f3a..b65e978 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvSplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvSplitScreenController.java
@@ -33,7 +33,6 @@
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.SystemWindows;
 import com.android.wm.shell.common.TransactionPool;
-import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.splitscreen.StageCoordinator;
@@ -88,7 +87,8 @@
                 syncQueue, rootTDAOrganizer, displayController, displayImeController,
                 displayInsetsController, null, transitions, transactionPool,
                 iconProvider, recentTasks, launchAdjacentController, Optional.empty(),
-                Optional.empty(), null /* stageCoordinator */, multiInstanceHelper, mainExecutor);
+                Optional.empty(), null /* stageCoordinator */, multiInstanceHelper, mainExecutor,
+                mainHandler);
 
         mTaskOrganizer = shellTaskOrganizer;
         mSyncQueue = syncQueue;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvStageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvStageCoordinator.java
index 7947691..81ca48f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvStageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/tv/TvStageCoordinator.java
@@ -56,7 +56,7 @@
             SystemWindows systemWindows) {
         super(context, displayId, syncQueue, taskOrganizer, displayController, displayImeController,
                 displayInsetsController, transitions, transactionPool, iconProvider,
-                mainExecutor, recentTasks, launchAdjacentController, Optional.empty());
+                mainExecutor, mainHandler, recentTasks, launchAdjacentController, Optional.empty());
 
         mTvSplitMenuController = new TvSplitMenuController(context, this,
                 systemWindows, mainHandler);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
index 5b95b15..1c5d5e9 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
@@ -50,7 +50,7 @@
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
 import android.os.Bundle;
-import android.os.IBinder;
+import android.os.Handler;
 import android.window.IWindowContainerToken;
 import android.window.WindowContainerToken;
 
@@ -104,6 +104,7 @@
     @Mock SyncTransactionQueue mSyncQueue;
     @Mock RootTaskDisplayAreaOrganizer mRootTDAOrganizer;
     @Mock ShellExecutor mMainExecutor;
+    @Mock Handler mMainHandler;
     @Mock DisplayController mDisplayController;
     @Mock DisplayImeController mDisplayImeController;
     @Mock DisplayInsetsController mDisplayInsetsController;
@@ -134,7 +135,7 @@
                 mDisplayInsetsController, mDragAndDropController, mTransitions, mTransactionPool,
                 mIconProvider, Optional.of(mRecentTasks), mLaunchAdjacentController,
                 Optional.of(mWindowDecorViewModel), Optional.of(mDesktopTasksController),
-                mStageCoordinator, mMultiInstanceHelper, mMainExecutor));
+                mStageCoordinator, mMultiInstanceHelper, mMainExecutor, mMainHandler));
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
index a3009a5..29d3fb4 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTestUtils.java
@@ -22,6 +22,7 @@
 import android.app.ActivityManager;
 import android.content.Context;
 import android.graphics.Rect;
+import android.os.Handler;
 import android.view.SurfaceControl;
 import android.view.SurfaceSession;
 
@@ -77,13 +78,13 @@
                 DisplayController displayController, DisplayImeController imeController,
                 DisplayInsetsController insetsController, SplitLayout splitLayout,
                 Transitions transitions, TransactionPool transactionPool,
-                ShellExecutor mainExecutor,
+                ShellExecutor mainExecutor, Handler mainHandler,
                 Optional<RecentTasksController> recentTasks,
                 LaunchAdjacentController launchAdjacentController,
                 Optional<WindowDecorViewModel> windowDecorViewModel) {
             super(context, displayId, syncQueue, taskOrganizer, mainStage,
                     sideStage, displayController, imeController, insetsController, splitLayout,
-                    transitions, transactionPool, mainExecutor, recentTasks,
+                    transitions, transactionPool, mainExecutor, mainHandler, recentTasks,
                     launchAdjacentController, windowDecorViewModel);
 
             // Prepare root task for testing.
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
index 34b2eeb..37ef788 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitTransitionTests.java
@@ -49,6 +49,7 @@
 
 import android.annotation.NonNull;
 import android.app.ActivityManager;
+import android.os.Handler;
 import android.os.IBinder;
 import android.view.SurfaceControl;
 import android.view.SurfaceSession;
@@ -107,6 +108,7 @@
     @Mock private IconProvider mIconProvider;
     @Mock private WindowDecorViewModel mWindowDecorViewModel;
     @Mock private ShellExecutor mMainExecutor;
+    @Mock private Handler mMainHandler;
     @Mock private LaunchAdjacentController mLaunchAdjacentController;
     @Mock private DefaultMixedHandler mMixedHandler;
     @Mock private SplitScreen.SplitInvocationListener mInvocationListener;
@@ -140,7 +142,7 @@
         mStageCoordinator = new SplitTestUtils.TestStageCoordinator(mContext, DEFAULT_DISPLAY,
                 mSyncQueue, mTaskOrganizer, mMainStage, mSideStage, mDisplayController,
                 mDisplayImeController, mDisplayInsetsController, mSplitLayout, mTransitions,
-                mTransactionPool, mMainExecutor, Optional.empty(),
+                mTransactionPool, mMainExecutor, mMainHandler, Optional.empty(),
                 mLaunchAdjacentController, Optional.empty());
         mStageCoordinator.setMixedHandler(mMixedHandler);
         mSplitScreenTransitions = mStageCoordinator.getSplitTransitions();
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index d18fec2..315332f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -138,7 +138,8 @@
         mStageCoordinator = spy(new StageCoordinator(mContext, DEFAULT_DISPLAY, mSyncQueue,
                 mTaskOrganizer, mMainStage, mSideStage, mDisplayController, mDisplayImeController,
                 mDisplayInsetsController, mSplitLayout, mTransitions, mTransactionPool,
-                mMainExecutor, Optional.empty(), mLaunchAdjacentController, Optional.empty()));
+                mMainExecutor, mMainHandler, Optional.empty(), mLaunchAdjacentController,
+                Optional.empty()));
         mDividerLeash = new SurfaceControl.Builder(mSurfaceSession).setName("fakeDivider").build();
 
         when(mSplitLayout.getBounds1()).thenReturn(mBounds1);
@@ -359,10 +360,11 @@
         mMainStage.mRootTaskInfo = new TestRunningTaskInfoBuilder().setVisible(true).build();
         when(mStageCoordinator.isSplitActive()).thenReturn(true);
         when(mStageCoordinator.isSplitScreenVisible()).thenReturn(true);
+        when(mStageCoordinator.willSleepOnFold()).thenReturn(true);
 
         mStageCoordinator.onFoldedStateChanged(true);
 
-        assertEquals(mStageCoordinator.mTopStageAfterFoldDismiss, STAGE_TYPE_MAIN);
+        assertEquals(mStageCoordinator.mLastActiveStage, STAGE_TYPE_MAIN);
 
         mStageCoordinator.onFinishedWakingUp();
 
diff --git a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
index ec9b5cf..a714351 100644
--- a/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
+++ b/packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java
@@ -263,6 +263,11 @@
             public void onFinishedWakingUp() {
                 splitScreen.onFinishedWakingUp();
             }
+
+            @Override
+            public void onStartedGoingToSleep() {
+                splitScreen.onStartedGoingToSleep();
+            }
         });
         mCommandQueue.addCallback(new CommandQueue.Callbacks() {
             @Override