Add support for multiple fill contexts

When saving data filled by the user the platform provides to
an autofill provider the state of the UI allowing the provider
to interpret this state and store relevant information.

A limitation of the current design is that the fill provider
needs to interpret the screen content twice, once handling a
fill request and once handling a save request. To address this
we are introducing a id for each fill request allowing the
autofill provider to associate arbitrary state with each fill
request and store it in the client data bundle later passed
to save.

Another limitation of the current design is that if the screen
changes dynamically while the user interacts with the app the
UI state passed on save represents a static snapshot, therefore
it is not possible to the autofill provider to determine the
context in which the data in the UI was filled. We could
keep the views and have deltas for views being removed/added
/moved/changed but this is not enough as the fill provider
needs to know not only what changed but what changed for every
fill request and in one session there could be multiple fill
requests. To address this we provide a list of fill contexts
on save each of which has the id of the corresponding fill
request. This allows the fill provider to know the exact context
in which the data was popuplated and also use its custom client
state for this fill request if desired.

This change deprecates the old APIs and the new ones delegate
to the old ones. Once the clients migrate to the new APIs we
will remove the old ones.

Test: all autofill CTS tests pass

Change-Id: Idcebcc671aa3c078a305d8c358e225274fccc588
diff --git a/api/current.txt b/api/current.txt
index 501ca61..0b5a991 100644
--- a/api/current.txt
+++ b/api/current.txt
@@ -37030,8 +37030,10 @@
     method public final android.os.IBinder onBind(android.content.Intent);
     method public void onConnected();
     method public void onDisconnected();
-    method public abstract void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
-    method public abstract void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
+    method public void onFillRequest(android.service.autofill.FillRequest, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public abstract deprecated void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public void onSaveRequest(android.service.autofill.SaveRequest, android.service.autofill.SaveCallback);
+    method public abstract deprecated void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
     field public static final java.lang.String SERVICE_INTERFACE = "android.service.autofill.AutofillService";
     field public static final java.lang.String SERVICE_META_DATA = "android.autofill";
   }
@@ -37056,6 +37058,25 @@
     method public void onSuccess(android.service.autofill.FillResponse);
   }
 
+  public final class FillContext implements android.os.Parcelable {
+    method public int describeContents();
+    method public int getRequestId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillContext> CREATOR;
+  }
+
+  public final class FillRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public int getFlags();
+    method public int getId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillRequest> CREATOR;
+    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+  }
+
   public final class FillResponse implements android.os.Parcelable {
     method public int describeContents();
     method public void writeToParcel(android.os.Parcel, int);
@@ -37067,7 +37088,8 @@
     method public android.service.autofill.FillResponse.Builder addDataset(android.service.autofill.Dataset);
     method public android.service.autofill.FillResponse build();
     method public android.service.autofill.FillResponse.Builder setAuthentication(android.view.autofill.AutofillId[], android.content.IntentSender, android.widget.RemoteViews);
-    method public android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
+    method public android.service.autofill.FillResponse.Builder setClientState(android.os.Bundle);
+    method public deprecated android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
     method public android.service.autofill.FillResponse.Builder setSaveInfo(android.service.autofill.SaveInfo);
   }
 
@@ -37096,6 +37118,14 @@
     method public android.service.autofill.SaveInfo.Builder setOptionalIds(android.view.autofill.AutofillId[]);
   }
 
+  public final class SaveRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public java.util.List<android.service.autofill.FillContext> getFillContexts();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.SaveRequest> CREATOR;
+  }
+
 }
 
 package android.service.carrier {
@@ -47798,7 +47828,7 @@
     field public static final java.lang.String EXTRA_ASSIST_STRUCTURE = "android.view.autofill.extra.ASSIST_STRUCTURE";
     field public static final java.lang.String EXTRA_AUTHENTICATION_RESULT = "android.view.autofill.extra.AUTHENTICATION_RESULT";
     field public static final java.lang.String EXTRA_DATA_EXTRAS = "android.view.autofill.extra.DATA_EXTRAS";
-    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+    field public static final deprecated int FLAG_MANUAL_REQUEST = 1; // 0x1
   }
 
   public static abstract class AutofillManager.AutofillCallback {
diff --git a/api/system-current.txt b/api/system-current.txt
index 99a50d8..87a2fcf 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -40143,8 +40143,10 @@
     method public final android.os.IBinder onBind(android.content.Intent);
     method public void onConnected();
     method public void onDisconnected();
-    method public abstract void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
-    method public abstract void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
+    method public void onFillRequest(android.service.autofill.FillRequest, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public abstract deprecated void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public void onSaveRequest(android.service.autofill.SaveRequest, android.service.autofill.SaveCallback);
+    method public abstract deprecated void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
     field public static final java.lang.String SERVICE_INTERFACE = "android.service.autofill.AutofillService";
     field public static final java.lang.String SERVICE_META_DATA = "android.autofill";
   }
@@ -40169,6 +40171,25 @@
     method public void onSuccess(android.service.autofill.FillResponse);
   }
 
+  public final class FillContext implements android.os.Parcelable {
+    method public int describeContents();
+    method public int getRequestId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillContext> CREATOR;
+  }
+
+  public final class FillRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public int getFlags();
+    method public int getId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillRequest> CREATOR;
+    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+  }
+
   public final class FillResponse implements android.os.Parcelable {
     method public int describeContents();
     method public void writeToParcel(android.os.Parcel, int);
@@ -40180,7 +40201,8 @@
     method public android.service.autofill.FillResponse.Builder addDataset(android.service.autofill.Dataset);
     method public android.service.autofill.FillResponse build();
     method public android.service.autofill.FillResponse.Builder setAuthentication(android.view.autofill.AutofillId[], android.content.IntentSender, android.widget.RemoteViews);
-    method public android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
+    method public android.service.autofill.FillResponse.Builder setClientState(android.os.Bundle);
+    method public deprecated android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
     method public android.service.autofill.FillResponse.Builder setSaveInfo(android.service.autofill.SaveInfo);
   }
 
@@ -40209,6 +40231,14 @@
     method public android.service.autofill.SaveInfo.Builder setOptionalIds(android.view.autofill.AutofillId[]);
   }
 
