diff --git a/core/java/android/app/Activity.java b/core/java/android/app/Activity.java
index 6fc60e9eaaf22..a820da43a6ba1 100644
--- a/core/java/android/app/Activity.java
+++ b/core/java/android/app/Activity.java
@@ -844,6 +844,8 @@ public class Activity extends ContextThemeWrapper
 
     private boolean mHasCurrentPermissionsRequest;
 
+    private boolean mAutoFillResetNeeded;
+
     private static native String getDlWarning();
 
     /** Return the intent that started this activity. */
@@ -1780,7 +1782,7 @@ public class Activity extends ContextThemeWrapper
         getApplication().dispatchActivityStopped(this);
         mTranslucentCallback = null;
         mCalled = true;
-        if (isFinishing() && AutoFillManager.isClientActive(getActivityToken())) {
+        if (isFinishing() && mAutoFillResetNeeded) {
             getSystemService(AutoFillManager.class).reset();
         }
     }
@@ -6746,8 +6748,6 @@ public class Activity extends ContextThemeWrapper
         mCurrentConfig = config;
 
         mWindow.setColorMode(info.colorMode);
-
-        AutoFillManager.addClient(token, this);
     }
 
     /** @hide */
@@ -7214,6 +7214,12 @@ public class Activity extends ContextThemeWrapper
         }
     }
 
+    /** @hide */
+    @Override
+    public void resetableStateAvailable() {
+        mAutoFillResetNeeded = true;
+    }
+
     class HostCallbacks extends FragmentHostCallback<Activity> {
         public HostCallbacks() {
             super(Activity.this /*activity*/);
diff --git a/core/java/android/view/autofill/AutoFillManager.java b/core/java/android/view/autofill/AutoFillManager.java
index 2168444b9b58f..baba3895ccca0 100644
--- a/core/java/android/view/autofill/AutoFillManager.java
+++ b/core/java/android/view/autofill/AutoFillManager.java
@@ -23,10 +23,8 @@ import android.content.Intent;
 import android.content.IntentSender;
 import android.graphics.Rect;
 import android.os.Bundle;
-import android.os.IBinder;
 import android.os.Parcelable;
 import android.os.RemoteException;
-import android.util.ArrayMap;
 import android.util.Log;
 import android.view.View;
 
@@ -70,13 +68,6 @@ public final class AutoFillManager {
     /** @hide */ public static final int FLAG_FOCUS_LOST = 0x4;
     /** @hide */ public static final int FLAG_VALUE_CHANGED = 0x8;
 
-    // These are activities that may have auto-fill UI which are keyed off their tokens.
-    // This is done instead of the activity setting the client in the auto-fill manager
-    // to avoid unnecessary instantiation of the manager and do this only if there is an
-    // auto-fillable focused. This has only the cost of loading the class vs creating an
-    // auto-fill manager for every activity even one that cannot be filled.
-    private static final ArrayMap<IBinder, AutoFillClient> sPendingClients = new ArrayMap<>();
-
     private final Rect mTempRect = new Rect();
 
     private final IAutoFillManager mService;
@@ -84,8 +75,6 @@ public final class AutoFillManager {
 
     private Context mContext;
 
-    private AutoFillClient mClient;
-
     private boolean mHasSession;
     private boolean mEnabled;
 
@@ -106,6 +95,11 @@ public final class AutoFillManager {
          * @param fillInIntent The authentication fill-in intent.
          */
         void authenticate(IntentSender intent, Intent fillInIntent);
+
+        /**
+         * Tells the client this manager has state to be reset.
+         */
+        void resetableStateAvailable();
     }
 
     /**
@@ -219,25 +213,11 @@ public final class AutoFillManager {
         finishSession();
     }
 
-    /** @hide */
-    public static void addClient(IBinder token, AutoFillClient client) {
-        sPendingClients.put(token, client);
-    }
-
-    /** @hide */
-    public static boolean isClientActive(IBinder token) {
-        return !sPendingClients.containsKey(token);
-    }
-
-    private void activateClient() {
-        mClient = sPendingClients.remove(mContext.getActivityToken());
-    }
-
     private AutoFillClient getClient() {
-        if (mClient == null) {
-            return sPendingClients.get(mContext.getActivityToken());
+        if (mContext instanceof AutoFillClient) {
+            return (AutoFillClient) mContext;
         }
-        return mClient;
+        return null;
     }
 
     /** @hide */
@@ -278,8 +258,11 @@ public final class AutoFillManager {
         try {
             mService.startSession(mContext.getActivityToken(), mServiceClient.asBinder(),
                     id, bounds, value, mContext.getUserId());
+            AutoFillClient client = getClient();
+            if (client != null) {
+                client.resetableStateAvailable();
+            }
             mHasSession = true;
-            activateClient();
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }