Pass FillResponse bundle on subsequent onFillRequest() calls.

BUG: 35707731

Test: existing CtsAutoFillServiceTestCases pass
Test: more tests on PartitionedActivityTest

Change-Id: I510aff54d0d9d7e0cce90f9786e3e1f6cab180f9
diff --git a/core/java/android/service/autofill/AutofillService.java b/core/java/android/service/autofill/AutofillService.java
index 9f8a621..ef14095 100644
--- a/core/java/android/service/autofill/AutofillService.java
+++ b/core/java/android/service/autofill/AutofillService.java
@@ -204,11 +204,12 @@
      * to notify the result of the request.
      *
      * @param structure {@link Activity}'s view structure.
-     * @param data bundle containing data passed by the service on previous calls to fill.
-     *     This bundle allows your service to keep state between fill and save requests
-     *     as well as when filling different sections of the UI as the system will try to
-     *     aggressively unbind from the service to conserve resources. See {@link
-     *     FillResponse} Javadoc for examples of multiple-sections requests.
+     * @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
+     *        service to keep state between fill and save requests as well as when filling different
+     *        sections of the UI as the system will try to aggressively unbind from the service to
+     *        conserve resources.
+     *        See {@link FillResponse} for examples of multiple-sections requests.
      * @param flags either {@code 0} or {@link AutofillManager#FLAG_MANUAL_REQUEST}.
      * @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
@@ -242,11 +243,12 @@
      * to notify the result of the request.
      *
      * @param structure {@link Activity}'s view structure.
-     * @param data bundle containing data passed by the service on previous calls to fill.
-     *     This bundle allows your service to keep state between fill and save requests
-     *     as well as when filling different sections of the UI as the system will try to
-     *     aggressively unbind from the service to conserve resources. See {@link
-     *     FillResponse} Javadoc for examples of multiple-sections requests.
+     * @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
+     *        service to keep state between fill and save requests as well as when filling different
+     *        sections of the UI as the system will try to aggressively unbind from the service to
+     *        conserve resources.
+     *        See {@link FillResponse} for examples of multiple-sections requests.
      * @param callback object used to notify the result of the request.
      */
     public abstract void onSaveRequest(@NonNull AssistStructure structure, @Nullable Bundle data,
diff --git a/core/java/android/service/autofill/FillResponse.java b/core/java/android/service/autofill/FillResponse.java
index 717312b..3117f98 100644
--- a/core/java/android/service/autofill/FillResponse.java
+++ b/core/java/android/service/autofill/FillResponse.java
@@ -297,6 +297,9 @@
          * android.os.CancellationSignal, FillCallback)} and {@link AutofillService#onSaveRequest(
          * android.app.assist.AssistStructure, Bundle, SaveCallback)}.
          *
+         * <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.
          * @return This builder.
          */
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
index f8b8e76..c8a5780 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
@@ -148,7 +148,7 @@
                 // TODO(b/33197203): since service is fetching the data (to use for save later),
                 // we should optimize what's sent (for example, remove layout containers,
                 // color / font info, etc...)
-                session.mStructure = structure;
+                session.setStructureLocked(structure);
             }
 
 
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index 1ffc82f..a3da50e 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -133,7 +133,8 @@
      * Assist structure sent by the app; it will be updated (sanitized, change values for save)
      * before sent to {@link AutofillService}.
      */
-    @GuardedBy("mLock") AssistStructure mStructure;
+    @GuardedBy("mLock")
+    private AssistStructure mStructure;
 
     /**
      * Whether the client has an {@link android.view.autofill.AutofillManager.AutofillCallback}.
@@ -141,6 +142,13 @@
     private boolean mHasCallback;
 
     /**
+     * Extras sent by service on {@code onFillRequest()} calls; the first non-null extra is saved
+     * and used on subsequent {@code onFillRequest()} and {@code onSaveRequest()} calls.
+     */
+    @GuardedBy("mLock")
+    private Bundle mExtras;
+
+    /**
      * Flags used to start the session.
      */
     int mFlags;
@@ -353,6 +361,10 @@
         mHasCallback = hasIt;
     }
 
+    public void setStructureLocked(AssistStructure structure) {
+        mStructure = structure;
+    }
+
     /**
      * Shows the save UI, when session can be saved.
      *
@@ -475,9 +487,6 @@
             Slog.d(TAG, "callSaveLocked(): mViewStates=" + mViewStates);
         }
 
-        // TODO(b/33197203 , b/35707731): decide how to handle bundle in multiple partitions
-        final Bundle extras = mResponses != null ? mResponses.get(0).getExtras() : null;
-
         for (Entry<AutofillId, ViewState> entry : mViewStates.entrySet()) {
             final AutofillValue value = entry.getValue().getCurrentValue();
             if (value == null) {
@@ -507,7 +516,7 @@
             mStructure.dump();
         }
 
-        mRemoteFillService.onSaveRequest(mStructure, extras);
+        mRemoteFillService.onSaveRequest(mStructure, mExtras);
     }
 
     void updateLocked(AutofillId id, Rect virtualBounds, AutofillValue value, int flags) {
@@ -593,13 +602,6 @@
                 new ViewState(this, id, value, this,ViewState.STATE_STARTED_PARTITION);
         mViewStates.put(id, newViewState);
 
-        /*
-         * TODO(b/33197203 , b/35707731): when start a new partition, it should
-         *
-         * - pass the first onFillRequest() bundle
-         * - optional: perhaps add a new flag onFilLRequest() to indicate it's a new partition?
-         */
-
         // Must update value of nodes so:
         // - proper node is focused
         // - autofillValue is sent back to service when it was previously autofilled
@@ -621,7 +623,7 @@
             overlay.focused = id.equals(viewState.id);
             node.setAutofillOverlay(overlay);
         }
-        mRemoteFillService.onFillRequest(mStructure, null, 0);
+        mRemoteFillService.onFillRequest(mStructure, mExtras, 0);
 
         return newViewState;
     }
@@ -671,6 +673,9 @@
             mResponses = new ArrayList<>(4);
         }
         mResponses.add(response);
+        if (response != null) {
+            mExtras = response.getExtras();
+        }
 
         setViewStatesLocked(response, ViewState.STATE_FILLABLE);
 
@@ -812,6 +817,7 @@
             }
         }
         pw.print(prefix); pw.print("mHasCallback: "); pw.println(mHasCallback);
+        pw.print(prefix); pw.print("mExtras: "); pw.println(Helper.bundleToString(mExtras));
         mRemoteFillService.dump(prefix, pw);
     }