+  public final class SaveRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public java.util.List<android.service.autofill.FillContext> getFillContexts();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.SaveRequest> CREATOR;
+  }
+
 }
 
 package android.service.carrier {
@@ -51367,7 +51397,7 @@
     field public static final java.lang.String EXTRA_ASSIST_STRUCTURE = "android.view.autofill.extra.ASSIST_STRUCTURE";
     field public static final java.lang.String EXTRA_AUTHENTICATION_RESULT = "android.view.autofill.extra.AUTHENTICATION_RESULT";
     field public static final java.lang.String EXTRA_DATA_EXTRAS = "android.view.autofill.extra.DATA_EXTRAS";
-    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+    field public static final deprecated int FLAG_MANUAL_REQUEST = 1; // 0x1
   }
 
   public static abstract class AutofillManager.AutofillCallback {
diff --git a/api/test-current.txt b/api/test-current.txt
index a113cc1..8698bbd 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -37183,8 +37183,10 @@
     method public final android.os.IBinder onBind(android.content.Intent);
     method public void onConnected();
     method public void onDisconnected();
-    method public abstract void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
-    method public abstract void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
+    method public void onFillRequest(android.service.autofill.FillRequest, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public abstract deprecated void onFillRequest(android.app.assist.AssistStructure, android.os.Bundle, int, android.os.CancellationSignal, android.service.autofill.FillCallback);
+    method public void onSaveRequest(android.service.autofill.SaveRequest, android.service.autofill.SaveCallback);
+    method public abstract deprecated void onSaveRequest(android.app.assist.AssistStructure, android.os.Bundle, android.service.autofill.SaveCallback);
     field public static final java.lang.String SERVICE_INTERFACE = "android.service.autofill.AutofillService";
     field public static final java.lang.String SERVICE_META_DATA = "android.autofill";
   }
@@ -37209,6 +37211,25 @@
     method public void onSuccess(android.service.autofill.FillResponse);
   }
 
+  public final class FillContext implements android.os.Parcelable {
+    method public int describeContents();
+    method public int getRequestId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillContext> CREATOR;
+  }
+
+  public final class FillRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public int getFlags();
+    method public int getId();
+    method public android.app.assist.AssistStructure getStructure();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.FillRequest> CREATOR;
+    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+  }
+
   public final class FillResponse implements android.os.Parcelable {
     method public int describeContents();
     method public void writeToParcel(android.os.Parcel, int);
@@ -37220,7 +37241,8 @@
     method public android.service.autofill.FillResponse.Builder addDataset(android.service.autofill.Dataset);
     method public android.service.autofill.FillResponse build();
     method public android.service.autofill.FillResponse.Builder setAuthentication(android.view.autofill.AutofillId[], android.content.IntentSender, android.widget.RemoteViews);
-    method public android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
+    method public android.service.autofill.FillResponse.Builder setClientState(android.os.Bundle);
+    method public deprecated android.service.autofill.FillResponse.Builder setExtras(android.os.Bundle);
     method public android.service.autofill.FillResponse.Builder setSaveInfo(android.service.autofill.SaveInfo);
   }
 
@@ -37249,6 +37271,14 @@
     method public android.service.autofill.SaveInfo.Builder setOptionalIds(android.view.autofill.AutofillId[]);
   }
 
+  public final class SaveRequest implements android.os.Parcelable {
+    method public int describeContents();
+    method public android.os.Bundle getClientState();
+    method public java.util.List<android.service.autofill.FillContext> getFillContexts();
+    method public void writeToParcel(android.os.Parcel, int);
+    field public static final android.os.Parcelable.Creator<android.service.autofill.SaveRequest> CREATOR;
+  }
+
 }
 
 package android.service.carrier {
@@ -48176,7 +48206,7 @@
     field public static final java.lang.String EXTRA_ASSIST_STRUCTURE = "android.view.autofill.extra.ASSIST_STRUCTURE";
     field public static final java.lang.String EXTRA_AUTHENTICATION_RESULT = "android.view.autofill.extra.AUTHENTICATION_RESULT";
     field public static final java.lang.String EXTRA_DATA_EXTRAS = "android.view.autofill.extra.DATA_EXTRAS";
-    field public static final int FLAG_MANUAL_REQUEST = 1; // 0x1
+    field public static final deprecated int FLAG_MANUAL_REQUEST = 1; // 0x1
   }
 
   public static abstract class AutofillManager.AutofillCallback {
diff --git a/core/java/android/service/autofill/AutofillService.java b/core/java/android/service/autofill/AutofillService.java
index a4d3fb2..5e49b8f 100644
--- a/core/java/android/service/autofill/AutofillService.java
+++ b/core/java/android/service/autofill/AutofillService.java
@@ -34,6 +34,8 @@
 
 import com.android.internal.os.SomeArgs;
 
+import java.util.List;
+
 //TODO(b/33197203): improve javadoc (of both class and methods); in particular, make sure the
 //life-cycle (and how state could be maintained on server-side) is well documented.
 
@@ -103,24 +105,22 @@
         }
 
         @Override
