Merge changes I7c2c9411,Ibc72c84d,Ib3968644 into oc-dev

* changes:
  No need to deal with windowTokens
  Persistable accessibility ID from ContextWrappers
  Check callbacks when operting on UI
diff --git a/core/java/android/app/Activity.java b/core/java/android/app/Activity.java
index abe5dc3..37c287e 100644
--- a/core/java/android/app/Activity.java
+++ b/core/java/android/app/Activity.java
@@ -25,6 +25,7 @@
 import com.android.internal.app.IVoiceInteractor;
 import com.android.internal.app.ToolbarActionBar;
 import com.android.internal.app.WindowDecorActionBar;
+import com.android.internal.policy.DecorView;
 import com.android.internal.policy.PhoneWindow;
 
 import android.annotation.CallSuper;
@@ -3147,19 +3148,6 @@
     }
 
     /**
-     * Called before {@link #onAttachedToWindow}.
-     *
-     * @hide
-     */
-    @CallSuper
-    public void onBeforeAttachedToWindow() {
-        if (mAutoFillResetNeeded) {
-            getAutofillManager().onAttachedToWindow(
-                    getWindow().getDecorView().getRootView().getWindowToken());
-        }
-    }
-
-    /**
      * Called when the main window associated with the activity has been
      * attached to the window manager.
      * See {@link View#onAttachedToWindow() View.onAttachedToWindow()}
@@ -7471,45 +7459,62 @@
     }
 
     /** @hide */
+    @NonNull public View[] findViewsByAccessibilityIdTraversal(@NonNull int[] viewIds) {
+        final View[] views = new View[viewIds.length];
+        final ArrayList<ViewRootImpl> roots =
+                WindowManagerGlobal.getInstance().getRootViews(getActivityToken());
+
+        for (int rootNum = 0; rootNum < roots.size(); rootNum++) {
+            final View rootView = roots.get(rootNum).getView();
+
+            if (rootView != null) {
+                for (int viewNum = 0; viewNum < viewIds.length; viewNum++) {
+                    if (views[viewNum] == null) {
+                        views[viewNum] = rootView.findViewByAccessibilityIdTraversal(
+                                viewIds[viewNum]);
+                    }
+                }
+            }
+        }
+
+        return views;
+    }
+
+    /** @hide */
     @Override
