diff --git a/core/java/android/provider/ContactsContract.java b/core/java/android/provider/ContactsContract.java
index f8aa98e..5036abc 100644
--- a/core/java/android/provider/ContactsContract.java
+++ b/core/java/android/provider/ContactsContract.java
@@ -8788,7 +8788,12 @@
         public static final String KEY_DEFAULT_ACCOUNT = "key_default_account";
 
         /**
-         * Return the account that was set to default account for new contacts.
+         * Get the account that is set as the default account for new contacts, which should be
+         * initially selected when creating a new contact on contact management apps.
+         *
+         * @param resolver the ContentResolver to query.
+         * @return the default account for new contacts, or null if it's not set or set to NULL
+         * account.
          */
         @Nullable
         public static Account getDefaultAccount(@NonNull ContentResolver resolver) {
@@ -8798,8 +8803,10 @@
         }
 
         /**
-         * Set the account to be the default account for new contacts.
+         * Sets the account as the default account that should be initially selected
+         * when creating a new contact on contact management apps.
          *
+         * @param resolver the ContentResolver to query.
          * @param account the account to be set to default.
          * @hide
          */
diff --git a/core/java/android/provider/Settings.java b/core/java/android/provider/Settings.java
index dfc4fa7..17272d3 100644
--- a/core/java/android/provider/Settings.java
+++ b/core/java/android/provider/Settings.java
@@ -13587,7 +13587,7 @@
          * Whether of not to send keycode sleep for ungaze when Home is the foreground activity on
          * watch type devices.
          * Type: int (0 for false, 1 for true)
-         * Default: 0
+         * Default: 1
          * @hide
          */
         @Readable
diff --git a/core/java/android/util/Log.java b/core/java/android/util/Log.java
index 12bcd8b..b5fe4f5 100644
--- a/core/java/android/util/Log.java
+++ b/core/java/android/util/Log.java
@@ -62,6 +62,10 @@
  * another buffer allocation and copy, and even more pressure on the gc.
  * That means that if your log message is filtered out, you might be doing
  * significant work and incurring significant overhead.