-        public void onFillRequest(AssistStructure structure, Bundle extras,
-                IFillCallback callback, int flags) {
+        public void onFillRequest(FillRequest request, IFillCallback callback) {
             ICancellationSignal transport = CancellationSignal.createTransport();
             try {
                 callback.onCancellable(transport);
             } catch (RemoteException e) {
                 e.rethrowFromSystemServer();
             }
-            mHandlerCaller.obtainMessageIIOOOO(MSG_ON_FILL_REQUEST, flags, UNUSED_ARG, structure,
-                    CancellationSignal.fromTransport(transport), extras, callback)
+            mHandlerCaller.obtainMessageOOO(MSG_ON_FILL_REQUEST, request,
+                    CancellationSignal.fromTransport(transport), callback)
                     .sendToTarget();
         }
 
         @Override
-        public void onSaveRequest(AssistStructure structure, Bundle extras,
-                ISaveCallback callback) {
-            mHandlerCaller.obtainMessageOOO(MSG_ON_SAVE_REQUEST, structure,
-                    extras, callback).sendToTarget();
+        public void onSaveRequest(SaveRequest request, ISaveCallback callback) {
+            mHandlerCaller.obtainMessageOO(MSG_ON_SAVE_REQUEST, request,
+                    callback).sendToTarget();
         }
     };
 
@@ -131,23 +131,20 @@
                 break;
             } case MSG_ON_FILL_REQUEST: {
                 final SomeArgs args = (SomeArgs) msg.obj;
-                final AssistStructure structure = (AssistStructure) args.arg1;
+                final FillRequest request = (FillRequest) args.arg1;
                 final CancellationSignal cancellation = (CancellationSignal) args.arg2;
-                final Bundle extras = (Bundle) args.arg3;
-                final IFillCallback callback = (IFillCallback) args.arg4;
-                final FillCallback fillCallback = new FillCallback(callback);
-                final int flags = msg.arg1;
+                final IFillCallback callback = (IFillCallback) args.arg3;
+                final FillCallback fillCallback = new FillCallback(callback, request.getId());
                 args.recycle();
-                onFillRequest(structure, extras, flags, cancellation, fillCallback);
+                onFillRequest(request, cancellation, fillCallback);
                 break;
             } case MSG_ON_SAVE_REQUEST: {
                 final SomeArgs args = (SomeArgs) msg.obj;
-                final AssistStructure structure = (AssistStructure) args.arg1;
-                final Bundle extras = (Bundle) args.arg2;
-                final ISaveCallback callback = (ISaveCallback) args.arg3;
+                final SaveRequest request = (SaveRequest) args.arg1;
+                final ISaveCallback callback = (ISaveCallback) args.arg2;
                 final SaveCallback saveCallback = new SaveCallback(callback);
                 args.recycle();
-                onSaveRequest(structure, extras, saveCallback);
+                onSaveRequest(request, saveCallback);
                 break;
             } case MSG_DISCONNECT: {
                 onDisconnected();
@@ -198,6 +195,28 @@
      * or {@link FillCallback#onFailure(CharSequence)})
      * to notify the result of the request.
      *
+     * @param request the {@link FillRequest request} to handle.
+     *        See {@link FillResponse} for examples of multiple-sections requests.
+     * @param cancellationSignal signal for observing cancellation requests. The system will use
+     *     this to notify you that the fill result is no longer needed and you should stop
+     *     handling this fill request in order to save resources.
+     * @param callback object used to notify the result of the request.
+     */
+    public void onFillRequest(@NonNull FillRequest request,
+            @NonNull CancellationSignal cancellationSignal, @NonNull FillCallback callback) {
+        onFillRequest(request.getStructure(), request.getClientState(), request.getFlags(),
+                cancellationSignal, callback);
+    }
+
+    /**
+     * Called by the Android system do decide if an {@link Activity} can be autofilled by the
+     * service.
+     *
+     * <p>Service must call one of the {@link FillCallback} methods (like
+     * {@link FillCallback#onSuccess(FillResponse)}
+     * or {@link FillCallback#onFailure(CharSequence)})
+     * to notify the result of the request.
+     *
      * @param structure {@link Activity}'s view structure.
      * @param data bundle containing data passed by the service in a last call to
      *        {@link FillResponse.Builder#setExtras(Bundle)}, if any. This bundle allows your
@@ -211,6 +230,7 @@
      *     handling this fill request in order to save resources.
      * @param callback object used to notify the result of the request.
      */
+    @Deprecated
     public abstract void onFillRequest(@NonNull AssistStructure structure, @Nullable Bundle data,
             int flags, @NonNull CancellationSignal cancellationSignal,
             @NonNull FillCallback callback);
@@ -222,6 +242,23 @@
      * {@link SaveCallback#onSuccess()} or {@link SaveCallback#onFailure(CharSequence)})
      * to notify the result of the request.
      *
+     * @param request the {@link SaveRequest request} to handle.
+     *        See {@link FillResponse} for examples of multiple-sections requests.
+     * @param callback object used to notify the result of the request.
+     */
+    public void onSaveRequest(@NonNull SaveRequest request, @NonNull SaveCallback callback) {
+        List<FillContext> contexts = request.getFillContexts();
+        onSaveRequest(contexts.get(contexts.size() - 1).getStructure(),
+                request.getClientState(), callback);
+    }
+
+    /**
+     * Called when user requests service to save the fields of an {@link Activity}.
+     *
+     * <p>Service must call one of the {@link SaveCallback} methods (like
+     * {@link SaveCallback#onSuccess()} or {@link SaveCallback#onFailure(CharSequence)})
+     * to notify the result of the request.
+     *
      * @param structure {@link Activity}'s view structure.
      * @param data bundle containing data passed by the service in a last call to
      *        {@link FillResponse.Builder#setExtras(Bundle)}, if any. This bundle allows your
@@ -231,6 +268,7 @@
      *        See {@link FillResponse} for examples of multiple-sections requests.
      * @param callback object used to notify the result of the request.
      */
+    @Deprecated
     public abstract void onSaveRequest(@NonNull AssistStructure structure, @Nullable Bundle data,
             @NonNull SaveCallback callback);
 
diff --git a/core/java/android/service/autofill/FillCallback.java b/core/java/android/service/autofill/FillCallback.java
index e8ad14f..a009be8 100644
--- a/core/java/android/service/autofill/FillCallback.java
+++ b/core/java/android/service/autofill/FillCallback.java
@@ -27,11 +27,13 @@
  */
 public final class FillCallback {
     private final IFillCallback mCallback;
+    private final int mRequestId;
     private boolean mCalled;
 
     /** @hide */
-    public FillCallback(IFillCallback callback) {
+    public FillCallback(IFillCallback callback, int requestId) {
         mCallback = callback;
+        mRequestId = requestId;
     }
 
     /**
@@ -47,7 +49,7 @@
         assertNotCalled();
         mCalled = true;
         try {
-            mCallback.onSuccess(response);
+            mCallback.onSuccess(response, mRequestId);
         } catch (RemoteException e) {
             e.rethrowAsRuntimeException();
         }
diff --git a/core/java/android/service/autofill/FillContext.java b/core/java/android/service/autofill/FillContext.java
new file mode 100644
index 0000000..2efa08c
--- /dev/null
+++ b/core/java/android/service/autofill/FillContext.java
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2017 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 android.service.autofill;
+
+import android.annotation.NonNull;
+import android.app.assist.AssistStructure;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.Parcel;
+import android.os.Parcelable;
+
+/**
+ * This class represents a context for each fill request made via {@link
+ * AutofillService#onFillRequest(FillRequest, CancellationSignal, FillCallback)}.
+ * It contains a snapshot of the UI state, the view ids that were returned by
+ * the {@link AutofillService autofill service} as both required to trigger a save
+ * and optional that can be saved, and the id of the corresponding {@link
+ * FillRequest}.
+ * <p>
+ * This context allows you to inspect the values for the interesting views
+ * in the context they appeared. Also a reference to the corresponding fill
+ * request is useful to store meta-data in the client state bundle passed
+ * to {@link FillResponse.Builder#setClientState(Bundle)} to avoid interpreting
+ * the UI state again while saving.
+ */
+public final class FillContext implements Parcelable {
+    private final int mRequestId;
+    private final @NonNull AssistStructure mStructure;
+
+    /** @hide */
+    public FillContext(int requestId, @NonNull AssistStructure structure) {
+        mRequestId = requestId;
+        mStructure = structure;
+    }
+
+    private FillContext(Parcel parcel) {
+        this(parcel.readInt(), parcel.readParcelable(null));
+    }
+
+    /**
+     * Gets the id of the {@link FillRequest fill request} this context
+     * corresponds to. This is useful to associate your custom client
+     * state with every request to avoid reinterpreting the UI when saving
+     * user data.
+     *
+     * @return The request id.
+     */
+    public int getRequestId() {
+        return mRequestId;
+    }
+
+    /**
+     * @return The screen content.
+     */
+    public AssistStructure getStructure() {
+        return mStructure;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel parcel, int flags) {
+        parcel.writeInt(mRequestId);
+        parcel.writeParcelable(mStructure, flags);
+    }
+
+    public static final Parcelable.Creator<FillContext> CREATOR =
+            new Parcelable.Creator<FillContext>() {
+        @Override
+        public FillContext createFromParcel(Parcel parcel) {
+            return new FillContext(parcel);
+        }
+
+        @Override
+        public FillContext[] newArray(int size) {
+            return new FillContext[size];
+        }
+    };
+}
diff --git a/core/java/android/service/autofill/FillRequest.aidl b/core/java/android/service/autofill/FillRequest.aidl
new file mode 100644
index 0000000..2b1a8fe
--- /dev/null
+++ b/core/java/android/service/autofill/FillRequest.aidl
@@ -0,0 +1,19 @@
+/**
+ * Copyright (c) 2017, 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 android.service.autofill;
+
+parcelable FillRequest;
diff --git a/core/java/android/service/autofill/FillRequest.java b/core/java/android/service/autofill/FillRequest.java
new file mode 100644
index 0000000..aa6db4d
--- /dev/null
+++ b/core/java/android/service/autofill/FillRequest.java
@@ -0,0 +1,144 @@
+/*
+ * Copyright (C) 2017 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 android.service.autofill;
+
+import android.annotation.IntDef;
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.assist.AssistStructure;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.Parcel;
+import android.os.Parcelable;
+import com.android.internal.util.Preconditions;
+
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * This class represents a request to an {@link AutofillService autofill provider}
+ * to interpret the screen and provide information to the system which views are
+ * interesting for saving and what are the possible ways to fill the inputs on
+ * the screen if applicable.
+ *
+ * @see AutofillService#onFillRequest(FillRequest, CancellationSignal, FillCallback)
+ */
+public final class FillRequest implements Parcelable {
+    private static AtomicInteger sIdCounter = new AtomicInteger();
+
+    /**
+     * Indicates autofill was explicitly requested by the user.
+     */
+    public static final int FLAG_MANUAL_REQUEST = 0x1;
+
+    /** @hide */
+    @IntDef(
+        flag = true,
+        value = {FLAG_MANUAL_REQUEST})
+    @Retention(RetentionPolicy.SOURCE)
+    @interface RequestFlags{}
+
+    private final int mId;
+    private final @RequestFlags int mFlags;
+    private final @NonNull AssistStructure mStructure;
+    private final @Nullable Bundle mClientState;
+
+    /** @hide */
+    public FillRequest(@NonNull AssistStructure structure,
+            @Nullable Bundle clientState, @RequestFlags int flags) {
+        this(sIdCounter.incrementAndGet(), structure, clientState, flags);
+    }
+
+    private FillRequest(@NonNull Parcel parcel) {
+        mId = parcel.readInt();
+        mStructure = parcel.readParcelable(null);
+        mClientState = parcel.readBundle();
+        mFlags = parcel.readInt();
+    }
+
+    private FillRequest(int id, @NonNull AssistStructure structure,
+            @Nullable Bundle clientState, @RequestFlags int flags) {
+        mId = id;
+        mFlags = Preconditions.checkFlagsArgument(flags, FLAG_MANUAL_REQUEST);
+        mStructure = Preconditions.checkNotNull(structure, "structure");
+        mClientState = clientState;
+    }
+
+    /**
+     * @return The unique id of this request.
+     */
+    public int getId() {
+        return mId;
+    }
+
+    /**
+     * @return The flags associated with this request.
+     *
+     * @see #FLAG_MANUAL_REQUEST
+     */
+    public @RequestFlags int getFlags() {
+        return mFlags;
+    }
+
+    /**
+     * @return The structure capturing the UI state.
+     */
+    public @NonNull AssistStructure getStructure() {
+        return mStructure;
+    }
+
+    /**
+     * Gets the extra client state returned from the last {@link
+     * AutofillService#onFillRequest(FillRequest, CancellationSignal, FillCallback)
+     * fill request}.
+     * <p>
+     * Once a {@link AutofillService#onSaveRequest(SaveRequest, SaveCallback)
+     * save request} is made the client state is cleared.
+     *
+     * @return The client state.
+     */
+    public @Nullable Bundle getClientState() {
+        return mClientState;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel parcel, int flags) {
+        parcel.writeInt(mId);
+        parcel.writeParcelable(mStructure, flags);
+        parcel.writeBundle(mClientState);
+        parcel.writeInt(mFlags);
+    }
+
+    public static final Parcelable.Creator<FillRequest> CREATOR =
+            new Parcelable.Creator<FillRequest>() {
+        @Override
+        public FillRequest createFromParcel(Parcel parcel) {
+            return new FillRequest(parcel);
+        }
+
+        @Override
+        public FillRequest[] newArray(int size) {
+            return new FillRequest[size];
+        }
+    };
+}
diff --git a/core/java/android/service/autofill/FillResponse.java b/core/java/android/service/autofill/FillResponse.java
index eab0d4c..8c8060a 100644
--- a/core/java/android/service/autofill/FillResponse.java
+++ b/core/java/android/service/autofill/FillResponse.java
@@ -132,25 +132,25 @@
  */
 public final class FillResponse implements Parcelable {
 
-    private final ArrayList<Dataset> mDatasets;
-    private final SaveInfo mSaveInfo;
-    private final Bundle mExtras;
-    private final RemoteViews mPresentation;
-    private final IntentSender mAuthentication;
-    private AutofillId[] mAuthenticationIds;
+    private final @Nullable ArrayList<Dataset> mDatasets;
+    private final @Nullable SaveInfo mSaveInfo;
+    private final @Nullable Bundle mClientState;
+    private final @Nullable RemoteViews mPresentation;
+    private final @Nullable IntentSender mAuthentication;
+    private final @Nullable AutofillId[] mAuthenticationIds;
 
     private FillResponse(@NonNull Builder builder) {
         mDatasets = builder.mDatasets;
         mSaveInfo = builder.mSaveInfo;
-        mExtras = builder.mExtras;
+        mClientState = builder.mCLientState;
         mPresentation = builder.mPresentation;
         mAuthentication = builder.mAuthentication;
         mAuthenticationIds = builder.mAuthenticationIds;
     }
 
     /** @hide */
-    public @Nullable Bundle getExtras() {
-        return mExtras;
+    public @Nullable Bundle getClientState() {
+        return mClientState;
     }
 
     /** @hide */
@@ -185,7 +185,7 @@
     public static final class Builder {
         private ArrayList<Dataset> mDatasets;
         private SaveInfo mSaveInfo;
-        private Bundle mExtras;
+        private Bundle mCLientState;
         private RemoteViews mPresentation;
         private IntentSender mAuthentication;
         private AutofillId[] mAuthenticationIds;
@@ -289,23 +289,35 @@
             return this;
         }
 
+        @Deprecated
+        public Builder setExtras(@Nullable Bundle extras) {
+            throwIfDestroyed();
+            mCLientState = extras;
+            return this;
+        }
+
         /**
-         * Sets a {@link Bundle} that will be passed to subsequent APIs that
+         * Sets a {@link Bundle state} that will be passed to subsequent APIs that
          * manipulate this response. For example, they are passed to subsequent
          * calls to {@link AutofillService#onFillRequest(
          * android.app.assist.AssistStructure, Bundle, int,
          * android.os.CancellationSignal, FillCallback)} and {@link AutofillService#onSaveRequest(
-         * android.app.assist.AssistStructure, Bundle, SaveCallback)}.
+         * android.app.assist.AssistStructure, Bundle, SaveCallback)}. You can use
+         * this to store intermediate state that is persistent across multiple
+         * fill requests and the subsequent save request.
          *
          * <p>If this method is called on multiple {@link FillResponse} objects for the same
          * activity, just the latest bundle is passed back to the service.
          *
-         * @param extras The response extras.
+         * <p>Once a {@link AutofillService#onSaveRequest(SaveRequest, SaveCallback)
+         * save request} is made the client state is cleared.
+         *
+         * @param clientState The custom client state.
          * @return This builder.
          */
-        public Builder setExtras(Bundle extras) {
+        public Builder setClientState(@Nullable Bundle clientState) {
             throwIfDestroyed();
-            mExtras = extras;
+            mCLientState = clientState;
             return this;
         }
 
@@ -344,7 +356,7 @@
         return new StringBuilder(
                 "FillResponse: [datasets=").append(mDatasets)
                 .append(", saveInfo=").append(mSaveInfo)
-                .append(", hasExtras=").append(mExtras != null)
+                .append(", clientState=").append(mClientState != null)
                 .append(", hasPresentation=").append(mPresentation != null)
                 .append(", hasAuthentication=").append(mAuthentication != null)
                 .append(", authenticationSize=").append(mAuthenticationIds != null
@@ -365,7 +377,7 @@
     public void writeToParcel(Parcel parcel, int flags) {
         parcel.writeTypedArrayList(mDatasets, flags);
         parcel.writeParcelable(mSaveInfo, flags);
-        parcel.writeParcelable(mExtras, flags);
+        parcel.writeParcelable(mClientState, flags);
         parcel.writeParcelableArray(mAuthenticationIds, flags);
         parcel.writeParcelable(mAuthentication, flags);
         parcel.writeParcelable(mPresentation, flags);
diff --git a/core/java/android/service/autofill/IAutoFillService.aidl b/core/java/android/service/autofill/IAutoFillService.aidl
index a8d86ca..23a1a3f 100644
--- a/core/java/android/service/autofill/IAutoFillService.aidl
+++ b/core/java/android/service/autofill/IAutoFillService.aidl
@@ -16,10 +16,10 @@
 
 package android.service.autofill;
 
-import android.app.assist.AssistStructure;
-import android.os.Bundle;
+import android.service.autofill.FillRequest;
 import android.service.autofill.IFillCallback;
 import android.service.autofill.ISaveCallback;
+import android.service.autofill.SaveRequest;
 import com.android.internal.os.IResultReceiver;
 
 /**
@@ -29,8 +29,6 @@
  */
 oneway interface IAutoFillService {
     void onConnectedStateChanged(boolean connected);
-    void onFillRequest(in AssistStructure structure, in Bundle extras,
-            in IFillCallback callback, int flags);
-    void onSaveRequest(in AssistStructure structure, in Bundle extras,
-            in ISaveCallback callback);
+    void onFillRequest(in FillRequest request, in IFillCallback callback);
+    void onSaveRequest(in SaveRequest request, in ISaveCallback callback);
 }
diff --git a/core/java/android/service/autofill/IFillCallback.aidl b/core/java/android/service/autofill/IFillCallback.aidl
index 2bb3e9a..688ac84 100644
--- a/core/java/android/service/autofill/IFillCallback.aidl
+++ b/core/java/android/service/autofill/IFillCallback.aidl
@@ -27,6 +27,6 @@
  */
 interface IFillCallback {
     void onCancellable(in ICancellationSignal cancellation);
-    void onSuccess(in FillResponse response);
+    void onSuccess(in FillResponse response, int requestId);
     void onFailure(CharSequence message);
 }
diff --git a/core/java/android/service/autofill/SaveInfo.java b/core/java/android/service/autofill/SaveInfo.java
index 4ad0f08..258d257 100644
--- a/core/java/android/service/autofill/SaveInfo.java
+++ b/core/java/android/service/autofill/SaveInfo.java
@@ -140,7 +140,19 @@
      */
     public static final int SAVE_DATA_TYPE_EMAIL_ADDRESS = 0x10;
 
-    private final int mType;
+    /** @hide */
+    @IntDef(
+       flag = true,
+       value = {
+               SAVE_DATA_TYPE_GENERIC,
+               SAVE_DATA_TYPE_PASSWORD,
+               SAVE_DATA_TYPE_ADDRESS,
+               SAVE_DATA_TYPE_CREDIT_CARD,
+               SAVE_DATA_TYPE_EMAIL_ADDRESS})
+    @Retention(RetentionPolicy.SOURCE)
+    @interface SaveDataType{}
+
+    private final @SaveDataType int mType;
     private final CharSequence mNegativeActionTitle;
     private final IntentSender mNegativeActionListener;
     private final AutofillId[] mRequiredIds;
@@ -177,7 +189,7 @@
     }
 
     /** @hide */
-    public int getType() {
+    public @SaveDataType int getType() {
         return mType;
     }
 
@@ -191,7 +203,7 @@
      */
     public static final class Builder {
 
-        private final int mType;
+        private final @SaveDataType int mType;
         private CharSequence mNegativeActionTitle;
         private IntentSender mNegativeActionListener;
         // TODO(b/33197203): make mRequiredIds final once addSavableIds() is gone
@@ -215,7 +227,7 @@
          *
          * @throws IllegalArgumentException if {@code requiredIds} is {@code null} or empty.
          */
-        public Builder(int type, @NonNull AutofillId[] requiredIds) {
+        public Builder(@SaveDataType int type, @NonNull AutofillId[] requiredIds) {
             if (false) {// TODO(b/33197203): re-move when clients use it
             Preconditions.checkArgument(requiredIds != null && requiredIds.length > 0,
                     "must have at least one required id: " + Arrays.toString(requiredIds));
@@ -230,7 +242,7 @@
          * // TODO(b/33197203): make sure is removed when clients migrated
          */
         @Deprecated
-        public Builder(int type) {
+        public Builder(@SaveDataType int type) {
             this(type, null);
         }
 
diff --git a/core/java/android/service/autofill/SaveRequest.aidl b/core/java/android/service/autofill/SaveRequest.aidl
new file mode 100644
index 0000000..7789b577
--- /dev/null
+++ b/core/java/android/service/autofill/SaveRequest.aidl
@@ -0,0 +1,19 @@
+/**
+ * Copyright (c) 2017, 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 android.service.autofill;
+
+parcelable SaveRequest;
diff --git a/core/java/android/service/autofill/SaveRequest.java b/core/java/android/service/autofill/SaveRequest.java
new file mode 100644
index 0000000..9de9315
--- /dev/null
+++ b/core/java/android/service/autofill/SaveRequest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2017 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 android.service.autofill;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.os.Bundle;
+import android.os.CancellationSignal;
+import android.os.Parcel;
+import android.os.Parcelable;
+import com.android.internal.util.Preconditions;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * This class represents a request to an {@link AutofillService
+ * autofill provider} to save applicable data entered by the user.
+ *
+ * @see AutofillService#onSaveRequest(SaveRequest, SaveCallback)
+ */
+public final class SaveRequest implements Parcelable {
+    private final @NonNull ArrayList<FillContext> mFillContexts;
+    private final @Nullable Bundle mClientState;
+
+    /** @hide */
+    public SaveRequest(@NonNull ArrayList<FillContext> fillContexts,
+            @Nullable Bundle clientState) {
+        mFillContexts = Preconditions.checkNotNull(fillContexts, "fillContexts");
+        mClientState = clientState;
+    }
+
+    private SaveRequest(@NonNull Parcel parcel) {
+        this(parcel.readTypedArrayList(null), parcel.readBundle());
+    }
+
+    /**
+     * @return The contexts associated with each previous fill request.
+     */
+    public @NonNull List<FillContext> getFillContexts() {
+        return mFillContexts;
+    }
+
+    /**
+     * Gets the extra client state returned from the last {@link
+     * AutofillService#onFillRequest(FillRequest, CancellationSignal, FillCallback)}
+     * fill request}.
+     *
+     * @return The client state.
+     */
+    public @Nullable Bundle getClientState() {
+        return mClientState;
+    }
+
+    @Override
+    public int describeContents() {
+        return 0;
+    }
+
+    @Override
+    public void writeToParcel(Parcel parcel, int flags) {
+        parcel.writeTypedArrayList(mFillContexts, flags);
+        parcel.writeBundle(mClientState);
+    }
+
+    public static final Creator<SaveRequest> CREATOR =
+            new Creator<SaveRequest>() {
+        @Override
+        public SaveRequest createFromParcel(Parcel parcel) {
+            return new SaveRequest(parcel);
+        }
+
+        @Override
+        public SaveRequest[] newArray(int size) {
+            return new SaveRequest[size];
+        }
+    };
+}
diff --git a/core/java/android/view/autofill/AutofillManager.java b/core/java/android/view/autofill/AutofillManager.java
index 41c209c..327b242 100644
--- a/core/java/android/view/autofill/AutofillManager.java
+++ b/core/java/android/view/autofill/AutofillManager.java
@@ -101,7 +101,10 @@
     // Public flags start from the lowest bit
     /**
      * Indicates autofill was explicitly requested by the user.
+     *
+     * @deprecated Use {@link android.service.autofill.FillRequest#FLAG_MANUAL_REQUEST}
      */
+    @Deprecated
     public static final int FLAG_MANUAL_REQUEST = 0x1;
 
     // Private flags start from the highest bit
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
index 63bf373..57ef6d0 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
@@ -48,6 +48,7 @@
 import android.provider.Settings;
 import android.service.autofill.AutofillService;
 import android.service.autofill.AutofillServiceInfo;
+import android.service.autofill.FillRequest;
 import android.service.autofill.IAutoFillService;
 import android.text.TextUtils;
 import android.util.LocalLog;
@@ -169,7 +170,8 @@
             structure.sanitizeForParceling(true);
 
             // TODO(b/33197203): Need to pipe the bundle
-            session.mRemoteFillService.onFillRequest(structure, null, session.mFlags);
+            FillRequest request = new FillRequest(structure, null, session.mFlags);
+            session.mRemoteFillService.onFillRequest(request);
         }
     };
 
diff --git a/services/autofill/java/com/android/server/autofill/RemoteFillService.java b/services/autofill/java/com/android/server/autofill/RemoteFillService.java
index dd520ac..3badcfc 100644
--- a/services/autofill/java/com/android/server/autofill/RemoteFillService.java
+++ b/services/autofill/java/com/android/server/autofill/RemoteFillService.java
@@ -20,12 +20,10 @@
 
 import android.annotation.NonNull;
 import android.annotation.Nullable;
-import android.app.assist.AssistStructure;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.ServiceConnection;
-import android.os.Bundle;
 import android.os.IBinder;
 import android.os.IBinder.DeathRecipient;
 import android.os.ICancellationSignal;
@@ -33,10 +31,12 @@
 import android.os.RemoteException;
 import android.os.UserHandle;
 import android.service.autofill.AutofillService;
+import android.service.autofill.FillRequest;
 import android.service.autofill.FillResponse;
 import android.service.autofill.IAutoFillService;
 import android.service.autofill.IFillCallback;
 import android.service.autofill.ISaveCallback;
+import android.service.autofill.SaveRequest;
 import android.text.format.DateUtils;
 import android.util.Slog;
 
@@ -88,7 +88,7 @@
 
     public interface FillServiceCallbacks {
         void onFillRequestSuccess(@Nullable FillResponse response,
-                @NonNull String servicePackageName);
+                @NonNull String servicePackageName, int requestId);
         void onFillRequestFailure(@Nullable CharSequence message,
                 @NonNull String servicePackageName);
         void onSaveRequestSuccess(@NonNull String servicePackageName);
@@ -134,17 +134,16 @@
         mCallbacks.onServiceDied(this);
     }
 
-    public void onFillRequest(@NonNull AssistStructure structure, @Nullable Bundle extras,
-            int flags) {
+    public void onFillRequest(@NonNull FillRequest request) {
         cancelScheduledUnbind();
-        final PendingFillRequest request = new PendingFillRequest(structure, extras, this, flags);
-        mHandler.obtainMessageO(MyHandler.MSG_ON_PENDING_REQUEST, request).sendToTarget();
+        final PendingFillRequest pendingRequest = new PendingFillRequest(request, this);
+        mHandler.obtainMessageO(MyHandler.MSG_ON_PENDING_REQUEST, pendingRequest).sendToTarget();
     }
 
-    public void onSaveRequest(@NonNull AssistStructure structure, @Nullable Bundle extras) {
+    public void onSaveRequest(@NonNull SaveRequest request) {
         cancelScheduledUnbind();
-        final PendingSaveRequest request = new PendingSaveRequest(structure, extras, this);
-        mHandler.obtainMessageO(MyHandler.MSG_ON_PENDING_REQUEST, request).sendToTarget();
+        final PendingSaveRequest pendingRequest = new PendingSaveRequest(request, this);
+        mHandler.obtainMessageO(MyHandler.MSG_ON_PENDING_REQUEST, pendingRequest).sendToTarget();
     }
 
     // Note: we are dumping without a lock held so this is a bit racy but
@@ -253,10 +252,11 @@
     }
 
     private void dispatchOnFillRequestSuccess(PendingRequest pendingRequest,
-            FillResponse response) {
+            FillResponse response, int requestId) {
         mHandler.getHandler().post(() -> {
             if (handleResponseCallbackCommon(pendingRequest)) {
-                mCallbacks.onFillRequestSuccess(response, mComponentName.getPackageName());
+                mCallbacks.onFillRequestSuccess(response, mComponentName.getPackageName(),
+                        requestId);
             }
         });
     }
@@ -392,18 +392,13 @@
     private static final class PendingFillRequest extends PendingRequest {
         private final Object mLock = new Object();
         private final WeakReference<RemoteFillService> mWeakService;
-        private final AssistStructure mStructure;
-        private final Bundle mExtras;
+        private final FillRequest mRequest;
         private final IFillCallback mCallback;
         private ICancellationSignal mCancellation;
         private boolean mCancelled;
-        private int mFlags;
 
-        public PendingFillRequest(AssistStructure structure,
-                Bundle extras, RemoteFillService service, int flags) {
-            mStructure = structure;
-            mExtras = extras;
-            mFlags = flags;
+        public PendingFillRequest(FillRequest request, RemoteFillService service) {
+            mRequest = request;
             mWeakService = new WeakReference<>(service);
             mCallback = new IFillCallback.Stub() {
                 @Override
@@ -425,11 +420,11 @@
                 }
 
                 @Override
-                public void onSuccess(FillResponse response) {
+                public void onSuccess(FillResponse response, int requestId) {
                     RemoteFillService remoteService = mWeakService.get();
                     if (remoteService != null) {
                         remoteService.dispatchOnFillRequestSuccess(
-                                PendingFillRequest.this, response);
+                                PendingFillRequest.this, response, requestId);
                     }
                 }
 
@@ -449,8 +444,7 @@
             RemoteFillService remoteService = mWeakService.get();
             if (remoteService != null) {
                 try {
-                    remoteService.mAutoFillService.onFillRequest(mStructure,
-                            mExtras, mCallback, mFlags);
+                    remoteService.mAutoFillService.onFillRequest(mRequest, mCallback);
                 } catch (RemoteException e) {
                     Slog.e(LOG_TAG, "Error calling on fill request", e);
                     cancel();
@@ -481,14 +475,12 @@
 
     private static final class PendingSaveRequest extends PendingRequest {
         private final WeakReference<RemoteFillService> mWeakService;
-        private final AssistStructure mStructure;
-        private final Bundle mExtras;
+        private final SaveRequest mRequest;
         private final ISaveCallback mCallback;
 
-        public PendingSaveRequest(@NonNull AssistStructure structure, @Nullable Bundle extras,
+        public PendingSaveRequest(@NonNull SaveRequest request,
                 @NonNull RemoteFillService service) {
-            mStructure = structure;
-            mExtras = extras;
+            mRequest = request;
             mWeakService = new WeakReference<>(service);
             mCallback = new ISaveCallback.Stub() {
                 @Override
@@ -516,7 +508,7 @@
             final RemoteFillService service = mWeakService.get();
             if (service != null) {
                 try {
-                    service.mAutoFillService.onSaveRequest(mStructure, mExtras, mCallback);
+                    service.mAutoFillService.onSaveRequest(mRequest, mCallback);
                 } catch (RemoteException e) {
                     Slog.e(LOG_TAG, "Error calling on save request", e);
                 }
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 67b5282..1df34ea 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -44,11 +44,15 @@
 import android.os.RemoteException;
 import android.service.autofill.AutofillService;
 import android.service.autofill.Dataset;
+import android.service.autofill.FillContext;
+import android.service.autofill.FillRequest;
 import android.service.autofill.FillResponse;
 import android.service.autofill.SaveInfo;
+import android.service.autofill.SaveRequest;
 import android.util.ArrayMap;
 import android.util.DebugUtils;
 import android.util.Slog;
+import android.util.SparseArray;
 import android.view.autofill.AutofillId;
 import android.view.autofill.AutofillManager;
 import android.view.autofill.AutofillValue;
@@ -125,7 +129,7 @@
     RemoteFillService mRemoteFillService;
 
     @GuardedBy("mLock")
-    private ArrayList<FillResponse> mResponses;
+    private SparseArray<FillResponse> mResponses;
 
     /**
      * Response that requires a service authentitcation request.
@@ -156,7 +160,7 @@
      * and used on subsequent {@code onFillRequest()} and {@code onSaveRequest()} calls.
      */
     @GuardedBy("mLock")
-    private Bundle mExtras;
+    private Bundle mClientState;
 
     /**
      * Flags used to start the session.
@@ -222,7 +226,7 @@
     // FillServiceCallbacks
     @Override
     public void onFillRequestSuccess(@Nullable FillResponse response,
-            @NonNull String servicePackageName) {
+            @NonNull String servicePackageName, int requestId) {
         if (response == null) {
             if ((mFlags & FLAG_MANUAL_REQUEST) != 0) {
                 getUiForShowing().showError(R.string.autofill_error_cannot_autofill);
@@ -243,7 +247,7 @@
                 // TODO(b/33197203 , b/35707731): make sure it's ignored if there is one already
                 mResponseWaitingAuth = response;
             }
-            processResponseLocked(response);
+            processResponseLocked(response, requestId);
         }
 
         final LogMaker log = (new LogMaker(MetricsEvent.AUTOFILL_REQUEST))
@@ -395,12 +399,18 @@
                     AutofillManager.EXTRA_AUTHENTICATION_RESULT);
             if (result instanceof FillResponse) {
                 mMetricsLogger.action(MetricsEvent.AUTOFILL_AUTHENTICATED, mPackageName);
+                final int requestIndex = mResponses.indexOfValue(mResponseWaitingAuth);
                 mResponseWaitingAuth = null;
-                processResponseLocked((FillResponse) result);
+                if (requestIndex >= 0) {
+                    final int requestId = mResponses.keyAt(requestIndex);
+                    processResponseLocked((FillResponse) result, requestId);
+                } else {
+                    Slog.e(TAG, "Error cannot find id for auth response");
+                }
             } else if (result instanceof Dataset) {
                 final Dataset dataset = (Dataset) result;
                 for (int i = 0; i < mResponses.size(); i++) {
-                    final FillResponse response = mResponses.get(i);
+                    final FillResponse response = mResponses.valueAt(i);
                     final int index = response.getDatasets().indexOf(mDatasetWaitingAuth);
                     if (index >= 0) {
                         response.getDatasets().set(index, dataset);
@@ -441,12 +451,18 @@
             return true;
         }
 
-        final FillResponse response = mResponses.get(mResponses.size() - 1);
+        final int lastResponseIdx = getLastResponseIndex();
+        if (lastResponseIdx < 0) {
+            Slog.d(TAG, "showSaveLocked(): mResponses=" + mResponses
+                    + ", mViewStates=" + mViewStates);
+            return true;
+        }
 
+        final FillResponse response = mResponses.valueAt(lastResponseIdx);
         final SaveInfo saveInfo = response.getSaveInfo();
         if (DEBUG) {
-            Slog.d(TAG,
-                    "showSaveLocked(): mResponses=" + mResponses + ", mViewStates=" + mViewStates);
+            Slog.d(TAG, "showSaveLocked(): mResponses=" + mResponses
+                    + ", mViewStates=" + mViewStates);
         }
 
         /*
@@ -573,7 +589,15 @@
             mStructure.dump();
         }
 
-        mRemoteFillService.onSaveRequest(mStructure, mExtras);
+        // TODO(b/33197203): Implement partitioning properly
+        final int lastResponseIdx = getLastResponseIndex();
+        final int requestId = mResponses.keyAt(lastResponseIdx);
+        final FillContext fillContext = new FillContext(requestId, mStructure);
+        final ArrayList fillContexts = new ArrayList(1);
+        fillContexts.add(fillContext);
+
+        final SaveRequest saveRequest = new SaveRequest(fillContexts, mClientState);
+        mRemoteFillService.onSaveRequest(saveRequest);
     }
 
     void updateLocked(AutofillId id, Rect virtualBounds, AutofillValue value, int flags) {
@@ -687,7 +711,9 @@
             overlay.focused = id.equals(viewState.id);
             node.setAutofillOverlay(overlay);
         }
-        mRemoteFillService.onFillRequest(mStructure, mExtras, 0);
+
+        FillRequest request = new FillRequest(mStructure, mClientState, 0);
+        mRemoteFillService.onFillRequest(request);
 
         return newViewState;
     }
@@ -724,17 +750,17 @@
         }
     }
 
-    private void processResponseLocked(FillResponse response) {
+    private void processResponseLocked(FillResponse response, int requestId) {
         if (DEBUG) {
             Slog.d(TAG, "processResponseLocked(mCurrentViewId=" + mCurrentViewId + "):" + response);
         }
 
         if (mResponses == null) {
-            mResponses = new ArrayList<>(4);
+            mResponses = new SparseArray<>(4);
         }
-        mResponses.add(response);
+        mResponses.put(requestId, response);
         if (response != null) {
-            mExtras = response.getExtras();
+            mClientState = response.getClientState();
         }
 
         setViewStatesLocked(response, ViewState.STATE_FILLABLE);
@@ -885,7 +911,8 @@
             }
         }
         pw.print(prefix); pw.print("mHasCallback: "); pw.println(mHasCallback);
-        pw.print(prefix); pw.print("mExtras: "); pw.println(Helper.bundleToString(mExtras));
+        pw.print(prefix); pw.print("mClientState: "); pw.println(
+                Helper.bundleToString(mClientState));
         mRemoteFillService.dump(prefix, pw);
     }
 
@@ -962,4 +989,20 @@
         destroyLocked();
         mService.removeSessionLocked(id);
     }
+
+    private int getLastResponseIndex() {
+        // The response ids are monotonically increasing so
+        // we just find the largest id which is the last. We
+        // do not rely on the internal ordering in sparse
+        // array to avoid - wow this stopped working!?
+        int lastResponseIdx = -1;
+        int lastResponseId = -1;
+        final int responseCount = mResponses.size();
+        for (int i = 0; i < responseCount; i++) {
+            if (mResponses.keyAt(i) > lastResponseId) {
+                lastResponseIdx = i;
+            }
+        }
+        return lastResponseIdx;
+    }
 }
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 4449da9..ab6a3a7 100644
--- a/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
+++ b/services/autofill/java/com/android/server/autofill/ui/AutoFillUI.java
@@ -161,7 +161,8 @@
                     log.setType(MetricsProto.MetricsEvent.TYPE_DETAIL);
                     hideFillUiUiThread();
                     if (mCallback != null) {
-                        mCallback.authenticate(response.getAuthentication(), response.getExtras());
+                        mCallback.authenticate(response.getAuthentication(),
+                                response.getClientState());
                     }
                 }