Reusing the remote inline suggestion view during autofill filtering

* This gets rid of the flicker in the sample keyboard for most of
  tiem time, during filtering.
* We introduce a new object InlineFillUi to store the autofill
  inline candidates returned from the autofill provider. Filtering
  is applied internally to ensure a single place managing
  generating the InlineSuggestionsResponse that is sent to the IME.
  This enables us to reuse the backing remote view for filtering.
* The InlineFillUi is stored in the SessionController.
* The other classes within the UI package are marked as package
  private.
* Also implement the reference count to track the right time for
  releasing the remote view.

Test: atest android.autofillservice.cts.inline (sanity test)

Bug: 151467650
Change-Id: I023bd8cf8acf838088fba2d1d97c0e78d563222e

9 files changed, 492 insertions, 245 deletions

diff --git a/services/autofill/java/com/android/server/autofill/AutofillInlineSessionController.java b/services/autofill/java/com/android/server/autofill/AutofillInlineSessionController.java
index 3612e093c8bd..3282870fe281 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillInlineSessionController.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillInlineSessionController.java
@@ -17,17 +17,17 @@
 package com.android.server.autofill;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.content.ComponentName;
 import android.os.Bundle;
 import android.os.Handler;
 import android.view.autofill.AutofillId;
 import android.view.inputmethod.InlineSuggestionsRequest;
-import android.view.inputmethod.InlineSuggestionsResponse;
 
 import com.android.internal.annotations.GuardedBy;
+import com.android.server.autofill.ui.InlineFillUi;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 
-import java.util.Collections;
 import java.util.Optional;
 import java.util.function.Consumer;
 
