diff --git a/core/java/android/app/WallpaperManager.java b/core/java/android/app/WallpaperManager.java
index 62db90f..e6d51c6 100644
--- a/core/java/android/app/WallpaperManager.java
+++ b/core/java/android/app/WallpaperManager.java
@@ -255,6 +255,13 @@
     public static final String COMMAND_GOING_TO_SLEEP = "android.wallpaper.goingtosleep";
 
     /**
+     * Command for {@link #sendWallpaperCommand}: reported when a physical display switch event
+     * happens, e.g. fold and unfold.
+     * @hide
+     */
+    public static final String COMMAND_DISPLAY_SWITCH = "android.wallpaper.displayswitch";
+
+    /**
      * Command for {@link #sendWallpaperCommand}: reported when the wallpaper that was already
      * set is re-applied by the user.
      * @hide
diff --git a/core/java/android/service/wallpaper/WallpaperService.java b/core/java/android/service/wallpaper/WallpaperService.java
index 1a2be15..76e0c25 100644
--- a/core/java/android/service/wallpaper/WallpaperService.java
+++ b/core/java/android/service/wallpaper/WallpaperService.java
@@ -16,6 +16,7 @@
 
 package android.service.wallpaper;
 
+import static android.app.WallpaperManager.COMMAND_DISPLAY_SWITCH;
 import static android.app.WallpaperManager.COMMAND_FREEZE;
 import static android.app.WallpaperManager.COMMAND_UNFREEZE;
 import static android.app.WallpaperManager.SetWallpaperFlags;
@@ -153,6 +154,7 @@
     static final boolean DEBUG = false;
     static final float MIN_PAGE_ALLOWED_MARGIN = .05f;
     private static final int MIN_BITMAP_SCREENSHOT_WIDTH = 64;
+    private static final long PRESERVE_VISIBLE_TIMEOUT_MS = 1000;
     private static final long DEFAULT_UPDATE_SCREENSHOT_DURATION = 60 * 1000; //Once per minute
     private static final @NonNull RectF LOCAL_COLOR_BOUNDS =
             new RectF(0, 0, 1, 1);
@@ -165,6 +167,7 @@
 
     private static final int MSG_UPDATE_SURFACE = 10000;
     private static final int MSG_VISIBILITY_CHANGED = 10010;
+    private static final int MSG_REFRESH_VISIBILITY = 10011;
     private static final int MSG_WALLPAPER_OFFSETS = 10020;
     private static final int MSG_WALLPAPER_COMMAND = 10025;
     @UnsupportedAppUsage(maxTargetSdk = Build.VERSION_CODES.R, trackingBug = 170729553)
@@ -248,6 +251,11 @@
          */
         private boolean mIsScreenTurningOn;
         boolean mReportedVisible;
+        /**
+         * This is used with {@link #PRESERVE_VISIBLE_TIMEOUT_MS} to avoid intermediate visibility
+         * changes if the display may be toggled in a short time, e.g. display switch.
+         */
+        boolean mPreserveVisible;
         boolean mDestroyed;
         // Set to true after receiving WallpaperManager#COMMAND_FREEZE. It's reset back to false
         // after receiving WallpaperManager#COMMAND_UNFREEZE. COMMAND_FREEZE is fully applied once
@@ -1084,6 +1092,9 @@
             if (pendingCount != 0) {
                 out.print(prefix); out.print("mPendingResizeCount="); out.println(pendingCount);
             }