-    public boolean getViewVisibility(int viewId) {
-        Window window = getWindow();
-        if (window == null) {
-            Log.i(TAG, "no window");
-            return false;
-        }
+    @NonNull public boolean[] getViewVisibility(@NonNull int[] viewIds) {
+        final boolean[] isVisible = new boolean[viewIds.length];
+        final View views[] = findViewsByAccessibilityIdTraversal(viewIds);
 
-        View decorView = window.peekDecorView();
-        if (decorView == null) {
-            Log.i(TAG, "no decorView");
-            return false;
-        }
-
-        View view = decorView.findViewByAccessibilityIdTraversal(viewId);
-        if (view == null) {
-            Log.i(TAG, "cannot find view");
-            return false;
-        }
-
-        // Check if the view is visible by checking all parents
-        while (view != null) {
-            if (view == decorView) {
-                break;
+        for (int i = 0; i < viewIds.length; i++) {
+            View view = views[i];
+            if (view == null) {
+                isVisible[i] = false;
+                continue;
             }
 
-            if (view.getVisibility() != View.VISIBLE) {
-                Log.i(TAG, view + " is not visible");
-                return false;
-            }
+            isVisible[i] = true;
 
-            if (view.getParent() instanceof View) {
-                view = (View) view.getParent();
-            } else {
-                break;
+            // Check if the view is visible by checking all parents
+            while (true) {
+                if (view instanceof DecorView && view.getViewRootImpl() == view.getParent()) {
+                    break;
+                }
+
+                if (view.getVisibility() != View.VISIBLE) {
+                    isVisible[i] = false;
+                    break;
+                }
+
+                if (view.getParent() instanceof View) {
+                    view = (View) view.getParent();
+                } else {
+                    break;
+                }
             }
         }
 
-        return true;
+        return isVisible;
     }
 
     /** @hide */
diff --git a/core/java/android/content/ContextWrapper.java b/core/java/android/content/ContextWrapper.java
index 5264cd7..e127ca3 100644
--- a/core/java/android/content/ContextWrapper.java
+++ b/core/java/android/content/ContextWrapper.java
@@ -952,4 +952,12 @@
     public Handler getMainThreadHandler() {
         return mBase.getMainThreadHandler();
     }
+
+    /**
+     * @hide
+     */
+    @Override
+    public int getNextAccessibilityId() {
+        return mBase.getNextAccessibilityId();
+    }
 }
diff --git a/core/java/android/view/View.java b/core/java/android/view/View.java
index 7399f77..6c74980 100644
--- a/core/java/android/view/View.java
+++ b/core/java/android/view/View.java
@@ -7706,7 +7706,8 @@
     }
 
     private boolean isAutofillable() {
-        return getAutofillType() != AUTOFILL_TYPE_NONE && isImportantForAutofill();
+        return getAutofillType() != AUTOFILL_TYPE_NONE && isImportantForAutofill()
+                && getAccessibilityViewId() > LAST_APP_ACCESSIBILITY_ID;
     }
 
     private void populateVirtualStructure(ViewStructure structure,
diff --git a/core/java/android/view/Window.java b/core/java/android/view/Window.java
index a432d30..6dd8ecf 100644
--- a/core/java/android/view/Window.java
+++ b/core/java/android/view/Window.java
@@ -480,13 +480,6 @@
         public void onWindowFocusChanged(boolean hasFocus);
 
         /**
-         * @hide
-         */
-        default void onBeforeAttachedToWindow() {
-            // empty
-        }
-
-        /**
          * Called when the window has been attached to the window manager.
          * See {@link View#onAttachedToWindow() View.onAttachedToWindow()}
          * for more information.
diff --git a/core/java/android/view/WindowCallbackWrapper.java b/core/java/android/view/WindowCallbackWrapper.java
index 7018529..02c8945 100644
--- a/core/java/android/view/WindowCallbackWrapper.java
+++ b/core/java/android/view/WindowCallbackWrapper.java
@@ -109,11 +109,6 @@
     }
 
     @Override
-    public void onBeforeAttachedToWindow() {
-        mWrapped.onBeforeAttachedToWindow();
-    }
-
-    @Override
     public void onAttachedToWindow() {
         mWrapped.onAttachedToWindow();
     }
diff --git a/core/java/android/view/autofill/AutofillManager.java b/core/java/android/view/autofill/AutofillManager.java
index d0133ed..c7151db 100644
--- a/core/java/android/view/autofill/AutofillManager.java
+++ b/core/java/android/view/autofill/AutofillManager.java
@@ -28,7 +28,6 @@
 import android.graphics.Rect;
 import android.metrics.LogMaker;
 import android.os.Bundle;
-import android.os.IBinder;
 import android.os.Parcelable;
 import android.os.RemoteException;
 import android.service.autofill.AutofillService;
@@ -38,7 +37,6 @@
 import android.util.Log;
 import android.util.SparseArray;
 import android.view.View;
-import android.view.WindowManagerGlobal;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.logging.MetricsLogger;
@@ -47,6 +45,7 @@
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 
@@ -188,11 +187,11 @@
         boolean autofillCallbackRequestHideFillUi();
 
         /**
-         * Checks if the view is currently attached and visible.
+         * Checks if views are currently attached and visible.
          *
-         * @return {@code true} iff the view is attached or visible
+         * @return And array with {@code true} iff the view is attached or visible
          */
-        boolean getViewVisibility(int viewId);
+        @NonNull boolean[] getViewVisibility(@NonNull int[] viewId);
 
         /**
          * Checks is the client is currently visible as understood by autofill.
@@ -200,6 +199,15 @@
          * @return {@code true} if the client is currently visible
          */
         boolean isVisibleForAutofill();
+
+        /**
+         * Find views by traversing the hierarchies of the client.
+         *
+         * @param viewIds The accessibility ids of the views to find
+         *
+         * @return And array containing the views, or {@code null} if not found
+         */
+        @NonNull View[] findViewsByAccessibilityIdTraversal(@NonNull int[] viewIds);
     }
 
     /**
@@ -259,30 +267,6 @@
     }
 
     /**
-     * Set window future popup windows should be attached to.
-     *
-     * @param windowToken The window the popup windows should be attached to
-     *
-     * {@hide}
-     */
-    public void onAttachedToWindow(@NonNull IBinder windowToken) {
-        if (!hasAutofillFeature()) {
-            return;
-        }
-        synchronized (mLock) {
-            if (mSessionId == NO_SESSION) {
-                return;
-            }
-
-            try {
-                mService.setWindow(mSessionId, windowToken);
-            } catch (RemoteException e) {
-                Log.e(TAG, "Could not attach window to session " + mSessionId);
-            }
-        }
-    }
-
-    /**
      * Called once the client becomes visible.
      *
      * @see AutofillClient#isVisibleForAutofill()
@@ -292,7 +276,7 @@
     public void onVisibleForAutofill() {
         synchronized (mLock) {
             if (mEnabled && mSessionId != NO_SESSION && mTrackedViews != null) {
-                mTrackedViews.onVisibleForAutofill();
+                mTrackedViews.onVisibleForAutofillLocked();
             }
         }
     }
@@ -406,7 +390,7 @@
 
                 if (mSessionId == NO_SESSION) {
                     // Starts new session.
-                    startSessionLocked(id, view.getWindowToken(), null, value, flags);
+                    startSessionLocked(id, null, value, flags);
                 } else {
                     // Update focus on existing session.
                     updateSessionLocked(id, null, value, ACTION_VIEW_ENTERED, flags);
@@ -484,7 +468,7 @@
 
                 if (mSessionId == NO_SESSION) {
                     // Starts new session.
-                    startSessionLocked(id, view.getWindowToken(), bounds, null, flags);
+                    startSessionLocked(id, bounds, null, flags);
                 } else {
                     // Update focus on existing session.
                     updateSessionLocked(id, bounds, null, ACTION_VIEW_ENTERED, flags);
@@ -725,15 +709,15 @@
         return new AutofillId(parent.getAccessibilityViewId(), childId);
     }
 
-    private void startSessionLocked(@NonNull AutofillId id, @NonNull IBinder windowToken,
-            @NonNull Rect bounds, @NonNull AutofillValue value, int flags) {
+    private void startSessionLocked(@NonNull AutofillId id, @NonNull Rect bounds,
+            @NonNull AutofillValue value, int flags) {
         if (sVerbose) {
             Log.v(TAG, "startSessionLocked(): id=" + id + ", bounds=" + bounds + ", value=" + value
                     + ", flags=" + flags);
         }
 
         try {
-            mSessionId = mService.startSession(mContext.getActivityToken(), windowToken,
+            mSessionId = mService.startSession(mContext.getActivityToken(),
                     mServiceClient.asBinder(), id, bounds, value, mContext.getUserId(),
                     mCallback != null, flags, mContext.getOpPackageName());
             final AutofillClient client = getClientLocked();
@@ -855,9 +839,9 @@
         }
     }
 
-    private void requestShowFillUi(int sessionId, IBinder windowToken, AutofillId id, int width,
-            int height, Rect anchorBounds, IAutofillWindowPresenter presenter) {
-        final View anchor = findAchorView(windowToken, id);
+    private void requestShowFillUi(int sessionId, AutofillId id, int width, int height,
+            Rect anchorBounds, IAutofillWindowPresenter presenter) {
+        final View anchor = findView(id);
         if (anchor == null) {
             return;
         }
@@ -930,27 +914,27 @@
         }
     }
 
-    private void autofill(int sessionId, IBinder windowToken, List<AutofillId> ids,
-            List<AutofillValue> values) {
+    private void autofill(int sessionId, List<AutofillId> ids, List<AutofillValue> values) {
         synchronized (mLock) {
             if (sessionId != mSessionId) {
                 return;
             }
 
-            final View root = WindowManagerGlobal.getInstance().getWindowView(windowToken);
-            if (root == null) {
+            final AutofillClient client = getClientLocked();
+            if (client == null) {
                 return;
             }
 
             final int itemCount = ids.size();
             int numApplied = 0;
             ArrayMap<View, SparseArray<AutofillValue>> virtualValues = null;
+            final View[] views = client.findViewsByAccessibilityIdTraversal(getViewIds(ids));
 
             for (int i = 0; i < itemCount; i++) {
                 final AutofillId id = ids.get(i);
                 final AutofillValue value = values.get(i);
                 final int viewId = id.getViewId();
-                final View view = root.findViewByAccessibilityIdTraversal(viewId);
+                final View view = views[i];
                 if (view == null) {
                     Log.w(TAG, "autofill(): no View with id " + viewId);
                     continue;
@@ -1022,8 +1006,11 @@
         }
     }
 
-    private void requestHideFillUi(int sessionId, IBinder windowToken, AutofillId id) {
-        final View anchor = findAchorView(windowToken, id);
+    private void requestHideFillUi(int sessionId, AutofillId id) {
+        final View anchor = findView(id);
+        if (anchor == null) {
+            return;
+        }
 
         AutofillCallback callback = null;
         synchronized (mLock) {
@@ -1049,8 +1036,11 @@
         }
     }
 
-    private void notifyNoFillUi(int sessionId, IBinder windowToken, AutofillId id) {
-        final View anchor = findAchorView(windowToken, id);
+    private void notifyNoFillUi(int sessionId, AutofillId id) {
+        final View anchor = findView(id);
+        if (anchor == null) {
+            return;
+        }
 
         AutofillCallback callback = null;
         synchronized (mLock) {
@@ -1070,18 +1060,38 @@
         }
     }
 
-    private View findAchorView(IBinder windowToken, AutofillId id) {
-        final View root = WindowManagerGlobal.getInstance().getWindowView(windowToken);
-        if (root == null) {
-            Log.w(TAG, "no window with token " + windowToken);
+    /**
+     * Get an array of viewIds from a List of {@link AutofillId}.
+     *
+     * @param autofillIds The autofill ids to convert
+     *
+     * @return The array of viewIds.
+     */
+    @NonNull private int[] getViewIds(@NonNull List<AutofillId> autofillIds) {
+        final int numIds = autofillIds.size();
+        final int[] viewIds = new int[numIds];
+        for (int i = 0; i < numIds; i++) {
+            viewIds[i] = autofillIds.get(i).getViewId();
+        }
+
+        return viewIds;
+    }
+
+    /**
+     * Find a single view by its id.
+     *
+     * @param autofillId The autofill id of the view
+     *
+     * @return The view or {@code null} if view was not found
+     */
+    private View findView(@NonNull AutofillId autofillId) {
+        final AutofillClient client = getClientLocked();
+
+        if (client == null) {
             return null;
         }
-        final View view = root.findViewByAccessibilityIdTraversal(id.getViewId());
-        if (view == null) {
-            Log.w(TAG, "no view with id " + id);
-            return null;
-        }
-        return view;
+
+        return client.findViewsByAccessibilityIdTraversal(new int[]{autofillId.getViewId()})[0];
     }
 
     /** @hide */
@@ -1160,22 +1170,26 @@
          *
          * @param trackedIds The views to be tracked
          */
-        TrackedViews(@NonNull List<AutofillId> trackedIds) {
+        TrackedViews(List<AutofillId> trackedIds) {
             mVisibleTrackedIds = null;
             mInvisibleTrackedIds = null;
 
             AutofillClient client = getClientLocked();
-            if (trackedIds != null) {
-                int numIds = trackedIds.size();
+            if (trackedIds != null && client != null) {
+                final boolean[] isVisible;
+
+                if (client.isVisibleForAutofill()) {
+                    isVisible = client.getViewVisibility(getViewIds(trackedIds));
+                } else {
+                    // All false
+                    isVisible = new boolean[trackedIds.size()];
+                }
+
+                final int numIds = trackedIds.size();
                 for (int i = 0; i < numIds; i++) {
-                    AutofillId id = trackedIds.get(i);
+                    final AutofillId id = trackedIds.get(i);
 
-                    boolean isVisible = true;
-                    if (client != null && client.isVisibleForAutofill()) {
-                        isVisible = client.getViewVisibility(id.getViewId());
-                    }
-
-                    if (isVisible) {
+                    if (isVisible[i]) {
                         mVisibleTrackedIds = addToSet(mVisibleTrackedIds, id);
                     } else {
                         mInvisibleTrackedIds = addToSet(mInvisibleTrackedIds, id);
@@ -1233,16 +1247,23 @@
          *
          * @see AutofillClient#isVisibleForAutofill()
          */
-        void onVisibleForAutofill() {
-            // The visibility of the views might have changed while the client was not started,
+        void onVisibleForAutofillLocked() {
+            // The visibility of the views might have changed while the client was not be visible,
             // hence update the visibility state for all views.
             AutofillClient client = getClientLocked();
             ArraySet<AutofillId> updatedVisibleTrackedIds = null;
             ArraySet<AutofillId> updatedInvisibleTrackedIds = null;
             if (client != null) {
                 if (mInvisibleTrackedIds != null) {
-                    for (AutofillId id : mInvisibleTrackedIds) {
-                        if (client.getViewVisibility(id.getViewId())) {
+                    final ArrayList<AutofillId> orderedInvisibleIds =
+                            new ArrayList<>(mInvisibleTrackedIds);
+                    final boolean[] isVisible = client.getViewVisibility(
+                            getViewIds(orderedInvisibleIds));
+
+                    final int numInvisibleTrackedIds = orderedInvisibleIds.size();
+                    for (int i = 0; i < numInvisibleTrackedIds; i++) {
+                        final AutofillId id = orderedInvisibleIds.get(i);
+                        if (isVisible[i]) {
                             updatedVisibleTrackedIds = addToSet(updatedVisibleTrackedIds, id);
 
                             if (sDebug) {
@@ -1255,8 +1276,16 @@
                 }
 
                 if (mVisibleTrackedIds != null) {
-                    for (AutofillId id : mVisibleTrackedIds) {
-                        if (client.getViewVisibility(id.getViewId())) {
+                    final ArrayList<AutofillId> orderedVisibleIds =
+                            new ArrayList<>(mVisibleTrackedIds);
+                    final boolean[] isVisible = client.getViewVisibility(
+                            getViewIds(orderedVisibleIds));
+
+                    final int numVisibleTrackedIds = orderedVisibleIds.size();
+                    for (int i = 0; i < numVisibleTrackedIds; i++) {
+                        final AutofillId id = orderedVisibleIds.get(i);
+
+                        if (isVisible[i]) {
                             updatedVisibleTrackedIds = addToSet(updatedVisibleTrackedIds, id);
                         } else {
                             updatedInvisibleTrackedIds = addToSet(updatedInvisibleTrackedIds, id);
@@ -1355,12 +1384,11 @@
         }
 
         @Override
-        public void autofill(int sessionId, IBinder windowToken, List<AutofillId> ids,
-                List<AutofillValue> values) {
+        public void autofill(int sessionId, List<AutofillId> ids, List<AutofillValue> values) {
             final AutofillManager afm = mAfm.get();
             if (afm != null) {
                 afm.mContext.getMainThreadHandler().post(
-                        () -> afm.autofill(sessionId, windowToken, ids, values));
+                        () -> afm.autofill(sessionId, ids, values));
             }
         }
 
@@ -1374,31 +1402,30 @@
         }
 
         @Override
-        public void requestShowFillUi(int sessionId, IBinder windowToken, AutofillId id,
-                int width, int height, Rect anchorBounds, IAutofillWindowPresenter presenter) {
+        public void requestShowFillUi(int sessionId, AutofillId id, int width, int height,
+                Rect anchorBounds, IAutofillWindowPresenter presenter) {
             final AutofillManager afm = mAfm.get();
             if (afm != null) {
                 afm.mContext.getMainThreadHandler().post(
-                        () -> afm.requestShowFillUi(sessionId, windowToken, id, width, height,
-                                anchorBounds, presenter));
+                        () -> afm.requestShowFillUi(sessionId, id, width, height, anchorBounds,
+                                presenter));
             }
         }
 
         @Override
-        public void requestHideFillUi(int sessionId, IBinder windowToken, AutofillId id) {
+        public void requestHideFillUi(int sessionId, AutofillId id) {
             final AutofillManager afm = mAfm.get();
             if (afm != null) {
                 afm.mContext.getMainThreadHandler().post(
-                        () -> afm.requestHideFillUi(sessionId, windowToken, id));
+                        () -> afm.requestHideFillUi(sessionId, id));
             }
         }
 
         @Override
-        public void notifyNoFillUi(int sessionId, IBinder windowToken, AutofillId id) {
+        public void notifyNoFillUi(int sessionId, AutofillId id) {
             final AutofillManager afm = mAfm.get();
             if (afm != null) {
-                afm.mContext.getMainThreadHandler().post(
-                        () -> afm.notifyNoFillUi(sessionId, windowToken, id));
+                afm.mContext.getMainThreadHandler().post(() -> afm.notifyNoFillUi(sessionId, id));
             }
         }
 
diff --git a/core/java/android/view/autofill/IAutoFillManager.aidl b/core/java/android/view/autofill/IAutoFillManager.aidl
index f28d8ba..4193a3c 100644
--- a/core/java/android/view/autofill/IAutoFillManager.aidl
+++ b/core/java/android/view/autofill/IAutoFillManager.aidl
@@ -32,12 +32,11 @@
 interface IAutoFillManager {
     // Returns flags: FLAG_ADD_CLIENT_ENABLED | FLAG_ADD_CLIENT_DEBUG | FLAG_ADD_CLIENT_VERBOSE
     int addClient(in IAutoFillManagerClient client, int userId);
-    int startSession(IBinder activityToken, IBinder windowToken, in IBinder appCallback,
-            in AutofillId autoFillId, in Rect bounds, in AutofillValue value, int userId,
-            boolean hasCallback, int flags, String packageName);
+    int startSession(IBinder activityToken, in IBinder appCallback, in AutofillId autoFillId,
+            in Rect bounds, in AutofillValue value, int userId, boolean hasCallback, int flags,
+            String packageName);
     FillEventHistory getFillEventHistory();
     boolean restoreSession(int sessionId, in IBinder activityToken, in IBinder appCallback);
-    void setWindow(int sessionId, in IBinder windowToken);
     void updateSession(int sessionId, in AutofillId id, in Rect bounds,
             in AutofillValue value, int action, int flags, int userId);
     void finishSession(int sessionId, int userId);
diff --git a/core/java/android/view/autofill/IAutoFillManagerClient.aidl b/core/java/android/view/autofill/IAutoFillManagerClient.aidl
index aef98b7..825d311 100644
--- a/core/java/android/view/autofill/IAutoFillManagerClient.aidl
+++ b/core/java/android/view/autofill/IAutoFillManagerClient.aidl
@@ -40,8 +40,7 @@
     /**
       * Autofills the activity with the contents of a dataset.
       */
-    void autofill(int sessionId, in IBinder windowToken, in List<AutofillId> ids,
-            in List<AutofillValue> values);
+    void autofill(int sessionId, in List<AutofillId> ids, in List<AutofillValue> values);
 
     /**
       * Authenticates a fill response or a data set.
@@ -58,18 +57,18 @@
     /**
      * Requests showing the fill UI.
      */
-    void requestShowFillUi(int sessionId, in IBinder windowToken, in AutofillId id, int width,
-            int height, in Rect anchorBounds, in IAutofillWindowPresenter presenter);
+    void requestShowFillUi(int sessionId, in AutofillId id, int width, int height,
+    in Rect anchorBounds, in IAutofillWindowPresenter presenter);
 
     /**
      * Requests hiding the fill UI.
      */
-    void requestHideFillUi(int sessionId, in IBinder windowToken, in AutofillId id);
+    void requestHideFillUi(int sessionId, in AutofillId id);
 
     /**
      * Notifies no fill UI will be shown.
      */
-    void notifyNoFillUi(int sessionId, in IBinder windowToken, in AutofillId id);
+    void notifyNoFillUi(int sessionId, in AutofillId id);
 
     /**
      * Starts the provided intent sender
diff --git a/core/java/com/android/internal/policy/DecorView.java b/core/java/com/android/internal/policy/DecorView.java
index 80b6b08..1b83708 100644
--- a/core/java/com/android/internal/policy/DecorView.java
+++ b/core/java/com/android/internal/policy/DecorView.java
@@ -1498,7 +1498,6 @@
 
         final Window.Callback cb = mWindow.getCallback();
         if (cb != null && !mWindow.isDestroyed() && mFeatureId < 0) {
-            cb.onBeforeAttachedToWindow();
             cb.onAttachedToWindow();
         }
 
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerService.java b/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
index 9081746..e3398c9 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerService.java
@@ -118,7 +118,7 @@
         @Override
         public void onReceive(Context context, Intent intent) {
             if (Intent.ACTION_CLOSE_SYSTEM_DIALOGS.equals(intent.getAction())) {
-                mUi.hideAll();
+                mUi.hideAll(null);
             }
         }
     };
@@ -470,9 +470,9 @@
         }
 
         @Override
-        public int startSession(IBinder activityToken, IBinder windowToken, IBinder appCallback,
-                AutofillId autofillId, Rect bounds, AutofillValue value, int userId,
-                boolean hasCallback, int flags, String packageName) {
+        public int startSession(IBinder activityToken, IBinder appCallback, AutofillId autofillId,
+                Rect bounds, AutofillValue value, int userId, boolean hasCallback, int flags,
+                String packageName) {
 
             activityToken = Preconditions.checkNotNull(activityToken, "activityToken");
             appCallback = Preconditions.checkNotNull(appCallback, "appCallback");
@@ -489,8 +489,8 @@
 
             synchronized (mLock) {
                 final AutofillManagerServiceImpl service = getServiceForUserLocked(userId);
-                return service.startSessionLocked(activityToken, getCallingUid(), windowToken,
-                        appCallback, autofillId, bounds, value, hasCallback, flags, packageName);
+                return service.startSessionLocked(activityToken, getCallingUid(), appCallback,
+                        autofillId, bounds, value, hasCallback, flags, packageName);
             }
         }
 
@@ -528,19 +528,6 @@
         }
 
         @Override
-        public void setWindow(int sessionId, IBinder windowToken) throws RemoteException {
-            windowToken = Preconditions.checkNotNull(windowToken, "windowToken");
-
-            synchronized (mLock) {
-                final AutofillManagerServiceImpl service = mServicesCache.get(
-                        UserHandle.getCallingUserId());
-                if (service != null) {
-                    service.setWindow(sessionId, getCallingUid(), windowToken);
-                }
-            }
-        }
-
-        @Override
         public void updateSession(int sessionId, AutofillId id, Rect bounds,
                 AutofillValue value, int action, int flags, int userId) {
             synchronized (mLock) {
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
index 3ab8b8d..e315f9d 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
@@ -263,7 +263,7 @@
         }
     }
 
-    int startSessionLocked(@NonNull IBinder activityToken, int uid, @Nullable IBinder windowToken,
+    int startSessionLocked(@NonNull IBinder activityToken, int uid,
             @NonNull IBinder appCallbackToken, @NonNull AutofillId autofillId,
             @NonNull Rect virtualBounds, @Nullable AutofillValue value, boolean hasCallback,
             int flags, @NonNull String packageName) {
@@ -274,8 +274,8 @@
         // Occasionally clean up abandoned sessions
         pruneAbandonedSessionsLocked();
 
-        final Session newSession = createSessionByTokenLocked(activityToken, uid, windowToken,
-                appCallbackToken, hasCallback, flags, packageName);
+        final Session newSession = createSessionByTokenLocked(activityToken, uid, appCallbackToken,
+                hasCallback, flags, packageName);
         if (newSession == null) {
             return NO_SESSION;
         }
@@ -359,8 +359,8 @@
     }
 
     private Session createSessionByTokenLocked(@NonNull IBinder activityToken, int uid,
-            @Nullable IBinder windowToken, @NonNull IBinder appCallbackToken, boolean hasCallback,
-            int flags, @NonNull String packageName) {
+            @NonNull IBinder appCallbackToken, boolean hasCallback, int flags,
+            @NonNull String packageName) {
         // use random ids so that one app cannot know that another app creates sessions
         int sessionId;
         int tries = 0;
@@ -375,7 +375,7 @@
         } while (sessionId == NO_SESSION || mSessions.indexOfKey(sessionId) >= 0);
 
         final Session newSession = new Session(this, mUi, mContext, mHandlerCaller, mUserId, mLock,
-                sessionId, uid, activityToken, windowToken, appCallbackToken, hasCallback,
+                sessionId, uid, activityToken, appCallbackToken, hasCallback,
                 mInfo.getServiceInfo().getComponentName(), packageName);
         mSessions.put(newSession.id, newSession);
 
@@ -402,24 +402,6 @@
         }
     }
 
-    /**
-     * Set the window the UI should get attached to
-     *
-     * @param sessionId The id of the session to restore
-     * @param uid UID of the process that tries to restore the session
-     * @param windowToken The window the activity is now in
-     */
-    boolean setWindow(int sessionId, int uid, @NonNull IBinder windowToken) {
-        final Session session = mSessions.get(sessionId);
-
-        if (session == null || uid != session.uid) {
-            return false;
-        } else {
-            session.switchWindow(windowToken);
-            return true;
-        }
-    }
-
     void updateSessionLocked(int sessionId, int uid, AutofillId autofillId, Rect virtualBounds,
             AutofillValue value, int action, int flags) {
         final Session session = mSessions.get(sessionId);
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index a98821d..3f78fb8 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -119,9 +119,6 @@
     @GuardedBy("mLock")
     @NonNull private IBinder mActivityToken;
 
-    @GuardedBy("mLock")
-    @NonNull private IBinder mWindowToken;
-
     /** Package name of the app that is auto-filled */
     @NonNull private final String mPackageName;
 
@@ -341,7 +338,7 @@
     Session(@NonNull AutofillManagerServiceImpl service, @NonNull AutoFillUI ui,
             @NonNull Context context, @NonNull HandlerCaller handlerCaller, int userId,
             @NonNull Object lock, int sessionId, int uid, @NonNull IBinder activityToken,
-            @Nullable IBinder windowToken, @NonNull IBinder client, boolean hasCallback,
+            @NonNull IBinder client, boolean hasCallback,
             @NonNull ComponentName componentName, @NonNull String packageName) {
         id = sessionId;
         this.uid = uid;
@@ -351,7 +348,6 @@
         mHandlerCaller = handlerCaller;
         mRemoteFillService = new RemoteFillService(context, componentName, userId, this);
         mActivityToken = activityToken;
-        mWindowToken = windowToken;
         mHasCallback = hasCallback;
         mPackageName = packageName;
         mClient = IAutoFillManagerClient.Stub.asInterface(client);
@@ -369,23 +365,6 @@
     }
 
     /**
-     * Sets new window  for this session.
-     *
-     * @param newWindow The window the Ui should be attached to. Can be {@code null} if no
-     *                  further UI is needed.
-     */
-    void switchWindow(@NonNull IBinder newWindow) {
-        synchronized (mLock) {
-            if (mDestroyed) {
-                Slog.w(TAG, "Call to Session#switchWindow() rejected - session: "
-                        + id + " destroyed");
-                return;
-            }
-            mWindowToken = newWindow;
-        }
-    }
-
-    /**
      * Sets new activity and client for this session.
      *
      * @param newActivity The token of the new activity
@@ -419,7 +398,7 @@
         }
         if (response == null) {
             if ((requestFlags & FLAG_MANUAL_REQUEST) != 0) {
-                getUiForShowing().showError(R.string.autofill_error_cannot_autofill);
+                getUiForShowing().showError(R.string.autofill_error_cannot_autofill, this);
             }
             // Nothing to be done, but need to notify client.
             notifyUnavailableToClient();
@@ -469,7 +448,7 @@
                 .addTaggedData(MetricsEvent.FIELD_AUTOFILL_SERVICE, servicePackageName);
         mMetricsLogger.write(log);
 
-        getUiForShowing().showError(message);
+        getUiForShowing().showError(message, this);
         removeSelf();
     }
 
@@ -516,7 +495,7 @@
                 .addTaggedData(MetricsEvent.FIELD_AUTOFILL_SERVICE, servicePackageName);
         mMetricsLogger.write(log);
 
-        getUiForShowing().showError(message);
+        getUiForShowing().showError(message, this);
         removeSelf();
     }
 
@@ -627,8 +606,8 @@
             if (id.equals(mCurrentViewId)) {
                 try {
                     final ViewState view = mViewStates.get(id);
-                    mClient.requestShowFillUi(this.id, mWindowToken, id, width, height,
-                            view.getVirtualBounds(), presenter);
+                    mClient.requestShowFillUi(this.id, id, width, height, view.getVirtualBounds(),
+                            presenter);
                 } catch (RemoteException e) {
                     Slog.e(TAG, "Error requesting to show fill UI", e);
                 }
@@ -648,7 +627,7 @@
             // NOTE: We allow this call in a destroyed state as the UI is
             // asked to go away after we get destroyed, so let it do that.
             try {
-                mClient.requestHideFillUi(this.id, mWindowToken, id);
+                mClient.requestHideFillUi(this.id, id);
             } catch (RemoteException e) {
                 Slog.e(TAG, "Error requesting to hide fill UI", e);
             }
@@ -706,7 +685,7 @@
                         // refactor auth to support partitions; if it doesn't, we need to
                         // investigate it further (it can be reproduced by running
                         // LoginActivityTest.testFillResponseAuthServiceHasNoData())
-                        mUi.hideAll();
+                        mUi.hideAll(this);
                     }
                     processResponseLocked(response);
                 } else {
@@ -851,7 +830,8 @@
             }
             if (atLeastOneChanged) {
                 mService.setSaveShown();
-                getUiForShowing().showSaveUi(mService.getServiceLabel(), saveInfo, mPackageName);
+                getUiForShowing().showSaveUi(mService.getServiceLabel(), saveInfo, mPackageName,
+                        this);
 
                 mIsSaving = true;
                 return false;
@@ -1062,9 +1042,9 @@
 
                     //..and the UI
                     if (value.isText()) {
-                        getUiForShowing().filterFillUi(value.getTextValue().toString());
+                        getUiForShowing().filterFillUi(value.getTextValue().toString(), this);
                     } else {
-                        getUiForShowing().filterFillUi(null);
+                        getUiForShowing().filterFillUi(null, this);
                     }
                 }
                 break;
@@ -1076,7 +1056,7 @@
 
                 // Remove the UI if the ViewState has changed.
                 if (mCurrentViewId != viewState.id) {
-                    mUi.hideFillUi(mCurrentViewId != null ? mCurrentViewId : null);
+                    mUi.hideFillUi(this);
                     mCurrentViewId = viewState.id;
                 }
 
@@ -1086,7 +1066,7 @@
             case ACTION_VIEW_EXITED:
                 if (mCurrentViewId == viewState.id) {
                     if (sVerbose) Slog.d(TAG, "Exiting view " + id);
-                    mUi.hideFillUi(viewState.id);
+                    mUi.hideFillUi(this);
                     mCurrentViewId = null;
                 }
                 break;
@@ -1123,7 +1103,7 @@
             filterText = value.getTextValue().toString();
         }
 
-        getUiForShowing().showFillUi(filledId, response, filterText, mPackageName);
+        getUiForShowing().showFillUi(filledId, response, filterText, mPackageName, this);
     }
 
     boolean isDestroyed() {
@@ -1142,10 +1122,9 @@
         synchronized (mLock) {
             if (!mHasCallback) return;
             try {
-                mClient.notifyNoFillUi(id, mWindowToken, mCurrentViewId);
+                mClient.notifyNoFillUi(id, mCurrentViewId);
             } catch (RemoteException e) {
-                Slog.e(TAG, "Error notifying client no fill UI: windowToken=" + mWindowToken
-                        + " id=" + mCurrentViewId, e);
+                Slog.e(TAG, "Error notifying client no fill UI: id=" + mCurrentViewId, e);
             }
         }
     }
@@ -1413,8 +1392,7 @@
             }
             try {
                 if (sDebug) Slog.d(TAG, "autoFillApp(): the buck is on the app: " + dataset);
-                mClient.autofill(id, mWindowToken, dataset.getFieldIds(),
-                        dataset.getFieldValues());
+                mClient.autofill(id, dataset.getFieldIds(), dataset.getFieldValues());
                 setViewStatesLocked(null, dataset, ViewState.STATE_AUTOFILLED);
             } catch (RemoteException e) {
                 Slog.w(TAG, "Error autofilling activity: " + e);
@@ -1434,8 +1412,8 @@
             return;
         }
         mRemoteFillService.destroy();
-        mUi.hideAll();
-        mUi.setCallback(null);
+        mUi.hideAll(this);
+        mUi.clearCallback(this);
         mDestroyed = true;
         mMetricsLogger.action(MetricsEvent.AUTOFILL_SESSION_FINISHED, mPackageName);
     }
diff --git a/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java b/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
index 7428460..9eaabfe2 100644
--- a/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
+++ b/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
@@ -74,31 +74,43 @@
         mContext = context;
     }
 
-    public void setCallback(@Nullable AutoFillUiCallback callback) {
+    public void setCallback(@NonNull AutoFillUiCallback callback) {
         mHandler.post(() -> {
             if (mCallback != callback) {
-                hideAllUiThread();
+                if (mCallback != null) {
+                    hideAllUiThread(mCallback);
+                }
+
                 mCallback = callback;
             }
         });
     }
 
+    public void clearCallback(@NonNull AutoFillUiCallback callback) {
+        mHandler.post(() -> {
+            if (mCallback == callback) {
+                hideAllUiThread(callback);
+                mCallback = null;
+            }
+        });
+    }
+
     /**
      * Displays an error message to the user.
      */
-    public void showError(int resId) {
-        showError(mContext.getString(resId));
+    public void showError(int resId, @NonNull AutoFillUiCallback callback) {
+        showError(mContext.getString(resId), callback);
     }
 
     /**
      * Displays an error message to the user.
      */
-    public void showError(@Nullable CharSequence message) {
+    public void showError(@Nullable CharSequence message, @NonNull AutoFillUiCallback callback) {
         mHandler.post(() -> {
-            if (!hasCallback()) {
+            if (mCallback != callback) {
                 return;
             }
-            hideAllUiThread();
+            hideAllUiThread(callback);
             if (!TextUtils.isEmpty(message)) {
                 Toast.makeText(mContext, message, Toast.LENGTH_LONG).show();
             }
@@ -108,8 +120,8 @@
     /**
      * Hides the fill UI.
      */
-    public void hideFillUi(AutofillId id) {
-        mHandler.post(this::hideFillUiUiThread);
+    public void hideFillUi(@NonNull AutoFillUiCallback callback) {
+        mHandler.post(() -> hideFillUiUiThread(callback));
     }
 
     /**
@@ -117,12 +129,12 @@
      *
      * @param filterText The filter prefix.
      */
-    public void filterFillUi(@Nullable String filterText) {
+    public void filterFillUi(@Nullable String filterText, @NonNull AutoFillUiCallback callback) {
         mHandler.post(() -> {
-            if (!hasCallback()) {
+            if (callback != mCallback) {
                 return;
             }
-            hideSaveUiUiThread();
+            hideSaveUiUiThread(callback);
             if (mFillUi != null) {
                 mFillUi.setFilterText(filterText);
             }
@@ -136,9 +148,11 @@
      * @param response the current fill response
      * @param filterText text of the view to be filled
      * @param packageName package name of the activity that is filled
+     * @param callback Identifier for the caller
      */
     public void showFillUi(@NonNull AutofillId focusedId, @NonNull FillResponse response,
-            @Nullable String filterText, @NonNull String packageName) {
+            @Nullable String filterText, @NonNull String packageName,
+            @NonNull AutoFillUiCallback callback) {
         if (sDebug) {
             Slog.d(TAG, "showFillUi(): id=" + focusedId + ", filter=" + filterText);
         }
@@ -150,16 +164,16 @@
                         response.getDatasets() == null ? 0 : response.getDatasets().size());
 
         mHandler.post(() -> {
-            if (!hasCallback()) {
+            if (callback != mCallback) {
                 return;
             }
-            hideAllUiThread();
+            hideAllUiThread(callback);
             mFillUi = new FillUi(mContext, response, focusedId,
                     filterText, new FillUi.Callback() {
                 @Override
                 public void onResponsePicked(FillResponse response) {
                     log.setType(MetricsProto.MetricsEvent.TYPE_DETAIL);
-                    hideFillUiUiThread();
+                    hideFillUiUiThread(callback);
                     if (mCallback != null) {
                         mCallback.authenticate(response.getRequestId(),
                                 response.getAuthentication(), response.getClientState());
@@ -169,7 +183,7 @@
                 @Override
                 public void onDatasetPicked(Dataset dataset) {
                     log.setType(MetricsProto.MetricsEvent.TYPE_ACTION);
-                    hideFillUiUiThread();
+                    hideFillUiUiThread(callback);
                     if (mCallback != null) {
                         mCallback.fill(response.getRequestId(), dataset);
                     }
@@ -178,7 +192,7 @@
                 @Override
                 public void onCanceled() {
                     log.setType(MetricsProto.MetricsEvent.TYPE_DISMISS);
-                    hideFillUiUiThread();
+                    hideFillUiUiThread(callback);
                 }
 
                 @Override
@@ -218,7 +232,7 @@
      * Shows the UI asking the user to save for autofill.
      */
     public void showSaveUi(@NonNull CharSequence providerLabel, @NonNull SaveInfo info,
-            @NonNull String packageName) {
+            @NonNull String packageName, @NonNull AutoFillUiCallback callback) {
         int numIds = 0;
         numIds += info.getRequiredIds() == null ? 0 : info.getRequiredIds().length;
         numIds += info.getOptionalIds() == null ? 0 : info.getOptionalIds().length;
@@ -228,16 +242,16 @@
                         MetricsProto.MetricsEvent.FIELD_AUTOFILL_NUM_IDS, numIds);
 
         mHandler.post(() -> {
-            if (!hasCallback()) {
+            if (callback != mCallback) {
                 return;
             }
-            hideAllUiThread();
+            hideAllUiThread(callback);
             mSaveUi = new SaveUi(mContext, providerLabel, info,
                     new SaveUi.OnSaveListener() {
                 @Override
                 public void onSave() {
                     log.setType(MetricsProto.MetricsEvent.TYPE_ACTION);
-                    hideSaveUiUiThread();
+                    hideSaveUiUiThread(callback);
                     if (mCallback != null) {
                         mCallback.save();
                     }
@@ -246,7 +260,7 @@
                 @Override
                 public void onCancel(IntentSender listener) {
                     log.setType(MetricsProto.MetricsEvent.TYPE_DISMISS);
-                    hideSaveUiUiThread();
+                    hideSaveUiUiThread(callback);
                     if (listener != null) {
                         try {
                             listener.sendIntent(mContext, 0, null, null, null);
@@ -278,8 +292,8 @@
     /**
      * Hides all UI affordances.
      */
-    public void hideAll() {
-        mHandler.post(this::hideAllUiThread);
+    public void hideAll(@Nullable AutoFillUiCallback callback) {
+        mHandler.post(() -> hideAllUiThread(callback));
     }
 
     public void dump(PrintWriter pw) {
@@ -301,28 +315,24 @@
     }
 
     @android.annotation.UiThread
-    private void hideFillUiUiThread() {
-        if (mFillUi != null) {
+    private void hideFillUiUiThread(@Nullable AutoFillUiCallback callback) {
+        if (mFillUi != null && (callback == null || callback == mCallback)) {
             mFillUi.destroy();
             mFillUi = null;
         }
     }
 
     @android.annotation.UiThread
-    private void hideSaveUiUiThread() {
-        if (mSaveUi != null) {
+    private void hideSaveUiUiThread(@Nullable AutoFillUiCallback callback) {
+        if (mSaveUi != null && (callback == null || callback == mCallback)) {
             mSaveUi.destroy();
             mSaveUi = null;
         }
     }
 
     @android.annotation.UiThread
-    private void hideAllUiThread() {
-        hideFillUiUiThread();
-        hideSaveUiUiThread();
-    }
-
-    private boolean hasCallback() {
-        return mCallback != null;
+    private void hideAllUiThread(@Nullable AutoFillUiCallback callback) {
+        hideFillUiUiThread(callback);
+        hideSaveUiUiThread(callback);
     }
 }