diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index d7b3728e98033..963a9dcb1b01e 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -1795,6 +1795,12 @@ final class ActivityStack extends ConfigurationContainer {
         }
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int taskNdx = mTaskHistory.size() - 1; taskNdx >= 0; --taskNdx) {
+            mTaskHistory.get(taskNdx).addStartingWindowsForVisibleActivities(taskSwitch);
+        }
+    }
+
     /**
      * @return true if the top visible activity wants to occlude the Keyguard, false otherwise
      */
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index 14899b4d52561..2c1c3a191d11b 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -3229,6 +3229,17 @@ public class ActivityStackSupervisor extends ConfigurationContainer implements D
         }
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int displayNdx = mActivityDisplays.size() - 1; displayNdx >= 0; --displayNdx) {
+            final ArrayList<ActivityStack> stacks = mActivityDisplays.valueAt(displayNdx).mStacks;
+            final int topStackNdx = stacks.size() - 1;
+            for (int stackNdx = topStackNdx; stackNdx >= 0; --stackNdx) {
+                final ActivityStack stack = stacks.get(stackNdx);
+                stack.addStartingWindowsForVisibleActivities(taskSwitch);
+            }
+        }
+    }
+
     void invalidateTaskLayers() {
         mTaskLayersChanged = true;
     }
diff --git a/services/core/java/com/android/server/am/KeyguardController.java b/services/core/java/com/android/server/am/KeyguardController.java
index cfe2eb0755309..b0a4746b8ff1d 100644
--- a/services/core/java/com/android/server/am/KeyguardController.java
+++ b/services/core/java/com/android/server/am/KeyguardController.java
@@ -29,6 +29,7 @@ import static com.android.server.wm.AppTransition.TRANSIT_FLAG_KEYGUARD_GOING_AW
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_GOING_AWAY;
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_OCCLUDE;
 import static com.android.server.wm.AppTransition.TRANSIT_KEYGUARD_UNOCCLUDE;
+import static com.android.server.wm.AppTransition.TRANSIT_NONE;
 import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
 
 import android.os.IBinder;
@@ -120,6 +121,7 @@ class KeyguardController {
 
                 // Some stack visibility might change (e.g. docked stack)
                 mStackSupervisor.ensureActivitiesVisibleLocked(null, 0, !PRESERVE_WINDOWS);
+                mStackSupervisor.addStartingWindowsForVisibleActivities(true /* taskSwitch */);
                 mWindowManager.executeAppTransition();
             } finally {
                 mWindowManager.continueSurfaceLayout();
diff --git a/services/core/java/com/android/server/am/TaskRecord.java b/services/core/java/com/android/server/am/TaskRecord.java
index 1f5152a1606d4..2745fe5092687 100644
--- a/services/core/java/com/android/server/am/TaskRecord.java
+++ b/services/core/java/com/android/server/am/TaskRecord.java
@@ -1981,6 +1981,15 @@ final class TaskRecord extends ConfigurationContainer implements TaskWindowConta
         return rootAffinity != null && getStackId() != PINNED_STACK_ID;
     }
 
+    void addStartingWindowsForVisibleActivities(boolean taskSwitch) {
+        for (int activityNdx = mActivities.size() - 1; activityNdx >= 0; --activityNdx) {
+            final ActivityRecord r = mActivities.get(activityNdx);
+            if (r.visible) {
+                r.showStartingWindow(null /* prev */, false /* newTask */, taskSwitch);
+            }
+        }
+    }
+
     void dump(PrintWriter pw, String prefix) {
         pw.print(prefix); pw.print("userId="); pw.print(userId);
                 pw.print(" effectiveUid="); UserHandle.formatUid(pw, effectiveUid);
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index ac9859d42fcce..29c685d5fe97e 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -375,6 +375,7 @@ class AppWindowToken extends WindowToken implements WindowManagerService.AppFree
                 // affected.
                 mService.getDefaultDisplayContentLocked().getDockedDividerController()
                         .notifyAppVisibilityChanged();
+                mService.mTaskSnapshotController.notifyAppVisibilityChanged(this, visible);
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 3958510f59579..73e88c1958a95 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -2282,7 +2282,7 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
                 final boolean foundTargetWs =
                         (w.mAppToken != null && w.mAppToken.token == appToken)
                                 || (mScreenshotApplicationState.appWin != null && wallpaperOnly);
-                if (foundTargetWs && w.isDisplayedLw() && winAnim.getShown()) {
+                if (foundTargetWs && w.isVisible() && winAnim.getShown()) {
                     mScreenshotApplicationState.screenshotReady = true;
                 }
 
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 15878f6e414a3..2b74f84701b49 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -26,6 +26,8 @@ import android.os.Environment;
 import android.util.ArraySet;
 import android.view.WindowManagerPolicy.StartingSurface;
 
+import com.google.android.collect.Sets;
+
 import com.android.internal.annotations.VisibleForTesting;
 
 import java.io.PrintWriter;
@@ -66,10 +68,27 @@ class TaskSnapshotController {
         if (!ENABLE_TASK_SNAPSHOTS) {
             return;
         }
+        handleClosingApps(mService.mClosingApps);
+    }
+
+
+    /**
+     * Called when the visibility of an app changes outside of the regular app transition flow.
+     */
+    void notifyAppVisibilityChanged(AppWindowToken appWindowToken, boolean visible) {
+        if (!ENABLE_TASK_SNAPSHOTS) {
+            return;
+        }
+        if (!visible) {
+            handleClosingApps(Sets.newArraySet(appWindowToken));
+        }
+    }
+
+    private void handleClosingApps(ArraySet<AppWindowToken> closingApps) {
 
         // We need to take a snapshot of the task if and only if all activities of the task are
         // either closing or hidden.
-        getClosingTasks(mService.mClosingApps, mTmpTasks);
+        getClosingTasks(closingApps, mTmpTasks);
         for (int i = mTmpTasks.size() - 1; i >= 0; i--) {
             final Task task = mTmpTasks.valueAt(i);
             if (!canSnapshotTask(task)) {