+            if (mPreserveVisible) {
+                out.print(prefix); out.print("mPreserveVisible=true");
+            }
             synchronized (mLock) {
                 out.print(prefix); out.print("mPendingXOffset="); out.print(mPendingXOffset);
                         out.print(" mPendingXOffset="); out.println(mPendingXOffset);
@@ -1643,7 +1654,8 @@
                                 ? false
                                 : mIWallpaperEngine.mInfo.supportsAmbientMode();
                 // Report visibility only if display is fully on or wallpaper supports ambient mode.
-                boolean visible = mVisible && (displayFullyOn || supportsAmbientMode);
+                final boolean visible = (mVisible && (displayFullyOn || supportsAmbientMode))
+                        || mPreserveVisible;
                 if (DEBUG) {
                     Log.v(
                             TAG,
@@ -2080,6 +2092,9 @@
             if (!mDestroyed) {
                 if (COMMAND_FREEZE.equals(cmd.action) || COMMAND_UNFREEZE.equals(cmd.action)) {
                     updateFrozenState(/* frozenRequested= */ !COMMAND_UNFREEZE.equals(cmd.action));
+                } else if (COMMAND_DISPLAY_SWITCH.equals(cmd.action)) {
+                    handleDisplaySwitch(cmd.z == 1 /* startToSwitch */);
+                    return;
                 }
                 result = onCommand(cmd.action, cmd.x, cmd.y, cmd.z,
                         cmd.extras, cmd.sync);
@@ -2095,6 +2110,23 @@
             }
         }
 
+        private void handleDisplaySwitch(boolean startToSwitch) {
+            if (startToSwitch && mReportedVisible) {
+                // The display may be off/on in a short time when the display is switching.
+                // Keep the visible state until onScreenTurnedOn or !startToSwitch is received, so
+                // the rendering thread can be active to redraw in time when receiving size change.
+                mPreserveVisible = true;
+                mCaller.removeMessages(MSG_REFRESH_VISIBILITY);
+                mCaller.sendMessageDelayed(mCaller.obtainMessage(MSG_REFRESH_VISIBILITY),
+                        PRESERVE_VISIBLE_TIMEOUT_MS);
+            } else if (!startToSwitch && mPreserveVisible) {
+                // The switch is finished, so restore to actual visibility.
+                mPreserveVisible = false;
+                mCaller.removeMessages(MSG_REFRESH_VISIBILITY);
+                reportVisibility(false /* forceReport */);
+            }
+        }
+
         private void updateFrozenState(boolean frozenRequested) {
             if (mIWallpaperEngine.mInfo == null
                     // Procees the unfreeze command in case the wallaper became static while
@@ -2638,6 +2670,10 @@
                             + ": " + message.arg1);
                     mEngine.doVisibilityChanged(message.arg1 != 0);
                     break;
+                case MSG_REFRESH_VISIBILITY:
+                    mEngine.mPreserveVisible = false;
+                    mEngine.reportVisibility(false /* forceReport */);
+                    break;
                 case MSG_UPDATE_SCREEN_TURNING_ON:
                     if (DEBUG) {
                         Log.v(TAG,
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index bf669fb..0abf304 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -5634,7 +5634,7 @@
             Trace.asyncTraceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "screenTurningOn",
                     0 /* cookie */);
             updateScreenOffSleepToken(false /* acquire */, false /* isSwappingDisplay */);
-            mDefaultDisplayPolicy.screenTurnedOn(screenOnListener);
+            mDefaultDisplayPolicy.screenTurningOn(screenOnListener);
             mBootAnimationDismissable = false;
 
             synchronized (mLock) {
@@ -5676,6 +5676,7 @@
                 mKeyguardDelegate.onScreenTurnedOn();
             }
         }
+        mDefaultDisplayPolicy.screenTurnedOn();
         reportScreenStateToVrManager(true);
     }
 
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 63ca592..e2bc59b 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -794,6 +794,9 @@
             }
             mService.mAtmService.mKeyguardController.updateDeferTransitionForAod(
                     mAwake /* waiting */);
+            if (!awake) {
+                mDisplayContent.mWallpaperController.onDisplaySwitchFinished();
+            }
         }
     }
 
@@ -836,7 +839,8 @@
         mRemoteInsetsControllerControlsSystemBars = remoteInsetsControllerControlsSystemBars;
     }
 
