diff --git a/services/autofill/java/com/android/server/autofill/Helper.java b/services/autofill/java/com/android/server/autofill/Helper.java
index bc5d645..ac77043 100644
--- a/services/autofill/java/com/android/server/autofill/Helper.java
+++ b/services/autofill/java/com/android/server/autofill/Helper.java
@@ -16,13 +16,18 @@
 
 package com.android.server.autofill;
 
+import static com.android.server.autofill.Helper.sDebug;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.app.assist.AssistStructure;
 import android.app.assist.AssistStructure.ViewNode;
 import android.app.assist.AssistStructure.WindowNode;
 import android.content.ComponentName;
+import android.content.Context;
+import android.hardware.display.DisplayManager;
 import android.metrics.LogMaker;
+import android.os.UserManager;
 import android.service.autofill.Dataset;
 import android.service.autofill.InternalSanitizer;
 import android.service.autofill.SaveInfo;
@@ -30,6 +35,7 @@
 import android.util.ArrayMap;
 import android.util.ArraySet;
 import android.util.Slog;
+import android.view.Display;
 import android.view.View;
 import android.view.WindowManager;
 import android.view.autofill.AutofillId;
@@ -37,6 +43,7 @@
 
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.internal.util.ArrayUtils;
+import com.android.server.utils.Slogf;
 
 import java.io.PrintWriter;
 import java.util.ArrayDeque;
@@ -281,6 +288,38 @@
         return true;
     }
 
+    /**
+     * Gets a context with the proper display id.
+     *
+     * <p>For most cases it will return the provided context, but on devices that
+     * {@link UserManager#isVisibleBackgroundUsersEnabled() support visible background users}, it
+     * will return a context with the display pased as parameter.
+     */
+    static Context getDisplayContext(Context context, int displayId) {
+        if (!UserManager.isVisibleBackgroundUsersEnabled()) {
+            return context;
+        }
+        if (context.getDisplayId() == displayId) {
+            if (sDebug) {
+                Slogf.d(TAG, "getDisplayContext(): context %s already has displayId %d", context,
+                        displayId);
+            }
+            return context;
+        }
+        if (sDebug) {
+            Slogf.d(TAG, "Creating context for display %d", displayId);
+        }
+        Display display = context.getSystemService(DisplayManager.class).getDisplay(displayId);
+        if (display == null) {
+            Slogf.wtf(TAG, "Could not get context with displayId %d, Autofill operations will "
+                    + "probably fail)", displayId);
+            return context;
+        }
+
+        return context.createDisplayContext(display);
+    }
+
+
     private interface ViewNodeFilter {
         boolean matches(ViewNode node);
     }
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 3ab5bca..f83d734 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -164,10 +164,12 @@
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.internal.util.ArrayUtils;
+import com.android.server.LocalServices;
 import com.android.server.autofill.ui.AutoFillUI;
 import com.android.server.autofill.ui.InlineFillUi;
 import com.android.server.autofill.ui.PendingUi;
 import com.android.server.inputmethod.InputMethodManagerInternal;
+import com.android.server.wm.ActivityTaskManagerInternal;
 
 import java.io.PrintWriter;
 import java.lang.annotation.Retention;
@@ -214,7 +216,12 @@
     private final AutofillManagerServiceImpl mService;
     private final Handler mHandler;
     private final AutoFillUI mUi;
-    @NonNull private final Context mContext;
+
+    /**
+     * Context associated with the session, it has the same {@link Context#getDisplayId() displayId}
+     * of the activity being autofilled.
+     */
+    private final Context mContext;
 
     private final MetricsLogger mMetricsLogger = new MetricsLogger();
 
@@ -1352,7 +1359,9 @@
         mHasCallback = hasCallback;
         mUiLatencyHistory = uiLatencyHistory;
         mWtfHistory = wtfHistory;
-        mContext = context;
+        int displayId = LocalServices.getService(ActivityTaskManagerInternal.class)
+                .getDisplayId(activityToken);
+        mContext = Helper.getDisplayContext(context, displayId);
         mComponentName = componentName;
         mCompatMode = compatMode;
         mSessionState = STATE_ACTIVE;