+ *
+ * <p>When calling the log methods that take a Throwable parameter,
+ * if any of the throwables in the cause chain is an <code>UnknownHostException</code>,
+ * then the stack trace is not logged.
  */
 public final class Log {
     /** @hide */
@@ -341,6 +345,9 @@
 
     /**
      * Handy function to get a loggable stack trace from a Throwable
+
+     * <p>If any of the throwables in the cause chain is an <code>UnknownHostException</code>,
+     * this returns an empty string.
      * @param tr An exception to log
      */
     @NonNull
diff --git a/core/java/android/view/IWindow.aidl b/core/java/android/view/IWindow.aidl
index 2ee112b..fceda20 100644
--- a/core/java/android/view/IWindow.aidl
+++ b/core/java/android/view/IWindow.aidl
@@ -96,9 +96,9 @@
 
     /**
      * Tell the window that it is either gaining or losing focus.  Keep it up
-     * to date on the current state showing navigational focus (touch mode) too.
+     * to date on the current state showing navigational focus too.
      */
-    void windowFocusChanged(boolean hasFocus, boolean inTouchMode);
+    void windowFocusChanged(boolean hasFocus);
 
     void closeSystemDialogs(String reason);
 
diff --git a/core/java/android/view/InputEventReceiver.java b/core/java/android/view/InputEventReceiver.java
index 2165f55..c9abec9 100644
--- a/core/java/android/view/InputEventReceiver.java
+++ b/core/java/android/view/InputEventReceiver.java
@@ -139,11 +139,9 @@
      * @param hasFocus if true, the window associated with this input channel has just received
      *                 focus
      *                 if false, the window associated with this input channel has just lost focus
-     * @param inTouchMode if true, the device is in touch mode
-     *                    if false, the device is not in touch mode
      */
     // Called from native code.
-    public void onFocusEvent(boolean hasFocus, boolean inTouchMode) {
+    public void onFocusEvent(boolean hasFocus) {
     }
 
     /**
@@ -175,7 +173,7 @@
      * exited touch mode.
      *
      * @param inTouchMode {@code true} if the display showing the window associated with the
-     *                                  input channel entered touch mode.
+     *                    input channel entered touch mode or {@code false} if left touch mode
      */
     public void onTouchModeChanged(boolean inTouchMode) {
     }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index ce96eca..89b95c0 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -8705,8 +8705,8 @@
         }
 
         @Override
-        public void onFocusEvent(boolean hasFocus, boolean inTouchMode) {
-            windowFocusChanged(hasFocus, inTouchMode);
+        public void onFocusEvent(boolean hasFocus) {
+            windowFocusChanged(hasFocus);
         }
 
         @Override
@@ -8975,9 +8975,7 @@
     /**
      * Notifies this {@link ViewRootImpl} object that window focus has changed.
      */
-    public void windowFocusChanged(boolean hasFocus, boolean unusedInTouchMode) {
-        // TODO(b/193718270): Remove inTouchMode parameter from this method and update related code
-        //     accordingly.
+    public void windowFocusChanged(boolean hasFocus) {
         synchronized (this) {
             mWindowFocusChanged = true;
             mUpcomingWindowFocus = hasFocus;
@@ -9732,10 +9730,10 @@
         }
 
         @Override
-        public void windowFocusChanged(boolean hasFocus, boolean inTouchMode) {
+        public void windowFocusChanged(boolean hasFocus) {
             final ViewRootImpl viewAncestor = mViewAncestor.get();
             if (viewAncestor != null) {
-                viewAncestor.windowFocusChanged(hasFocus, inTouchMode);
+                viewAncestor.windowFocusChanged(hasFocus);
             }
         }
 
diff --git a/core/java/com/android/internal/policy/PhoneWindow.java b/core/java/com/android/internal/policy/PhoneWindow.java
index b428970..7755b69 100644
--- a/core/java/com/android/internal/policy/PhoneWindow.java
+++ b/core/java/com/android/internal/policy/PhoneWindow.java
@@ -1842,7 +1842,7 @@
     @Override
     public void setLocalFocus(boolean hasFocus, boolean inTouchMode) {
         ViewRootImpl viewRoot = getViewRootImpl();
-        viewRoot.windowFocusChanged(hasFocus, inTouchMode);
+        viewRoot.windowFocusChanged(hasFocus);
         viewRoot.touchModeChanged(inTouchMode);
     }
 
diff --git a/core/java/com/android/internal/view/BaseIWindow.java b/core/java/com/android/internal/view/BaseIWindow.java
index 9e741e2..f93e32f 100644
--- a/core/java/com/android/internal/view/BaseIWindow.java
+++ b/core/java/com/android/internal/view/BaseIWindow.java
@@ -90,7 +90,7 @@
     }
 
     @Override
-    public void windowFocusChanged(boolean hasFocus, boolean touchEnabled) {
+    public void windowFocusChanged(boolean hasFocus) {
     }
 
     @Override
diff --git a/core/jni/android_view_InputEventReceiver.cpp b/core/jni/android_view_InputEventReceiver.cpp
index 67ab30b..1159c8a 100644
--- a/core/jni/android_view_InputEventReceiver.cpp
+++ b/core/jni/android_view_InputEventReceiver.cpp
@@ -390,13 +390,11 @@
             case AINPUT_EVENT_TYPE_FOCUS: {
                 FocusEvent* focusEvent = static_cast<FocusEvent*>(inputEvent);
                 if (kDebugDispatchCycle) {
-                    ALOGD("channel '%s' ~ Received focus event: hasFocus=%s, inTouchMode=%s.",
-                          getInputChannelName().c_str(), toString(focusEvent->getHasFocus()),
-                          toString(focusEvent->getInTouchMode()));
+                    ALOGD("channel '%s' ~ Received focus event: hasFocus=%s.",
+                          getInputChannelName().c_str(), toString(focusEvent->getHasFocus()));
                 }
                 env->CallVoidMethod(receiverObj.get(), gInputEventReceiverClassInfo.onFocusEvent,
-                                    jboolean(focusEvent->getHasFocus()),
-                                    jboolean(focusEvent->getInTouchMode()));
+                                    jboolean(focusEvent->getHasFocus()));
                 finishInputEvent(seq, true /* handled */);
                 continue;
             }
@@ -615,7 +613,7 @@
             gInputEventReceiverClassInfo.clazz,
             "dispatchInputEvent", "(ILandroid/view/InputEvent;)V");
     gInputEventReceiverClassInfo.onFocusEvent =
-            GetMethodIDOrDie(env, gInputEventReceiverClassInfo.clazz, "onFocusEvent", "(ZZ)V");
+            GetMethodIDOrDie(env, gInputEventReceiverClassInfo.clazz, "onFocusEvent", "(Z)V");
     gInputEventReceiverClassInfo.onPointerCaptureEvent =
             GetMethodIDOrDie(env, gInputEventReceiverClassInfo.clazz, "onPointerCaptureEvent",
                              "(Z)V");
diff --git a/core/proto/android/server/windowmanagerservice.proto b/core/proto/android/server/windowmanagerservice.proto
index 8ef3825..211f78e 100644
--- a/core/proto/android/server/windowmanagerservice.proto
+++ b/core/proto/android/server/windowmanagerservice.proto
@@ -112,8 +112,9 @@
     option (.android.msg_privacy).dest = DEST_AUTOMATIC;
 
     optional bool keyguard_showing = 1;
-    repeated KeyguardOccludedProto keyguard_occluded_states = 2;
+    repeated KeyguardOccludedProto keyguard_occluded_states = 2 [deprecated=true];
     optional bool aod_showing = 3;
+    repeated KeyguardPerDisplayProto keyguard_per_display = 4;
 }
 
 message KeyguardOccludedProto {
@@ -123,6 +124,15 @@
     optional bool keyguard_occluded = 2;
 }
 
+message KeyguardPerDisplayProto {
+    option (.android.msg_privacy).dest = DEST_AUTOMATIC;
+
+    optional int32 display_id = 1;
+    optional bool keyguard_showing = 2;
+    optional bool aod_showing = 3;
+    optional bool keyguard_occluded = 4;
+}
+
 /* represents PhoneWindowManager */
 message WindowManagerPolicyProto {
     option (.android.msg_privacy).dest = DEST_AUTOMATIC;
diff --git a/core/tests/coretests/src/android/view/ViewRootImplTest.java b/core/tests/coretests/src/android/view/ViewRootImplTest.java
index 882a0da..e958a96 100644
--- a/core/tests/coretests/src/android/view/ViewRootImplTest.java
+++ b/core/tests/coretests/src/android/view/ViewRootImplTest.java
@@ -62,7 +62,7 @@
  * Tests for {@link ViewRootImpl}
  *
  * Build/Install/Run:
- * atest FrameworksCoreTests:ViewRootImplTest
+ *  atest FrameworksCoreTests:ViewRootImplTest
  */
 @Presubmit
 @SmallTest
@@ -300,7 +300,7 @@
     @Test
     public void whenWindowDoesNotHaveFocus_keysAreDropped() {
         checkKeyEvent(() -> {
-            mViewRootImpl.windowFocusChanged(false /*hasFocus*/, true /*inTouchMode*/);
+            mViewRootImpl.windowFocusChanged(false /*hasFocus*/);
         }, false /*shouldReceiveKey*/);
     }
 
@@ -310,7 +310,7 @@
     @Test
     public void whenWindowHasFocus_keysAreReceived() {
         checkKeyEvent(() -> {
-            mViewRootImpl.windowFocusChanged(true /*hasFocus*/, true /*inTouchMode*/);
+            mViewRootImpl.windowFocusChanged(true /*hasFocus*/);
         }, true /*shouldReceiveKey*/);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index caa5327..ec59fad 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -111,8 +111,7 @@
 
     private static final String TAG = TAG_WITH_CLASS_NAME ? "BubbleController" : TAG_BUBBLES;
 
-    // TODO(b/173386799) keep in sync with Launcher3 and also don't do a broadcast
-    public static final String TASKBAR_CHANGED_BROADCAST = "taskbarChanged";
+    // TODO(b/173386799) keep in sync with Launcher3, not hooked up to anything
     public static final String EXTRA_TASKBAR_CREATED = "taskbarCreated";
     public static final String EXTRA_BUBBLE_OVERFLOW_OPENED = "bubbleOverflowOpened";
     public static final String EXTRA_TASKBAR_VISIBLE = "taskbarVisible";
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/SystemWindows.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/SystemWindows.java
index 33ce383..c5713b4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/SystemWindows.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/SystemWindows.java
@@ -356,7 +356,7 @@
         public void dispatchGetNewSurface() {}
 
         @Override
-        public void windowFocusChanged(boolean hasFocus, boolean inTouchMode) {}
+        public void windowFocusChanged(boolean hasFocus) {}
 
         @Override
         public void executeCommand(String command, String parameters, ParcelFileDescriptor out) {}
diff --git a/media/java/android/media/tv/TvInputService.java b/media/java/android/media/tv/TvInputService.java
index f312f64..d285b13 100755
--- a/media/java/android/media/tv/TvInputService.java
+++ b/media/java/android/media/tv/TvInputService.java
@@ -53,6 +53,7 @@
 import android.view.MotionEvent;
 import android.view.Surface;
 import android.view.View;
+import android.view.ViewRootImpl;
 import android.view.WindowManager;
 import android.view.accessibility.CaptioningManager;
 import android.widget.FrameLayout;
@@ -1581,7 +1582,8 @@
                 return TvInputManager.Session.DISPATCH_NOT_HANDLED;
             }
             if (!mOverlayViewContainer.hasWindowFocus()) {
-                mOverlayViewContainer.getViewRootImpl().windowFocusChanged(true, true);
+                ViewRootImpl viewRoot = mOverlayViewContainer.getViewRootImpl();
+                viewRoot.windowFocusChanged(true);
             }
             if (isNavigationKey && mOverlayViewContainer.hasFocusable()) {
                 // If mOverlayView has focusable views, navigation key events should be always
diff --git a/media/jni/android_media_Utils.cpp b/media/jni/android_media_Utils.cpp
index ecd9cc1..39b560b 100644
--- a/media/jni/android_media_Utils.cpp
+++ b/media/jni/android_media_Utils.cpp
@@ -65,6 +65,7 @@
         case HAL_PIXEL_FORMAT_Y8:
         case HAL_PIXEL_FORMAT_Y16:
         case HAL_PIXEL_FORMAT_RAW16:
+        case HAL_PIXEL_FORMAT_RAW12:
         case HAL_PIXEL_FORMAT_RAW10:
         case HAL_PIXEL_FORMAT_RAW_OPAQUE:
         case HAL_PIXEL_FORMAT_BLOB:
diff --git a/media/jni/android_media_tv_Tuner.cpp b/media/jni/android_media_tv_Tuner.cpp
index 24eea30..c4dfee4 100644
--- a/media/jni/android_media_tv_Tuner.cpp
+++ b/media/jni/android_media_tv_Tuner.cpp
@@ -3881,8 +3881,8 @@
 
     Result r = filterClient->close();
     filterClient->decStrong(filter);
-       env->SetLongField(filter, gFields.sharedFilterContext, 0);
     if (shared) {
+        env->SetLongField(filter, gFields.sharedFilterContext, 0);
     } else {
         env->SetLongField(filter, gFields.filterContext, 0);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index dd44f72..022baf6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
@@ -507,15 +507,19 @@
 
         riv.setRevealParameters(cx, cy, r);
         riv.setPendingIntent(pendingIntent);
+        riv.getController().setPendingIntent(pendingIntent);
         riv.setRemoteInput(inputs, input, editedSuggestionInfo);
+        riv.getController().setRemoteInput(input);
+        riv.getController().setRemoteInputs(inputs);
         riv.focusAnimated();
         if (userMessageContent != null) {
             riv.setEditTextContent(userMessageContent);
         }
         if (deferBouncer) {
             final ExpandableNotificationRow finalRow = row;
-            riv.setBouncerChecker(() -> !authBypassCheck.canSendRemoteInputWithoutBouncer()
-                    && showBouncerForRemoteInput(view, pendingIntent, finalRow));
+            riv.getController().setBouncerChecker(() ->
+                    !authBypassCheck.canSendRemoteInputWithoutBouncer()
+                            && showBouncerForRemoteInput(view, pendingIntent, finalRow));
         }
 
         return true;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilder.kt
index 010b6f80..8182e73 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilder.kt
@@ -51,8 +51,10 @@
 
             // If this notif begins a new section, first add the section's header view
             if (section != currentSection) {
-                section.headerController?.let { headerController ->
-                    root.children.add(NodeSpecImpl(root, headerController))
+                if (section.headerController != currentSection?.headerController) {
+                    section.headerController?.let { headerController ->
+                        root.children.add(NodeSpecImpl(root, headerController))
+                    }
                 }
                 prevSections.add(currentSection)
                 currentSection = section
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
index 4568470..1b4d369 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
@@ -1289,8 +1289,8 @@
                     result.mView = riv;
                     // Create a new controller for the view. The lifetime of the controller is 1:1
                     // with that of the view.
-                    RemoteInputViewSubcomponent subcomponent =
-                            mRemoteInputSubcomponentFactory.create(result.mView);
+                    RemoteInputViewSubcomponent subcomponent = mRemoteInputSubcomponentFactory
+                            .create(result.mView, mRemoteInputController);
                     result.mController = subcomponent.getController();
                     result.mView.setController(result.mController);
                 } else {
@@ -1311,8 +1311,9 @@
                     Notification.Action[] actions = entry.getSbn().getNotification().actions;
                     if (existingPendingIntent != null) {
                         result.mView.setPendingIntent(existingPendingIntent);
+                        result.mController.setPendingIntent(existingPendingIntent);
                     }
-                    if (result.mView.updatePendingIntentFromActions(actions)) {
+                    if (result.mController.updatePendingIntentFromActions(actions)) {
                         if (!result.mView.isActive()) {
                             result.mView.focus();
                         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index b47e71a..06cf70d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -40,7 +40,6 @@
 import static com.android.systemui.statusbar.phone.BarTransitions.MODE_SEMI_TRANSPARENT;
 import static com.android.systemui.statusbar.phone.BarTransitions.MODE_TRANSPARENT;
 import static com.android.systemui.statusbar.phone.BarTransitions.TransitionMode;
-import static com.android.wm.shell.bubbles.BubbleController.TASKBAR_CHANGED_BROADCAST;
 
 import android.annotation.Nullable;
 import android.app.ActivityManager;
@@ -918,8 +917,6 @@
         mBypassHeadsUpNotifier.setUp();
         if (mBubblesOptional.isPresent()) {
             mBubblesOptional.get().setExpandListener(mBubbleExpandListener);
-            IntentFilter filter = new IntentFilter(TASKBAR_CHANGED_BROADCAST);
-            mBroadcastDispatcher.registerReceiver(mTaskbarChangeReceiver, filter);
         }
 
         mKeyguardIndicationController.init();
@@ -4270,13 +4267,6 @@
         }
     };
 
-    BroadcastReceiver mTaskbarChangeReceiver = new BroadcastReceiver() {
-        @Override
-        public void onReceive(Context context, Intent intent) {
-            mBubblesOptional.ifPresent(bubbles -> bubbles.onTaskbarChanged(intent.getExtras()));
-        }
-    };
-
     private final ConfigurationListener mConfigurationListener = new ConfigurationListener() {
         @Override
         public void onConfigChanged(Configuration newConfig) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
index d25013a..46fa20d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputView.java
@@ -25,9 +25,7 @@
 import android.app.PendingIntent;
 import android.app.RemoteInput;
 import android.content.Context;
-import android.content.Intent;
 import android.content.pm.PackageManager;
-import android.content.pm.ShortcutManager;
 import android.content.res.ColorStateList;
 import android.content.res.TypedArray;
 import android.graphics.BlendMode;
@@ -35,13 +33,9 @@
 import android.graphics.PorterDuff;
 import android.graphics.Rect;
 import android.graphics.drawable.GradientDrawable;
-import android.net.Uri;
-import android.os.Bundle;
-import android.os.SystemClock;
 import android.os.UserHandle;
 import android.text.Editable;
 import android.text.SpannedString;
-import android.text.TextUtils;
 import android.text.TextWatcher;
 import android.util.ArraySet;
 import android.util.AttributeSet;
@@ -78,7 +72,6 @@
 import com.android.internal.logging.UiEventLogger;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
-import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.RemoteInputController;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry.EditedSuggestionInfo;
@@ -89,7 +82,6 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.List;
 import java.util.function.Consumer;
 
@@ -107,7 +99,7 @@
 
     private final SendButtonTextWatcher mTextWatcher;
     private final TextView.OnEditorActionListener mEditorActionHandler;
-    private final ArrayList<OnSendRemoteInputListener> mOnSendListeners = new ArrayList<>();
+    private final ArrayList<Runnable> mOnSendListeners = new ArrayList<>();
     private final ArrayList<Consumer<Boolean>> mOnVisibilityChangedListeners = new ArrayList<>();
     private final ArrayList<OnFocusChangeListener> mEditTextFocusChangeListeners =
             new ArrayList<>();
@@ -133,7 +125,6 @@
     private PendingIntent mPendingIntent;
     private RemoteInput mRemoteInput;
     private RemoteInput[] mRemoteInputs;
-    private NotificationRemoteInputManager.BouncerChecker mBouncerChecker;
     private boolean mRemoved;
     private NotificationViewWrapper mWrapper;
 
@@ -178,6 +169,12 @@
         ta.recycle();
     }
 
+    // TODO(b/193539698): move to Controller, since we're just directly accessing a system service
+    /** Hide the IME, if visible. */
+    public void hideIme() {
+        mEditText.hideIme();
+    }
+
     private ColorStateList colorStateListWithDisabledAlpha(int color, int disabledAlpha) {
         return new ColorStateList(new int[][]{
                 new int[]{-com.android.internal.R.attr.state_enabled}, // disabled
@@ -303,6 +300,11 @@
         return mViewController;
     }
 
+    /** Clear the attachment, if present. */
+    public void clearAttachment() {
+        setAttachment(null);
+    }
+
     @VisibleForTesting
     protected void setAttachment(ContentInfo item) {
         if (mEntry.remoteInputAttachment != null && mEntry.remoteInputAttachment != item) {
@@ -339,121 +341,18 @@
         updateSendButton();
     }
 
-    /**
-     * Reply intent
-     * @return returns intent with granted URI permissions that should be used immediately
-     */
-    private Intent prepareRemoteInput() {
-        return mEntry.remoteInputAttachment == null
-                ? prepareRemoteInputFromText()
-                : prepareRemoteInputFromData(mEntry.remoteInputMimeType, mEntry.remoteInputUri);
-    }
-
-    private Intent prepareRemoteInputFromText() {
-        Bundle results = new Bundle();
-        results.putString(mRemoteInput.getResultKey(), mEditText.getText().toString());
-        Intent fillInIntent = new Intent().addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
-        RemoteInput.addResultsToIntent(mRemoteInputs, fillInIntent,
-                results);
-
-        mEntry.remoteInputText = mEditText.getText().toString();
-        setAttachment(null);
-        mEntry.remoteInputUri = null;
-        mEntry.remoteInputMimeType = null;
-
-        if (mEntry.editedSuggestionInfo == null) {
-            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_FREE_FORM_INPUT);
-        } else {
-            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_CHOICE);
-        }
-
-        return fillInIntent;
-    }
-
-    private Intent prepareRemoteInputFromData(String contentType, Uri data) {
-        HashMap<String, Uri> results = new HashMap<>();
-        results.put(contentType, data);
-        // grant for the target app.
-        mController.grantInlineReplyUriPermission(mEntry.getSbn(), data);
-        Intent fillInIntent = new Intent().addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
-        RemoteInput.addDataResultToIntent(mRemoteInput, fillInIntent, results);
-
-        Bundle bundle = new Bundle();
-        bundle.putString(mRemoteInput.getResultKey(), mEditText.getText().toString());
-        RemoteInput.addResultsToIntent(mRemoteInputs, fillInIntent,
-                bundle);
-
-        CharSequence attachmentText =
-                mEntry.remoteInputAttachment.getClip().getDescription().getLabel();
-
-        CharSequence attachmentLabel = TextUtils.isEmpty(attachmentText)
-                ? mContext.getString(R.string.remote_input_image_insertion_text)
-                : attachmentText;
-        // add content description to reply text for context
-        CharSequence fullText = TextUtils.isEmpty(mEditText.getText())
-                ? attachmentLabel
-                : "\"" + attachmentLabel + "\" " + mEditText.getText();
-
-        mEntry.remoteInputText = fullText;
-
-        // mirror prepareRemoteInputFromText for text input
-        if (mEntry.editedSuggestionInfo == null) {
-            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_FREE_FORM_INPUT);
-        } else if (mEntry.remoteInputAttachment == null) {
-            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_CHOICE);
-        }
-
-        return fillInIntent;
-    }
-
-    private void sendRemoteInput(Intent intent) {
-        if (mBouncerChecker != null && mBouncerChecker.showBouncerIfNecessary()) {
-            mEditText.hideIme();
-            for (OnSendRemoteInputListener listener : new ArrayList<>(mOnSendListeners)) {
-                listener.onSendRequestBounced();
-            }
-            return;
-        }
-
+    /** Show the "sending in-progress" UI. */
+    public void startSending() {
         mEditText.setEnabled(false);
         mSendButton.setVisibility(INVISIBLE);
         mProgressBar.setVisibility(VISIBLE);
-        mEntry.lastRemoteInputSent = SystemClock.elapsedRealtime();
-        mEntry.mRemoteEditImeAnimatingAway = true;
-        mController.addSpinning(mEntry.getKey(), mToken);
-        mController.removeRemoteInput(mEntry, mToken);
         mEditText.mShowImeOnInputConnection = false;
-        mController.remoteInputSent(mEntry);
-        mEntry.setHasSentReply();
+    }
 
-        for (OnSendRemoteInputListener listener : new ArrayList<>(mOnSendListeners)) {
-            listener.onSendRemoteInput();
+    private void sendRemoteInput() {
+        for (Runnable listener : new ArrayList<>(mOnSendListeners)) {
+            listener.run();
         }
-
-        // Tell ShortcutManager that this package has been "activated".  ShortcutManager
-        // will reset the throttling for this package.
-        // Strictly speaking, the intent receiver may be different from the notification publisher,
-        // but that's an edge case, and also because we can't always know which package will receive
-        // an intent, so we just reset for the publisher.
-        getContext().getSystemService(ShortcutManager.class).onApplicationActive(
-                mEntry.getSbn().getPackageName(),
-                mEntry.getSbn().getUser().getIdentifier());
-
-        mUiEventLogger.logWithInstanceId(
-                NotificationRemoteInputEvent.NOTIFICATION_REMOTE_INPUT_SEND,
-                mEntry.getSbn().getUid(), mEntry.getSbn().getPackageName(),
-                mEntry.getSbn().getInstanceId());
-        try {
-            mPendingIntent.send(mContext, 0, intent);
-        } catch (PendingIntent.CanceledException e) {
-            Log.i(TAG, "Unable to send remote input result", e);
-            mUiEventLogger.logWithInstanceId(
-                    NotificationRemoteInputEvent.NOTIFICATION_REMOTE_INPUT_FAILURE,
-                    mEntry.getSbn().getUid(), mEntry.getSbn().getPackageName(),
-                    mEntry.getSbn().getInstanceId());
-        }
-
-        setAttachment(null);
     }
 
     public CharSequence getText() {
@@ -478,7 +377,7 @@
     @Override
     public void onClick(View v) {
         if (v == mSendButton) {
-            sendRemoteInput(prepareRemoteInput());
+            sendRemoteInput();
         }
     }
 
@@ -687,55 +586,18 @@
         return mEditText.isFocused() && mEditText.isEnabled();
     }
 
+    // TODO(b/193539698): move this to the controller
     public void stealFocusFrom(RemoteInputView other) {
         other.close();
         setPendingIntent(other.mPendingIntent);
         setRemoteInput(other.mRemoteInputs, other.mRemoteInput, mEntry.editedSuggestionInfo);
         setRevealParameters(other.mRevealCx, other.mRevealCy, other.mRevealR);
+        getController().setPendingIntent(other.mPendingIntent);
+        getController().setRemoteInput(other.mRemoteInput);
+        getController().setRemoteInputs(other.mRemoteInputs);
         focus();
     }
 
-    /**
-     * Tries to find an action in {@param actions} that matches the current pending intent
-     * of this view and updates its state to that of the found action
-     *
-     * @return true if a matching action was found, false otherwise
-     */
-    public boolean updatePendingIntentFromActions(Notification.Action[] actions) {
-        if (mPendingIntent == null || actions == null) {
-            return false;
-        }
-        Intent current = mPendingIntent.getIntent();
-        if (current == null) {
-            return false;
-        }
-
-        for (Notification.Action a : actions) {
-            RemoteInput[] inputs = a.getRemoteInputs();
-            if (a.actionIntent == null || inputs == null) {
-                continue;
-            }
-            Intent candidate = a.actionIntent.getIntent();
-            if (!current.filterEquals(candidate)) {
-                continue;
-            }
-
-            RemoteInput input = null;
-            for (RemoteInput i : inputs) {
-                if (i.getAllowFreeFormInput()) {
-                    input = i;
-                }
-            }
-            if (input == null) {
-                continue;
-            }
-            setPendingIntent(a.actionIntent);
-            setRemoteInput(inputs, input, null /* editedSuggestionInfo*/);
-            return true;
-        }
-        return false;
-    }
-
     public PendingIntent getPendingIntent() {
         return mPendingIntent;
     }
@@ -814,16 +676,6 @@
         return getVisibility() == VISIBLE && mController.isSpinning(mEntry.getKey(), mToken);
     }
 
-    /**
-     * Sets a {@link com.android.systemui.statusbar.NotificationRemoteInputManager.BouncerChecker}
-     * that will be used to determine if the device needs to be unlocked before sending the
-     * RemoteInput.
-     */
-    public void setBouncerChecker(
-            @Nullable NotificationRemoteInputManager.BouncerChecker bouncerChecker) {
-        mBouncerChecker = bouncerChecker;
-    }
-
     /** Registers a listener for focus-change events on the EditText */
     public void addOnEditTextFocusChangedListener(View.OnFocusChangeListener listener) {
         mEditTextFocusChangeListeners.add(listener);
@@ -846,26 +698,15 @@
     }
 
     /** Registers a listener for send events on this RemoteInputView */
-    public void addOnSendRemoteInputListener(OnSendRemoteInputListener listener) {
+    public void addOnSendRemoteInputListener(Runnable listener) {
         mOnSendListeners.add(listener);
     }
 
     /** Removes a previously-added listener for send events on this RemoteInputView */
-    public void removeOnSendRemoteInputListener(OnSendRemoteInputListener listener) {
+    public void removeOnSendRemoteInputListener(Runnable listener) {
         mOnSendListeners.remove(listener);
     }
 
-    /** Listener for send events */
-    public interface OnSendRemoteInputListener {
-        /** Invoked when the remote input has been sent successfully. */
-        void onSendRemoteInput();
-        /**
-         * Invoked when the user had requested to send the remote input, but authentication was
-         * required and the bouncer was shown instead.
-         */
-        void onSendRequestBounced();
-    }
-
     /** Handler for button click on send action in IME. */
     private class EditorActionHandler implements TextView.OnEditorActionListener {
 
@@ -881,7 +722,7 @@
 
             if (isSoftImeEvent || isKeyboardEnterKey) {
                 if (mEditText.length() > 0 || mEntry.remoteInputAttachment != null) {
-                    sendRemoteInput(prepareRemoteInput());
+                    sendRemoteInput();
                 }
                 // Consume action to prevent IME from closing.
                 return true;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputViewController.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputViewController.kt
index 383170e..530da43 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputViewController.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/RemoteInputViewController.kt
@@ -16,28 +16,102 @@
 
 package com.android.systemui.statusbar.policy
 
+import android.app.Notification
+import android.app.PendingIntent
+import android.app.RemoteInput
+import android.content.Intent
+import android.content.pm.ShortcutManager
+import android.net.Uri
+import android.os.Bundle
+import android.os.SystemClock
+import android.text.TextUtils
+import android.util.ArraySet
+import android.util.Log
 import android.view.View
+import com.android.internal.logging.UiEventLogger
+import com.android.systemui.R
+import com.android.systemui.statusbar.NotificationRemoteInputManager
+import com.android.systemui.statusbar.RemoteInputController
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.policy.RemoteInputView.NotificationRemoteInputEvent
 import com.android.systemui.statusbar.policy.dagger.RemoteInputViewScope
 import javax.inject.Inject
 
 interface RemoteInputViewController {
     fun bind()
     fun unbind()
+
+    /**
+     * A [NotificationRemoteInputManager.BouncerChecker] that will be used to determine if the
+     * device needs to be unlocked before sending the RemoteInput.
+     */
+    var bouncerChecker: NotificationRemoteInputManager.BouncerChecker?
+
+    // TODO(b/193539698): these properties probably shouldn't be nullable
+    /** A [PendingIntent] to be used to send the RemoteInput. */
+    var pendingIntent: PendingIntent?
+    /** The [RemoteInput] data backing this Controller. */
+    var remoteInput: RemoteInput?
+    /** Other [RemoteInput]s from the notification associated with this Controller. */
+    var remoteInputs: Array<RemoteInput>?
+
+    /**
+     * Tries to find an action in {@param actions} that matches the current pending intent
+     * of this view and updates its state to that of the found action
+     *
+     * @return true if a matching action was found, false otherwise
+     */
+    fun updatePendingIntentFromActions(actions: Array<Notification.Action>?): Boolean
+
+    /** Registers a listener for send events. */
+    fun addOnSendRemoteInputListener(listener: OnSendRemoteInputListener)
+
+    /** Unregisters a listener previously registered via [addOnSendRemoteInputListener] */
+    fun removeOnSendRemoteInputListener(listener: OnSendRemoteInputListener)
 }
 
+/** Listener for send events  */
+interface OnSendRemoteInputListener {
+
+    /** Invoked when the remote input has been sent successfully.  */
+    fun onSendRemoteInput()
+
+    /**
+     * Invoked when the user had requested to send the remote input, but authentication was
+     * required and the bouncer was shown instead.
+     */
+    fun onSendRequestBounced()
+}
+
+private const val TAG = "RemoteInput"
+
 @RemoteInputViewScope
 class RemoteInputViewControllerImpl @Inject constructor(
     private val view: RemoteInputView,
-    private val remoteInputQuickSettingsDisabler: RemoteInputQuickSettingsDisabler
+    private val entry: NotificationEntry,
+    private val remoteInputQuickSettingsDisabler: RemoteInputQuickSettingsDisabler,
+    private val remoteInputController: RemoteInputController,
+    private val shortcutManager: ShortcutManager,
+    private val uiEventLogger: UiEventLogger
 ) : RemoteInputViewController {
 
+    private object Token
+    private val onSendListeners = ArraySet<OnSendRemoteInputListener>()
+    private val resources get() = view.resources
+
     private var isBound = false
 
+    override var pendingIntent: PendingIntent? = null
+    override var bouncerChecker: NotificationRemoteInputManager.BouncerChecker? = null
+    override var remoteInput: RemoteInput? = null
+    override var remoteInputs: Array<RemoteInput>? = null
+
     override fun bind() {
         if (isBound) return
         isBound = true
 
         view.addOnEditTextFocusChangedListener(onFocusChangeListener)
+        view.addOnSendRemoteInputListener(onSendRemoteInputListener)
     }
 
     override fun unbind() {
@@ -45,9 +119,159 @@
         isBound = false
 
         view.removeOnEditTextFocusChangedListener(onFocusChangeListener)
+        view.removeOnSendRemoteInputListener(onSendRemoteInputListener)
+    }
+
+    override fun updatePendingIntentFromActions(actions: Array<Notification.Action>?): Boolean {
+        actions ?: return false
+        val current: Intent = pendingIntent?.intent ?: return false
+        for (a in actions) {
+            val actionIntent = a.actionIntent ?: continue
+            val inputs = a.remoteInputs ?: continue
+            if (!current.filterEquals(actionIntent.intent)) continue
+            val input = inputs.firstOrNull { it.allowFreeFormInput } ?: continue
+            pendingIntent = actionIntent
+            remoteInput = input
+            remoteInputs = inputs
+            view.pendingIntent = actionIntent
+            view.setRemoteInput(inputs, input, null /* editedSuggestionInfo */)
+            return true
+        }
+        return false
+    }
+
+    override fun addOnSendRemoteInputListener(listener: OnSendRemoteInputListener) {
+        onSendListeners.add(listener)
+    }
+
+    /** Removes a previously-added listener for send events on this RemoteInputView  */
+    override fun removeOnSendRemoteInputListener(listener: OnSendRemoteInputListener) {
+        onSendListeners.remove(listener)
     }
 
     private val onFocusChangeListener = View.OnFocusChangeListener { _, hasFocus ->
         remoteInputQuickSettingsDisabler.setRemoteInputActive(hasFocus)
     }
-}
\ No newline at end of file
+
+    private val onSendRemoteInputListener = Runnable {
+        val remoteInput = remoteInput ?: run {
+            Log.e(TAG, "cannot send remote input, RemoteInput data is null")
+            return@Runnable
+        }
+        val pendingIntent = pendingIntent ?: run {
+            Log.e(TAG, "cannot send remote input, PendingIntent is null")
+            return@Runnable
+        }
+        val intent = prepareRemoteInput(remoteInput)
+        sendRemoteInput(pendingIntent, intent)
+    }
+
+    private fun sendRemoteInput(pendingIntent: PendingIntent, intent: Intent) {
+        if (bouncerChecker?.showBouncerIfNecessary() == true) {
+            view.hideIme()
+            for (listener in onSendListeners.toList()) {
+                listener.onSendRequestBounced()
+            }
+            return
+        }
+
+        view.startSending()
+
+        entry.lastRemoteInputSent = SystemClock.elapsedRealtime()
+        entry.mRemoteEditImeAnimatingAway = true
+        remoteInputController.addSpinning(entry.key, Token)
+        remoteInputController.removeRemoteInput(entry, Token)
+        remoteInputController.remoteInputSent(entry)
+        entry.setHasSentReply()
+
+        for (listener in onSendListeners.toList()) {
+            listener.onSendRemoteInput()
+        }
+
+        // Tell ShortcutManager that this package has been "activated". ShortcutManager will reset
+        // the throttling for this package.
+        // Strictly speaking, the intent receiver may be different from the notification publisher,
+        // but that's an edge case, and also because we can't always know which package will receive
+        // an intent, so we just reset for the publisher.
+        shortcutManager.onApplicationActive(entry.sbn.packageName, entry.sbn.user.identifier)
+
+        uiEventLogger.logWithInstanceId(
+                NotificationRemoteInputEvent.NOTIFICATION_REMOTE_INPUT_SEND,
+                entry.sbn.uid, entry.sbn.packageName,
+                entry.sbn.instanceId)
+
+        try {
+            pendingIntent.send(view.context, 0, intent)
+        } catch (e: PendingIntent.CanceledException) {
+            Log.i(TAG, "Unable to send remote input result", e)
+            uiEventLogger.logWithInstanceId(
+                    NotificationRemoteInputEvent.NOTIFICATION_REMOTE_INPUT_FAILURE,
+                    entry.sbn.uid, entry.sbn.packageName,
+                    entry.sbn.instanceId)
+        }
+
+        view.clearAttachment()
+    }
+
+    /**
+     * Reply intent
+     * @return returns intent with granted URI permissions that should be used immediately
+     */
+    private fun prepareRemoteInput(remoteInput: RemoteInput): Intent =
+            if (entry.remoteInputAttachment == null) prepareRemoteInputFromText(remoteInput)
+            else prepareRemoteInputFromData(
+                    remoteInput,
+                    entry.remoteInputMimeType,
+                    entry.remoteInputUri)
+
+    private fun prepareRemoteInputFromText(remoteInput: RemoteInput): Intent {
+        val results = Bundle()
+        results.putString(remoteInput.resultKey, view.text.toString())
+        val fillInIntent = Intent().addFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+        RemoteInput.addResultsToIntent(remoteInputs, fillInIntent, results)
+        entry.remoteInputText = view.text
+        view.clearAttachment()
+        entry.remoteInputUri = null
+        entry.remoteInputMimeType = null
+        if (entry.editedSuggestionInfo == null) {
+            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_FREE_FORM_INPUT)
+        } else {
+            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_CHOICE)
+        }
+        return fillInIntent
+    }
+
+    private fun prepareRemoteInputFromData(
+        remoteInput: RemoteInput,
+        contentType: String,
+        data: Uri
+    ): Intent {
+        val results = HashMap<String, Uri>()
+        results[contentType] = data
+        // grant for the target app.
+        remoteInputController.grantInlineReplyUriPermission(entry.sbn, data)
+        val fillInIntent = Intent().addFlags(Intent.FLAG_RECEIVER_FOREGROUND)
+        RemoteInput.addDataResultToIntent(remoteInput, fillInIntent, results)
+        val bundle = Bundle()
+        bundle.putString(remoteInput.resultKey, view.text.toString())
+        RemoteInput.addResultsToIntent(remoteInputs, fillInIntent, bundle)
+        val attachmentText: CharSequence = entry.remoteInputAttachment.clip.description.label
+        val attachmentLabel =
+                if (TextUtils.isEmpty(attachmentText))
+                    resources.getString(R.string.remote_input_image_insertion_text)
+                else attachmentText
+        // add content description to reply text for context
+        val fullText =
+                if (TextUtils.isEmpty(view.text)) attachmentLabel
+                else "\"" + attachmentLabel + "\" " + view.text
+        entry.remoteInputText = fullText
+
+        // mirror prepareRemoteInputFromText for text input
+        if (entry.editedSuggestionInfo == null) {
+            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_FREE_FORM_INPUT)
+        } else if (entry.remoteInputAttachment == null) {
+            RemoteInput.setResultsSource(fillInIntent, RemoteInput.SOURCE_CHOICE)
+        }
+        return fillInIntent
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/RemoteInput.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/RemoteInput.kt
index aff39ef..d4abc40 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/RemoteInput.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/dagger/RemoteInput.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.policy.dagger
 
+import com.android.systemui.statusbar.RemoteInputController
 import com.android.systemui.statusbar.policy.RemoteInputView
 import com.android.systemui.statusbar.policy.RemoteInputViewController
 import com.android.systemui.statusbar.policy.RemoteInputViewControllerImpl
@@ -32,7 +33,10 @@
 
     @Subcomponent.Factory
     interface Factory {
-        fun create(@BindsInstance view: RemoteInputView): RemoteInputViewSubcomponent
+        fun create(
+            @BindsInstance view: RemoteInputView,
+            @BindsInstance remoteInputController: RemoteInputController
+        ): RemoteInputViewSubcomponent
     }
 }
 
@@ -44,4 +48,4 @@
 
 @Qualifier
 @Retention(AnnotationRetention.BINARY)
-internal annotation class RemoteInputViewScope
\ No newline at end of file
+internal annotation class RemoteInputViewScope
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
index ed48452..6313d3a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/collection/render/NodeSpecBuilderTest.kt
@@ -58,6 +58,7 @@
     private val section1 = buildSection(1, section1Bucket, headerController1)
     private val section1NoHeader = buildSection(1, section1Bucket, null)
     private val section2 = buildSection(2, section2Bucket, headerController2)
+    private val section3 = buildSection(3, section2Bucket, headerController2)
 
     private val fakeViewBarn = FakeViewBarn()
 
@@ -75,6 +76,37 @@
     }
 
     @Test
+    fun testMultipleSectionsWithSameController() {
+        checkOutput(
+                listOf(
+                        notif(0, section0),
+                        notif(1, section2),
+                        notif(2, section3)
+                ),
+                tree(
+                        node(headerController0),
+                        notifNode(0),
+                        node(headerController2),
+                        notifNode(1),
+                        notifNode(2)
+                )
+        )
+    }
+
+    @Test(expected = RuntimeException::class)
+    fun testMultipleSectionsWithSameControllerNonConsecutive() {
+        checkOutput(
+                listOf(
+                        notif(0, section0),
+                        notif(1, section1),
+                        notif(2, section3),
+                        notif(3, section1)
+                ),
+                tree()
+        )
+    }
+
+    @Test
     fun testSimpleMapping() {
         checkOutput(
                 // GIVEN a simple flat list of notifications all in the same headerless section
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
index 27ddb36..ba7bbfe 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/RemoteInputViewTest.java
@@ -46,6 +46,7 @@
 import android.widget.EditText;
 import android.widget.ImageButton;
 
+import androidx.annotation.NonNull;
 import androidx.test.filters.SmallTest;
 
 import com.android.internal.logging.UiEventLogger;
@@ -55,6 +56,7 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.NotificationRemoteInputManager;
 import com.android.systemui.statusbar.RemoteInputController;
+import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
 import com.android.systemui.statusbar.phone.LightBarController;
@@ -85,7 +87,6 @@
     @Mock private LightBarController mLightBarController;
     private BlockingQueueIntentReceiver mReceiver;
     private final UiEventLoggerFake mUiEventLoggerFake = new UiEventLoggerFake();
-    private RemoteInputView mView;
 
     @Before
     public void setUp() throws Exception {
@@ -112,13 +113,17 @@
         mContext.unregisterReceiver(mReceiver);
     }
 
-    private void setTestPendingIntent(RemoteInputView view) {
+    private void setTestPendingIntent(RemoteInputView view, RemoteInputViewController controller) {
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0,
                 new Intent(TEST_ACTION), PendingIntent.FLAG_MUTABLE);
         RemoteInput input = new RemoteInput.Builder(TEST_RESULT_KEY).build();
+        RemoteInput[] inputs = {input};
 
         view.setPendingIntent(pendingIntent);
-        view.setRemoteInput(new RemoteInput[]{input}, input, null /* editedSuggestionInfo */);
+        controller.setPendingIntent(pendingIntent);
+        view.setRemoteInput(inputs, input, null /* editedSuggestionInfo */);
+        controller.setRemoteInput(input);
+        controller.setRemoteInputs(inputs);
     }
 
     @Test
@@ -129,8 +134,9 @@
                 TestableLooper.get(this));
         ExpandableNotificationRow row = helper.createRow();
         RemoteInputView view = RemoteInputView.inflate(mContext, null, row.getEntry(), mController);
+        RemoteInputViewController controller = bindController(view, row.getEntry());
 
-        setTestPendingIntent(view);
+        setTestPendingIntent(view, controller);
 
         view.focus();
 
@@ -140,6 +146,7 @@
         sendButton.performClick();
 
         Intent resultIntent = mReceiver.waitForIntent();
+        assertNotNull(resultIntent);
         assertEquals(TEST_REPLY,
                 RemoteInput.getResultsFromIntent(resultIntent).get(TEST_RESULT_KEY));
         assertEquals(RemoteInput.SOURCE_FREE_FORM_INPUT,
@@ -167,8 +174,9 @@
                 UserHandle.getUid(fromUser.getIdentifier(), DUMMY_MESSAGE_APP_ID),
                 toUser);
         RemoteInputView view = RemoteInputView.inflate(mContext, null, row.getEntry(), mController);
+        RemoteInputViewController controller = bindController(view, row.getEntry());
 
-        setTestPendingIntent(view);
+        setTestPendingIntent(view, controller);
 
         view.focus();
 
@@ -224,8 +232,9 @@
                 TestableLooper.get(this));
         ExpandableNotificationRow row = helper.createRow();
         RemoteInputView view = RemoteInputView.inflate(mContext, null, row.getEntry(), mController);
+        RemoteInputViewController controller = bindController(view, row.getEntry());
 
-        setTestPendingIntent(view);
+        setTestPendingIntent(view, controller);
 
         // Open view, send a reply
         view.focus();
@@ -253,8 +262,9 @@
                 TestableLooper.get(this));
         ExpandableNotificationRow row = helper.createRow();
         RemoteInputView view = RemoteInputView.inflate(mContext, null, row.getEntry(), mController);
+        RemoteInputViewController controller = bindController(view, row.getEntry());
 
-        setTestPendingIntent(view);
+        setTestPendingIntent(view, controller);
 
         // Open view, attach an image
         view.focus();
@@ -279,4 +289,21 @@
                         .NOTIFICATION_REMOTE_INPUT_ATTACH_IMAGE.getId(),
                 mUiEventLoggerFake.eventId(1));
     }
+
+    // NOTE: because we're refactoring the RemoteInputView and moving logic into the
+    //  RemoteInputViewController, it's easiest to just test the system of the two classes together.
+    @NonNull
+    private RemoteInputViewController bindController(
+            RemoteInputView view,
+            NotificationEntry entry) {
+        RemoteInputViewControllerImpl viewController = new RemoteInputViewControllerImpl(
+                view,
+                entry,
+                mRemoteInputQuickSettingsDisabler,
+                mController,
+                mShortcutManager,
+                mUiEventLoggerFake);
+        viewController.bind();
+        return viewController;
+    }
 }
diff --git a/services/core/java/com/android/server/location/provider/LocationProviderManager.java b/services/core/java/com/android/server/location/provider/LocationProviderManager.java
index cbca48c9..17efeb3 100644
--- a/services/core/java/com/android/server/location/provider/LocationProviderManager.java
+++ b/services/core/java/com/android/server/location/provider/LocationProviderManager.java
@@ -911,18 +911,21 @@
                 @Override
                 public void onPreExecute() {
                     mUseWakeLock = false;
-                    final int size = locationResult.size();
-                    for (int i = 0; i < size; ++i) {
-                        if (!locationResult.get(i).isMock()) {
-                            mUseWakeLock = true;
-                            break;
+
+                    // don't acquire a wakelock for passive requests or for mock locations
+                    if (getRequest().getIntervalMillis() != LocationRequest.PASSIVE_INTERVAL) {
+                        final int size = locationResult.size();
+                        for (int i = 0; i < size; ++i) {
+                            if (!locationResult.get(i).isMock()) {
+                                mUseWakeLock = true;
+                                break;
+                            }
                         }
                     }
 
                     // update last delivered location
                     setLastDeliveredLocation(locationResult.getLastLocation());
 
-                    // don't acquire a wakelock for mock locations to prevent abuse
                     if (mUseWakeLock) {
                         mWakeLock.acquire(WAKELOCK_TIMEOUT_MS);
                     }
diff --git a/services/core/java/com/android/server/pm/BackgroundDexOptService.java b/services/core/java/com/android/server/pm/BackgroundDexOptService.java
index af9c401..27db2f9 100644
--- a/services/core/java/com/android/server/pm/BackgroundDexOptService.java
+++ b/services/core/java/com/android/server/pm/BackgroundDexOptService.java
@@ -121,6 +121,8 @@
 
     private final Injector mInjector;
 
+    private final DexOptHelper mDexOptHelper;
+
     private final Object mLock = new Object();
 
     // Thread currently running dexopt. This will be null if dexopt is not running.
@@ -168,13 +170,15 @@
         void onPackagesUpdated(ArraySet<String> updatedPackages);
     }
 
-    public BackgroundDexOptService(Context context, DexManager dexManager) {
-        this(new Injector(context, dexManager));
+    public BackgroundDexOptService(Context context, DexManager dexManager,
+            PackageManagerService pm) {
+        this(new Injector(context, dexManager, pm));
     }
 
     @VisibleForTesting
     public BackgroundDexOptService(Injector injector) {
         mInjector = injector;
+        mDexOptHelper = mInjector.getDexOptHelper();
         LocalServices.addService(BackgroundDexOptService.class, this);
         mDowngradeUnusedAppsThresholdInMillis = mInjector.getDowngradeUnusedAppsThresholdInMillis();
     }
@@ -251,15 +255,13 @@
                 resetStatesForNewDexOptRunLocked(Thread.currentThread());
             }
             PackageManagerService pm = mInjector.getPackageManagerService();
-            DexOptHelper dexOptHelper = new DexOptHelper(pm);
             ArraySet<String> packagesToOptimize;
             if (packageNames == null) {
-                packagesToOptimize = dexOptHelper.getOptimizablePackages();
+                packagesToOptimize = mDexOptHelper.getOptimizablePackages();
             } else {
                 packagesToOptimize = new ArraySet<>(packageNames);
             }
-            return runIdleOptimization(pm, dexOptHelper, packagesToOptimize,
-                    /* isPostBootUpdate= */ false);
+            return runIdleOptimization(pm, packagesToOptimize, /* isPostBootUpdate= */ false);
         } finally {
             Binder.restoreCallingIdentity(identity);
             markDexOptCompleted();
@@ -320,8 +322,7 @@
             return false;
         }
 
-        DexOptHelper dexOptHelper = new DexOptHelper(pm);
-        ArraySet<String> pkgs = dexOptHelper.getOptimizablePackages();
+        ArraySet<String> pkgs = mDexOptHelper.getOptimizablePackages();
         if (pkgs.isEmpty()) {
             Slog.i(TAG, "No packages to optimize");
             markPostBootUpdateCompleted(params);
@@ -347,7 +348,7 @@
                         tr.traceBegin("jobExecution");
                         boolean completed = false;
                         try {
-                            completed = runIdleOptimization(pm, dexOptHelper, pkgs,
+                            completed = runIdleOptimization(pm, pkgs,
                                     params.getJobId() == JOB_POST_BOOT_UPDATE);
                         } finally { // Those cleanup should be done always.
                             tr.traceEnd();
@@ -461,7 +462,7 @@
     @GuardedBy("mLock")
     private void controlDexOptBlockingLocked(boolean block) {
         PackageManagerService pm = mInjector.getPackageManagerService();
-        new DexOptHelper(pm).controlDexOptBlocking(block);
+        mDexOptHelper.controlDexOptBlocking(block);
     }
 
     private void scheduleAJob(int jobId) {
@@ -511,10 +512,10 @@
     }
 
     /** Returns true if completed */
-    private boolean runIdleOptimization(PackageManagerService pm, DexOptHelper dexOptHelper,
-            ArraySet<String> pkgs, boolean isPostBootUpdate) {
+    private boolean runIdleOptimization(PackageManagerService pm, ArraySet<String> pkgs,
+            boolean isPostBootUpdate) {
         long lowStorageThreshold = getLowStorageThreshold();
-        int status = idleOptimizePackages(pm, dexOptHelper, pkgs, lowStorageThreshold,
+        int status = idleOptimizePackages(pm, pkgs, lowStorageThreshold,
                 isPostBootUpdate);
         logStatus(status);
         synchronized (mLock) {
@@ -562,8 +563,8 @@
     }
 
     @Status
-    private int idleOptimizePackages(PackageManagerService pm, DexOptHelper dexOptHelper,
-            ArraySet<String> pkgs, long lowStorageThreshold, boolean isPostBootUpdate) {
+    private int idleOptimizePackages(PackageManagerService pm, ArraySet<String> pkgs,
+            long lowStorageThreshold, boolean isPostBootUpdate) {
         ArraySet<String> updatedPackages = new ArraySet<>();
         ArraySet<String> updatedPackagesDueToSecondaryDex = new ArraySet<>();
 
@@ -600,7 +601,7 @@
                             // Should be aborted by the scheduler.
                             return abortCode;
                         }
-                        @DexOptResult int downgradeResult = downgradePackage(pm, dexOptHelper, pkg,
+                        @DexOptResult int downgradeResult = downgradePackage(pm, pkg,
                                 /* isForPrimaryDex= */ true, isPostBootUpdate);
                         if (downgradeResult == PackageDexOptimizer.DEX_OPT_PERFORMED) {
                             updatedPackages.add(pkg);
@@ -611,7 +612,7 @@
                             return status;
                         }
                         if (supportSecondaryDex) {
-                            downgradeResult = downgradePackage(pm, dexOptHelper, pkg,
+                            downgradeResult = downgradePackage(pm, pkg,
                                     /* isForPrimaryDex= */false, isPostBootUpdate);
                             status = convertPackageDexOptimizerStatusToInternal(downgradeResult);
                             if (status != STATUS_OK) {
@@ -625,9 +626,8 @@
                 }
             }
 
-            @Status int primaryResult = optimizePackages(dexOptHelper, pkgs,
-                    lowStorageThreshold, /*isForPrimaryDex=*/ true, updatedPackages,
-                    isPostBootUpdate);
+            @Status int primaryResult = optimizePackages(pkgs, lowStorageThreshold,
+                    /*isForPrimaryDex=*/ true, updatedPackages, isPostBootUpdate);
             if (primaryResult != STATUS_OK) {
                 return primaryResult;
             }
@@ -636,9 +636,9 @@
                 return STATUS_OK;
             }
 
-            @Status int secondaryResult = optimizePackages(dexOptHelper, pkgs,
-                    lowStorageThreshold, /*isForPrimaryDex*/ false,
-                    updatedPackagesDueToSecondaryDex, isPostBootUpdate);
+            @Status int secondaryResult = optimizePackages(pkgs, lowStorageThreshold,
+                    /*isForPrimaryDex*/ false, updatedPackagesDueToSecondaryDex,
+                    isPostBootUpdate);
             return secondaryResult;
         } finally {
             // Always let the pinner service know about changes.
@@ -651,9 +651,8 @@
     }
 
     @Status
-    private int optimizePackages(DexOptHelper dexOptHelper,
-            ArraySet<String> pkgs, long lowStorageThreshold, boolean isForPrimaryDex,
-            ArraySet<String> updatedPackages, boolean isPostBootUpdate) {
+    private int optimizePackages(ArraySet<String> pkgs, long lowStorageThreshold,
+            boolean isForPrimaryDex, ArraySet<String> updatedPackages, boolean isPostBootUpdate) {
         for (String pkg : pkgs) {
             int abortCode = abortIdleOptimizations(lowStorageThreshold);
             if (abortCode != STATUS_OK) {
@@ -661,8 +660,7 @@
                 return abortCode;
             }
 
-            @DexOptResult int result = optimizePackage(dexOptHelper, pkg, isForPrimaryDex,
-                    isPostBootUpdate);
+            @DexOptResult int result = optimizePackage(pkg, isForPrimaryDex, isPostBootUpdate);
             if (result == PackageDexOptimizer.DEX_OPT_PERFORMED) {
                 updatedPackages.add(pkg);
             } else if (result != PackageDexOptimizer.DEX_OPT_SKIPPED) {
@@ -681,8 +679,8 @@
      * @return PackageDexOptimizer.DEX_*
      */
     @DexOptResult
-    private int downgradePackage(PackageManagerService pm, DexOptHelper dexOptHelper, String pkg,
-            boolean isForPrimaryDex, boolean isPostBootUpdate) {
+    private int downgradePackage(PackageManagerService pm, String pkg, boolean isForPrimaryDex,
+            boolean isPostBootUpdate) {
         if (DEBUG) {
             Slog.d(TAG, "Downgrading " + pkg);
         }
@@ -704,10 +702,10 @@
                 // remove their compiler artifacts from dalvik cache.
                 pm.deleteOatArtifactsOfPackage(pkg);
             } else {
-                result = performDexOptPrimary(dexOptHelper, pkg, reason, dexoptFlags);
+                result = performDexOptPrimary(pkg, reason, dexoptFlags);
             }
         } else {
-            result = performDexOptSecondary(dexOptHelper, pkg, reason, dexoptFlags);
+            result = performDexOptSecondary(pkg, reason, dexoptFlags);
         }
 
         if (result == PackageDexOptimizer.DEX_OPT_PERFORMED) {
@@ -733,15 +731,13 @@
      *
      * Optimize package if needed. Note that there can be no race between
      * concurrent jobs because PackageDexOptimizer.performDexOpt is synchronized.
-     * @param dexOptHelper An instance of DexOptHelper
      * @param pkg The package to be downgraded.
      * @param isForPrimaryDex Apps can have several dex file, primary and secondary.
      * @param isPostBootUpdate is post boot update or not.
      * @return PackageDexOptimizer#DEX_OPT_*
      */
     @DexOptResult
-    private int optimizePackage(DexOptHelper dexOptHelper, String pkg,
-            boolean isForPrimaryDex, boolean isPostBootUpdate) {
+    private int optimizePackage(String pkg, boolean isForPrimaryDex, boolean isPostBootUpdate) {
         int reason = isPostBootUpdate ? PackageManagerService.REASON_POST_BOOT
                 : PackageManagerService.REASON_BACKGROUND_DEXOPT;
         int dexoptFlags = DexoptOptions.DEXOPT_BOOT_COMPLETE;
@@ -753,27 +749,27 @@
         // System server share the same code path as primary dex files.
         // PackageManagerService will select the right optimization path for it.
         if (isForPrimaryDex || PLATFORM_PACKAGE_NAME.equals(pkg)) {
-            return performDexOptPrimary(dexOptHelper, pkg, reason, dexoptFlags);
+            return performDexOptPrimary(pkg, reason, dexoptFlags);
         } else {
-            return performDexOptSecondary(dexOptHelper, pkg, reason, dexoptFlags);
+            return performDexOptSecondary(pkg, reason, dexoptFlags);
         }
     }
 
     @DexOptResult
-    private int performDexOptPrimary(DexOptHelper dexOptHelper, String pkg, int reason,
+    private int performDexOptPrimary(String pkg, int reason,
             int dexoptFlags) {
         return trackPerformDexOpt(pkg, /*isForPrimaryDex=*/ true,
-                () -> dexOptHelper.performDexOptWithStatus(
+                () -> mDexOptHelper.performDexOptWithStatus(
                         new DexoptOptions(pkg, reason, dexoptFlags)));
     }
 
     @DexOptResult
-    private int performDexOptSecondary(DexOptHelper dexOptHelper, String pkg, int reason,
+    private int performDexOptSecondary(String pkg, int reason,
             int dexoptFlags) {
         DexoptOptions dexoptOptions = new DexoptOptions(pkg, reason,
                 dexoptFlags | DexoptOptions.DEXOPT_ONLY_SECONDARY_DEX);
         return trackPerformDexOpt(pkg, /*isForPrimaryDex=*/ false,
-                () -> dexOptHelper.performDexOpt(dexoptOptions)
+                () -> mDexOptHelper.performDexOpt(dexoptOptions)
                     ? PackageDexOptimizer.DEX_OPT_PERFORMED : PackageDexOptimizer.DEX_OPT_FAILED
         );
     }
@@ -911,11 +907,13 @@
     static final class Injector {
         private final Context mContext;
         private final DexManager mDexManager;
+        private final PackageManagerService mPackageManagerService;
         private final File mDataDir = Environment.getDataDirectory();
 
-        Injector(Context context, DexManager dexManager) {
+        Injector(Context context, DexManager dexManager, PackageManagerService pm) {
             mContext = context;
             mDexManager = dexManager;
+            mPackageManagerService = pm;
         }
 
         Context getContext() {
@@ -923,7 +921,11 @@
         }
 
         PackageManagerService getPackageManagerService() {
-            return (PackageManagerService) ServiceManager.getService("package");
+            return mPackageManagerService;
+        }
+
+        DexOptHelper getDexOptHelper() {
+            return new DexOptHelper(getPackageManagerService());
         }
 
         JobScheduler getJobScheduler() {
diff --git a/services/core/java/com/android/server/pm/PackageManagerService.java b/services/core/java/com/android/server/pm/PackageManagerService.java
index accb7a0..5f56d4e 100644
--- a/services/core/java/com/android/server/pm/PackageManagerService.java
+++ b/services/core/java/com/android/server/pm/PackageManagerService.java
@@ -1486,7 +1486,7 @@
                 new DefaultSystemWrapper(),
                 LocalServices::getService,
                 context::getSystemService,
-                (i, pm) -> new BackgroundDexOptService(i.getContext(), i.getDexManager()));
+                (i, pm) -> new BackgroundDexOptService(i.getContext(), i.getDexManager(), pm));
 
         if (Build.VERSION.SDK_INT <= 0) {
             Slog.w(TAG, "**** ro.build.version.sdk not set!");
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index 8644d1e..5b062fc 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -2755,14 +2755,13 @@
             return false;
         }
 
-        boolean isKeyguardLocked = mAtmService.isKeyguardLocked();
         boolean isCurrentAppLocked =
                 mAtmService.getLockTaskModeState() != LOCK_TASK_MODE_NONE;
         final TaskDisplayArea taskDisplayArea = getDisplayArea();
         boolean hasRootPinnedTask = taskDisplayArea != null && taskDisplayArea.hasPinnedTask();
         // Don't return early if !isNotLocked, since we want to throw an exception if the activity
         // is in an incorrect state
-        boolean isNotLockedOrOnKeyguard = !isKeyguardLocked && !isCurrentAppLocked;
+        boolean isNotLockedOrOnKeyguard = !isKeyguardLocked() && !isCurrentAppLocked;
 
         // We don't allow auto-PiP when something else is already pipped.
         if (beforeStopping && hasRootPinnedTask) {
@@ -3161,7 +3160,7 @@
                 // If the current activity is not opaque, we need to make sure the visibilities of
                 // activities be updated, they may be seen by users.
                 ensureVisibility = true;
-            } else if (mTaskSupervisor.getKeyguardController().isKeyguardLocked()
+            } else if (isKeyguardLocked()
                     && mTaskSupervisor.getKeyguardController().topActivityOccludesKeyguard(this)) {
                 // Ensure activity visibilities and update lockscreen occluded/dismiss state when
                 // finishing the top activity that occluded keyguard. So that, the
@@ -3990,6 +3989,11 @@
         });
     }
 
+    boolean isKeyguardLocked() {
+        return (mDisplayContent != null) ? mDisplayContent.isKeyguardLocked() :
+                mRootWindowContainer.getDefaultDisplay().isKeyguardLocked();
+    }
+
     void checkKeyguardFlagsChanged() {
         final boolean containsDismissKeyguard = containsDismissKeyguardWindow();
         final boolean containsShowWhenLocked = containsShowWhenLockedWindow();
@@ -5179,7 +5183,7 @@
     void notifyUnknownVisibilityLaunchedForKeyguardTransition() {
         // No display activities never add a window, so there is no point in waiting them for
         // relayout.
-        if (noDisplay || !mTaskSupervisor.getKeyguardController().isKeyguardLocked()) {
+        if (noDisplay || !isKeyguardLocked()) {
             return;
         }
 
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index 9455ce6..0a85ba1 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -2799,7 +2799,10 @@
                 mH.sendMessage(msg);
             }
             try {
-                mKeyguardController.setKeyguardShown(keyguardShowing, aodShowing);
+                mRootWindowContainer.forAllDisplays(displayContent -> {
+                    mKeyguardController.setKeyguardShown(displayContent.getDisplayId(),
+                            keyguardShowing, aodShowing);
+                });
             } finally {
                 Binder.restoreCallingIdentity(ident);
             }
@@ -3359,7 +3362,9 @@
         final long token = Binder.clearCallingIdentity();
         try {
             synchronized (mGlobalLock) {
-                mKeyguardController.keyguardGoingAway(flags);
+                mRootWindowContainer.forAllDisplays(displayContent -> {
+                    mKeyguardController.keyguardGoingAway(displayContent.getDisplayId(), flags);
+                });
             }
         } finally {
             Binder.restoreCallingIdentity(token);
@@ -3442,7 +3447,7 @@
             }
         };
 
-        if (isKeyguardLocked()) {
+        if (r.isKeyguardLocked()) {
             // If the keyguard is showing or occluded, then try and dismiss it before
             // entering picture-in-picture (this will prompt the user to authenticate if the
             // device is currently locked).
@@ -3787,8 +3792,8 @@
         mRecentTasks.notifyTaskPersisterLocked(task, flush);
     }
 
-    boolean isKeyguardLocked() {
-        return mKeyguardController.isKeyguardLocked();
+    boolean isKeyguardLocked(int displayId) {
+        return mKeyguardController.isKeyguardLocked(displayId);
     }
 
     /**
@@ -6035,7 +6040,7 @@
                 final long ident = Binder.clearCallingIdentity();
                 try {
                     if (mAmInternal.shouldConfirmCredentials(userId)) {
-                        if (mKeyguardController.isKeyguardLocked()) {
+                        if (mKeyguardController.isKeyguardLocked(DEFAULT_DISPLAY)) {
                             // Showing launcher to avoid user entering credential twice.
                             startHomeActivity(currentUserId, "notifyLockedProfile");
                         }
diff --git a/services/core/java/com/android/server/wm/DisplayArea.java b/services/core/java/com/android/server/wm/DisplayArea.java
index 7485a1e..63fb793 100644
--- a/services/core/java/com/android/server/wm/DisplayArea.java
+++ b/services/core/java/com/android/server/wm/DisplayArea.java
@@ -560,7 +560,7 @@
                 // Ignore the orientation of keyguard if it is going away or is not showing while
                 // the device is fully awake. In other words, use the orientation of keyguard if
                 // its window is visible while the device is going to sleep or is sleeping.
-                if (!mWmService.mAtmService.isKeyguardLocked()
+                if (!mDisplayContent.isKeyguardLocked()
                         && mDisplayContent.getDisplayPolicy().isAwake()
                         // Device is not going to sleep.
                         && policy.okToAnimate(true /* ignoreScreenOn */)) {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 7a38554..4cb70fd 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -5841,6 +5841,30 @@
         return (flags & FLAG_CAN_SHOW_WITH_INSECURE_KEYGUARD) != 0;
     }
 
+    /**
+     * @return whether keyguard is locked for this display
+     */
+    boolean isKeyguardLocked() {
+        return mRootWindowContainer.mTaskSupervisor
+                .getKeyguardController().isKeyguardLocked(mDisplayId);
+    }
+
+    /**
+     * @return whether keyguard is going away on this display
+     */
+    boolean isKeyguardGoingAway() {
+        return mRootWindowContainer.mTaskSupervisor
+                .getKeyguardController().isKeyguardGoingAway(mDisplayId);
+    }
+
+    /**
+     * @return whether AOD is showing on this display
+     */
+    boolean isAodShowing() {
+        return mRootWindowContainer.mTaskSupervisor
+                .getKeyguardController().isAodShowing(mDisplayId);
+    }
+
     @VisibleForTesting
     void removeAllTasks() {
         forAllTasks((t) -> { t.getRootTask().removeChild(t, "removeAllTasks"); });
diff --git a/services/core/java/com/android/server/wm/DisplayPolicy.java b/services/core/java/com/android/server/wm/DisplayPolicy.java
index 2ea6334..e0dc4e4 100644
--- a/services/core/java/com/android/server/wm/DisplayPolicy.java
+++ b/services/core/java/com/android/server/wm/DisplayPolicy.java
@@ -1730,7 +1730,7 @@
         // drawing. This way, we know that the IME can be safely shown since the other windows are
         // now shown.
         final boolean hideIme = win.mIsImWindow
-                && (mService.mAtmService.mKeyguardController.isAodShowing()
+                && (mDisplayContent.isAodShowing()
                         || (mDisplayContent.isDefaultDisplay && !mWindowManagerDrawComplete));
         if (hideIme) {
             return true;
diff --git a/services/core/java/com/android/server/wm/EventLogTags.logtags b/services/core/java/com/android/server/wm/EventLogTags.logtags
index 40b80f7..e228a4b 100644
--- a/services/core/java/com/android/server/wm/EventLogTags.logtags
+++ b/services/core/java/com/android/server/wm/EventLogTags.logtags
@@ -52,7 +52,7 @@
 30066 wm_add_to_stopping (User|1|5),(Token|1|5),(Component Name|3),(Reason|3)
 
 # Keyguard status changed
-30067 wm_set_keyguard_shown (keyguardShowing|1),(aodShowing|1),(keyguardGoingAway|1),(Reason|3)
+30067 wm_set_keyguard_shown (Display Id|1|5),(keyguardShowing|1),(aodShowing|1),(keyguardGoingAway|1),(Reason|3)
 
 # Out of memory for surfaces.
 31000 wm_no_surface_memory (Window|3),(PID|1|5),(Operation|3)
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index d17c109..83fd3fa 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -497,12 +497,16 @@
 
             resetInputConsumers(mInputTransaction);
             // Update recents input consumer layer if active
-            if (mAddRecentsAnimationInputConsumerHandle
-                    && getWeak(mActiveRecentsActivity) != null) {
-                final WindowContainer layer = getWeak(mActiveRecentsLayerRef);
-                mRecentsAnimationInputConsumer.show(mInputTransaction,
-                        layer != null ? layer : getWeak(mActiveRecentsActivity));
-                mAddRecentsAnimationInputConsumerHandle = false;
+            final ActivityRecord activeRecents = getWeak(mActiveRecentsActivity);
+            if (mAddRecentsAnimationInputConsumerHandle && activeRecents != null
+                    && activeRecents.getSurfaceControl() != null) {
+                WindowContainer layer = getWeak(mActiveRecentsLayerRef);
+                layer = layer != null ? layer : activeRecents;
+                // Handle edge-case for SUW where windows don't exist yet
+                if (layer.getSurfaceControl() != null) {
+                    mRecentsAnimationInputConsumer.show(mInputTransaction, layer);
+                    mAddRecentsAnimationInputConsumerHandle = false;
+                }
             }
             mDisplayContent.forAllWindows(this, true /* traverseTopToBottom */);
             updateInputFocusRequest(mRecentsAnimationInputConsumer);
diff --git a/services/core/java/com/android/server/wm/KeyguardController.java b/services/core/java/com/android/server/wm/KeyguardController.java
index bd41de3..fee9884 100644
--- a/services/core/java/com/android/server/wm/KeyguardController.java
+++ b/services/core/java/com/android/server/wm/KeyguardController.java
@@ -38,10 +38,8 @@
 import static com.android.server.wm.ActivityTaskManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.ActivityTaskSupervisor.PRESERVE_WINDOWS;
 import static com.android.server.wm.KeyguardControllerProto.AOD_SHOWING;
-import static com.android.server.wm.KeyguardControllerProto.KEYGUARD_OCCLUDED_STATES;
+import static com.android.server.wm.KeyguardControllerProto.KEYGUARD_PER_DISPLAY;
 import static com.android.server.wm.KeyguardControllerProto.KEYGUARD_SHOWING;
-import static com.android.server.wm.KeyguardOccludedProto.DISPLAY_ID;
-import static com.android.server.wm.KeyguardOccludedProto.KEYGUARD_OCCLUDED;
 
 import android.annotation.Nullable;
 import android.os.IBinder;
@@ -73,10 +71,7 @@
 
     private final ActivityTaskSupervisor mTaskSupervisor;
     private WindowManagerService mWindowManager;
-    private boolean mKeyguardShowing;
-    private boolean mAodShowing;
-    private boolean mKeyguardGoingAway;
-    private boolean mDismissalRequested;
+
     private final SparseArray<KeyguardDisplayState> mDisplayStates = new SparseArray<>();
     private final ActivityTaskManagerService mService;
     private RootWindowContainer mRootWindowContainer;
@@ -95,8 +90,8 @@
         mRootWindowContainer = mService.mRootWindowContainer;
     }
 
-    boolean isAodShowing() {
-        return mAodShowing;
+    boolean isAodShowing(int displayId) {
+        return getDisplayState(displayId).mAodShowing;
     }
 
     /**
@@ -104,7 +99,9 @@
      *         on the given display, false otherwise.
      */
     boolean isKeyguardOrAodShowing(int displayId) {
-        return (mKeyguardShowing || mAodShowing) && !mKeyguardGoingAway
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        return (state.mKeyguardShowing || state.mAodShowing)
+                && !state.mKeyguardGoingAway
                 && !isDisplayOccluded(displayId);
     }
 
@@ -114,8 +111,9 @@
      * TODO(b/125198167): Replace isKeyguardOrAodShowing() by this logic.
      */
     boolean isKeyguardUnoccludedOrAodShowing(int displayId) {
-        if (displayId == DEFAULT_DISPLAY && mAodShowing) {
-            return !mKeyguardGoingAway;
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        if (displayId == DEFAULT_DISPLAY && state.mAodShowing) {
+            return !state.mKeyguardGoingAway;
         }
         return isKeyguardOrAodShowing(displayId);
     }
@@ -125,14 +123,17 @@
      *         display, false otherwise
      */
     boolean isKeyguardShowing(int displayId) {
-        return mKeyguardShowing && !mKeyguardGoingAway && !isDisplayOccluded(displayId);
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        return state.mKeyguardShowing && !state.mKeyguardGoingAway
+                && !isDisplayOccluded(displayId);
     }
 
     /**
      * @return true if Keyguard is either showing or occluded, but not going away
      */
-    boolean isKeyguardLocked() {
-        return mKeyguardShowing && !mKeyguardGoingAway;
+    boolean isKeyguardLocked(int displayId) {
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        return state.mKeyguardShowing && !state.mKeyguardGoingAway;
     }
 
     /**
@@ -146,28 +147,31 @@
     /**
      * @return {@code true} if the keyguard is going away, {@code false} otherwise.
      */
-    boolean isKeyguardGoingAway() {
+    boolean isKeyguardGoingAway(int displayId) {
+        final KeyguardDisplayState state = getDisplayState(displayId);
         // Also check keyguard showing in case value is stale.
-        return mKeyguardGoingAway && mKeyguardShowing;
+        return state.mKeyguardGoingAway && state.mKeyguardShowing;
     }
 
     /**
      * Update the Keyguard showing state.
      */
-    void setKeyguardShown(boolean keyguardShowing, boolean aodShowing) {
-        final boolean aodChanged = aodShowing != mAodShowing;
+    void setKeyguardShown(int displayId, boolean keyguardShowing, boolean aodShowing) {
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        final boolean aodChanged = aodShowing != state.mAodShowing;
         // If keyguard is going away, but SystemUI aborted the transition, need to reset state.
         // Do not reset keyguardChanged status if this is aodChanged.
-        final boolean keyguardChanged = (keyguardShowing != mKeyguardShowing)
-                || (mKeyguardGoingAway && keyguardShowing && !aodChanged);
+        final boolean keyguardChanged = (keyguardShowing != state.mKeyguardShowing)
+                || (state.mKeyguardGoingAway && keyguardShowing && !aodChanged);
         if (!keyguardChanged && !aodChanged) {
             setWakeTransitionReady();
             return;
         }
         EventLogTags.writeWmSetKeyguardShown(
+                displayId,
                 keyguardShowing ? 1 : 0,
                 aodShowing ? 1 : 0,
-                mKeyguardGoingAway ? 1 : 0,
+                state.mKeyguardGoingAway ? 1 : 0,
                 "setKeyguardShown");
 
         // Update the task snapshot if the screen will not be turned off. To make sure that the
@@ -180,13 +184,13 @@
         // - The display state is ON. Because if AOD is not on or pulsing, the display state will
         //   be OFF or DOZE (the path of screen off may have handled it).
         if (((aodShowing ^ keyguardShowing) || (aodShowing && aodChanged && keyguardChanged))
-                && !mKeyguardGoingAway && Display.isOnState(
+                && !state.mKeyguardGoingAway && Display.isOnState(
                         mRootWindowContainer.getDefaultDisplay().getDisplayInfo().state)) {
             mWindowManager.mTaskSnapshotController.snapshotForSleeping(DEFAULT_DISPLAY);
         }
 
-        mKeyguardShowing = keyguardShowing;
-        mAodShowing = aodShowing;
+        state.mKeyguardShowing = keyguardShowing;
+        state.mAodShowing = aodShowing;
         if (aodChanged) {
             // Ensure the new state takes effect.
             mWindowManager.mWindowPlacerLocked.performSurfacePlacement();
@@ -194,10 +198,11 @@
 
         if (keyguardChanged) {
             // Irrelevant to AOD.
-            dismissMultiWindowModeForTaskIfNeeded(null /* currentTaskControllsingOcclusion */);
-            mKeyguardGoingAway = false;
+            dismissMultiWindowModeForTaskIfNeeded(displayId,
+                    null /* currentTaskControllsingOcclusion */);
+            state.mKeyguardGoingAway = false;
             if (keyguardShowing) {
-                mDismissalRequested = false;
+                state.mDismissalRequested = false;
             }
         }
 
@@ -223,17 +228,19 @@
      * @param flags See {@link WindowManagerPolicy#KEYGUARD_GOING_AWAY_FLAG_TO_SHADE}
      *              etc.
      */
-    void keyguardGoingAway(int flags) {
-        if (!mKeyguardShowing) {
+    void keyguardGoingAway(int displayId, int flags) {
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        if (!state.mKeyguardShowing) {
             return;
         }
         Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "keyguardGoingAway");
         mService.deferWindowLayout();
-        mKeyguardGoingAway = true;
+        state.mKeyguardGoingAway = true;
         try {
             EventLogTags.writeWmSetKeyguardShown(
+                    displayId,
                     1 /* keyguardShowing */,
-                    mAodShowing ? 1 : 0,
+                    state.mAodShowing ? 1 : 0,
                     1 /* keyguardGoingAway */,
                     "keyguardGoingAway");
             final int transitFlags = convertTransitFlags(flags);
@@ -307,10 +314,10 @@
         // Allow to show it when we are about to dismiss Keyguard. This isn't allowed if r is
         // already the dismissing activity, in which case we don't allow it to repeatedly dismiss
         // Keyguard.
-        return r.containsDismissKeyguardWindow() && canDismissKeyguard() && !mAodShowing
-                && (mDismissalRequested
-                || (r.canShowWhenLocked()
-                        && getDisplayState(r.getDisplayId()).mDismissingKeyguardActivity != r));
+        final KeyguardDisplayState state = getDisplayState(r.getDisplayId());
+        return r.containsDismissKeyguardWindow() && canDismissKeyguard() && !state.mAodShowing
+                && (state.mDismissalRequested
+                || (r.canShowWhenLocked() && state.mDismissingKeyguardActivity != r));
     }
 
     /**
@@ -335,7 +342,7 @@
             // If keyguard is showing, nothing is visible, except if we are able to dismiss Keyguard
             // right away and AOD isn't visible.
             return canShowActivityWhileKeyguardShowing(r);
-        } else if (isKeyguardLocked()) {
+        } else if (isKeyguardLocked(r.getDisplayId())) {
             return canShowWhileOccluded(r.containsDismissKeyguardWindow(), r.canShowWhenLocked());
         } else {
             return true;
@@ -348,19 +355,15 @@
      * ({@link ActivityTaskSupervisor#beginActivityVisibilityUpdate}).
      */
     void updateVisibility() {
-        boolean requestDismissKeyguard = false;
         for (int displayNdx = mRootWindowContainer.getChildCount() - 1;
              displayNdx >= 0; displayNdx--) {
             final DisplayContent display = mRootWindowContainer.getChildAt(displayNdx);
             if (display.isRemoving() || display.isRemoved()) continue;
             final KeyguardDisplayState state = getDisplayState(display.mDisplayId);
             state.updateVisibility(this, display);
-            requestDismissKeyguard |= state.mRequestDismissKeyguard;
-        }
-
-        // Dismissing Keyguard happens globally using the information from all displays.
-        if (requestDismissKeyguard) {
-            handleDismissKeyguard();
+            if (state.mRequestDismissKeyguard) {
+                handleDismissKeyguard(display.getDisplayId());
+            }
         }
     }
 
@@ -381,7 +384,7 @@
         }
 
         mWindowManager.mPolicy.onKeyguardOccludedChangedLw(isDisplayOccluded(DEFAULT_DISPLAY));
-        if (isKeyguardLocked()) {
+        if (isKeyguardLocked(displayId)) {
             mService.deferWindowLayout();
             try {
                 mRootWindowContainer.getDefaultDisplay()
@@ -395,14 +398,14 @@
                 mService.continueWindowLayout();
             }
         }
-        dismissMultiWindowModeForTaskIfNeeded(topActivity != null
+        dismissMultiWindowModeForTaskIfNeeded(displayId, topActivity != null
                 ? topActivity.getRootTask() : null);
     }
 
     /**
      * Called when somebody wants to dismiss the Keyguard via the flag.
      */
-    private void handleDismissKeyguard() {
+    private void handleDismissKeyguard(int displayId) {
         // We only allow dismissing Keyguard via the flag when Keyguard is secure for legacy
         // reasons, because that's how apps used to dismiss Keyguard in the secure case. In the
         // insecure case, we actually show it on top of the lockscreen. See #canShowWhileOccluded.
@@ -411,12 +414,13 @@
         }
 
         mWindowManager.dismissKeyguard(null /* callback */, null /* message */);
-        mDismissalRequested = true;
+        final KeyguardDisplayState state = getDisplayState(displayId);
+        state.mDismissalRequested = true;
 
         // If we are about to unocclude the Keyguard, but we can dismiss it without security,
         // we immediately dismiss the Keyguard so the activity gets shown without a flicker.
         final DisplayContent dc = mRootWindowContainer.getDefaultDisplay();
-        if (mKeyguardShowing && canDismissKeyguard()
+        if (state.mKeyguardShowing && canDismissKeyguard()
                 && dc.mAppTransition.containsTransitRequest(TRANSIT_KEYGUARD_UNOCCLUDE)) {
             mWindowManager.executeAppTransition();
         }
@@ -434,10 +438,10 @@
                 || !mWindowManager.isKeyguardSecure(mService.getCurrentUserId());
     }
 
-    private void dismissMultiWindowModeForTaskIfNeeded(
+    private void dismissMultiWindowModeForTaskIfNeeded(int displayId,
             @Nullable Task currentTaskControllingOcclusion) {
         // TODO(b/113840485): Handle docked stack for individual display.
-        if (!mKeyguardShowing || !isDisplayOccluded(DEFAULT_DISPLAY)) {
+        if (!getDisplayState(displayId).mKeyguardShowing || !isDisplayOccluded(DEFAULT_DISPLAY)) {
             return;
         }
 
@@ -497,6 +501,10 @@
     /** Represents Keyguard state per individual display. */
     private static class KeyguardDisplayState {
         private final int mDisplayId;
+        private boolean mKeyguardShowing;
+        private boolean mAodShowing;
+        private boolean mKeyguardGoingAway;
+        private boolean mDismissalRequested;
         private boolean mOccluded;
 
         private ActivityRecord mTopOccludesActivity;
@@ -604,7 +612,16 @@
         void dumpStatus(PrintWriter pw, String prefix) {
             final StringBuilder sb = new StringBuilder();
             sb.append(prefix);
-            sb.append("  Occluded=").append(mOccluded)
+            sb.append(" KeyguardShowing=")
+                    .append(mKeyguardShowing)
+                    .append(" AodShowing=")
+                    .append(mAodShowing)
+                    .append(" KeyguardGoingAway=")
+                    .append(mKeyguardGoingAway)
+                    .append(" DismissalRequested=")
+                    .append(mDismissalRequested)
+                    .append("  Occluded=")
+                    .append(mOccluded)
                     .append(" DismissingKeyguardActivity=")
                     .append(mDismissingKeyguardActivity)
                     .append(" TurnScreenOnActivity=")
@@ -616,27 +633,31 @@
 
         void dumpDebug(ProtoOutputStream proto, long fieldId) {
             final long token = proto.start(fieldId);
-            proto.write(DISPLAY_ID, mDisplayId);
-            proto.write(KEYGUARD_OCCLUDED, mOccluded);
+            proto.write(KeyguardPerDisplayProto.DISPLAY_ID, mDisplayId);
+            proto.write(KeyguardPerDisplayProto.KEYGUARD_SHOWING, mKeyguardShowing);
+            proto.write(KeyguardPerDisplayProto.AOD_SHOWING, mAodShowing);
+            proto.write(KeyguardPerDisplayProto.KEYGUARD_OCCLUDED, mOccluded);
             proto.end(token);
         }
     }
 
     void dump(PrintWriter pw, String prefix) {
+        final KeyguardDisplayState default_state = getDisplayState(DEFAULT_DISPLAY);
         pw.println(prefix + "KeyguardController:");
-        pw.println(prefix + "  mKeyguardShowing=" + mKeyguardShowing);
-        pw.println(prefix + "  mAodShowing=" + mAodShowing);
-        pw.println(prefix + "  mKeyguardGoingAway=" + mKeyguardGoingAway);
+        pw.println(prefix + "  mKeyguardShowing=" + default_state.mKeyguardShowing);
+        pw.println(prefix + "  mAodShowing=" + default_state.mAodShowing);
+        pw.println(prefix + "  mKeyguardGoingAway=" + default_state.mKeyguardGoingAway);
         dumpDisplayStates(pw, prefix);
-        pw.println(prefix + "  mDismissalRequested=" + mDismissalRequested);
+        pw.println(prefix + "  mDismissalRequested=" + default_state.mDismissalRequested);
         pw.println();
     }
 
     void dumpDebug(ProtoOutputStream proto, long fieldId) {
+        final KeyguardDisplayState default_state = getDisplayState(DEFAULT_DISPLAY);
         final long token = proto.start(fieldId);
-        proto.write(AOD_SHOWING, mAodShowing);
-        proto.write(KEYGUARD_SHOWING, mKeyguardShowing);
-        writeDisplayStatesToProto(proto, KEYGUARD_OCCLUDED_STATES);
+        proto.write(AOD_SHOWING, default_state.mAodShowing);
+        proto.write(KEYGUARD_SHOWING, default_state.mKeyguardShowing);
+        writeDisplayStatesToProto(proto, KEYGUARD_PER_DISPLAY);
         proto.end(token);
     }
 
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index c28d089..ccee458 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -3480,7 +3480,9 @@
         // avoid power mode from being cleared before that, add a special reason to consider whether
         // the unknown visibility is resolved. The case from SystemUI is excluded because it should
         // rely on keyguard-going-away.
-        if (mService.mKeyguardController.isKeyguardLocked() && targetActivity != null
+        final boolean isKeyguardLocked = (targetActivity != null)
+                ? targetActivity.isKeyguardLocked() : mDefaultDisplay.isKeyguardLocked();
+        if (isKeyguardLocked && targetActivity != null
                 && !targetActivity.isLaunchSourceType(ActivityRecord.LAUNCH_SOURCE_TYPE_SYSTEMUI)) {
             final ActivityOptions opts = targetActivity.getOptions();
             if (opts == null || opts.getSourceInfo() == null
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 35d93f6..bfb1a8e 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -6080,11 +6080,13 @@
 
     boolean shouldSleepActivities() {
         final DisplayContent display = mDisplayContent;
+        final boolean isKeyguardGoingAway = (mDisplayContent != null)
+                ? mDisplayContent.isKeyguardGoingAway()
+                : mRootWindowContainer.getDefaultDisplay().isKeyguardGoingAway();
 
         // Do not sleep activities in this root task if we're marked as focused and the keyguard
         // is in the process of going away.
-        if (mTaskSupervisor.getKeyguardController().isKeyguardGoingAway()
-                && isFocusedRootTaskOnDisplay()
+        if (isKeyguardGoingAway && isFocusedRootTaskOnDisplay()
                 // Avoid resuming activities on secondary displays since we don't want bubble
                 // activities to be resumed while bubble is still collapsed.
                 // TODO(b/113840485): Having keyguard going away state for secondary displays.
diff --git a/services/core/java/com/android/server/wm/TaskDisplayArea.java b/services/core/java/com/android/server/wm/TaskDisplayArea.java
index 3387e37..600545e 100644
--- a/services/core/java/com/android/server/wm/TaskDisplayArea.java
+++ b/services/core/java/com/android/server/wm/TaskDisplayArea.java
@@ -1728,7 +1728,7 @@
         // the locked state, the keyguard isn't locked, or we can show when locked.
         if (topRunning != null && considerKeyguardState
                 && mRootWindowContainer.mTaskSupervisor.getKeyguardController()
-                .isKeyguardLocked()
+                .isKeyguardLocked(topRunning.getDisplayId())
                 && !topRunning.canShowWhenLocked()) {
             return null;
         }
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 4db8ef4..f175eec 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -513,7 +513,7 @@
         mState = STATE_PLAYING;
         mController.moveToPlaying(this);
 
-        if (mController.mAtm.mTaskSupervisor.getKeyguardController().isKeyguardLocked()) {
+        if (mController.mAtm.mTaskSupervisor.getKeyguardController().isKeyguardLocked(displayId)) {
             mFlags |= TRANSIT_FLAG_KEYGUARD_LOCKED;
         }
 
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundDexOptServiceUnitTest.java b/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundDexOptServiceUnitTest.java
new file mode 100644
index 0000000..8223b8c
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/BackgroundDexOptServiceUnitTest.java
@@ -0,0 +1,480 @@
+/*
+ * Copyright (C) 2021 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.server.pm;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.atLeastOnce;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertThrows;
+
+import android.app.job.JobInfo;
+import android.app.job.JobParameters;
+import android.app.job.JobScheduler;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.os.HandlerThread;
+import android.os.PowerManager;
+import android.util.ArraySet;
+
+import com.android.server.LocalServices;
+import com.android.server.PinnerService;
+import com.android.server.pm.dex.DexManager;
+import com.android.server.pm.dex.DexoptOptions;
+
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnitRunner;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.stream.Collectors;
+
+@RunWith(MockitoJUnitRunner.class)
+public final class BackgroundDexOptServiceUnitTest {
+    private static final String TAG = BackgroundDexOptServiceUnitTest.class.getSimpleName();
+
+    private static final long USABLE_SPACE_NORMAL = 1_000_000_000;
+    private static final long STORAGE_LOW_BYTES = 1_000_000;
+
+    private static final long TEST_WAIT_TIMEOUT_MS = 10_000;
+
+    private static final ArraySet<String> DEFAULT_PACKAGE_LIST = new ArraySet<>(
+            Arrays.asList("aaa", "bbb"));
+    private static final ArraySet<String> EMPTY_PACKAGE_LIST = new ArraySet<>();
+
+    @Mock
+    private Context mContext;
+    @Mock
+    private PackageManagerService mPackageManager;
+    @Mock
+    private DexOptHelper mDexOptHelper;
+    @Mock
+    private DexManager mDexManager;
+    @Mock
+    private PinnerService mPinnerService;
+    @Mock
+    private JobScheduler mJobScheduler;
+    @Mock
+    private BackgroundDexOptService.Injector mInjector;
+    @Mock
+    private BackgroundDexOptJobService mJobServiceForPostBoot;
+    @Mock
+    private BackgroundDexOptJobService mJobServiceForIdle;
+
+    private final JobParameters mJobParametersForPostBoot = new JobParameters(null,
+            BackgroundDexOptService.JOB_POST_BOOT_UPDATE, null, null, null,
+            0, false, false, null, null, null);
+    private final JobParameters mJobParametersForIdle = new JobParameters(null,
+            BackgroundDexOptService.JOB_IDLE_OPTIMIZE, null, null, null,
+            0, false, false, null, null, null);
+
+    private BackgroundDexOptService mService;
+
+    private StartAndWaitThread mDexOptThread;
+    private StartAndWaitThread mCancelThread;
+
+    @Before
+    public void setUp() throws Exception {
+        when(mInjector.getContext()).thenReturn(mContext);
+        when(mInjector.getDexOptHelper()).thenReturn(mDexOptHelper);
+        when(mInjector.getDexManager()).thenReturn(mDexManager);
+        when(mInjector.getPinnerService()).thenReturn(mPinnerService);
+        when(mInjector.getJobScheduler()).thenReturn(mJobScheduler);
+        when(mInjector.getPackageManagerService()).thenReturn(mPackageManager);
+
+        // These mocking can be overwritten in some tests but still keep it here as alternative
+        // takes too many repetitive codes.
+        when(mInjector.getDataDirUsableSpace()).thenReturn(USABLE_SPACE_NORMAL);
+        when(mInjector.getDataDirStorageLowBytes()).thenReturn(STORAGE_LOW_BYTES);
+        when(mInjector.getDexOptThermalCutoff()).thenReturn(PowerManager.THERMAL_STATUS_CRITICAL);
+        when(mInjector.getCurrentThermalStatus()).thenReturn(PowerManager.THERMAL_STATUS_NONE);
+        when(mDexOptHelper.getOptimizablePackages()).thenReturn(DEFAULT_PACKAGE_LIST);
+        when(mDexOptHelper.performDexOptWithStatus(any())).thenReturn(
+                PackageDexOptimizer.DEX_OPT_PERFORMED);
+
+        mService = new BackgroundDexOptService(mInjector);
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        LocalServices.removeServiceForTest(BackgroundDexOptService.class);
+    }
+
+    @Test
+    public void testGetService() {
+        assertThat(BackgroundDexOptService.getService()).isEqualTo(mService);
+    }
+
+    @Test
+    public void testBootCompleted() {
+        initUntilBootCompleted();
+    }
+
+    @Test
+    public void testNoExecutionForIdleJobBeforePostBootUpdate() {
+        initUntilBootCompleted();
+
+        assertThat(mService.onStartJob(mJobServiceForIdle, mJobParametersForIdle)).isFalse();
+    }
+
+    @Test
+    public void testNoExecutionForLowStorage() {
+        initUntilBootCompleted();
+        when(mPackageManager.isStorageLow()).thenReturn(true);
+
+        assertThat(mService.onStartJob(mJobServiceForPostBoot,
+                mJobParametersForPostBoot)).isFalse();
+        verify(mDexOptHelper, never()).performDexOpt(any());
+    }
+
+    @Test
+    public void testNoExecutionForNoOptimizablePackages() {
+        initUntilBootCompleted();
+        when(mDexOptHelper.getOptimizablePackages()).thenReturn(EMPTY_PACKAGE_LIST);
+
+        assertThat(mService.onStartJob(mJobServiceForPostBoot,
+                mJobParametersForPostBoot)).isFalse();
+        verify(mDexOptHelper, never()).performDexOpt(any());
+    }
+
+    @Test
+    public void testPostBootUpdateFullRun() {
+        initUntilBootCompleted();
+
+        runFullJob(mJobServiceForPostBoot, mJobParametersForPostBoot, false, 1);
+    }
+
+    @Test
+    public void testIdleJobFullRun() {
+        initUntilBootCompleted();
+
+        runFullJob(mJobServiceForPostBoot, mJobParametersForPostBoot, false, 1);
+        runFullJob(mJobServiceForIdle, mJobParametersForIdle, true, 2);
+    }
+
+    @Test
+    public void testSystemReadyWhenDisabled() {
+        when(mInjector.isBackgroundDexOptDisabled()).thenReturn(true);
+
+        mService.systemReady();
+
+        verify(mContext, never()).registerReceiver(any(), any());
+    }
+
+    @Test
+    public void testStopByCancelFlag() {
+        when(mInjector.createAndStartThread(any(), any())).thenReturn(Thread.currentThread());
+        initUntilBootCompleted();
+
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+
+        ArgumentCaptor<Runnable> argDexOptThreadRunnable = ArgumentCaptor.forClass(Runnable.class);
+        verify(mInjector, atLeastOnce()).createAndStartThread(any(),
+                argDexOptThreadRunnable.capture());
+
+        // Stopping requires a separate thread
+        HandlerThread cancelThread = new HandlerThread("Stopping");
+        cancelThread.start();
+        when(mInjector.createAndStartThread(any(), any())).thenReturn(cancelThread);
+
+        // Cancel
+        assertThat(mService.onStopJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+
+        // Capture Runnable for cancel
+        ArgumentCaptor<Runnable> argCancelThreadRunnable = ArgumentCaptor.forClass(Runnable.class);
+        verify(mInjector, atLeastOnce()).createAndStartThread(any(),
+                argCancelThreadRunnable.capture());
+
+        // Execute cancelling part
+        cancelThread.getThreadHandler().post(argCancelThreadRunnable.getValue());
+
+        verify(mDexOptHelper, timeout(TEST_WAIT_TIMEOUT_MS)).controlDexOptBlocking(true);
+
+        // Dexopt thread run and cancelled
+        argDexOptThreadRunnable.getValue().run();
+
+        // Wait until cancellation Runnable is completed.
+        assertThat(cancelThread.getThreadHandler().runWithScissors(
+                argCancelThreadRunnable.getValue(), TEST_WAIT_TIMEOUT_MS)).isTrue();
+
+        // Now cancel completed
+        verify(mJobServiceForPostBoot).jobFinished(mJobParametersForPostBoot, true);
+        verifyLastControlDexOptBlockingCall(false);
+    }
+
+    @Test
+    public void testPostUpdateCancelFirst() throws Exception {
+        initUntilBootCompleted();
+        when(mInjector.createAndStartThread(any(), any())).thenAnswer(
+                i -> createAndStartExecutionThread(i.getArgument(0), i.getArgument(1)));
+
+        // Start
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+        // Cancel
+        assertThat(mService.onStopJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+
+        mCancelThread.runActualRunnable();
+
+        // Wait until cancel has set the flag.
+        verify(mDexOptHelper, timeout(TEST_WAIT_TIMEOUT_MS)).controlDexOptBlocking(
+                true);
+
+        mDexOptThread.runActualRunnable();
+
+        // All threads should finish.
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+        mCancelThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Retry later if post boot job was cancelled
+        verify(mJobServiceForPostBoot).jobFinished(mJobParametersForPostBoot, true);
+        verifyLastControlDexOptBlockingCall(false);
+    }
+
+    @Test
+    public void testPostUpdateCancelLater() throws Exception {
+        initUntilBootCompleted();
+        when(mInjector.createAndStartThread(any(), any())).thenAnswer(
+                i -> createAndStartExecutionThread(i.getArgument(0), i.getArgument(1)));
+
+        // Start
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+        // Cancel
+        assertThat(mService.onStopJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+
+        // Dexopt thread runs and finishes
+        mDexOptThread.runActualRunnable();
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        mCancelThread.runActualRunnable();
+        mCancelThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Already completed before cancel, so no rescheduling.
+        verify(mJobServiceForPostBoot).jobFinished(mJobParametersForPostBoot, false);
+        verify(mDexOptHelper, never()).controlDexOptBlocking(true);
+    }
+
+    @Test
+    public void testPeriodicJobCancelFirst() throws Exception {
+        initUntilBootCompleted();
+        when(mInjector.createAndStartThread(any(), any())).thenAnswer(
+                i -> createAndStartExecutionThread(i.getArgument(0), i.getArgument(1)));
+
+        // Start and finish post boot job
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+        mDexOptThread.runActualRunnable();
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Start
+        assertThat(mService.onStartJob(mJobServiceForIdle, mJobParametersForIdle)).isTrue();
+        // Cancel
+        assertThat(mService.onStopJob(mJobServiceForIdle, mJobParametersForIdle)).isTrue();
+
+        mCancelThread.runActualRunnable();
+
+        // Wait until cancel has set the flag.
+        verify(mDexOptHelper, timeout(TEST_WAIT_TIMEOUT_MS)).controlDexOptBlocking(
+                true);
+
+        mDexOptThread.runActualRunnable();
+
+        // All threads should finish.
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+        mCancelThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Always reschedule for periodic job
+        verify(mJobServiceForIdle).jobFinished(mJobParametersForIdle, true);
+        verifyLastControlDexOptBlockingCall(false);
+    }
+
+    @Test
+    public void testPeriodicJobCancelLater() throws Exception {
+        initUntilBootCompleted();
+        when(mInjector.createAndStartThread(any(), any())).thenAnswer(
+                i -> createAndStartExecutionThread(i.getArgument(0), i.getArgument(1)));
+
+        // Start and finish post boot job
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+        mDexOptThread.runActualRunnable();
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Start
+        assertThat(mService.onStartJob(mJobServiceForIdle, mJobParametersForIdle)).isTrue();
+        // Cancel
+        assertThat(mService.onStopJob(mJobServiceForIdle, mJobParametersForIdle)).isTrue();
+
+        // Dexopt thread finishes first.
+        mDexOptThread.runActualRunnable();
+        mDexOptThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        mCancelThread.runActualRunnable();
+        mCancelThread.join(TEST_WAIT_TIMEOUT_MS);
+
+        // Always reschedule for periodic job
+        verify(mJobServiceForIdle).jobFinished(mJobParametersForIdle, true);
+        verify(mDexOptHelper, never()).controlDexOptBlocking(true);
+    }
+
+    @Test
+    public void testStopByThermal() {
+        when(mInjector.createAndStartThread(any(), any())).thenReturn(Thread.currentThread());
+        initUntilBootCompleted();
+
+        assertThat(mService.onStartJob(mJobServiceForPostBoot, mJobParametersForPostBoot)).isTrue();
+
+        ArgumentCaptor<Runnable> argThreadRunnable = ArgumentCaptor.forClass(Runnable.class);
+        verify(mInjector, atLeastOnce()).createAndStartThread(any(), argThreadRunnable.capture());
+
+        // Thermal cancel level
+        when(mInjector.getCurrentThermalStatus()).thenReturn(PowerManager.THERMAL_STATUS_CRITICAL);
+
+        argThreadRunnable.getValue().run();
+
+        verify(mJobServiceForPostBoot).jobFinished(mJobParametersForPostBoot, true);
+        verifyLastControlDexOptBlockingCall(false);
+    }
+
+    @Test
+    public void testRunShellCommandWithInvalidUid() {
+        // Test uid cannot execute the command APIs
+        assertThrows(SecurityException.class, () -> mService.runBackgroundDexoptJob(null));
+    }
+
+    @Test
+    public void testCancelShellCommandWithInvalidUid() {
+        // Test uid cannot execute the command APIs
+        assertThrows(SecurityException.class, () -> mService.cancelBackgroundDexoptJob());
+    }
+
+    private void initUntilBootCompleted() {
+        ArgumentCaptor<BroadcastReceiver> argReceiver = ArgumentCaptor.forClass(
+                BroadcastReceiver.class);
+        ArgumentCaptor<IntentFilter> argIntentFilter = ArgumentCaptor.forClass(IntentFilter.class);
+
+        mService.systemReady();
+
+        verify(mContext).registerReceiver(argReceiver.capture(), argIntentFilter.capture());
+        assertThat(argIntentFilter.getValue().getAction(0)).isEqualTo(Intent.ACTION_BOOT_COMPLETED);
+
+        argReceiver.getValue().onReceive(mContext, null);
+
+        verify(mContext).unregisterReceiver(argReceiver.getValue());
+        ArgumentCaptor<JobInfo> argJobs = ArgumentCaptor.forClass(JobInfo.class);
+        verify(mJobScheduler, times(2)).schedule(argJobs.capture());
+
+        List<Integer> expectedJobIds = Arrays.asList(BackgroundDexOptService.JOB_IDLE_OPTIMIZE,
+                BackgroundDexOptService.JOB_POST_BOOT_UPDATE);
+        List<Integer> jobIds = argJobs.getAllValues().stream().map(job -> job.getId()).collect(
+                Collectors.toList());
+        assertThat(jobIds).containsExactlyElementsIn(expectedJobIds);
+    }
+
+    private void verifyLastControlDexOptBlockingCall(boolean expected) {
+        ArgumentCaptor<Boolean> argDexOptBlock = ArgumentCaptor.forClass(Boolean.class);
+        verify(mDexOptHelper, atLeastOnce()).controlDexOptBlocking(argDexOptBlock.capture());
+        assertThat(argDexOptBlock.getValue()).isEqualTo(expected);
+    }
+
+    private void runFullJob(BackgroundDexOptJobService jobService, JobParameters params,
+            boolean expectedReschedule, int totalJobRuns) {
+        when(mInjector.createAndStartThread(any(), any())).thenReturn(Thread.currentThread());
+        assertThat(mService.onStartJob(jobService, params)).isTrue();
+
+        ArgumentCaptor<Runnable> argThreadRunnable = ArgumentCaptor.forClass(Runnable.class);
+        verify(mInjector, atLeastOnce()).createAndStartThread(any(), argThreadRunnable.capture());
+
+        argThreadRunnable.getValue().run();
+
+        verify(jobService).jobFinished(params, expectedReschedule);
+        // Never block
+        verify(mDexOptHelper, never()).controlDexOptBlocking(true);
+        verifyPerformDexOpt(DEFAULT_PACKAGE_LIST, totalJobRuns);
+    }
+
+    private void verifyPerformDexOpt(ArraySet<String> pkgs, int expectedRuns) {
+        ArgumentCaptor<DexoptOptions> dexOptOptions = ArgumentCaptor.forClass(DexoptOptions.class);
+        verify(mDexOptHelper, atLeastOnce()).performDexOptWithStatus(dexOptOptions.capture());
+        HashMap<String, Integer> primaryPkgs = new HashMap<>(); // K: pkg, V: dexopt runs left
+        for (String pkg : pkgs) {
+            primaryPkgs.put(pkg, expectedRuns);
+        }
+
+        for (DexoptOptions opt : dexOptOptions.getAllValues()) {
+            assertThat(pkgs).contains(opt.getPackageName());
+            assertThat(opt.isDexoptOnlySecondaryDex()).isFalse();
+            Integer count = primaryPkgs.get(opt.getPackageName());
+            assertThat(count).isNotNull();
+            if (count == 1) {
+                primaryPkgs.remove(opt.getPackageName());
+            } else {
+                primaryPkgs.put(opt.getPackageName(), count - 1);
+            }
+        }
+        assertThat(primaryPkgs).isEmpty();
+    }
+
+    private static class StartAndWaitThread extends Thread {
+        private final Runnable mActualRunnable;
+        private final CountDownLatch mLatch = new CountDownLatch(1);
+
+        private StartAndWaitThread(String name, Runnable runnable) {
+            super(name);
+            mActualRunnable = runnable;
+        }
+
+        private void runActualRunnable() {
+            mLatch.countDown();
+        }
+
+        @Override
+        public void run() {
+            // Thread is started but does not run actual code. This is for controlling the execution
+            // order while still meeting Thread.isAlive() check.
+            try {
+                mLatch.await();
+            } catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+            mActualRunnable.run();
+        }
+    }
+
+    private Thread createAndStartExecutionThread(String name, Runnable runnable) {
+        final boolean isDexOptThread = !name.equals("DexOptCancel");
+        StartAndWaitThread thread = new StartAndWaitThread(name, runnable);
+        if (isDexOptThread) {
+            mDexOptThread = thread;
+        } else {
+            mCancelThread = thread;
+        }
+        thread.start();
+        return thread;
+    }
+}
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
index bc0cc1f..ae24785 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityRecordTests.java
@@ -1389,7 +1389,8 @@
         final Task task = activity.getTask();
         // Make keyguard locked and set the top activity show-when-locked.
         KeyguardController keyguardController = activity.mTaskSupervisor.getKeyguardController();
-        doReturn(true).when(keyguardController).isKeyguardLocked();
+        int displayId = activity.getDisplayId();
+        doReturn(true).when(keyguardController).isKeyguardLocked(eq(displayId));
         final ActivityRecord topActivity = new ActivityBuilder(mAtm).setTask(task).build();
         topActivity.mVisibleRequested = true;
         topActivity.nowVisible = true;
diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
index abd43a4..1266b2e 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
@@ -39,6 +39,7 @@
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.doCallRealMethod;
 import static org.mockito.Mockito.times;
@@ -246,7 +247,7 @@
         doReturn(true).when(record)
                 .checkEnterPictureInPictureState("enterPictureInPictureMode", false);
         doReturn(false).when(record).inPinnedWindowingMode();
-        doReturn(false).when(mAtm).isKeyguardLocked();
+        doReturn(false).when(mAtm).isKeyguardLocked(anyInt());
 
         //to simulate NPE
         doReturn(null).when(record).getParent();
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 046328d..562e958 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -827,12 +827,12 @@
                 SCREEN_ORIENTATION_LANDSCAPE, dc.getOrientation());
 
         keyguard.mAttrs.screenOrientation = SCREEN_ORIENTATION_PORTRAIT;
-        mAtm.mKeyguardController.setKeyguardShown(true /* keyguardShowing */,
+        mAtm.mKeyguardController.setKeyguardShown(window.getDisplayId(), true /* keyguardShowing */,
                 false /* aodShowing */);
         assertEquals("Visible keyguard must influence device orientation",
                 SCREEN_ORIENTATION_PORTRAIT, dc.getOrientation());
 
-        mAtm.mKeyguardController.keyguardGoingAway(0 /* flags */);
+        mAtm.mKeyguardController.keyguardGoingAway(window.getDisplayId(), 0 /* flags */);
         assertEquals("Keyguard that is going away must not influence device orientation",
                 SCREEN_ORIENTATION_LANDSCAPE, dc.getOrientation());
     }
@@ -2137,7 +2137,7 @@
         assertTopRunningActivity(activity, display);
 
         // Check to make sure activity not reported when it cannot show on lock and lock is on.
-        doReturn(true).when(keyguard).isKeyguardLocked();
+        doReturn(true).when(keyguard).isKeyguardLocked(anyInt());
         assertEquals(activity, display.topRunningActivity());
         assertNull(display.topRunningActivity(true /* considerKeyguardState */));
 
@@ -2175,11 +2175,11 @@
         final WindowState appWin = createWindow(null, TYPE_APPLICATION, mDisplayContent, "appWin");
         final ActivityRecord activity = appWin.mActivityRecord;
 
-        mAtm.mKeyguardController.setKeyguardShown(true /* keyguardShowing */,
+        mAtm.mKeyguardController.setKeyguardShown(appWin.getDisplayId(), true /* keyguardShowing */,
                 true /* aodShowing */);
         assertFalse(activity.isVisibleRequested());
 
-        mAtm.mKeyguardController.keyguardGoingAway(0 /* flags */);
+        mAtm.mKeyguardController.keyguardGoingAway(appWin.getDisplayId(), 0 /* flags */);
         assertTrue(activity.isVisibleRequested());
     }
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootTaskTests.java b/services/tests/wmtests/src/com/android/server/wm/RootTaskTests.java
index 030733b..bb49cd2 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootTaskTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootTaskTests.java
@@ -29,6 +29,7 @@
 import static android.content.pm.ActivityInfo.FLAG_RESUME_WHILE_PAUSING;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE;
 import static android.content.pm.ActivityInfo.SCREEN_ORIENTATION_PORTRAIT;
+import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doNothing;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn;
@@ -1505,7 +1506,7 @@
         final UnknownAppVisibilityController unknownAppVisibilityController =
                 mDefaultTaskDisplayArea.mDisplayContent.mUnknownAppVisibilityController;
         final KeyguardController keyguardController = mSupervisor.getKeyguardController();
-        doReturn(true).when(keyguardController).isKeyguardLocked();
+        doReturn(true).when(keyguardController).isKeyguardLocked(eq(DEFAULT_DISPLAY));
 
         // Start 2 activities that their processes have not yet started.
         final ActivityRecord[] activities = new ActivityRecord[2];
@@ -1571,7 +1572,7 @@
         display.isDefaultDisplay = isDefaultDisplay;
 
         task.mDisplayContent = display;
-        doReturn(keyguardGoingAway).when(keyguardController).isKeyguardGoingAway();
+        doReturn(keyguardGoingAway).when(display).isKeyguardGoingAway();
         doReturn(displaySleeping).when(display).isSleeping();
         doReturn(focusedRootTask).when(task).isFocusedRootTaskOnDisplay();
 
diff --git a/services/tests/wmtests/src/com/android/server/wm/TestIWindow.java b/services/tests/wmtests/src/com/android/server/wm/TestIWindow.java
index b6701dd..c742e56 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TestIWindow.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TestIWindow.java
@@ -69,7 +69,7 @@
     }
 
     @Override
-    public void windowFocusChanged(boolean hasFocus, boolean inTouchMode) throws RemoteException {
+    public void windowFocusChanged(boolean hasFocus) throws RemoteException {
     }
 
     @Override
diff --git a/telephony/java/android/telephony/ImsManager.java b/telephony/java/android/telephony/ImsManager.java
index a25abc9..fc76f99 100644
--- a/telephony/java/android/telephony/ImsManager.java
+++ b/telephony/java/android/telephony/ImsManager.java
@@ -135,7 +135,7 @@
             throw new IllegalArgumentException("Invalid subscription ID: " + subscriptionId);
         }
 
-        return new ImsMmTelManager(subscriptionId, sTelephonyCache);
+        return new ImsMmTelManager(mContext, subscriptionId, sTelephonyCache);
     }
 
     /**
diff --git a/telephony/java/android/telephony/ims/ImsMmTelManager.java b/telephony/java/android/telephony/ims/ImsMmTelManager.java
index 7a27ed7..683bb92 100644
--- a/telephony/java/android/telephony/ims/ImsMmTelManager.java
+++ b/telephony/java/android/telephony/ims/ImsMmTelManager.java
@@ -25,6 +25,7 @@
 import android.annotation.SuppressAutoDoc;
 import android.annotation.SuppressLint;
 import android.annotation.SystemApi;
+import android.content.Context;
 import android.os.Binder;
 import android.os.RemoteException;
 import android.os.ServiceSpecificException;
@@ -215,6 +216,7 @@
         }
     }
 
+    private final Context mContext;
     private final int mSubId;
     private final BinderCacheManager<ITelephony> mBinderCache;
 
@@ -256,6 +258,16 @@
      */
     @VisibleForTesting
     public ImsMmTelManager(int subId, BinderCacheManager<ITelephony> binderCache) {
+        this(null, subId, binderCache);
+    }
+
+    /**
+     * Only visible for testing, use {@link ImsManager#getImsMmTelManager(int)} instead.
+     * @hide
+     */
+    @VisibleForTesting
+    public ImsMmTelManager(Context context, int subId, BinderCacheManager<ITelephony> binderCache) {
+        mContext = context;
         mSubId = subId;
         mBinderCache = binderCache;
     }
@@ -1516,7 +1528,8 @@
 
         try {
             telephony.registerImsStateCallback(
-                    mSubId, ImsFeature.FEATURE_MMTEL, callback.getCallbackBinder());
+                    mSubId, ImsFeature.FEATURE_MMTEL,
+                    callback.getCallbackBinder(), getOpPackageName());
         } catch (ServiceSpecificException e) {
             throw new ImsException(e.getMessage(), e.errorCode);
         } catch (RemoteException | IllegalStateException e) {
@@ -1542,6 +1555,14 @@
         }
     }
 
+    private String getOpPackageName() {
+        if (mContext != null) {
+            return mContext.getOpPackageName();
+        } else {
+            return null;
+        }
+    }
+
     private ITelephony getITelephony() {
         return mBinderCache.getBinder();
     }
diff --git a/telephony/java/android/telephony/ims/ImsRcsManager.java b/telephony/java/android/telephony/ims/ImsRcsManager.java
index 42af025..1b047c7 100644
--- a/telephony/java/android/telephony/ims/ImsRcsManager.java
+++ b/telephony/java/android/telephony/ims/ImsRcsManager.java
@@ -572,7 +572,8 @@
 
         try {
             telephony.registerImsStateCallback(
-                    mSubId, ImsFeature.FEATURE_RCS, callback.getCallbackBinder());
+                    mSubId, ImsFeature.FEATURE_RCS,
+                    callback.getCallbackBinder(), mContext.getOpPackageName());
         } catch (ServiceSpecificException e) {
             throw new ImsException(e.getMessage(), e.errorCode);
         } catch (RemoteException | IllegalStateException e) {
diff --git a/telephony/java/android/telephony/ims/SipDelegateManager.java b/telephony/java/android/telephony/ims/SipDelegateManager.java
index 48e3d45..f913df5 100644
--- a/telephony/java/android/telephony/ims/SipDelegateManager.java
+++ b/telephony/java/android/telephony/ims/SipDelegateManager.java
@@ -484,7 +484,8 @@
 
         try {
             telephony.registerImsStateCallback(
-                    mSubId, ImsFeature.FEATURE_RCS, callback.getCallbackBinder());
+                    mSubId, ImsFeature.FEATURE_RCS,
+                    callback.getCallbackBinder(), mContext.getOpPackageName());
         } catch (ServiceSpecificException e) {
             throw new ImsException(e.getMessage(), e.errorCode);
         } catch (RemoteException | IllegalStateException e) {
diff --git a/telephony/java/com/android/internal/telephony/ITelephony.aidl b/telephony/java/com/android/internal/telephony/ITelephony.aidl
index b3c5d042..2e59fc7 100644
--- a/telephony/java/com/android/internal/telephony/ITelephony.aidl
+++ b/telephony/java/com/android/internal/telephony/ITelephony.aidl
@@ -2497,7 +2497,8 @@
     /**
      * Register an IMS connection state callback
      */
-    void registerImsStateCallback(int subId, int feature, in IImsStateCallback cb);
+    void registerImsStateCallback(int subId, int feature, in IImsStateCallback cb,
+            in String callingPackage);
 
     /**
      * Unregister an IMS connection state callback