-    public void screenTurnedOn(ScreenOnListener screenOnListener) {
+    /** Prepares to turn on screen. The given listener is used to notify that it is ready. */
+    public void screenTurningOn(ScreenOnListener screenOnListener) {
         WindowProcessController visibleDozeUiProcess = null;
         synchronized (mLock) {
             mScreenOnEarly = true;
@@ -858,6 +862,11 @@
         }
     }
 
+    /** It is called after {@link #finishScreenTurningOn}. This runs on PowerManager's thread. */
+    public void screenTurnedOn() {
+        mDisplayContent.mWallpaperController.onDisplaySwitchFinished();
+    }
+
     public void screenTurnedOff() {
         synchronized (mLock) {
             mScreenOnEarly = false;
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 6033220..a9f35ac 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2784,6 +2784,9 @@
         } else {
             throw new RuntimeException("Create the same sleep token twice: " + token);
         }
+        if (isSwappingDisplay) {
+            display.mWallpaperController.onDisplaySwitchStarted();
+        }
         return token;
     }
 
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index a9f0554..4ed3dbc 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -16,6 +16,7 @@
 
 package com.android.server.wm;
 
+import static android.app.WallpaperManager.COMMAND_DISPLAY_SWITCH;
 import static android.app.WallpaperManager.COMMAND_FREEZE;
 import static android.app.WallpaperManager.COMMAND_UNFREEZE;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
@@ -116,6 +117,11 @@
 
     private boolean mShouldOffsetWallpaperCenter;
 
+    /**
+     * Whether the wallpaper has been notified about a physical display switch event is started.
+     */
+    private volatile boolean mIsWallpaperNotifiedOnDisplaySwitch;
+
     private final Consumer<WindowState> mFindWallpapers = w -> {
         if (w.mAttrs.type == TYPE_WALLPAPER) {
             WallpaperWindowToken token = w.mToken.asWallpaperToken();
@@ -1010,6 +1016,52 @@
     }
 
     /**
+     * Notifies the wallpaper that the display turns off when switching physical device. If the
+     * wallpaper is currently visible, its client visibility will be preserved until the display is
+     * confirmed to be off or on.
+     */
+    void onDisplaySwitchStarted() {
+        mIsWallpaperNotifiedOnDisplaySwitch = notifyDisplaySwitch(true /* start */);
+    }
+
+    /**
+     * Called when the screen has finished turning on or the device goes to sleep. This is no-op if
+     * the operation is not part of a display switch.
+     */
+    void onDisplaySwitchFinished() {
+        // The method can be called outside WM lock (turned on), so only acquire lock if needed.
+        // This is to optimize the common cases that regular devices don't have display switch.
+        if (mIsWallpaperNotifiedOnDisplaySwitch) {
+            synchronized (mService.mGlobalLock) {
+                mIsWallpaperNotifiedOnDisplaySwitch = false;
+                notifyDisplaySwitch(false /* start */);
+            }
+        }
+    }
+
+    private boolean notifyDisplaySwitch(boolean start) {
+        boolean notified = false;
+        for (int curTokenNdx = mWallpaperTokens.size() - 1; curTokenNdx >= 0; curTokenNdx--) {
+            final WallpaperWindowToken token = mWallpaperTokens.get(curTokenNdx);
+            for (int i = token.getChildCount() - 1; i >= 0; i--) {
+                final WindowState w = token.getChildAt(i);
+                if (start && !w.mWinAnimator.getShown()) {
+                    continue;
+                }
+                try {
+                    w.mClient.dispatchWallpaperCommand(COMMAND_DISPLAY_SWITCH, 0 /* x */, 0 /* y */,
+                            start ? 1 : 0 /* use z as start or finish */,
+                            null /* bundle */, false /* sync */);
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Failed to dispatch COMMAND_DISPLAY_SWITCH " + e);
+                }
+                notified = true;
+            }
+        }
+        return notified;
+    }
+
+    /**
      * Each window can request a zoom, example:
      * - User is in overview, zoomed out.
      * - User also pulls down the shade.
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
index be96e60..9e00f92 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayPolicyTests.java
@@ -283,12 +283,12 @@
 
         policy.screenTurnedOff();
         policy.setAwake(false);
-        policy.screenTurnedOn(null /* screenOnListener */);
+        policy.screenTurningOn(null /* screenOnListener */);
         assertTrue(wpc.isShowingUiWhileDozing());
         policy.screenTurnedOff();
         assertFalse(wpc.isShowingUiWhileDozing());
 
-        policy.screenTurnedOn(null /* screenOnListener */);
+        policy.screenTurningOn(null /* screenOnListener */);
         assertTrue(wpc.isShowingUiWhileDozing());
         policy.setAwake(true);
         assertFalse(wpc.isShowingUiWhileDozing());
diff --git a/services/tests/wmtests/src/com/android/server/wm/TestDisplayContent.java b/services/tests/wmtests/src/com/android/server/wm/TestDisplayContent.java
index 45ecc3f..00ecd00 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TestDisplayContent.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TestDisplayContent.java
@@ -215,7 +215,7 @@
                 doReturn(false).when(newDisplay).supportsSystemDecorations();
             }
             // Update the display policy to make the screen fully turned on so animation is allowed
-            displayPolicy.screenTurnedOn(null /* screenOnListener */);
+            displayPolicy.screenTurningOn(null /* screenOnListener */);
             displayPolicy.finishKeyguardDrawn();
             displayPolicy.finishWindowsDrawn();
             displayPolicy.finishScreenTurningOn();
diff --git a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
index 60e84b0..04a252a 100644
--- a/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/WindowTestsBase.java
@@ -225,7 +225,7 @@
         mDefaultDisplay = mWm.mRoot.getDefaultDisplay();
         // Update the display policy to make the screen fully turned on so animation is allowed
         final DisplayPolicy displayPolicy = mDefaultDisplay.getDisplayPolicy();
-        displayPolicy.screenTurnedOn(null /* screenOnListener */);
+        displayPolicy.screenTurningOn(null /* screenOnListener */);
         displayPolicy.finishKeyguardDrawn();
         displayPolicy.finishWindowsDrawn();
         displayPolicy.finishScreenTurningOn();