Use context for correct display in the renderer service

bug:149958394

Test: atest android.autofillservice.cts.inline.InlineLoginActivityTest

Change-Id: I4bcdca8188b9b97076c4146f0f9eddffe3e1b3b1
(cherry picked from commit 2aa4f109961966c49e9cb1be4b3e96605a996c54)

6 files changed, 157 insertions, 55 deletions

diff --git a/core/java/android/service/autofill/IInlineSuggestionRenderService.aidl b/core/java/android/service/autofill/IInlineSuggestionRenderService.aidl
index c389b1a1a5ff..dd434b440af4 100644
--- a/core/java/android/service/autofill/IInlineSuggestionRenderService.aidl
+++ b/core/java/android/service/autofill/IInlineSuggestionRenderService.aidl
@@ -28,5 +28,5 @@ import android.service.autofill.InlinePresentation;
 oneway interface IInlineSuggestionRenderService {
     void renderSuggestion(in IInlineSuggestionUiCallback callback,
                           in InlinePresentation presentation, int width, int height,
-                          in IBinder hostInputToken);
+                          in IBinder hostInputToken, int displayId);
 }
diff --git a/core/java/android/service/autofill/InlineSuggestionRenderService.java b/core/java/android/service/autofill/InlineSuggestionRenderService.java
index 29069e7035f9..17e0456df156 100644
--- a/core/java/android/service/autofill/InlineSuggestionRenderService.java
+++ b/core/java/android/service/autofill/InlineSuggestionRenderService.java
@@ -23,14 +23,18 @@ import android.annotation.SystemApi;
 import android.annotation.TestApi;
 import android.app.Service;
 import android.app.slice.Slice;
+import android.content.Context;
 import android.content.Intent;
 import android.graphics.PixelFormat;
+import android.hardware.display.DisplayManager;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.IBinder;
 import android.os.Looper;
 import android.os.RemoteException;
+import android.util.DisplayMetrics;
 import android.util.Log;
+import android.view.Display;
 import android.view.SurfaceControl;
 import android.view.SurfaceControlViewHost;
 import android.view.View;