@@ -46,8 +46,12 @@ final class AutofillInlineSessionController {
     @NonNull
     private final Handler mHandler;
 
+    @Nullable
     @GuardedBy("mLock")
     private AutofillInlineSuggestionsRequestSession mSession;
+    @Nullable
+    @GuardedBy("mLock")
+    private InlineFillUi mInlineFillUi;
 
     AutofillInlineSessionController(InputMethodManagerInternal inputMethodManagerInternal,
             int userId, ComponentName componentName, Handler handler, Object lock) {
@@ -72,16 +76,16 @@ final class AutofillInlineSessionController {
         // TODO(b/151123764): rename the method to better reflect what it does.
         if (mSession != null) {
             // Send an empty response to IME and destroy the existing session.
-            mSession.onInlineSuggestionsResponseLocked(mSession.getAutofillIdLocked(),
-                    new InlineSuggestionsResponse(Collections.EMPTY_LIST));
+            mSession.onInlineSuggestionsResponseLocked(
+                    InlineFillUi.emptyUi(mSession.getAutofillIdLocked()));
             mSession.destroySessionLocked();
+            mInlineFillUi = null;
         }
         // TODO(b/151123764): consider reusing the same AutofillInlineSession object for the
         // same field.
         mSession = new AutofillInlineSuggestionsRequestSession(mInputMethodManagerInternal, mUserId,
                 mComponentName, mHandler, mLock, autofillId, requestConsumer, uiExtras);
         mSession.onCreateInlineSuggestionsRequestLocked();
-
     }
 
     /**
@@ -101,30 +105,63 @@ final class AutofillInlineSessionController {
 
     /**
      * Requests the IME to hide the current suggestions, if any. Returns true if the message is sent
-     * to the IME.
+     * to the IME. This only hides the UI temporarily. For example if user starts typing/deleting
+     * characters, new filterText will kick in and may revive the suggestion UI.
      */
     @GuardedBy("mLock")
     boolean hideInlineSuggestionsUiLocked(@NonNull AutofillId autofillId) {
         if (mSession != null) {
-            return mSession.onInlineSuggestionsResponseLocked(autofillId,
-                    new InlineSuggestionsResponse(Collections.EMPTY_LIST));
+            return mSession.onInlineSuggestionsResponseLocked(InlineFillUi.emptyUi(autofillId));
+        }
+        return false;
+    }
+
+    /**
+     * Permanently delete the current inline fill UI. Notify the IME to hide the suggestions as
+     * well.
+     */
+    @GuardedBy("mLock")
+    boolean deleteInlineFillUiLocked(@NonNull AutofillId autofillId) {
+        mInlineFillUi = null;
+        return hideInlineSuggestionsUiLocked(autofillId);
+    }
+
+    /**
+     * Updates the inline fill UI with the filter text. It'll send updated inline suggestions to
+     * the IME.
+     */
+    @GuardedBy("mLock")
+    boolean filterInlineFillUiLocked(@NonNull AutofillId autofillId, @Nullable String filterText) {
+        if (mInlineFillUi != null && mInlineFillUi.getAutofillId().equals(autofillId)) {
+            mInlineFillUi.setFilterText(filterText);
+            return requestImeToShowInlineSuggestionsLocked();
         }
         return false;
     }
 
     /**
-     * Requests showing the inline suggestion in the IME when the IME becomes visible and is focused
-     * on the {@code autofillId}.
+     * Set the current inline fill UI. It'll request the IME to show the inline suggestions when
+     * the IME becomes visible and is focused on the {@code autofillId}.
      *
-     * @return false if there is no session, or if the IME callback is not available in the session.
+     * @return false if the suggestions are not sent to IME because there is no session, or if the
+     * IME callback is not available in the session.
      */
     @GuardedBy("mLock")
-    boolean onInlineSuggestionsResponseLocked(@NonNull AutofillId autofillId,
-            @NonNull InlineSuggestionsResponse inlineSuggestionsResponse) {
-        // TODO(b/151123764): rename the method to better reflect what it does.
-        if (mSession != null) {
-            return mSession.onInlineSuggestionsResponseLocked(autofillId,
-                    inlineSuggestionsResponse);
+    boolean setInlineFillUiLocked(@NonNull InlineFillUi inlineFillUi) {
+        mInlineFillUi = inlineFillUi;
+        return requestImeToShowInlineSuggestionsLocked();
+    }
+
+    /**
+     * Sends the suggestions from the current inline fill UI to the IME.
+     *
+     * @return false if the suggestions are not sent to IME because there is no session, or if the
+     * IME callback is not available in the session.
+     */
+    @GuardedBy("mLock")
+    private boolean requestImeToShowInlineSuggestionsLocked() {
+        if (mSession != null && mInlineFillUi != null) {
+            return mSession.onInlineSuggestionsResponseLocked(mInlineFillUi);
         }
         return false;
     }
diff --git a/services/autofill/java/com/android/server/autofill/AutofillInlineSuggestionsRequestSession.java b/services/autofill/java/com/android/server/autofill/AutofillInlineSuggestionsRequestSession.java
index 22451e1d992e..0bf89936f2ce 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillInlineSuggestionsRequestSession.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillInlineSuggestionsRequestSession.java
@@ -27,7 +27,6 @@ import android.content.ComponentName;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.RemoteException;
-import android.util.Log;
 import android.util.Slog;
 import android.view.autofill.AutofillId;
 import android.view.inputmethod.InlineSuggestionsRequest;
@@ -37,7 +36,7 @@ import com.android.internal.annotations.GuardedBy;
 import com.android.internal.view.IInlineSuggestionsRequestCallback;
 import com.android.internal.view.IInlineSuggestionsResponseCallback;
 import com.android.internal.view.InlineSuggestionsRequestInfo;
-import com.android.server.autofill.ui.InlineSuggestionFactory;
+import com.android.server.autofill.ui.InlineFillUi;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 
 import java.lang.ref.WeakReference;
@@ -106,7 +105,7 @@ final class AutofillInlineSuggestionsRequestSession {
     private boolean mImeInputViewStarted;
     @GuardedBy("mLock")
     @Nullable
-    private InlineSuggestionsResponse mInlineSuggestionsResponse;
+    private InlineFillUi mInlineFillUi;
     @GuardedBy("mLock")
     private boolean mPreviousResponseIsNotEmpty;
 
@@ -156,18 +155,20 @@ final class AutofillInlineSuggestionsRequestSession {
      * @return false if the IME callback is not available.
      */
     @GuardedBy("mLock")
-    boolean onInlineSuggestionsResponseLocked(@NonNull AutofillId autofillId,
-            @NonNull InlineSuggestionsResponse inlineSuggestionsResponse) {
+    boolean onInlineSuggestionsResponseLocked(@NonNull InlineFillUi inlineFillUi) {
         if (mDestroyed) {
             return false;
         }
-        if (sDebug) Log.d(TAG, "onInlineSuggestionsResponseLocked called for:" + autofillId);
+        if (sDebug) {
+            Slog.d(TAG,
+                    "onInlineSuggestionsResponseLocked called for:" + inlineFillUi.getAutofillId());
+        }
         if (mImeRequest == null || mResponseCallback == null) {
             return false;
         }
         // TODO(b/151123764): each session should only correspond to one field.
-        mAutofillId = autofillId;
-        mInlineSuggestionsResponse = inlineSuggestionsResponse;
+        mAutofillId = inlineFillUi.getAutofillId();
+        mInlineFillUi = inlineFillUi;
         maybeUpdateResponseToImeLocked();
         return true;
     }
@@ -191,12 +192,12 @@ final class AutofillInlineSuggestionsRequestSession {
         if (mDestroyed) {
             return;
         }
-        if (sDebug) Log.d(TAG, "onCreateInlineSuggestionsRequestLocked called: " + mAutofillId);
+        if (sDebug) Slog.d(TAG, "onCreateInlineSuggestionsRequestLocked called: " + mAutofillId);
         mInputMethodManagerInternal.onCreateInlineSuggestionsRequest(mUserId,
                 new InlineSuggestionsRequestInfo(mComponentName, mAutofillId, mUiExtras),
                 new InlineSuggestionsRequestCallbackImpl(this));
         mTimeoutCallback = () -> {
-            Log.w(TAG, "Timed out waiting for IME callback InlineSuggestionsRequest.");
+            Slog.w(TAG, "Timed out waiting for IME callback InlineSuggestionsRequest.");
             handleOnReceiveImeRequest(null, null);
         };
         mHandler.postDelayed(mTimeoutCallback, CREATE_INLINE_SUGGESTIONS_REQUEST_TIMEOUT_MS);
@@ -207,7 +208,7 @@ final class AutofillInlineSuggestionsRequestSession {
      */
     @GuardedBy("mLock")
     private void maybeUpdateResponseToImeLocked() {
-        if (sVerbose) Log.v(TAG, "maybeUpdateResponseToImeLocked called");
+        if (sVerbose) Slog.v(TAG, "maybeUpdateResponseToImeLocked called");
         if (mDestroyed || mResponseCallback == null) {
             return;
         }
@@ -217,18 +218,19 @@ final class AutofillInlineSuggestionsRequestSession {
             // Although the inline suggestions should disappear when IME hides which removes them
             // from the view hierarchy, but we still send an empty response to be extra safe.
 
-            if (sVerbose) Log.v(TAG, "Send empty inline response");
+            if (sVerbose) Slog.v(TAG, "Send empty inline response");
             updateResponseToImeUncheckLocked(new InlineSuggestionsResponse(Collections.EMPTY_LIST));
             mPreviousResponseIsNotEmpty = false;
-        } else if (mImeInputViewStarted && mInlineSuggestionsResponse != null && match(mAutofillId,
+        } else if (mImeInputViewStarted && mInlineFillUi != null && match(mAutofillId,
                 mImeCurrentFieldId)) {
             // 2. if IME is visible, and response is not null, send the response
-            boolean isEmptyResponse = mInlineSuggestionsResponse.getInlineSuggestions().isEmpty();
+            InlineSuggestionsResponse response = mInlineFillUi.getInlineSuggestionsResponse();
+            boolean isEmptyResponse = response.getInlineSuggestions().isEmpty();
             if (isEmptyResponse && !mPreviousResponseIsNotEmpty) {
                 // No-op if both the previous response and current response are empty.
                 return;
             }
-            updateResponseToImeUncheckLocked(mInlineSuggestionsResponse);
+            updateResponseToImeUncheckLocked(response);
             mPreviousResponseIsNotEmpty = !isEmptyResponse;
         }
     }
@@ -241,10 +243,9 @@ final class AutofillInlineSuggestionsRequestSession {
         if (mDestroyed) {
             return;
         }
-        if (sDebug) Log.d(TAG, "Send inline response: " + response.getInlineSuggestions().size());
+        if (sDebug) Slog.d(TAG, "Send inline response: " + response.getInlineSuggestions().size());
         try {
-            mResponseCallback.onInlineSuggestionsResponse(mAutofillId,
-                    InlineSuggestionFactory.copy(response));
+            mResponseCallback.onInlineSuggestionsResponse(mAutofillId, response);
         } catch (RemoteException e) {
             Slog.e(TAG, "RemoteException sending InlineSuggestionsResponse to IME");
         }
@@ -264,7 +265,7 @@ final class AutofillInlineSuggestionsRequestSession {
             mImeRequestReceived = true;
 
             if (mTimeoutCallback != null) {
-                if (sVerbose) Log.v(TAG, "removing timeout callback");
+                if (sVerbose) Slog.v(TAG, "removing timeout callback");
                 mHandler.removeCallbacks(mTimeoutCallback);
                 mTimeoutCallback = null;
             }
@@ -335,7 +336,7 @@ final class AutofillInlineSuggestionsRequestSession {
         @BinderThread
         @Override
         public void onInlineSuggestionsUnsupported() throws RemoteException {
-            if (sDebug) Log.d(TAG, "onInlineSuggestionsUnsupported() called.");
+            if (sDebug) Slog.d(TAG, "onInlineSuggestionsUnsupported() called.");
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
@@ -348,7 +349,7 @@ final class AutofillInlineSuggestionsRequestSession {
         @Override
         public void onInlineSuggestionsRequest(InlineSuggestionsRequest request,
                 IInlineSuggestionsResponseCallback callback) {
-            if (sDebug) Log.d(TAG, "onInlineSuggestionsRequest() received: " + request);
+            if (sDebug) Slog.d(TAG, "onInlineSuggestionsRequest() received: " + request);
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
@@ -359,7 +360,7 @@ final class AutofillInlineSuggestionsRequestSession {
 
         @Override
         public void onInputMethodStartInput(AutofillId imeFieldId) throws RemoteException {
-            if (sVerbose) Log.v(TAG, "onInputMethodStartInput() received on " + imeFieldId);
+            if (sVerbose) Slog.v(TAG, "onInputMethodStartInput() received on " + imeFieldId);
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
@@ -371,14 +372,14 @@ final class AutofillInlineSuggestionsRequestSession {
         @Override
         public void onInputMethodShowInputRequested(boolean requestResult) throws RemoteException {
             if (sVerbose) {
-                Log.v(TAG, "onInputMethodShowInputRequested() received: " + requestResult);
+                Slog.v(TAG, "onInputMethodShowInputRequested() received: " + requestResult);
             }
         }
 
         @BinderThread
         @Override
         public void onInputMethodStartInputView() {
-            if (sVerbose) Log.v(TAG, "onInputMethodStartInputView() received");
+            if (sVerbose) Slog.v(TAG, "onInputMethodStartInputView() received");
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
@@ -390,7 +391,7 @@ final class AutofillInlineSuggestionsRequestSession {
         @BinderThread
         @Override
         public void onInputMethodFinishInputView() {
-            if (sVerbose) Log.v(TAG, "onInputMethodFinishInputView() received");
+            if (sVerbose) Slog.v(TAG, "onInputMethodFinishInputView() received");
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
@@ -401,7 +402,7 @@ final class AutofillInlineSuggestionsRequestSession {
 
         @Override
         public void onInputMethodFinishInput() throws RemoteException {
-            if (sVerbose) Log.v(TAG, "onInputMethodFinishInput() received");
+            if (sVerbose) Slog.v(TAG, "onInputMethodFinishInput() received");
             final AutofillInlineSuggestionsRequestSession session = mSession.get();
             if (session != null) {
                 session.mHandler.sendMessage(obtainMessage(
diff --git a/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java b/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
index 6cec8d82f9d4..851e4cc0bfd1 100644
--- a/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
+++ b/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
@@ -48,17 +48,15 @@ import android.view.autofill.AutofillManager;
 import android.view.autofill.AutofillValue;
 import android.view.autofill.IAutoFillManagerClient;
 import android.view.inputmethod.InlineSuggestionsRequest;
-import android.view.inputmethod.InlineSuggestionsResponse;
 
 import com.android.internal.infra.AbstractRemoteService;
 import com.android.internal.infra.AndroidFuture;
 import com.android.internal.infra.ServiceConnector;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.internal.os.IResultReceiver;
-import com.android.server.autofill.ui.InlineSuggestionFactory;
+import com.android.server.autofill.ui.InlineFillUi;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.TimeUnit;
@@ -149,7 +147,7 @@ final class RemoteAugmentedAutofillService
             int taskId, @NonNull ComponentName activityComponent, @NonNull AutofillId focusedId,
             @Nullable AutofillValue focusedValue,
             @Nullable InlineSuggestionsRequest inlineSuggestionsRequest,
-            @Nullable Function<InlineSuggestionsResponse, Boolean> inlineSuggestionsCallback,
+            @Nullable Function<InlineFillUi, Boolean> inlineSuggestionsCallback,
             @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
         long requestTime = SystemClock.elapsedRealtime();
@@ -173,7 +171,8 @@ final class RemoteAugmentedAutofillService
                                     mCallbacks.resetLastResponse();
                                     maybeRequestShowInlineSuggestions(sessionId,
                                             inlineSuggestionsRequest, inlineSuggestionsData,
-                                            clientState, focusedId, inlineSuggestionsCallback,
+                                            clientState, focusedId, focusedValue,
+                                            inlineSuggestionsCallback,
                                             client, onErrorCallback, remoteRenderService);
                                     requestAutofill.complete(null);
                                 }
@@ -239,8 +238,8 @@ final class RemoteAugmentedAutofillService
     private void maybeRequestShowInlineSuggestions(int sessionId,
             @Nullable InlineSuggestionsRequest request,
             @Nullable List<Dataset> inlineSuggestionsData, @Nullable Bundle clientState,
-            @NonNull AutofillId focusedId,
-            @Nullable Function<InlineSuggestionsResponse, Boolean> inlineSuggestionsCallback,
+            @NonNull AutofillId focusedId, @Nullable AutofillValue focusedValue,
+            @Nullable Function<InlineFillUi, Boolean> inlineSuggestionsCallback,
             @NonNull IAutoFillManagerClient client, @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
         if (inlineSuggestionsData == null || inlineSuggestionsData.isEmpty()
@@ -250,10 +249,14 @@ final class RemoteAugmentedAutofillService
         }
         mCallbacks.setLastResponse(sessionId);
 
-        final InlineSuggestionsResponse inlineSuggestionsResponse =
-                InlineSuggestionFactory.createAugmentedInlineSuggestionsResponse(
-                        request, inlineSuggestionsData, focusedId,
-                        new InlineSuggestionFactory.InlineSuggestionUiCallback() {
+        final String filterText =
+                focusedValue != null && focusedValue.isText()
+                        ? focusedValue.getTextValue().toString() : null;
+
+        final InlineFillUi inlineFillUi =
+                InlineFillUi.forAugmentedAutofill(
+                        request, inlineSuggestionsData, focusedId, filterText,
+                        new InlineFillUi.InlineSuggestionUiCallback() {
                             @Override
                             public void autofill(Dataset dataset) {
                                 mCallbacks.logAugmentedAutofillSelected(sessionId,
@@ -265,8 +268,8 @@ final class RemoteAugmentedAutofillService
                                             && fieldIds.get(0).equals(focusedId);
                                     client.autofill(sessionId, fieldIds, dataset.getFieldValues(),
                                             hideHighlight);
-                                    inlineSuggestionsCallback.apply(new InlineSuggestionsResponse(
-                                            Collections.EMPTY_LIST));
+                                    inlineSuggestionsCallback.apply(
+                                            InlineFillUi.emptyUi(focusedId));
                                 } catch (RemoteException e) {
                                     Slog.w(TAG, "Encounter exception autofilling the values");
                                 }
@@ -283,11 +286,7 @@ final class RemoteAugmentedAutofillService
                             }
                         }, onErrorCallback, remoteRenderService);
 
-        if (inlineSuggestionsResponse == null) {
-            Slog.w(TAG, "InlineSuggestionFactory created null response");
-            return;
-        }
-        if (inlineSuggestionsCallback.apply(inlineSuggestionsResponse)) {
+        if (inlineSuggestionsCallback.apply(inlineFillUi)) {
             mCallbacks.logAugmentedAutofillShown(sessionId, clientState);
         }
     }
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 20d1b98f8647..ff4e7bac6cae 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -94,7 +94,6 @@ import android.view.autofill.AutofillValue;
 import android.view.autofill.IAutoFillManagerClient;
 import android.view.autofill.IAutofillWindowPresenter;
 import android.view.inputmethod.InlineSuggestionsRequest;
-import android.view.inputmethod.InlineSuggestionsResponse;
 
 import com.android.internal.R;
 import com.android.internal.annotations.GuardedBy;
@@ -102,7 +101,7 @@ import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.internal.util.ArrayUtils;
 import com.android.server.autofill.ui.AutoFillUI;
-import com.android.server.autofill.ui.InlineSuggestionFactory;
+import com.android.server.autofill.ui.InlineFillUi;
 import com.android.server.autofill.ui.PendingUi;
 import com.android.server.inputmethod.InputMethodManagerInternal;
 
@@ -2662,10 +2661,20 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
             }
         } else if (viewState.id.equals(this.mCurrentViewId)
                 && (viewState.getState() & ViewState.STATE_INLINE_SHOWN) != 0) {
-            requestShowInlineSuggestionsLocked(viewState.getResponse(), filterText);
+            if ((viewState.getState() & ViewState.STATE_INLINE_DISABLED) != 0) {
+                final FillResponse response = viewState.getResponse();
+                if (response != null) {
+                    response.getDatasets().clear();
+                }
+                mInlineSessionController.deleteInlineFillUiLocked(viewState.id);
+            } else {
+                mInlineSessionController.filterInlineFillUiLocked(mCurrentViewId, filterText);
+            }
         } else if (viewState.id.equals(this.mCurrentViewId)
                 && (viewState.getState() & ViewState.STATE_TRIGGERED_AUGMENTED_AUTOFILL) != 0) {
             if (!TextUtils.isEmpty(filterText)) {
+                // TODO: we should be able to replace this with controller#filterInlineFillUiLocked
+                // to accomplish filtering for augmented autofill.
                 mInlineSessionController.hideInlineSuggestionsUiLocked(mCurrentViewId);
             }
         }
@@ -2816,26 +2825,15 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
             return false;
         }
 
-        final ViewState currentView = mViewStates.get(focusedId);
-        if ((currentView.getState() & ViewState.STATE_INLINE_DISABLED) != 0) {
-            response.getDatasets().clear();
-        }
-        InlineSuggestionsResponse inlineSuggestionsResponse =
-                InlineSuggestionFactory.createInlineSuggestionsResponse(
-                        inlineSuggestionsRequest.get(), response, filterText, focusedId,
-                        this, () -> {
-                            synchronized (mLock) {
-                                mInlineSessionController.hideInlineSuggestionsUiLocked(
-                                        focusedId);
-                            }
-                        }, remoteRenderService);
-        if (inlineSuggestionsResponse == null) {
-            Slog.w(TAG, "InlineSuggestionFactory created null response");
-            return false;
-        }
-
-        return mInlineSessionController.onInlineSuggestionsResponseLocked(focusedId,
-                inlineSuggestionsResponse);
+        InlineFillUi inlineFillUi = InlineFillUi.forAutofill(
+                inlineSuggestionsRequest.get(), response, focusedId, filterText,
+                /*uiCallback*/this, /*onErrorCallback*/ () -> {
+                    synchronized (mLock) {
+                        mInlineSessionController.hideInlineSuggestionsUiLocked(
+                                focusedId);
+                    }
+                }, remoteRenderService);
+        return mInlineSessionController.setInlineFillUiLocked(inlineFillUi);
     }
 
     boolean isDestroyed() {
@@ -3119,11 +3117,10 @@ final class Session implements RemoteFillService.FillServiceCallbacks, ViewState
 
         final AutofillId focusedId = mCurrentViewId;
 
-        final Function<InlineSuggestionsResponse, Boolean> inlineSuggestionsResponseCallback =
+        final Function<InlineFillUi, Boolean> inlineSuggestionsResponseCallback =
                 response -> {
                     synchronized (mLock) {
-                        return mInlineSessionController.onInlineSuggestionsResponseLocked(
-                                focusedId, response);
+                        return mInlineSessionController.setInlineFillUiLocked(response);
                     }
                 };
         final Consumer<InlineSuggestionsRequest> requestAugmentedAutofill =
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 344b92f43089..890208720f97 100644
--- a/services/autofill/java/com/android/server/autofill/ui/FillUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/FillUi.java
@@ -130,9 +130,9 @@ final class FillUi {
     }
 
     FillUi(@NonNull Context context, @NonNull FillResponse response,
-           @NonNull AutofillId focusedViewId, @NonNull @Nullable String filterText,
-           @NonNull OverlayControl overlayControl, @NonNull CharSequence serviceLabel,
-           @NonNull Drawable serviceIcon, boolean nightMode, @NonNull Callback callback) {
+            @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);
         mThemeId = nightMode ? THEME_ID_DARK : THEME_ID_LIGHT;
         mCallback = callback;
diff --git a/services/autofill/java/com/android/server/autofill/ui/InlineContentProviderImpl.java b/services/autofill/java/com/android/server/autofill/ui/InlineContentProviderImpl.java
index 819f2b813a5e..7fbf4b9c590c 100644
--- a/services/autofill/java/com/android/server/autofill/ui/InlineContentProviderImpl.java
+++ b/services/autofill/java/com/android/server/autofill/ui/InlineContentProviderImpl.java
@@ -47,7 +47,7 @@ import com.android.server.FgThread;
  *
  * See also {@link RemoteInlineSuggestionUi} for relevant information.
  */
-public final class InlineContentProviderImpl extends IInlineContentProvider.Stub {
+final class InlineContentProviderImpl extends IInlineContentProvider.Stub {
 
     // TODO(b/153615023): consider not holding strong reference to heavy objects in this stub, to
     //  avoid memory leak in case the client app is holding the remote reference for a longer
diff --git a/services/autofill/java/com/android/server/autofill/ui/InlineFillUi.java b/services/autofill/java/com/android/server/autofill/ui/InlineFillUi.java
new file mode 100644
index 000000000000..652252220c1a
--- /dev/null
+++ b/services/autofill/java/com/android/server/autofill/ui/InlineFillUi.java
@@ -0,0 +1,282 @@
+/*
+ * Copyright (C) 2020 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.sVerbose;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.content.Intent;
+import android.content.IntentSender;
+import android.service.autofill.Dataset;
+import android.service.autofill.FillResponse;
+import android.service.autofill.InlinePresentation;
+import android.text.TextUtils;
+import android.util.Pair;
+import android.util.Slog;
+import android.util.SparseArray;
+import android.view.autofill.AutofillId;
+import android.view.autofill.AutofillValue;
+import android.view.inputmethod.InlineSuggestion;
+import android.view.inputmethod.InlineSuggestionsRequest;
+import android.view.inputmethod.InlineSuggestionsResponse;
+
+import com.android.internal.view.inline.IInlineContentProvider;
+import com.android.server.autofill.RemoteInlineSuggestionRenderService;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.regex.Pattern;
+
+
+/**
+ * UI for a particular field (i.e. {@link AutofillId}) based on an inline autofill response from
+ * the autofill service or the augmented autofill service. It wraps multiple inline suggestions.
+ *
+ * <p> This class is responsible for filtering the suggestions based on the filtered text.
+ * It'll create {@link InlineSuggestion} instances by reusing the backing remote views (from the
+ * renderer service) if possible.
+ */
+public final class InlineFillUi {
+
+    private static final String TAG = "InlineFillUi";
+
+    /**
+     * The id of the field which the current Ui is for.
+     */
+    @NonNull
+    final AutofillId mAutofillId;
+
+    /**
+     * The list of inline suggestions, before applying any filtering
+     */
+    @NonNull
+    private final ArrayList<InlineSuggestion> mInlineSuggestions;
+
+    /**
+     * The corresponding data sets for the inline suggestions. The list may be null if the current
+     * Ui is the authentication UI for the response. If non-null, the size of data sets should equal
+     * that of  inline suggestions.
+     */
+    @Nullable
+    private final ArrayList<Dataset> mDatasets;
+
+    /**
+     * The filter text which will be applied on the inline suggestion list before they are returned
+     * as a response.
+     */
+    @Nullable
+    private String mFilterText;
+
+    /**
+     * Returns an empty inline autofill UI.
+     */
+    @NonNull
+    public static InlineFillUi emptyUi(@NonNull AutofillId autofillId) {
+        return new InlineFillUi(autofillId, new SparseArray<>(), null);
+    }
+
+    /**
+     * Returns an inline autofill UI for a field based on an Autofilll response.
+     */
+    @NonNull
+    public static InlineFillUi forAutofill(@NonNull InlineSuggestionsRequest request,
+            @NonNull FillResponse response,
+            @NonNull AutofillId focusedViewId, @Nullable String filterText,
+            @NonNull AutoFillUI.AutoFillUiCallback uiCallback,
+            @NonNull Runnable onErrorCallback,
+            @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
+
+        if (InlineSuggestionFactory.responseNeedAuthentication(response)) {
+            InlineSuggestion inlineAuthentication =
+                    InlineSuggestionFactory.createInlineAuthentication(request, response,
+                            focusedViewId, uiCallback, onErrorCallback, remoteRenderService);
+            return new InlineFillUi(focusedViewId, inlineAuthentication, filterText);
+        } else if (response.getDatasets() != null) {
+            SparseArray<Pair<Dataset, InlineSuggestion>> inlineSuggestions =
+                    InlineSuggestionFactory.createAutofillInlineSuggestions(request,
+                            response.getRequestId(),
+                            response.getDatasets(), focusedViewId, uiCallback, onErrorCallback,
+                            remoteRenderService);
+            return new InlineFillUi(focusedViewId, inlineSuggestions, filterText);
+        }
+        return new InlineFillUi(focusedViewId, new SparseArray<>(), filterText);
+    }
+
+    /**
+     * Returns an inline autofill UI for a field based on an Autofilll response.
+     */
+    @NonNull
+    public static InlineFillUi forAugmentedAutofill(@NonNull InlineSuggestionsRequest request,
+            @NonNull List<Dataset> datasets,
+            @NonNull AutofillId focusedViewId, @Nullable String filterText,
+            @NonNull InlineSuggestionUiCallback uiCallback,
+            @NonNull Runnable onErrorCallback,
+            @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
+        SparseArray<Pair<Dataset, InlineSuggestion>> inlineSuggestions =
+                InlineSuggestionFactory.createAugmentedAutofillInlineSuggestions(request, datasets,
+                        focusedViewId,
+                        uiCallback, onErrorCallback, remoteRenderService);
+        return new InlineFillUi(focusedViewId, inlineSuggestions, filterText);
+    }
+
+    InlineFillUi(@NonNull AutofillId autofillId,
+            @NonNull SparseArray<Pair<Dataset, InlineSuggestion>> inlineSuggestions,
+            @Nullable String filterText) {
+        mAutofillId = autofillId;
+        int size = inlineSuggestions.size();
+        mDatasets = new ArrayList<>(size);
+        mInlineSuggestions = new ArrayList<>(size);
+        for (int i = 0; i < size; i++) {
+            Pair<Dataset, InlineSuggestion> value = inlineSuggestions.valueAt(i);
+            mDatasets.add(value.first);
+            mInlineSuggestions.add(value.second);
+        }
+        mFilterText = filterText;
+    }
+
+    InlineFillUi(@NonNull AutofillId autofillId, InlineSuggestion inlineSuggestion,
+            @Nullable String filterText) {
+        mAutofillId = autofillId;
+        mDatasets = null;
+        mInlineSuggestions = new ArrayList<>();
+        mInlineSuggestions.add(inlineSuggestion);
+        mFilterText = filterText;
+    }
+
+    @NonNull
+    public AutofillId getAutofillId() {
+        return mAutofillId;
+    }
+
+    public void setFilterText(@Nullable String filterText) {
+        mFilterText = filterText;
+    }
+
+    /**
+     * Returns the list of filtered inline suggestions suitable for being sent to the IME.
+     */
+    @NonNull
+    public InlineSuggestionsResponse getInlineSuggestionsResponse() {
+        final int size = mInlineSuggestions.size();
+        if (size == 0) {
+            return new InlineSuggestionsResponse(Collections.emptyList());
+        }
+        final List<InlineSuggestion> inlineSuggestions = new ArrayList<>();
+        if (mDatasets == null || mDatasets.size() != size) {
+            // authentication case
+            for (int i = 0; i < size; i++) {
+                inlineSuggestions.add(copy(i, mInlineSuggestions.get(i)));
+            }
+            return new InlineSuggestionsResponse(inlineSuggestions);
+        }
+        for (int i = 0; i < size; i++) {
+            final Dataset dataset = mDatasets.get(i);
+            final int fieldIndex = dataset.getFieldIds().indexOf(mAutofillId);
+            if (fieldIndex < 0) {
+                Slog.w(TAG, "AutofillId=" + mAutofillId + " not found in dataset");
+                continue;
+            }
+            final InlinePresentation inlinePresentation = dataset.getFieldInlinePresentation(
+                    fieldIndex);
+            if (inlinePresentation == null) {
+                Slog.w(TAG, "InlinePresentation not found in dataset");
+                continue;
+            }
+            if (!inlinePresentation.isPinned()  // don't filter pinned suggestions
+                    && !includeDataset(dataset, fieldIndex, mFilterText)) {
+                continue;
+            }
+            inlineSuggestions.add(copy(i, mInlineSuggestions.get(i)));
+        }
+        return new InlineSuggestionsResponse(inlineSuggestions);
+    }
+
+    /**
+     * Returns a copy of the suggestion, that internally copies the {@link IInlineContentProvider}
+     * so that it's not reused by the remote IME process across different inline suggestions.
+     * See {@link InlineContentProviderImpl} for why this is needed.
+     *
+     * <p>Note that although it copies the {@link IInlineContentProvider}, the underlying remote
+     * view (in the renderer service) is still reused.
+     */
+    @NonNull
+    private InlineSuggestion copy(int index, @NonNull InlineSuggestion inlineSuggestion) {
+        final IInlineContentProvider contentProvider = inlineSuggestion.getContentProvider();
+        if (contentProvider instanceof InlineContentProviderImpl) {
+            // We have to create a new inline suggestion instance to ensure we don't reuse the
+            // same {@link IInlineContentProvider}, but the underlying views are reused when
+            // calling {@link InlineContentProviderImpl#copy()}.
+            InlineSuggestion newInlineSuggestion = new InlineSuggestion(inlineSuggestion
+                    .getInfo(), ((InlineContentProviderImpl) contentProvider).copy());
+            // The remote view is only set when the content provider is called to inflate the view,
+            // which happens after it's sent to the IME (i.e. not now), so we keep the latest
+            // content provider (through newInlineSuggestion) to make sure the next time we copy it,
+            // we get to reuse the view.
+            mInlineSuggestions.set(index, newInlineSuggestion);
+            return newInlineSuggestion;
+        }
+        return inlineSuggestion;
+    }
+
+    // TODO: Extract the shared filtering logic here and in FillUi to a common method.
+    private static boolean includeDataset(Dataset dataset, int fieldIndex,
+            @Nullable String filterText) {
+        // Show everything when the user input is empty.
+        if (TextUtils.isEmpty(filterText)) {
+            return true;
+        }
+
+        final String constraintLowerCase = filterText.toString().toLowerCase();
+
+        // Use the filter provided by the service, if available.
+        final Dataset.DatasetFieldFilter filter = dataset.getFilter(fieldIndex);
+        if (filter != null) {
+            Pattern filterPattern = filter.pattern;
+            if (filterPattern == null) {
+                if (sVerbose) {
+                    Slog.v(TAG, "Explicitly disabling filter for dataset id" + dataset.getId());
+                }
+                return true;
+            }
+            return filterPattern.matcher(constraintLowerCase).matches();
+        }
+
+        final AutofillValue value = dataset.getFieldValues().get(fieldIndex);
+        if (value == null || !value.isText()) {
+            return dataset.getAuthentication() != null;
+        }
+        final String valueText = value.getTextValue().toString().toLowerCase();
+        return valueText.toLowerCase().startsWith(constraintLowerCase);
+    }
+
+    /**
+     * Callback from the inline suggestion Ui.
+     */
+    public interface InlineSuggestionUiCallback {
+        /**
+         * Callback to autofill a dataset to the client app.
+         */
+        void autofill(@NonNull Dataset dataset);
+
+        /**
+         * Callback to start Intent in client app.
+         */
+        void startIntentSender(@NonNull IntentSender intentSender, @NonNull Intent intent);
+    }
+}
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 e74463a8584b..089eeb2aa086 100644
--- a/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java
+++ b/services/autofill/java/com/android/server/autofill/ui/InlineSuggestionFactory.java
@@ -17,7 +17,6 @@
 package com.android.server.autofill.ui;
 
 import static com.android.server.autofill.Helper.sDebug;
-import static com.android.server.autofill.Helper.sVerbose;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -27,11 +26,11 @@ import android.os.IBinder;
 import android.service.autofill.Dataset;
 import android.service.autofill.FillResponse;
 import android.service.autofill.InlinePresentation;
-import android.text.TextUtils;
+import android.util.Pair;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.view.autofill.AutofillId;
 import android.view.autofill.AutofillManager;
-import android.view.autofill.AutofillValue;
 import android.view.inputmethod.InlineSuggestion;
 import android.view.inputmethod.InlineSuggestionInfo;
 import android.view.inputmethod.InlineSuggestionsRequest;
@@ -41,49 +40,34 @@ import android.widget.inline.InlinePresentationSpec;
 import com.android.internal.view.inline.IInlineContentProvider;
 import com.android.server.autofill.RemoteInlineSuggestionRenderService;
 
-import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
-import java.util.regex.Pattern;
 
-public final class InlineSuggestionFactory {
+final class InlineSuggestionFactory {
     private static final String TAG = "InlineSuggestionFactory";
 
-    /**
-     * Callback from the inline suggestion Ui.
-     */
-    public interface InlineSuggestionUiCallback {
-        /**
-         * Callback to autofill a dataset to the client app.
-         */
-        void autofill(@NonNull Dataset dataset);
-
-        /**
-         * Callback to start Intent in client app.
-         */
-        void startIntentSender(@NonNull IntentSender intentSender, @NonNull Intent intent);
+    public static boolean responseNeedAuthentication(@NonNull FillResponse response) {
+        return response.getAuthentication() != null && response.getInlinePresentation() != null;
     }
 
-    /**
-     * Returns a copy of the response, that internally copies the {@link IInlineContentProvider}
-     * so that it's not reused by the remote IME process across different inline suggestions.
-     * See {@link InlineContentProviderImpl} for why this is needed.
-     */
-    @NonNull
-    public static InlineSuggestionsResponse copy(@NonNull InlineSuggestionsResponse response) {
-        final ArrayList<InlineSuggestion> copiedInlineSuggestions = new ArrayList<>();
-        for (InlineSuggestion inlineSuggestion : response.getInlineSuggestions()) {
-            final IInlineContentProvider contentProvider = inlineSuggestion.getContentProvider();
-            if (contentProvider instanceof InlineContentProviderImpl) {
-                copiedInlineSuggestions.add(new
-                        InlineSuggestion(inlineSuggestion.getInfo(),
-                        ((InlineContentProviderImpl) contentProvider).copy()));
-            } else {
-                copiedInlineSuggestions.add(inlineSuggestion);
-            }
-        }
-        return new InlineSuggestionsResponse(copiedInlineSuggestions);
+    public static InlineSuggestion createInlineAuthentication(
+            @NonNull InlineSuggestionsRequest request, @NonNull FillResponse response,
+            @NonNull AutofillId autofillId,
+            @NonNull AutoFillUI.AutoFillUiCallback client, @NonNull Runnable onErrorCallback,
+            @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
+        final BiConsumer<Dataset, Integer> onClickFactory = (dataset, datasetIndex) -> {
+            client.requestHideFillUi(autofillId);
+            client.authenticate(response.getRequestId(),
+                    datasetIndex, response.getAuthentication(), response.getClientState(),
+                    /* authenticateInline= */ true);
+        };
+        final Consumer<IntentSender> intentSenderConsumer = (intentSender) ->
+                client.startIntentSender(intentSender, new Intent());
+        InlinePresentation inlineAuthentication = response.getInlinePresentation();
+        return createInlineAuthSuggestion(inlineAuthentication,
+                remoteRenderService, onClickFactory, onErrorCallback, intentSenderConsumer,
+                request.getHostInputToken(), request.getHostDisplayId());
     }
 
     /**
@@ -91,33 +75,23 @@ public final class InlineSuggestionFactory {
      * autofill service, potentially filtering the datasets.
      */
     @Nullable
-    public static InlineSuggestionsResponse createInlineSuggestionsResponse(
-            @NonNull InlineSuggestionsRequest request, @NonNull FillResponse response,
-            @Nullable String filterText, @NonNull AutofillId autofillId,
+    public static SparseArray<Pair<Dataset, InlineSuggestion>> createAutofillInlineSuggestions(
+            @NonNull InlineSuggestionsRequest request, int requestId,
+            @NonNull List<Dataset> datasets,
+            @NonNull AutofillId autofillId,
             @NonNull AutoFillUI.AutoFillUiCallback client, @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
         if (sDebug) Slog.d(TAG, "createInlineSuggestionsResponse called");
-        final BiConsumer<Dataset, Integer> onClickFactory;
-        if (response.getAuthentication() != null) {
-            onClickFactory = (dataset, datasetIndex) -> {
-                client.requestHideFillUi(autofillId);
-                client.authenticate(response.getRequestId(),
-                        datasetIndex, response.getAuthentication(), response.getClientState(),
-                        /* authenticateInline= */ true);
-            };
-        } else {
-            onClickFactory = (dataset, datasetIndex) -> {
-                client.requestHideFillUi(autofillId);
-                client.fill(response.getRequestId(), datasetIndex, dataset);
-            };
-        }
-
-        final InlinePresentation inlineAuthentication =
-                response.getAuthentication() == null ? null : response.getInlinePresentation();
-        return createInlineSuggestionsResponseInternal(/* isAugmented= */ false, request,
-                response.getDatasets(), filterText, inlineAuthentication, autofillId,
-                onErrorCallback, onClickFactory, (intentSender) ->
-                        client.startIntentSender(intentSender, new Intent()), remoteRenderService);
+        final Consumer<IntentSender> intentSenderConsumer = (intentSender) ->
+                client.startIntentSender(intentSender, new Intent());
+        final BiConsumer<Dataset, Integer> onClickFactory = (dataset, datasetIndex) -> {
+            client.requestHideFillUi(autofillId);
+            client.fill(requestId, datasetIndex, dataset);
+        };
+
+        return createInlineSuggestionsInternal(/* isAugmented= */ false, request,
+                datasets, autofillId,
+                onErrorCallback, onClickFactory, intentSenderConsumer, remoteRenderService);
     }
 
     /**
@@ -125,16 +99,16 @@ public final class InlineSuggestionFactory {
      * autofill service.
      */
     @Nullable
-    public static InlineSuggestionsResponse createAugmentedInlineSuggestionsResponse(
+    public static SparseArray<Pair<Dataset, InlineSuggestion>>
+            createAugmentedAutofillInlineSuggestions(
             @NonNull InlineSuggestionsRequest request, @NonNull List<Dataset> datasets,
             @NonNull AutofillId autofillId,
-            @NonNull InlineSuggestionUiCallback inlineSuggestionUiCallback,
+            @NonNull InlineFillUi.InlineSuggestionUiCallback inlineSuggestionUiCallback,
             @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
         if (sDebug) Slog.d(TAG, "createAugmentedInlineSuggestionsResponse called");
-        return createInlineSuggestionsResponseInternal(/* isAugmented= */ true, request,
-                datasets, /* filterText= */ null, /* inlineAuthentication= */ null,
-                autofillId, onErrorCallback,
+        return createInlineSuggestionsInternal(/* isAugmented= */ true, request,
+                datasets, autofillId, onErrorCallback,
                 (dataset, datasetIndex) ->
                         inlineSuggestionUiCallback.autofill(dataset),
                 (intentSender) ->
@@ -143,29 +117,13 @@ public final class InlineSuggestionFactory {
     }
 
     @Nullable
-    private static InlineSuggestionsResponse createInlineSuggestionsResponseInternal(
+    private static SparseArray<Pair<Dataset, InlineSuggestion>> createInlineSuggestionsInternal(
             boolean isAugmented, @NonNull InlineSuggestionsRequest request,
-            @Nullable List<Dataset> datasets, @Nullable String filterText,
-            @Nullable InlinePresentation inlineAuthentication, @NonNull AutofillId autofillId,
+            @NonNull List<Dataset> datasets, @NonNull AutofillId autofillId,
             @NonNull Runnable onErrorCallback, @NonNull BiConsumer<Dataset, Integer> onClickFactory,
             @NonNull Consumer<IntentSender> intentSenderConsumer,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
-
-        final ArrayList<InlineSuggestion> inlineSuggestions = new ArrayList<>();
-        if (inlineAuthentication != null) {
-            InlineSuggestion inlineAuthSuggestion = createInlineAuthSuggestion(inlineAuthentication,
-                    remoteRenderService, onClickFactory, onErrorCallback, intentSenderConsumer,
-                    request.getHostInputToken(), request.getHostDisplayId());
-            inlineSuggestions.add(inlineAuthSuggestion);
-
-            return new InlineSuggestionsResponse(inlineSuggestions);
-        }
-
-        if (datasets == null) {
-            Slog.w(TAG, "Datasets should not be null here");
-            return null;
-        }
-
+        SparseArray<Pair<Dataset, InlineSuggestion>> response = new SparseArray<>(datasets.size());
         for (int datasetIndex = 0; datasetIndex < datasets.size(); datasetIndex++) {
             final Dataset dataset = datasets.get(datasetIndex);
             final int fieldIndex = dataset.getFieldIds().indexOf(autofillId);
@@ -179,50 +137,14 @@ public final class InlineSuggestionFactory {
                 Slog.w(TAG, "InlinePresentation not found in dataset");
                 continue;
             }
-            if (!inlinePresentation.isPinned()  // don't filter pinned suggestions
-                    && !includeDataset(dataset, fieldIndex, filterText)) {
-                continue;
-            }
             InlineSuggestion inlineSuggestion = createInlineSuggestion(isAugmented, dataset,
                     datasetIndex,
                     mergedInlinePresentation(request, datasetIndex, inlinePresentation),
                     onClickFactory, remoteRenderService, onErrorCallback, intentSenderConsumer,
                     request.getHostInputToken(), request.getHostDisplayId());
-
-            inlineSuggestions.add(inlineSuggestion);
-        }
-        return new InlineSuggestionsResponse(inlineSuggestions);
-    }
-
-    // TODO: Extract the shared filtering logic here and in FillUi to a common method.
-    private static boolean includeDataset(Dataset dataset, int fieldIndex,
-            @Nullable String filterText) {
-        // Show everything when the user input is empty.
-        if (TextUtils.isEmpty(filterText)) {
-            return true;
-        }
-
-        final String constraintLowerCase = filterText.toString().toLowerCase();
-
-        // Use the filter provided by the service, if available.
-        final Dataset.DatasetFieldFilter filter = dataset.getFilter(fieldIndex);
-        if (filter != null) {
-            Pattern filterPattern = filter.pattern;
-            if (filterPattern == null) {
-                if (sVerbose) {
-                    Slog.v(TAG, "Explicitly disabling filter for dataset id" + dataset.getId());
-                }
-                return true;
-            }
-            return filterPattern.matcher(constraintLowerCase).matches();
-        }
-
-        final AutofillValue value = dataset.getFieldValues().get(fieldIndex);
-        if (value == null || !value.isText()) {
-            return dataset.getAuthentication() == null;
+            response.append(datasetIndex, Pair.create(dataset, inlineSuggestion));
         }
-        final String valueText = value.getTextValue().toString().toLowerCase();
-        return valueText.toLowerCase().startsWith(constraintLowerCase);
+        return response;
     }
 
     private static InlineSuggestion createInlineSuggestion(boolean isAugmented,
diff --git a/services/autofill/java/com/android/server/autofill/ui/RemoteInlineSuggestionUi.java b/services/autofill/java/com/android/server/autofill/ui/RemoteInlineSuggestionUi.java
index 00a5283c9b1f..368f71760b0d 100644
--- a/services/autofill/java/com/android/server/autofill/ui/RemoteInlineSuggestionUi.java
+++ b/services/autofill/java/com/android/server/autofill/ui/RemoteInlineSuggestionUi.java
@@ -17,6 +17,7 @@
 package com.android.server.autofill.ui;
 
 import static com.android.server.autofill.Helper.sDebug;
+import static com.android.server.autofill.Helper.sVerbose;
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
@@ -42,13 +43,8 @@ import com.android.internal.view.inline.IInlineContentCallback;
  * {@link InlineContentProviderImpl}s, each of which wraps a callback from the IME. But at any
  * given time, there is only one active IME callback which this class will callback into.
  *
- * <p>This class is thread safe, because all the outside calls are piped into the same single
- * thread handler to be processed.
- *
- * TODO(b/154683107): implement the reference counting in case there are multiple active
- * SurfacePackages at the same time. This will not happen for now since all the InlineSuggestions
- * sharing the same UI will be sent to the same IME window, so the previous view will be detached
- * before the new view are attached to the window.
+ * <p>This class is thread safe, because all the outside calls are piped into a single handler
+ * thread to be processed.
  */
 final class RemoteInlineSuggestionUi {
 
@@ -83,6 +79,7 @@ final class RemoteInlineSuggestionUi {
      */
     @Nullable
     private IInlineSuggestionUi mInlineSuggestionUi;
+    private int mRefCount = 0;
     private boolean mWaitingForUiCreation = false;
     private int mActualWidth;
     private int mActualHeight;
@@ -124,7 +121,7 @@ final class RemoteInlineSuggestionUi {
      * released.
      */
     void surfacePackageReleased() {
-        mHandler.post(this::handleSurfacePackageReleased);
+        mHandler.post(() -> handleUpdateRefCount(-1));
     }
 
     /**
@@ -134,24 +131,6 @@ final class RemoteInlineSuggestionUi {
         return mWidth == width && mHeight == height;
     }
 
-    private void handleSurfacePackageReleased() {
-        cancelPendingReleaseViewRequest();
-
-        // Schedule a delayed release view request
-        mDelayedReleaseViewRunnable = () -> {
-            if (mInlineSuggestionUi != null) {
-                try {
-                    mInlineSuggestionUi.releaseSurfaceControlViewHost();
-                    mInlineSuggestionUi = null;
-                } catch (RemoteException e) {
-                    Slog.w(TAG, "RemoteException calling releaseSurfaceControlViewHost");
-                }
-            }
-            mDelayedReleaseViewRunnable = null;
-        };
-        mHandler.postDelayed(mDelayedReleaseViewRunnable, RELEASE_REMOTE_VIEW_HOST_DELAY_MS);
-    }
-
     private void handleRequestSurfacePackage() {
         cancelPendingReleaseViewRequest();
 
@@ -174,10 +153,17 @@ final class RemoteInlineSuggestionUi {
             try {
                 mInlineSuggestionUi.getSurfacePackage(new ISurfacePackageResultCallback.Stub() {
                     @Override
-                    public void onResult(SurfaceControlViewHost.SurfacePackage result)
-                            throws RemoteException {
-                        if (sDebug) Slog.d(TAG, "Sending new SurfacePackage to IME");
-                        mInlineContentCallback.onContent(result, mActualWidth, mActualHeight);
+                    public void onResult(SurfaceControlViewHost.SurfacePackage result) {
+                        mHandler.post(() -> {
+                            if (sVerbose) Slog.v(TAG, "Sending refreshed SurfacePackage to IME");
+                            try {
+                                mInlineContentCallback.onContent(result, mActualWidth,
+                                        mActualHeight);
+                                handleUpdateRefCount(1);
+                            } catch (RemoteException e) {
+                                Slog.w(TAG, "RemoteException calling onContent");
+                            }
+                        });
                     }
                 });
             } catch (RemoteException e) {
@@ -186,6 +172,26 @@ final class RemoteInlineSuggestionUi {
         }
     }
 
+    private void handleUpdateRefCount(int delta) {
+        cancelPendingReleaseViewRequest();
+        mRefCount += delta;
+        if (mRefCount <= 0) {
+            mDelayedReleaseViewRunnable = () -> {
+                if (mInlineSuggestionUi != null) {
+                    try {
+                        if (sVerbose) Slog.v(TAG, "releasing the host");
+                        mInlineSuggestionUi.releaseSurfaceControlViewHost();
+                        mInlineSuggestionUi = null;
+                    } catch (RemoteException e) {
+                        Slog.w(TAG, "RemoteException calling releaseSurfaceControlViewHost");
+                    }
+                }
+                mDelayedReleaseViewRunnable = null;
+            };
+            mHandler.postDelayed(mDelayedReleaseViewRunnable, RELEASE_REMOTE_VIEW_HOST_DELAY_MS);
+        }
+    }
+
     private void cancelPendingReleaseViewRequest() {
         if (mDelayedReleaseViewRunnable != null) {
             mHandler.removeCallbacks(mDelayedReleaseViewRunnable);
@@ -199,11 +205,14 @@ final class RemoteInlineSuggestionUi {
     private void handleInlineSuggestionUiReady(IInlineSuggestionUi content,
             SurfaceControlViewHost.SurfacePackage surfacePackage, int width, int height) {
         mInlineSuggestionUi = content;
+        mRefCount = 0;
         mWaitingForUiCreation = false;
         mActualWidth = width;
         mActualHeight = height;
         if (mInlineContentCallback != null) {
             try {
+                if (sVerbose) Slog.v(TAG, "Sending new UI content to IME");
+                handleUpdateRefCount(1);
                 mInlineContentCallback.onContent(surfacePackage, mActualWidth, mActualHeight);
             } catch (RemoteException e) {
                 Slog.w(TAG, "RemoteException calling onContent");