diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationViewHierarchyManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationViewHierarchyManager.java
index 467f27f640f9e..396d86bab825e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationViewHierarchyManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationViewHierarchyManager.java
@@ -435,7 +435,8 @@ public class NotificationViewHierarchyManager implements DynamicPrivacyControlle
         final int N = mListContainer.getContainerChildCount();
 
         int visibleNotifications = 0;
-        boolean onKeyguard = mStatusBarStateController.getState() == StatusBarState.KEYGUARD;
+        boolean onKeyguard =
+                mStatusBarStateController.getCurrentOrUpcomingState() == StatusBarState.KEYGUARD;
         Stack<ExpandableNotificationRow> stack = new Stack<>();
         for (int i = N - 1; i >= 0; i--) {
             View child = mListContainer.getContainerChildAt(i);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarStateControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarStateControllerImpl.java
index f8a1ff879e721..0725bf961e137 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarStateControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarStateControllerImpl.java
@@ -82,6 +82,7 @@ public class StatusBarStateControllerImpl implements SysuiStatusBarStateControll
     private final UiEventLogger mUiEventLogger;
     private int mState;
     private int mLastState;
+    private int mUpcomingState;
     private boolean mLeaveOpenOnKeyguardHide;
     private boolean mKeyguardRequested;
 
@@ -169,6 +170,7 @@ public class StatusBarStateControllerImpl implements SysuiStatusBarStateControll
             }
             mLastState = mState;
             mState = state;
+            mUpcomingState = state;
             mUiEventLogger.log(StatusBarStateEvent.fromState(mState));
             for (RankedListener rl : new ArrayList<>(mListeners)) {
                 rl.mListener.onStateChanged(mState);
@@ -183,6 +185,16 @@ public class StatusBarStateControllerImpl implements SysuiStatusBarStateControll
         return true;
     }
 
+    @Override
+    public void setUpcomingState(int nextState) {
+        mUpcomingState = nextState;
+    }
+
+    @Override
+    public int getCurrentOrUpcomingState() {
+        return mUpcomingState;
+    }
+
     @Override
     public boolean isDozing() {
         return mIsDozing;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SysuiStatusBarStateController.java b/packages/SystemUI/src/com/android/systemui/statusbar/SysuiStatusBarStateController.java
index 73f3d90bd4f86..25200501a9165 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SysuiStatusBarStateController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SysuiStatusBarStateController.java
@@ -73,6 +73,20 @@ public interface SysuiStatusBarStateController extends StatusBarStateController
      */
     boolean setState(int state, boolean force);
 
+    /**
+     * Provides a hint that the status bar has started to transition to another
+     * {@link StatusBarState}. This suggests that a matching call to setState() with the same value
+     * will happen in the near future, although that may not happen if the animation is canceled,
+     * etc.
+     */
+    void setUpcomingState(int state);
+
+    /**
+     * If the status bar is in the process of transitioning to a new state, returns that state.
+     * Otherwise, returns the current state.
+     */
+    int getCurrentOrUpcomingState();
+
     /**
      * Update the dozing state from {@link StatusBar}'s perspective
      * @param isDozing well, are we dozing?
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
index e36c38edaa88e..1bf70e2a9c6d0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelViewController.java
@@ -4398,6 +4398,8 @@ public class NotificationPanelViewController extends PanelViewController {
      */
     public void showAodUi() {
         setDozing(true /* dozing */, false /* animate */, null);
+        mStatusBarStateController.setUpcomingState(KEYGUARD);
+        mEntryManager.updateNotifications("showAodUi");
         mStatusBarStateListener.onStateChanged(KEYGUARD);
         mStatusBarStateListener.onDozeAmountChanged(1f, 1f);
         setExpandedFraction(1f);