@@ -61,7 +65,8 @@ public abstract class InlineSuggestionRenderService extends Service {
     private final Handler mHandler = new Handler(Looper.getMainLooper(), null, true);
 
     private void handleRenderSuggestion(IInlineSuggestionUiCallback callback,
-            InlinePresentation presentation, int width, int height, IBinder hostInputToken) {
+            InlinePresentation presentation, int width, int height, IBinder hostInputToken,
+            int displayId) {
         if (hostInputToken == null) {
             try {
                 callback.onError();
@@ -70,8 +75,17 @@ public abstract class InlineSuggestionRenderService extends Service {
             }
             return;
         }
-        final SurfaceControlViewHost host = new SurfaceControlViewHost(this, this.getDisplay(),
-                hostInputToken);
+
+        final DisplayManager displayManager = getSystemService(DisplayManager.class);
+        final Display targetDisplay = displayManager.getDisplay(displayId);
+        if (targetDisplay == null) {
+            sendResult(callback, /*surface*/ null);
+            return;
+        }
+        final Context displayContext = createDisplayContext(targetDisplay);
+
+        final SurfaceControlViewHost host = new SurfaceControlViewHost(displayContext,
+                displayContext.getDisplay(), hostInputToken);
         final SurfaceControl surface = host.getSurfacePackage().getSurfaceControl();
 
         final View suggestionView = onRenderSuggestion(presentation, width, height);
@@ -90,6 +104,11 @@ public abstract class InlineSuggestionRenderService extends Service {
                 new WindowManager.LayoutParams(width, height,
                         WindowManager.LayoutParams.TYPE_APPLICATION, 0, PixelFormat.TRANSPARENT);
         host.addView(suggestionRoot, lp);
+        sendResult(callback, surface);
+    }
+
+    private void sendResult(@NonNull IInlineSuggestionUiCallback callback,
+            @Nullable SurfaceControl surface) {
         try {
             callback.onContent(surface);
         } catch (RemoteException e) {
@@ -105,11 +124,11 @@ public abstract class InlineSuggestionRenderService extends Service {
                 @Override
                 public void renderSuggestion(@NonNull IInlineSuggestionUiCallback callback,
                         @NonNull InlinePresentation presentation, int width, int height,
-                        @Nullable IBinder hostInputToken) {
+                        @Nullable IBinder hostInputToken, int displayId) {
                     mHandler.sendMessage(obtainMessage(
                             InlineSuggestionRenderService::handleRenderSuggestion,
                             InlineSuggestionRenderService.this, callback, presentation,
-                            width, height, hostInputToken));
+                            width, height, hostInputToken, displayId));
                 }
             }.asBinder();
         }
diff --git a/core/java/android/view/inputmethod/InlineSuggestionsRequest.java b/core/java/android/view/inputmethod/InlineSuggestionsRequest.java
index 5700dda3a2e3..e50da40cdc16 100644
--- a/core/java/android/view/inputmethod/InlineSuggestionsRequest.java
+++ b/core/java/android/view/inputmethod/InlineSuggestionsRequest.java
@@ -22,7 +22,9 @@ import android.app.ActivityThread;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.os.LocaleList;
+import android.os.Parcel;
 import android.os.Parcelable;
+import android.view.Display;
 import android.view.inline.InlinePresentationSpec;
 
 import com.android.internal.util.DataClass;
@@ -67,7 +69,11 @@ public final class InlineSuggestionsRequest implements Parcelable {
      */
     private @NonNull LocaleList mSupportedLocales;
 
-    // TODO(b/149609075): the generated code needs to be manually fixed due to the bug.
+    /**
+     * The extras state propagated from the IME to pass extra data.
+     */
+    private @Nullable Bundle mExtras;
+
     /**
      * The host input token of the IME that made the request. This will be set by the system for
      * safety reasons.
@@ -77,9 +83,12 @@ public final class InlineSuggestionsRequest implements Parcelable {
     private @Nullable IBinder mHostInputToken;
 
     /**
-     * The extras state propagated from the IME to pass extra data.
+     * The host display id of the IME that made the request. This will be set by the system for
+     * safety reasons.
+     *
+     * @hide
      */
-    private @Nullable Bundle mExtras;
+    private int mHostDisplayId;
 
     /**
      * @hide
@@ -89,6 +98,24 @@ public final class InlineSuggestionsRequest implements Parcelable {
         mHostInputToken = hostInputToken;
     }
 
+    // TODO(b/149609075): remove once IBinder parcelling is natively supported
+    private void parcelHostInputToken(@NonNull Parcel parcel, int flags) {
+        parcel.writeStrongBinder(mHostInputToken);
+    }
+
+    // TODO(b/149609075): remove once IBinder parcelling is natively supported
+    private @Nullable IBinder unparcelHostInputToken(Parcel parcel) {
+        return parcel.readStrongBinder();
+    }
+
+    /**
+     * @hide
+     * @see {@link #mHostDisplayId}.
+     */
+    public void setHostDisplayId(int hostDisplayId) {
+        mHostDisplayId = hostDisplayId;
+    }
+
     private void onConstructed() {
         Preconditions.checkState(mMaxSuggestionCount >= mPresentationSpecs.size());
     }
@@ -111,10 +138,17 @@ public final class InlineSuggestionsRequest implements Parcelable {
     }
 
     @Nullable
+    private static int defaultHostDisplayId() {
+        return Display.INVALID_DISPLAY;
+    }
+
+    @Nullable
     private static Bundle defaultExtras() {
         return null;
     }
 
+
+
     /** @hide */
     abstract static class BaseBuilder {
         abstract Builder setPresentationSpecs(@NonNull List<InlinePresentationSpec> value);
@@ -122,6 +156,8 @@ public final class InlineSuggestionsRequest implements Parcelable {
         abstract Builder setHostPackageName(@Nullable String value);
 
         abstract Builder setHostInputToken(IBinder hostInputToken);
+
+        abstract Builder setHostDisplayId(int value);
     }
 
 
@@ -145,8 +181,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
             @NonNull List<InlinePresentationSpec> presentationSpecs,
             @NonNull String hostPackageName,
             @NonNull LocaleList supportedLocales,
+            @Nullable Bundle extras,
             @Nullable IBinder hostInputToken,
-            @Nullable Bundle extras) {
+            int hostDisplayId) {
         this.mMaxSuggestionCount = maxSuggestionCount;
         this.mPresentationSpecs = presentationSpecs;
         com.android.internal.util.AnnotationValidations.validate(
@@ -157,8 +194,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
         this.mSupportedLocales = supportedLocales;
         com.android.internal.util.AnnotationValidations.validate(
                 NonNull.class, null, mSupportedLocales);
-        this.mHostInputToken = hostInputToken;
         this.mExtras = extras;
+        this.mHostInputToken = hostInputToken;
+        this.mHostDisplayId = hostDisplayId;
 
         onConstructed();
     }
@@ -202,6 +240,14 @@ public final class InlineSuggestionsRequest implements Parcelable {
     }
 
     /**
+     * The extras state propagated from the IME to pass extra data.
+     */
+    @DataClass.Generated.Member
+    public @Nullable Bundle getExtras() {
+        return mExtras;
+    }
+
+    /**
      * The host input token of the IME that made the request. This will be set by the system for
      * safety reasons.
      *
@@ -213,11 +259,14 @@ public final class InlineSuggestionsRequest implements Parcelable {
     }
 
     /**
-     * The extras state propagated from the IME to pass extra data.
+     * The host display id of the IME that made the request. This will be set by the system for
+     * safety reasons.
+     *
+     * @hide
      */
     @DataClass.Generated.Member
-    public @Nullable Bundle getExtras() {
-        return mExtras;
+    public int getHostDisplayId() {
+        return mHostDisplayId;
     }
 
     @Override
@@ -231,8 +280,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
                 "presentationSpecs = " + mPresentationSpecs + ", " +
                 "hostPackageName = " + mHostPackageName + ", " +
                 "supportedLocales = " + mSupportedLocales + ", " +
+                "extras = " + mExtras + ", " +
                 "hostInputToken = " + mHostInputToken + ", " +
-                "extras = " + mExtras +
+                "hostDisplayId = " + mHostDisplayId +
         " }";
     }
 
@@ -253,8 +303,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
                 && java.util.Objects.equals(mPresentationSpecs, that.mPresentationSpecs)
                 && java.util.Objects.equals(mHostPackageName, that.mHostPackageName)
                 && java.util.Objects.equals(mSupportedLocales, that.mSupportedLocales)
+                && java.util.Objects.equals(mExtras, that.mExtras)
                 && java.util.Objects.equals(mHostInputToken, that.mHostInputToken)
-                && java.util.Objects.equals(mExtras, that.mExtras);
+                && mHostDisplayId == that.mHostDisplayId;
     }
 
     @Override
@@ -268,27 +319,29 @@ public final class InlineSuggestionsRequest implements Parcelable {
         _hash = 31 * _hash + java.util.Objects.hashCode(mPresentationSpecs);
         _hash = 31 * _hash + java.util.Objects.hashCode(mHostPackageName);
         _hash = 31 * _hash + java.util.Objects.hashCode(mSupportedLocales);
-        _hash = 31 * _hash + java.util.Objects.hashCode(mHostInputToken);
         _hash = 31 * _hash + java.util.Objects.hashCode(mExtras);
+        _hash = 31 * _hash + java.util.Objects.hashCode(mHostInputToken);
+        _hash = 31 * _hash + mHostDisplayId;
         return _hash;
     }
 
     @Override
     @DataClass.Generated.Member
-    public void writeToParcel(@NonNull android.os.Parcel dest, int flags) {
+    public void writeToParcel(@NonNull Parcel dest, int flags) {
         // You can override field parcelling by defining methods like:
         // void parcelFieldName(Parcel dest, int flags) { ... }
 
         byte flg = 0;
-        if (mHostInputToken != null) flg |= 0x10;
-        if (mExtras != null) flg |= 0x20;
+        if (mExtras != null) flg |= 0x10;
+        if (mHostInputToken != null) flg |= 0x20;
         dest.writeByte(flg);
         dest.writeInt(mMaxSuggestionCount);
         dest.writeParcelableList(mPresentationSpecs, flags);
         dest.writeString(mHostPackageName);
         dest.writeTypedObject(mSupportedLocales, flags);
-        if (mHostInputToken != null) dest.writeStrongBinder(mHostInputToken);
         if (mExtras != null) dest.writeBundle(mExtras);
+        parcelHostInputToken(dest, flags);
+        dest.writeInt(mHostDisplayId);
     }
 
     @Override
@@ -298,7 +351,7 @@ public final class InlineSuggestionsRequest implements Parcelable {
     /** @hide */
     @SuppressWarnings({"unchecked", "RedundantCast"})
     @DataClass.Generated.Member
-    /* package-private */ InlineSuggestionsRequest(@NonNull android.os.Parcel in) {
+    /* package-private */ InlineSuggestionsRequest(@NonNull Parcel in) {
         // You can override field unparcelling by defining methods like:
         // static FieldType unparcelFieldName(Parcel in) { ... }
 
@@ -308,8 +361,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
         in.readParcelableList(presentationSpecs, InlinePresentationSpec.class.getClassLoader());
         String hostPackageName = in.readString();
         LocaleList supportedLocales = (LocaleList) in.readTypedObject(LocaleList.CREATOR);
-        IBinder hostInputToken = (flg & 0x10) == 0 ? null : in.readStrongBinder();
-        Bundle extras = (flg & 0x20) == 0 ? null : in.readBundle();
+        Bundle extras = (flg & 0x10) == 0 ? null : in.readBundle();
+        IBinder hostInputToken = unparcelHostInputToken(in);
+        int hostDisplayId = in.readInt();
 
         this.mMaxSuggestionCount = maxSuggestionCount;
         this.mPresentationSpecs = presentationSpecs;
@@ -321,8 +375,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
         this.mSupportedLocales = supportedLocales;
         com.android.internal.util.AnnotationValidations.validate(
                 NonNull.class, null, mSupportedLocales);
-        this.mHostInputToken = hostInputToken;
         this.mExtras = extras;
+        this.mHostInputToken = hostInputToken;
+        this.mHostDisplayId = hostDisplayId;
 
         onConstructed();
     }
@@ -336,7 +391,7 @@ public final class InlineSuggestionsRequest implements Parcelable {
         }
 
         @Override
-        public InlineSuggestionsRequest createFromParcel(@NonNull android.os.Parcel in) {
+        public InlineSuggestionsRequest createFromParcel(@NonNull Parcel in) {
             return new InlineSuggestionsRequest(in);
         }
     };
@@ -352,8 +407,9 @@ public final class InlineSuggestionsRequest implements Parcelable {
         private @NonNull List<InlinePresentationSpec> mPresentationSpecs;
         private @NonNull String mHostPackageName;
         private @NonNull LocaleList mSupportedLocales;
-        private @Nullable IBinder mHostInputToken;
         private @Nullable Bundle mExtras;
+        private @Nullable IBinder mHostInputToken;
+        private int mHostDisplayId;
 
         private long mBuilderFieldsSet = 0L;
 
@@ -436,6 +492,17 @@ public final class InlineSuggestionsRequest implements Parcelable {
         }
 
         /**
+         * The extras state propagated from the IME to pass extra data.
+         */
+        @DataClass.Generated.Member
+        public @NonNull Builder setExtras(@Nullable Bundle value) {
+            checkNotUsed();
+            mBuilderFieldsSet |= 0x10;
+            mExtras = value;
+            return this;
+        }
+
+        /**
          * The host input token of the IME that made the request. This will be set by the system for
          * safety reasons.
          *
@@ -445,26 +512,30 @@ public final class InlineSuggestionsRequest implements Parcelable {
         @Override
         @NonNull Builder setHostInputToken(@Nullable IBinder value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x10;
+            mBuilderFieldsSet |= 0x20;
             mHostInputToken = value;
             return this;
         }
 
         /**
-         * The extras state propagated from the IME to pass extra data.
+         * The host display id of the IME that made the request. This will be set by the system for
+         * safety reasons.
+         *
+         * @hide
          */
         @DataClass.Generated.Member
-        public @NonNull Builder setExtras(@Nullable Bundle value) {
+        @Override
+        @NonNull Builder setHostDisplayId(int value) {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x20;
-            mExtras = value;
+            mBuilderFieldsSet |= 0x40;
+            mHostDisplayId = value;
             return this;
         }
 
         /** Builds the instance. This builder should not be touched after calling this! */
         public @NonNull InlineSuggestionsRequest build() {
             checkNotUsed();
-            mBuilderFieldsSet |= 0x40; // Mark builder used
+            mBuilderFieldsSet |= 0x80; // Mark builder used
 
             if ((mBuilderFieldsSet & 0x1) == 0) {
                 mMaxSuggestionCount = defaultMaxSuggestionCount();
@@ -476,23 +547,27 @@ public final class InlineSuggestionsRequest implements Parcelable {
                 mSupportedLocales = defaultSupportedLocales();
             }
             if ((mBuilderFieldsSet & 0x10) == 0) {
-                mHostInputToken = defaultHostInputToken();
+                mExtras = defaultExtras();
             }
             if ((mBuilderFieldsSet & 0x20) == 0) {
-                mExtras = defaultExtras();
+                mHostInputToken = defaultHostInputToken();
+            }
+            if ((mBuilderFieldsSet & 0x40) == 0) {
+                mHostDisplayId = defaultHostDisplayId();
             }
             InlineSuggestionsRequest o = new InlineSuggestionsRequest(
                     mMaxSuggestionCount,
                     mPresentationSpecs,
                     mHostPackageName,
                     mSupportedLocales,
+                    mExtras,
                     mHostInputToken,
-                    mExtras);
+                    mHostDisplayId);
             return o;
         }
 
         private void checkNotUsed() {
-            if ((mBuilderFieldsSet & 0x40) != 0) {
+            if ((mBuilderFieldsSet & 0x80) != 0) {
                 throw new IllegalStateException(
                         "This Builder should not be reused. Use a new Builder instance instead");
             }
@@ -500,10 +575,10 @@ public final class InlineSuggestionsRequest implements Parcelable {
     }
 
     @DataClass.Generated(
-            time = 1581747892762L,
+            time = 1582339908980L,
             codegenVersion = "1.0.14",
             sourceFile = "frameworks/base/core/java/android/view/inputmethod/InlineSuggestionsRequest.java",
-            inputSignatures = "public static final  int SUGGESTION_COUNT_UNLIMITED\nprivate final  int mMaxSuggestionCount\nprivate final @android.annotation.NonNull java.util.List<android.view.inline.InlinePresentationSpec> mPresentationSpecs\nprivate @android.annotation.NonNull java.lang.String mHostPackageName\nprivate @android.annotation.NonNull android.os.LocaleList mSupportedLocales\nprivate @android.annotation.Nullable android.os.IBinder mHostInputToken\nprivate @android.annotation.Nullable android.os.Bundle mExtras\npublic  void setHostInputToken(android.os.IBinder)\nprivate  void onConstructed()\nprivate static  int defaultMaxSuggestionCount()\nprivate static  java.lang.String defaultHostPackageName()\nprivate static  android.os.LocaleList defaultSupportedLocales()\nprivate static @android.annotation.Nullable android.os.IBinder defaultHostInputToken()\nprivate static @android.annotation.Nullable android.os.Bundle defaultExtras()\nclass InlineSuggestionsRequest extends java.lang.Object implements [android.os.Parcelable]\n@com.android.internal.util.DataClass(genEqualsHashCode=true, genToString=true, genBuilder=true)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setPresentationSpecs(java.util.List<android.view.inline.InlinePresentationSpec>)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setHostPackageName(java.lang.String)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setHostInputToken(android.os.IBinder)\nclass BaseBuilder extends java.lang.Object implements []")
+            inputSignatures = "public static final  int SUGGESTION_COUNT_UNLIMITED\nprivate final  int mMaxSuggestionCount\nprivate final @android.annotation.NonNull java.util.List<android.view.inline.InlinePresentationSpec> mPresentationSpecs\nprivate @android.annotation.NonNull java.lang.String mHostPackageName\nprivate @android.annotation.NonNull android.os.LocaleList mSupportedLocales\nprivate @android.annotation.Nullable android.os.Bundle mExtras\nprivate @android.annotation.Nullable android.os.IBinder mHostInputToken\nprivate  int mHostDisplayId\npublic  void setHostInputToken(android.os.IBinder)\nprivate  void parcelHostInputToken(android.os.Parcel,int)\nprivate @android.annotation.Nullable android.os.IBinder unparcelHostInputToken(android.os.Parcel)\npublic  void setHostDisplayId(int)\nprivate  void onConstructed()\nprivate static  int defaultMaxSuggestionCount()\nprivate static  java.lang.String defaultHostPackageName()\nprivate static  android.os.LocaleList defaultSupportedLocales()\nprivate static @android.annotation.Nullable android.os.IBinder defaultHostInputToken()\nprivate static @android.annotation.Nullable int defaultHostDisplayId()\nprivate static @android.annotation.Nullable android.os.Bundle defaultExtras()\nclass InlineSuggestionsRequest extends java.lang.Object implements [android.os.Parcelable]\n@com.android.internal.util.DataClass(genEqualsHashCode=true, genToString=true, genBuilder=true)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setPresentationSpecs(java.util.List<android.view.inline.InlinePresentationSpec>)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setHostPackageName(java.lang.String)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setHostInputToken(android.os.IBinder)\nabstract  android.view.inputmethod.InlineSuggestionsRequest.Builder setHostDisplayId(int)\nclass BaseBuilder extends java.lang.Object implements []")
     @Deprecated
     private void __metadata() {}
 
diff --git a/services/autofill/java/com/android/server/autofill/RemoteInlineSuggestionRenderService.java b/services/autofill/java/com/android/server/autofill/RemoteInlineSuggestionRenderService.java
index 5d5af535920b..7ad5632dd70e 100644
--- a/services/autofill/java/com/android/server/autofill/RemoteInlineSuggestionRenderService.java
+++ b/services/autofill/java/com/android/server/autofill/RemoteInlineSuggestionRenderService.java
@@ -86,9 +86,9 @@ public final class RemoteInlineSuggestionRenderService extends
      */
     public void renderSuggestion(@NonNull IInlineSuggestionUiCallback callback,
             @NonNull InlinePresentation presentation, int width, int height,
-            @Nullable IBinder hostInputToken) {
-        scheduleAsyncRequest(
-                (s) -> s.renderSuggestion(callback, presentation, width, height, hostInputToken));
+            @Nullable IBinder hostInputToken, int displayId) {
+        scheduleAsyncRequest((s) -> s.renderSuggestion(callback, presentation, width, height,
+                hostInputToken, displayId));
     }
 
     @Nullable
diff --git a/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java b/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java
index 0d8c89b4124d..56e8707840ec 100644
--- a/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java
+++ b/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java
@@ -134,7 +134,7 @@ public final class InlineSuggestionFactory {
         if (inlineAuthentication != null) {
             InlineSuggestion inlineAuthSuggestion = createInlineAuthSuggestion(inlineAuthentication,
                     remoteRenderService, onClickFactory, onErrorCallback,
-                    request.getHostInputToken());
+                    request.getHostInputToken(), request.getHostDisplayId());
             inlineSuggestions.add(inlineAuthSuggestion);
 
             return new InlineSuggestionsResponse(inlineSuggestions);
@@ -164,7 +164,7 @@ public final class InlineSuggestionFactory {
             InlineSuggestion inlineSuggestion = createInlineSuggestion(isAugmented, dataset,
                     fieldIndex, mergedInlinePresentation(request, datasetIndex, inlinePresentation),
                     onClickFactory, remoteRenderService, onErrorCallback,
-                    request.getHostInputToken());
+                    request.getHostInputToken(), request.getHostDisplayId());
 
             inlineSuggestions.add(inlineSuggestion);
         }
@@ -172,7 +172,8 @@ public final class InlineSuggestionFactory {
             for (InlinePresentation inlinePresentation : inlineActions) {
                 final InlineSuggestion inlineAction = createInlineAction(isAugmented, context,
                         mergedInlinePresentation(request, 0, inlinePresentation),
-                        remoteRenderService, onErrorCallback, request.getHostInputToken());
+                        remoteRenderService, onErrorCallback, request.getHostInputToken(),
+                        request.getHostDisplayId());
                 inlineSuggestions.add(inlineAction);
             }
         }
@@ -215,7 +216,8 @@ public final class InlineSuggestionFactory {
             @NonNull Context context,
             @NonNull InlinePresentation inlinePresentation,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService,
-            @NonNull Runnable onErrorCallback, @Nullable IBinder hostInputToken) {
+            @NonNull Runnable onErrorCallback, @Nullable IBinder hostInputToken,
+            int displayId) {
         // TODO(b/146453195): fill in the autofill hint properly.
         final InlineSuggestionInfo inlineSuggestionInfo = new InlineSuggestionInfo(
                 inlinePresentation.getInlinePresentationSpec(),
@@ -227,7 +229,7 @@ public final class InlineSuggestionFactory {
         };
         return new InlineSuggestion(inlineSuggestionInfo,
                 createInlineContentProvider(inlinePresentation, onClickAction, onErrorCallback,
-                        remoteRenderService, hostInputToken));
+                        remoteRenderService, hostInputToken, displayId));
     }
 
     private static InlineSuggestion createInlineSuggestion(boolean isAugmented,
@@ -235,7 +237,8 @@ public final class InlineSuggestionFactory {
             @NonNull InlinePresentation inlinePresentation,
             @NonNull BiConsumer<Dataset, Integer> onClickFactory,
             @NonNull RemoteInlineSuggestionRenderService remoteRenderService,
-            @NonNull Runnable onErrorCallback, @Nullable IBinder hostInputToken) {
+            @NonNull Runnable onErrorCallback, @Nullable IBinder hostInputToken,
+            int displayId) {
         // TODO(b/146453195): fill in the autofill hint properly.
         final InlineSuggestionInfo inlineSuggestionInfo = new InlineSuggestionInfo(
                 inlinePresentation.getInlinePresentationSpec(),
@@ -246,7 +249,7 @@ public final class InlineSuggestionFactory {
         final InlineSuggestion inlineSuggestion = new InlineSuggestion(inlineSuggestionInfo,
                 createInlineContentProvider(inlinePresentation,
                         () -> onClickFactory.accept(dataset, datasetIndex), onErrorCallback,
-                        remoteRenderService, hostInputToken));
+                        remoteRenderService, hostInputToken, displayId));
 
         return inlineSuggestion;
     }
@@ -255,7 +258,7 @@ public final class InlineSuggestionFactory {
             @NonNull InlinePresentation inlinePresentation,
             @NonNull RemoteInlineSuggestionRenderService remoteRenderService,
             @NonNull BiConsumer<Dataset, Integer> onClickFactory, @NonNull Runnable onErrorCallback,
-            @Nullable IBinder hostInputToken) {
+            @Nullable IBinder hostInputToken, int displayId) {
         final InlineSuggestionInfo inlineSuggestionInfo = new InlineSuggestionInfo(
                 inlinePresentation.getInlinePresentationSpec(),
                 InlineSuggestionInfo.SOURCE_AUTOFILL, null, InlineSuggestionInfo.TYPE_SUGGESTION);
@@ -264,7 +267,7 @@ public final class InlineSuggestionFactory {
                 createInlineContentProvider(inlinePresentation,
                         () -> onClickFactory.accept(null,
                                 AutofillManager.AUTHENTICATION_ID_DATASET_ID_UNDEFINED),
-                        onErrorCallback, remoteRenderService, hostInputToken));
+                        onErrorCallback, remoteRenderService, hostInputToken, displayId));
     }
 
     /**
@@ -291,7 +294,8 @@ public final class InlineSuggestionFactory {
             @NonNull InlinePresentation inlinePresentation, @Nullable Runnable onClickAction,
             @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService,
-            @Nullable IBinder hostInputToken) {
+            @Nullable IBinder hostInputToken,
+            int displayId) {
         return new IInlineContentProvider.Stub() {
             @Override
             public void provideContent(int width, int height, IInlineContentCallback callback) {
@@ -305,7 +309,7 @@ public final class InlineSuggestionFactory {
                     }
 
                     remoteRenderService.renderSuggestion(uiCallback, inlinePresentation,
-                            width, height, hostInputToken);
+                            width, height, hostInputToken, displayId);
                 });
             }
         };
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index 87262a844b18..6f04af608e14 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -1882,7 +1882,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                 executeOrSendMessage(mCurMethod,
                         mCaller.obtainMessageOOO(MSG_INLINE_SUGGESTIONS_REQUEST, mCurMethod,
                                 requestInfo, new InlineSuggestionsRequestCallbackDecorator(callback,
-                                        imi.getPackageName())));
+                                        imi.getPackageName(), mCurTokenDisplayId)));
             } else {
                 callback.onInlineSuggestionsUnsupported();
             }
@@ -1902,11 +1902,14 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         @NonNull
         private final String mImePackageName;
 
+        private final int mImeDisplayId;
+
         InlineSuggestionsRequestCallbackDecorator(
                 @NonNull IInlineSuggestionsRequestCallback callback,
-                @NonNull String imePackageName) {
+                @NonNull String imePackageName, int displayId) {
             mCallback = callback;
             mImePackageName = imePackageName;
+            mImeDisplayId = displayId;
         }
 
         @Override
@@ -1923,6 +1926,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                                 + "] doesn't match the IME package name=[" + mImePackageName
                                 + "].");
             }
+            request.setHostDisplayId(mImeDisplayId);
             mCallback.onInlineSuggestionsRequest(request, callback);
         }
     }