Merge "Don't tint smart actions in shelf UI" into 24D1-dev
diff --git a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayView.java b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayView.java
index e345836..ba236ba 100644
--- a/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayView.java
+++ b/packages/SystemUI/src/com/android/systemui/clipboardoverlay/ClipboardOverlayView.java
@@ -109,6 +109,7 @@
     private View mDismissButton;
     private LinearLayout mActionContainer;
     private ClipboardOverlayCallbacks mClipboardCallbacks;
+    private ActionButtonViewBinder mActionButtonViewBinder = new ActionButtonViewBinder();
 
     public ClipboardOverlayView(Context context) {
         this(context, null);
@@ -152,7 +153,7 @@
 
     private void bindDefaultActionChips() {
         if (screenshotShelfUi2()) {
-            ActionButtonViewBinder.INSTANCE.bind(mRemoteCopyChip,
+            mActionButtonViewBinder.bind(mRemoteCopyChip,
                     ActionButtonViewModel.Companion.withNextId(
                             new ActionButtonAppearance(
                                     Icon.createWithResource(mContext,
@@ -170,7 +171,7 @@
                                     return null;
                                 }
                             }));
-            ActionButtonViewBinder.INSTANCE.bind(mShareChip,
+            mActionButtonViewBinder.bind(mShareChip,
                     ActionButtonViewModel.Companion.withNextId(
                             new ActionButtonAppearance(
                                     Icon.createWithResource(mContext,
@@ -515,7 +516,7 @@
     private View constructShelfActionChip(RemoteAction action, Runnable onFinish) {
         View chip = LayoutInflater.from(mContext).inflate(
                 R.layout.shelf_action_chip, mActionContainer, false);
-        ActionButtonViewBinder.INSTANCE.bind(chip, ActionButtonViewModel.Companion.withNextId(
+        mActionButtonViewBinder.bind(chip, ActionButtonViewModel.Companion.withNextId(
                 new ActionButtonAppearance(action.getIcon().loadDrawable(mContext),
                         action.getTitle(), action.getTitle(), false), new Function0<>() {
                     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
index dfd742c..71363b2 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotController.java
@@ -47,7 +47,6 @@
 import android.graphics.Bitmap;
 import android.graphics.Insets;
 import android.graphics.Rect;
-import android.hardware.display.DisplayManager;
 import android.net.Uri;
 import android.os.Process;
 import android.os.UserHandle;
@@ -209,8 +208,7 @@
     @Nullable
     private final ScreenshotSoundController mScreenshotSoundController;
     private final PhoneWindow mWindow;
-    private final DisplayManager mDisplayManager;
-    private final int mDisplayId;
+    private final Display mDisplay;
     private final ScrollCaptureExecutor mScrollCaptureExecutor;
     private final ScreenshotNotificationSmartActionsProvider
             mScreenshotNotificationSmartActionsProvider;
@@ -250,7 +248,6 @@
     @AssistedInject
     ScreenshotController(
             Context context,
-            DisplayManager displayManager,
             WindowManager windowManager,
             FeatureFlags flags,
             ScreenshotViewProxy.Factory viewProxyFactory,
@@ -272,12 +269,14 @@
             AssistContentRequester assistContentRequester,
             MessageContainerController messageContainerController,
             Provider<ScreenshotSoundController> screenshotSoundController,
-            @Assisted int displayId,
+            @Assisted Display display,
             @Assisted boolean showUIOnExternalDisplay
     ) {
         mScreenshotSmartActions = screenshotSmartActions;
+        mWindowManager = windowManager;
         mActionsProviderFactory = actionsProviderFactory;
-        mNotificationsController = screenshotNotificationsControllerFactory.create(displayId);
+        mNotificationsController = screenshotNotificationsControllerFactory.create(
+                display.getDisplayId());
         mUiEventLogger = uiEventLogger;
         mImageExporter = imageExporter;
         mImageCapture = imageCapture;
@@ -290,12 +289,8 @@
 
         mScreenshotHandler = timeoutHandler;
         mScreenshotHandler.setDefaultTimeoutMillis(SCREENSHOT_CORNER_DEFAULT_TIMEOUT_MILLIS);
-
-
-        mDisplayId = displayId;
-        mDisplayManager = displayManager;
-        mWindowManager = windowManager;
-        final Context displayContext = context.createDisplayContext(getDisplay());
+        mDisplay = display;
+        final Context displayContext = context.createDisplayContext(display);
         mContext = (WindowContext) displayContext.createWindowContext(TYPE_SCREENSHOT, null);
         mFlags = flags;
         mActionIntentExecutor = actionIntentExecutor;
@@ -303,7 +298,7 @@
         mMessageContainerController = messageContainerController;
         mAssistContentRequester = assistContentRequester;
 
-        mViewProxy = viewProxyFactory.getProxy(mContext, mDisplayId);
+        mViewProxy = viewProxyFactory.getProxy(mContext, mDisplay.getDisplayId());
 
         mScreenshotHandler.setOnTimeoutRunnable(() -> {
             if (DEBUG_UI) {
@@ -329,7 +324,7 @@
                 });
 
         // Sound is only reproduced from the controller of the default display.
-        if (displayId == Display.DEFAULT_DISPLAY) {
+        if (display.getDisplayId() == Display.DEFAULT_DISPLAY) {
             mScreenshotSoundController = screenshotSoundController.get();
         } else {
             mScreenshotSoundController = null;
@@ -357,7 +352,7 @@
         if (screenshot.getType() == WindowManager.TAKE_SCREENSHOT_FULLSCREEN) {
             Rect bounds = getFullScreenRect();
             screenshot.setBitmap(
-                    mImageCapture.captureDisplay(mDisplayId, bounds));
+                    mImageCapture.captureDisplay(mDisplay.getDisplayId(), bounds));
             screenshot.setScreenBounds(bounds);
         }
 
@@ -460,7 +455,7 @@
     }
 
     private boolean shouldShowUi() {
-        return mDisplayId == Display.DEFAULT_DISPLAY || mShowUIOnExternalDisplay;
+        return mDisplay.getDisplayId() == Display.DEFAULT_DISPLAY || mShowUIOnExternalDisplay;
     }
 
     void prepareViewForNewScreenshot(@NonNull ScreenshotData screenshot, String oldPackageName) {
@@ -623,7 +618,7 @@
 
     private void requestScrollCapture(UserHandle owner) {
         mScrollCaptureExecutor.requestScrollCapture(
-                mDisplayId,
+                mDisplay.getDisplayId(),
                 mWindow.getDecorView().getWindowToken(),
                 (response) -> {
                     mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_IMPRESSION,
@@ -646,7 +641,8 @@
         }
         mUiEventLogger.log(ScreenshotEvent.SCREENSHOT_LONG_SCREENSHOT_REQUESTED, 0,
                 response.getPackageName());
-        Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplayId, getFullScreenRect());
+        Bitmap newScreenshot = mImageCapture.captureDisplay(mDisplay.getDisplayId(),
+                getFullScreenRect());
         if (newScreenshot == null) {
             Log.e(TAG, "Failed to capture current screenshot for scroll transition!");
             return;
@@ -824,7 +820,8 @@
     private void saveScreenshotInBackground(
             ScreenshotData screenshot, UUID requestId, Consumer<Uri> finisher) {
         ListenableFuture<ImageExporter.Result> future = mImageExporter.export(mBgExecutor,
-                requestId, screenshot.getBitmap(), screenshot.getUserOrDefault(), mDisplayId);
+                requestId, screenshot.getBitmap(), screenshot.getUserOrDefault(),
+                mDisplay.getDisplayId());
         future.addListener(() -> {
             try {
                 ImageExporter.Result result = future.get();
@@ -866,7 +863,7 @@
         data.mActionsReadyListener = actionsReadyListener;
         data.mQuickShareActionsReadyListener = quickShareActionsReadyListener;
         data.owner = owner;
-        data.displayId = mDisplayId;
+        data.displayId = mDisplay.getDisplayId();
 
         if (mSaveInBgTask != null) {
             // just log success/failure for the pre-existing screenshot
@@ -991,13 +988,9 @@
         }
     }
 
-    private Display getDisplay() {
-        return mDisplayManager.getDisplay(mDisplayId);
-    }
-
     private Rect getFullScreenRect() {
         DisplayMetrics displayMetrics = new DisplayMetrics();
-        getDisplay().getRealMetrics(displayMetrics);
+        mDisplay.getRealMetrics(displayMetrics);
         return new Rect(0, 0, displayMetrics.widthPixels, displayMetrics.heightPixels);
     }
 
@@ -1033,10 +1026,10 @@
         /**
          * Creates an instance of the controller for that specific displayId.
          *
-         * @param displayId:               display to capture
-         * @param showUIOnExternalDisplay: Whether the UI should be shown if this is an external
-         *                                 display.
+         * @param display                 Display to capture.
+         * @param showUIOnExternalDisplay Whether the UI should be shown if this is an external
+         *                                display.
          */
-        ScreenshotController create(int displayId, boolean showUIOnExternalDisplay);
+        ScreenshotController create(Display display, boolean showUIOnExternalDisplay);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
index 1e66cd1..846884f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ScreenshotShelfViewProxy.kt
@@ -31,6 +31,7 @@
 import android.view.WindowManager
 import android.window.OnBackInvokedCallback
 import android.window.OnBackInvokedDispatcher
+import androidx.appcompat.content.res.AppCompatResources
 import androidx.core.animation.doOnEnd
 import androidx.core.animation.doOnStart
 import com.android.internal.logging.UiEventLogger
@@ -58,6 +59,7 @@
     private val logger: UiEventLogger,
     private val viewModel: ScreenshotViewModel,
     private val windowManager: WindowManager,
+    shelfViewBinder: ScreenshotShelfViewBinder,
     private val thumbnailObserver: ThumbnailObserver,
     @Assisted private val context: Context,
     @Assisted private val displayId: Int
@@ -69,7 +71,17 @@
     override var callbacks: ScreenshotView.ScreenshotViewCallback? = null
     override var screenshot: ScreenshotData? = null
         set(value) {
-            viewModel.setScreenshotBitmap(value?.bitmap)
+            value?.let {
+                viewModel.setScreenshotBitmap(it.bitmap)
+                val badgeBg =
+                    AppCompatResources.getDrawable(context, R.drawable.overlay_badge_background)
+                val user = it.userHandle
+                if (badgeBg != null && user != null) {
+                    viewModel.setScreenshotBadge(
+                        context.packageManager.getUserBadgedIcon(badgeBg, user)
+                    )
+                }
+            }
             field = value
         }
 
@@ -81,7 +93,7 @@
     private val animationController = ScreenshotAnimationController(view)
 
     init {
-        ScreenshotShelfViewBinder.bind(
+        shelfViewBinder.bind(
             view,
             viewModel,
             LayoutInflater.from(context),
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
index a9179bf..5feac80 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotExecutor.kt
@@ -52,11 +52,13 @@
         onSaved: (Uri?) -> Unit,
         requestCallback: RequestCallback
     ) {
-        val displayIds = getDisplaysToScreenshot(screenshotRequest.type)
+        val displays = getDisplaysToScreenshot(screenshotRequest.type)
         val resultCallbackWrapper = MultiResultCallbackWrapper(requestCallback)
-        displayIds.forEach { displayId: Int ->
+        displays.forEach { display: Display ->
+            val displayId = display.displayId
             Log.d(TAG, "Executing screenshot for display $displayId")
             dispatchToController(
+                display,
                 rawScreenshotData = ScreenshotData.fromRequest(screenshotRequest, displayId),
                 onSaved =
                     if (displayId == Display.DEFAULT_DISPLAY) {
@@ -69,6 +71,7 @@
 
     /** All logging should be triggered only by this method. */
     private suspend fun dispatchToController(
+        display: Display,
         rawScreenshotData: ScreenshotData,
         onSaved: (Uri?) -> Unit,
         callback: RequestCallback
@@ -88,8 +91,7 @@
         logScreenshotRequested(screenshotData)
         Log.d(TAG, "Screenshot request: $screenshotData")
         try {
-            getScreenshotController(screenshotData.displayId)
-                .handleScreenshot(screenshotData, onSaved, callback)
+            getScreenshotController(display).handleScreenshot(screenshotData, onSaved, callback)
         } catch (e: IllegalStateException) {
             Log.e(TAG, "Error while ScreenshotController was handling ScreenshotData!", e)
             onFailedScreenshotRequest(screenshotData, callback)
@@ -119,12 +121,13 @@
         callback.reportError()
     }
 
-    private suspend fun getDisplaysToScreenshot(requestType: Int): List<Int> {
+    private suspend fun getDisplaysToScreenshot(requestType: Int): List<Display> {
+        val allDisplays = displays.first()
         return if (requestType == TAKE_SCREENSHOT_PROVIDED_IMAGE) {
             // If this is a provided image, let's show the UI on the default display only.
-            listOf(Display.DEFAULT_DISPLAY)
+            allDisplays.filter { it.displayId == Display.DEFAULT_DISPLAY }
         } else {
-            displays.first().filter { it.type in ALLOWED_DISPLAY_TYPES }.map { it.displayId }
+            allDisplays.filter { it.type in ALLOWED_DISPLAY_TYPES }
         }
     }
 
@@ -158,9 +161,9 @@
         screenshotControllers.clear()
     }
 
-    private fun getScreenshotController(id: Int): ScreenshotController {
-        return screenshotControllers.computeIfAbsent(id) {
-            screenshotControllerFactory.create(id, /* showUIOnExternalDisplay= */ false)
+    private fun getScreenshotController(display: Display): ScreenshotController {
+        return screenshotControllers.computeIfAbsent(display.displayId) {
+            screenshotControllerFactory.create(display, /* showUIOnExternalDisplay= */ false)
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
index 9cf347b..c03ba65 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/TakeScreenshotService.java
@@ -37,6 +37,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
+import android.hardware.display.DisplayManager;
 import android.net.Uri;
 import android.os.Handler;
 import android.os.IBinder;
@@ -116,7 +117,8 @@
             UiEventLogger uiEventLogger,
             ScreenshotNotificationsController.Factory notificationsControllerFactory,
             Context context, @Background Executor bgExecutor, FeatureFlags featureFlags,
-            RequestProcessor processor, Provider<TakeScreenshotExecutor> takeScreenshotExecutor) {
+            RequestProcessor processor, Provider<TakeScreenshotExecutor> takeScreenshotExecutor,
+            DisplayManager displayManager) {
         if (DEBUG_SERVICE) {
             Log.d(TAG, "new " + this);
         }
@@ -134,7 +136,8 @@
             mScreenshot = null;
         } else {
             mScreenshot = screenshotControllerFactory.create(
-                    Display.DEFAULT_DISPLAY, /* showUIOnExternalDisplay= */ false);
+                    displayManager.getDisplay(
+                            Display.DEFAULT_DISPLAY), /* showUIOnExternalDisplay= */ false);
         }
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
index ccd7a68..36aa39f 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ActionButtonViewBinder.kt
@@ -23,8 +23,9 @@
 import com.android.systemui.res.R
 import com.android.systemui.screenshot.ui.TransitioningIconDrawable
 import com.android.systemui.screenshot.ui.viewmodel.ActionButtonViewModel
+import javax.inject.Inject
 
-object ActionButtonViewBinder {
+class ActionButtonViewBinder @Inject constructor() {
     /** Binds the given view to the given view-model */
     fun bind(view: View, viewModel: ActionButtonViewModel) {
         val iconView = view.requireViewById<ImageView>(R.id.overlay_action_chip_icon)
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
index 89f904a..c7bc50cb 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/binder/ScreenshotShelfViewBinder.kt
@@ -35,10 +35,13 @@
 import com.android.systemui.screenshot.ui.viewmodel.AnimationState
 import com.android.systemui.screenshot.ui.viewmodel.ScreenshotViewModel
 import com.android.systemui.util.children
+import javax.inject.Inject
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
 
-object ScreenshotShelfViewBinder {
+class ScreenshotShelfViewBinder
+@Inject
+constructor(private val buttonViewBinder: ActionButtonViewBinder) {
     fun bind(
         view: ScreenshotShelfView,
         viewModel: ScreenshotViewModel,
@@ -68,6 +71,7 @@
         dismissButton.setOnClickListener {
             onDismissalRequested(ScreenshotEvent.SCREENSHOT_EXPLICIT_DISMISSAL, null)
         }
+        val badgeView = view.requireViewById<ImageView>(R.id.screenshot_badge)
 
         // use immediate dispatcher to ensure screenshot bitmap is set before animation
         view.repeatWhenAttached(Dispatchers.Main.immediate) {
@@ -87,6 +91,12 @@
                         }
                     }
                     launch {
+                        viewModel.badge.collect { badge ->
+                            badgeView.setImageDrawable(badge)
+                            badgeView.visibility = if (badge != null) View.VISIBLE else View.GONE
+                        }
+                    }
+                    launch {
                         viewModel.previewAction.collect { onClick ->
                             previewView.setOnClickListener { onClick?.invoke() }
                         }
@@ -151,14 +161,14 @@
             val currentView: View? = actionsContainer.getChildAt(index)
             if (action.id == currentView?.tag) {
                 // Same ID, update the display
-                ActionButtonViewBinder.bind(currentView, action)
+                buttonViewBinder.bind(currentView, action)
             } else {
                 // Different ID. Removals have already happened so this must
                 // mean that the new action must be inserted here.
                 val actionButton =
                     layoutInflater.inflate(R.layout.shelf_action_chip, actionsContainer, false)
                 actionsContainer.addView(actionButton, index)
-                ActionButtonViewBinder.bind(actionButton, action)
+                buttonViewBinder.bind(actionButton, action)
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
index 5f36f73..81bc281 100644
--- a/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/screenshot/ui/viewmodel/ScreenshotViewModel.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.screenshot.ui.viewmodel
 
 import android.graphics.Bitmap
+import android.graphics.drawable.Drawable
 import android.util.Log
 import android.view.accessibility.AccessibilityManager
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -25,6 +26,8 @@
 class ScreenshotViewModel(private val accessibilityManager: AccessibilityManager) {
     private val _preview = MutableStateFlow<Bitmap?>(null)
     val preview: StateFlow<Bitmap?> = _preview
+    private val _badge = MutableStateFlow<Drawable?>(null)
+    val badge: StateFlow<Drawable?> = _badge
     private val _previewAction = MutableStateFlow<(() -> Unit)?>(null)
     val previewAction: StateFlow<(() -> Unit)?> = _previewAction
     private val _actions = MutableStateFlow(emptyList<ActionButtonViewModel>())
@@ -39,6 +42,10 @@
         _preview.value = bitmap
     }
 
+    fun setScreenshotBadge(badge: Drawable?) {
+        _badge.value = badge
+    }
+
     fun setPreviewAction(onClick: () -> Unit) {
         _previewAction.value = onClick
     }
@@ -109,6 +116,7 @@
 
     fun reset() {
         _preview.value = null
+        _badge.value = null
         _previewAction.value = null
         _actions.value = listOf()
         _animationState.value = AnimationState.NOT_STARTED
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index 5f3a83a..1bb5a77 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -1765,9 +1765,7 @@
      */
     public ExpandableNotificationRow(Context context, AttributeSet attrs) {
         this(context, attrs, context);
-        if (com.android.systemui.Flags.notificationRowUserContext()) {
-            Log.wtf(TAG, "This constructor shouldn't be called");
-        }
+        Log.wtf(TAG, "This constructor shouldn't be called");
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolver.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolver.java
index 609b15e..3e932aa 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolver.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationInlineImageResolver.java
@@ -31,7 +31,6 @@
 import com.android.internal.widget.ImageResolver;
 import com.android.internal.widget.LocalImageResolver;
 import com.android.internal.widget.MessagingMessage;
-import com.android.systemui.Flags;
 
 import java.util.HashSet;
 import java.util.List;
@@ -67,11 +66,7 @@
      * @param imageCache The implementation of internal cache.
      */
     public NotificationInlineImageResolver(Context context, ImageCache imageCache) {
-        if (Flags.notificationRowUserContext()) {
-            mContext = context;
-        } else {
-            mContext = context.getApplicationContext();
-        }
+        mContext = context;
         mImageCache = imageCache;
 
         if (mImageCache != null) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
index 9445d56..9a85b2f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/RowInflaterTask.java
@@ -60,9 +60,8 @@
             mInflateOrigin = new Throwable("inflate requested here");
         }
         mListener = listener;
-        AsyncLayoutInflater inflater = com.android.systemui.Flags.notificationRowUserContext()
-                ? new AsyncLayoutInflater(context, new RowAsyncLayoutInflater(entry))
-                : new AsyncLayoutInflater(context);
+        AsyncLayoutInflater inflater = new AsyncLayoutInflater(context,
+                new RowAsyncLayoutInflater(entry));
         mEntry = entry;
         entry.setInflationTask(this);
         inflater.inflate(R.layout.status_bar_notification_row, parent, this);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
index 22312f4..0dda41f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotExecutorTest.kt
@@ -69,8 +69,9 @@
 
     @Before
     fun setUp() {
-        whenever(controllerFactory.create(eq(0), any())).thenReturn(controller0)
-        whenever(controllerFactory.create(eq(1), any())).thenReturn(controller1)
+        whenever(controllerFactory.create(any(), any())).thenAnswer {
+            if (it.getArgument<Display>(0).displayId == 0) controller0 else controller1
+        }
         whenever(notificationControllerFactory.create(eq(0))).thenReturn(notificationsController0)
         whenever(notificationControllerFactory.create(eq(1))).thenReturn(notificationsController1)
     }
@@ -78,12 +79,14 @@
     @Test
     fun executeScreenshots_severalDisplays_callsControllerForEachOne() =
         testScope.runTest {
-            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val internalDisplay = display(TYPE_INTERNAL, id = 0)
+            val externalDisplay = display(TYPE_EXTERNAL, id = 1)
+            setDisplays(internalDisplay, externalDisplay)
             val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(createScreenshotRequest(), onSaved, callback)
 
-            verify(controllerFactory).create(eq(0), any())
-            verify(controllerFactory).create(eq(1), any())
+            verify(controllerFactory).create(eq(internalDisplay), any())
+            verify(controllerFactory).create(eq(externalDisplay), any())
 
             val capturer = ArgumentCaptor<ScreenshotData>()
 
@@ -107,7 +110,9 @@
     @Test
     fun executeScreenshots_providedImageType_callsOnlyDefaultDisplayController() =
         testScope.runTest {
-            setDisplays(display(TYPE_INTERNAL, id = 0), display(TYPE_EXTERNAL, id = 1))
+            val internalDisplay = display(TYPE_INTERNAL, id = 0)
+            val externalDisplay = display(TYPE_EXTERNAL, id = 1)
+            setDisplays(internalDisplay, externalDisplay)
             val onSaved = { _: Uri? -> }
             screenshotExecutor.executeScreenshots(
                 createScreenshotRequest(TAKE_SCREENSHOT_PROVIDED_IMAGE),
@@ -115,8 +120,8 @@
                 callback
             )
 
-            verify(controllerFactory).create(eq(0), any())
-            verify(controllerFactory, never()).create(eq(1), any())
+            verify(controllerFactory).create(eq(internalDisplay), any())
+            verify(controllerFactory, never()).create(eq(externalDisplay), any())
 
             val capturer = ArgumentCaptor<ScreenshotData>()
 
@@ -473,6 +478,7 @@
         var processed: ScreenshotData? = null
         var toReturn: ScreenshotData? = null
         var shouldThrowException = false
+
         override suspend fun process(screenshot: ScreenshotData): ScreenshotData {
             if (shouldThrowException) throw RequestProcessorException("")
             processed = screenshot
diff --git a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
index f3809aa..b6163512 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/screenshot/TakeScreenshotServiceTest.kt
@@ -20,6 +20,7 @@
 import android.app.admin.DevicePolicyResources.Strings.SystemUi.SCREENSHOT_BLOCKED_BY_ADMIN
 import android.app.admin.DevicePolicyResourcesManager
 import android.content.ComponentName
+import android.hardware.display.DisplayManager
 import android.os.UserHandle
 import android.os.UserManager
 import android.testing.AndroidTestingRunner
@@ -38,6 +39,7 @@
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
 import com.android.systemui.util.mockito.whenever
 import java.util.function.Consumer
 import org.junit.Assert.assertEquals
@@ -68,6 +70,7 @@
     private val notificationsControllerFactory = mock<ScreenshotNotificationsController.Factory>()
     private val notificationsController = mock<ScreenshotNotificationsController>()
     private val callback = mock<RequestCallback>()
+    private val displayManager = mock<DisplayManager>()
 
     private val eventLogger = UiEventLoggerFake()
     private val flags = FakeFeatureFlags()
@@ -87,7 +90,7 @@
             )
             .thenReturn(false)
         whenever(userManager.isUserUnlocked).thenReturn(true)
-        whenever(controllerFactory.create(any(), any())).thenReturn(controller)
+        whenever(controllerFactory.create(nullable<Display>(), any())).thenReturn(controller)
         whenever(notificationsControllerFactory.create(any())).thenReturn(notificationsController)
 
         // Stub request processor as a synchronous no-op for tests with the flag enabled
@@ -331,6 +334,7 @@
                 flags,
                 requestProcessor,
                 { takeScreenshotExecutor },
+                displayManager,
             )
         service.attach(
             mContext,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
index 42a6924..9c31bb2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowTest.java
@@ -873,7 +873,6 @@
     }
 
     @Test
-    @EnableFlags(com.android.systemui.Flags.FLAG_NOTIFICATION_ROW_USER_CONTEXT)
     public void imageResolver_differentNotificationUser_createsUserContext() throws Exception {
         UserHandle user = new UserHandle(33);
         Context userContext = new SysuiTestableContext(mContext);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
index e78081f..84139da 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
@@ -571,9 +571,7 @@
 
         LayoutInflater inflater = (LayoutInflater) mContext.getSystemService(
                 Context.LAYOUT_INFLATER_SERVICE);
-        if (com.android.systemui.Flags.notificationRowUserContext()) {
-            inflater.setFactory2(new RowInflaterTask.RowAsyncLayoutInflater(entry));
-        }
+        inflater.setFactory2(new RowInflaterTask.RowAsyncLayoutInflater(entry));
         mRow = (ExpandableNotificationRow) inflater.inflate(
                 R.layout.status_bar_notification_row,
                 null /* root */,
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
index d8098b7..133719c 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/display/data/repository/FakeDisplayRepository.kt
@@ -22,6 +22,7 @@
 import org.mockito.Mockito.`when` as whenever
 
 /** Creates a mock display. */
+@JvmOverloads
 fun display(
     type: Int,
     flags: Int = 0,
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index fb62785..29b28e5 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -1364,6 +1364,9 @@
 
     @GuardedBy("this")
     private void packageRemovedLocked(int uid, String packageName) {
+        mHandler.post(PooledLambda.obtainRunnable(HistoricalRegistry::clearHistory,
+                mHistoricalRegistry, uid, packageName));
+
         UidState uidState = mUidStates.get(uid);
         if (uidState == null) {
             return;
@@ -1398,9 +1401,6 @@
                 }
             }
         }
-
-        mHandler.post(PooledLambda.obtainRunnable(HistoricalRegistry::clearHistory,
-                    mHistoricalRegistry, uid, packageName));
     }
 
     public void uidRemoved(int uid) {