@@ -3401,7 +3410,7 @@
                 final long saveUiDisplayStartTimestamp = SystemClock.elapsedRealtime();
                 getUiForShowing().showSaveUi(serviceLabel, serviceIcon,
                         mService.getServicePackageName(), saveInfo, this,
-                        mComponentName, this, userId, mPendingSaveUi, isUpdate, mCompatMode,
+                        mComponentName, this, mContext,  mPendingSaveUi, isUpdate, mCompatMode,
                         response.getShowSaveDialogIcon());
                 mSaveEventLogger.maybeSetLatencySaveUiDisplayMillis(
                     SystemClock.elapsedRealtime()- saveUiDisplayStartTimestamp);
@@ -4272,7 +4281,7 @@
 
         getUiForShowing().showFillUi(filledId, response, filterText,
                 mService.getServicePackageName(), mComponentName,
-                targetLabel, targetIcon, this, userId, id, mCompatMode);
+                targetLabel, targetIcon, this, mContext, id, mCompatMode);
 
         synchronized (mLock) {
             mPresentationStatsEventLogger.maybeSetCountShown(
@@ -5453,6 +5462,7 @@
         pw.print(prefix); pw.print("uid: "); pw.println(uid);
         pw.print(prefix); pw.print("taskId: "); pw.println(taskId);
         pw.print(prefix); pw.print("flags: "); pw.println(mFlags);
+        pw.print(prefix); pw.print("displayId: "); pw.println(mContext.getDisplayId());
         pw.print(prefix); pw.print("state: "); pw.println(sessionStateAsString(mSessionState));
         pw.print(prefix); pw.print("mComponentName: "); pw.println(mComponentName);
         pw.print(prefix); pw.print("mActivityToken: "); pw.println(mActivityToken);
diff --git a/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java b/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
index cfd66f1..b3cbe45 100644
--- a/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
+++ b/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
@@ -23,7 +23,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.annotation.UserIdInt;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -53,6 +52,7 @@
 import com.android.server.UiModeManagerInternal;
 import com.android.server.UiThread;
 import com.android.server.autofill.Helper;
+import com.android.server.utils.Slogf;
 
 import java.io.PrintWriter;
 
@@ -198,6 +198,7 @@
      * @param serviceIcon icon of autofill service
      * @param callback identifier for the caller
      * @param userId the user associated wit the session
+     * @param context context with the proper state (like display id) to show the UI
      * @param sessionId id of the autofill session
      * @param compatMode whether the app is being autofilled in compatibility mode.
      */
@@ -205,11 +206,11 @@
             @Nullable String filterText, @Nullable String servicePackageName,
             @NonNull ComponentName componentName, @NonNull CharSequence serviceLabel,
             @NonNull Drawable serviceIcon, @NonNull AutoFillUiCallback callback,
-            @UserIdInt int userId, int sessionId, boolean compatMode) {
+            @NonNull Context context, int sessionId, boolean compatMode) {
         if (sDebug) {
             final int size = filterText == null ? 0 : filterText.length();
-            Slog.d(TAG, "showFillUi(): id=" + focusedId + ", filter=" + size + " chars, userId="
-                    + userId);
+            Slogf.d(TAG, "showFillUi(): id=%s, filter=%d chars, displayId=%d", focusedId, size,
+                    context.getDisplayId());
         }
         final LogMaker log = Helper
                 .newLogMaker(MetricsEvent.AUTOFILL_FILL_UI, componentName, servicePackageName,
@@ -224,10 +225,8 @@
                 return;
             }
             hideAllUiThread(callback);
-            mFillUi = new FillUi(mContext, userId, response, focusedId,
-                    filterText, mOverlayControl, serviceLabel, serviceIcon,
-                    mUiModeMgr.isNightMode(),
-                    new FillUi.Callback() {
+            mFillUi = new FillUi(context, response, focusedId, filterText, mOverlayControl,
+                    serviceLabel, serviceIcon, mUiModeMgr.isNightMode(), new FillUi.Callback() {
                 @Override
                 public void onResponsePicked(FillResponse response) {
                     log.setType(MetricsEvent.TYPE_DETAIL);
@@ -325,12 +324,12 @@
     public void showSaveUi(@NonNull CharSequence serviceLabel, @NonNull Drawable serviceIcon,
             @Nullable String servicePackageName, @NonNull SaveInfo info,
             @NonNull ValueFinder valueFinder, @NonNull ComponentName componentName,
-            @NonNull AutoFillUiCallback callback, @UserIdInt int userId,
+            @NonNull AutoFillUiCallback callback, @NonNull Context context,
             @NonNull PendingUi pendingSaveUi, boolean isUpdate, boolean compatMode,
             boolean showServiceIcon) {
         if (sVerbose) {
-            Slog.v(TAG, "showSaveUi(update=" + isUpdate + ") for " + componentName.toShortString()
-                    + " and user " + userId + ": " + info);
+            Slogf.v(TAG, "showSaveUi(update=%b) for %s and display %d: %s", isUpdate,
+                    componentName.toShortString(), context.getDisplayId(), info);
         }
         int numIds = 0;
         numIds += info.getRequiredIds() == null ? 0 : info.getRequiredIds().length;
@@ -350,7 +349,7 @@
             }
             hideAllUiThread(callback);
             mSaveUiCallback = callback;
-            mSaveUi = new SaveUi(mContext, userId, pendingSaveUi, serviceLabel, serviceIcon,
+            mSaveUi = new SaveUi(context, pendingSaveUi, serviceLabel, serviceIcon,
                     servicePackageName, componentName, info, valueFinder, mOverlayControl,
                     new SaveUi.OnSaveListener() {
                 @Override
diff --git a/services/autofill/java/com/android/server/autofill/ui/DisplayHelper.java b/services/autofill/java/com/android/server/autofill/ui/DisplayHelper.java
deleted file mode 100644
index 5353480..0000000
--- a/services/autofill/java/com/android/server/autofill/ui/DisplayHelper.java
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright (C) 2023 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.autofill.ui;
-
-import static com.android.server.autofill.Helper.sDebug;
-
-import android.annotation.UserIdInt;
-import android.content.Context;
-import android.hardware.display.DisplayManager;
-import android.os.UserManager;
-import android.view.Display;
-
-import com.android.server.LocalServices;
-import com.android.server.pm.UserManagerInternal;
-import com.android.server.utils.Slogf;
-
-/**
- * Helper for display-related needs.
- */
-final class DisplayHelper {
-
-    private static final String TAG = "AutofillDisplayHelper";
-
-    private static final UserManagerInternal sUmi = LocalServices
-            .getService(UserManagerInternal.class);
-
-    /**
-     * Gets a context with the proper display id set for the given user.
-     *
-     * <p>For most cases it will return the provided context, but on devices that
-     * {@link UserManager#isVisibleBackgroundUsersEnabled() support visible background users}, it
-     * will return a context with the display the user started visible on.
-     */
-    static Context getDisplayContext(Context context, @UserIdInt int userId) {
-        if (!UserManager.isVisibleBackgroundUsersEnabled()) {
-            return context;
-        }
-        int displayId = sUmi.getMainDisplayAssignedToUser(userId);
-        if (sDebug) {
-            Slogf.d(TAG, "Creating context for display %d for user %d", displayId, userId);
-        }
-        Display display = context.getSystemService(DisplayManager.class).getDisplay(displayId);
-        if (display == null) {
-            Slogf.wtf(TAG, "Could not get display with id %d (which is associated with user %d; "
-                    + "FillUi operations will probably fail", displayId, userId);
-            return context;
-        }
-
-        return context.createDisplayContext(display);
-    }
-
-    private DisplayHelper() {
-        throw new UnsupportedOperationException("Contains only static methods");
-    }
-}
diff --git a/services/autofill/java/com/android/server/autofill/ui/FillUi.java b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
index b651ae59..129ce72 100644
--- a/services/autofill/java/com/android/server/autofill/ui/FillUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
@@ -22,7 +22,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.annotation.UserIdInt;
 import android.content.Context;
 import android.content.IntentSender;
 import android.content.pm.PackageManager;
@@ -61,6 +60,7 @@
 import com.android.server.UiThread;
 import com.android.server.autofill.AutofillManagerService;
 import com.android.server.autofill.Helper;
+import com.android.server.utils.Slogf;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -134,14 +134,15 @@
         return context.getPackageManager().hasSystemFeature(PackageManager.FEATURE_LEANBACK);
     }
 
-    FillUi(@NonNull Context systemContext, @UserIdInt int userId, @NonNull FillResponse response,
+    FillUi(@NonNull Context context, @NonNull FillResponse response,
             @NonNull AutofillId focusedViewId, @Nullable String filterText,
             @NonNull OverlayControl overlayControl, @NonNull CharSequence serviceLabel,
             @NonNull Drawable serviceIcon, boolean nightMode, @NonNull Callback callback) {
-        if (sVerbose) Slog.v(TAG, "nightMode: " + nightMode);
+        if (sVerbose) {
+            Slogf.v(TAG, "nightMode: %b displayId: %d", nightMode, context.getDisplayId());
+        }
         mThemeId = nightMode ? THEME_ID_DARK : THEME_ID_LIGHT;
         mCallback = callback;
-        Context context = DisplayHelper.getDisplayContext(systemContext, userId);
         mFullScreen = isFullScreen(context);
         mContext = new ContextThemeWrapper(context, mThemeId);
 
diff --git a/services/autofill/java/com/android/server/autofill/ui/SaveUi.java b/services/autofill/java/com/android/server/autofill/ui/SaveUi.java
index aec0bdf..1204259 100644
--- a/services/autofill/java/com/android/server/autofill/ui/SaveUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/SaveUi.java
@@ -21,7 +21,6 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.annotation.UserIdInt;
 import android.app.Dialog;
 import android.app.PendingIntent;
 import android.content.ComponentName;
@@ -71,6 +70,7 @@
 import com.android.internal.util.ArrayUtils;
 import com.android.server.UiThread;
 import com.android.server.autofill.Helper;
+import com.android.server.utils.Slogf;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -173,14 +173,15 @@
 
     private boolean mDestroyed;
 
-    SaveUi(@NonNull Context systemContext, @UserIdInt int userId, @NonNull PendingUi pendingUi,
+    SaveUi(@NonNull Context context, @NonNull PendingUi pendingUi,
            @NonNull CharSequence serviceLabel, @NonNull Drawable serviceIcon,
            @Nullable String servicePackageName, @NonNull ComponentName componentName,
            @NonNull SaveInfo info, @NonNull ValueFinder valueFinder,
            @NonNull OverlayControl overlayControl, @NonNull OnSaveListener listener,
            boolean nightMode, boolean isUpdate, boolean compatMode, boolean showServiceIcon) {
-        Context context = DisplayHelper.getDisplayContext(systemContext, userId);
-        if (sVerbose) Slog.v(TAG, "nightMode: " + nightMode);
+        if (sVerbose) {
+            Slogf.v(TAG, "nightMode: %b displayId: %d", nightMode, context.getDisplayId());
+        }
         mThemeId = nightMode ? THEME_ID_DARK : THEME_ID_LIGHT;
         mPendingUi = pendingUi;
         mListener = new OneActionThenDestroyListener(listener);
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
index 4949ebc..3f4a775 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerInternal.java
@@ -755,4 +755,10 @@
 
     /** Unregister a task stack listener so that it stops receiving callbacks. */;
     public abstract void unregisterTaskStackListener(ITaskStackListener listener);
+
+    /**
+     * Gets the id of the display the activity was launched on.
+     * @param token The activity token.
+     */
+    public abstract int getDisplayId(IBinder token);
 }
diff --git a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
index b377032..1f4606b 100644
--- a/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/wm/ActivityTaskManagerService.java
@@ -5816,6 +5816,18 @@
         }
 
         @Override
+        public int getDisplayId(IBinder token) {
+            synchronized (mGlobalLock) {
+                ActivityRecord r = ActivityRecord.forTokenLocked(token);
+                if (r == null) {
+                    throw new IllegalArgumentException(
+                            "setFocusedActivity: No activity record matching token=" + token);
+                }
+                return r.getDisplayId();
+            }
+        }
+
+        @Override
         public void registerScreenObserver(ScreenObserver observer) {
             mScreenObservers.add(observer);
         }