diff --git a/core/java/android/service/autofill/augmented/AugmentedAutofillService.java b/core/java/android/service/autofill/augmented/AugmentedAutofillService.java
index 1efa3dd0c8655..b47b2b41bdd90 100644
--- a/core/java/android/service/autofill/augmented/AugmentedAutofillService.java
+++ b/core/java/android/service/autofill/augmented/AugmentedAutofillService.java
@@ -31,6 +31,7 @@ import android.content.ComponentName;
 import android.content.Intent;
 import android.graphics.Rect;
 import android.os.Build;
+import android.os.Bundle;
 import android.os.CancellationSignal;
 import android.os.Handler;
 import android.os.IBinder;
@@ -558,9 +559,10 @@ public abstract class AugmentedAutofillService extends Service {
             }
         }
 
-        void reportResult(@Nullable List<Dataset> inlineSuggestionsData) {
+        void reportResult(@Nullable List<Dataset> inlineSuggestionsData,
+                @Nullable Bundle clientState) {
             try {
-                mCallback.onSuccess(inlineSuggestionsData);
+                mCallback.onSuccess(inlineSuggestionsData, clientState);
             } catch (RemoteException e) {
                 Log.e(TAG, "Error calling back with the inline suggestions data: " + e);
             }
diff --git a/core/java/android/service/autofill/augmented/FillCallback.java b/core/java/android/service/autofill/augmented/FillCallback.java
index 526a749c95a10..21738d80f2d32 100644
--- a/core/java/android/service/autofill/augmented/FillCallback.java
+++ b/core/java/android/service/autofill/augmented/FillCallback.java
@@ -21,6 +21,7 @@ import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
+import android.os.Bundle;
 import android.service.autofill.Dataset;
 import android.service.autofill.augmented.AugmentedAutofillService.AutofillProxy;
 import android.util.Log;
@@ -55,14 +56,15 @@ public final class FillCallback {
 
         if (response == null) {
             mProxy.logEvent(AutofillProxy.REPORT_EVENT_NO_RESPONSE);
-            mProxy.reportResult(/* inlineSuggestionsData */ null);
+            mProxy.reportResult(/* inlineSuggestionsData */ null, /* clientState */ null);
             return;
         }
 
         List<Dataset> inlineSuggestions = response.getInlineSuggestions();
+        Bundle clientState = response.getClientState();
         if (inlineSuggestions != null && !inlineSuggestions.isEmpty()) {
             mProxy.logEvent(AutofillProxy.REPORT_EVENT_INLINE_RESPONSE);
-            mProxy.reportResult(inlineSuggestions);
+            mProxy.reportResult(inlineSuggestions, clientState);
             return;
         }
 
diff --git a/core/java/android/service/autofill/augmented/IFillCallback.aidl b/core/java/android/service/autofill/augmented/IFillCallback.aidl
index 24af1e51dd563..609e382e2b962 100644
--- a/core/java/android/service/autofill/augmented/IFillCallback.aidl
+++ b/core/java/android/service/autofill/augmented/IFillCallback.aidl
@@ -30,7 +30,7 @@ import java.util.List;
  */
 interface IFillCallback {
     void onCancellable(in ICancellationSignal cancellation);
-    void onSuccess(in @nullable List<Dataset> inlineSuggestionsData);
+    void onSuccess(in @nullable List<Dataset> inlineSuggestionsData, in @nullable Bundle clientState);
     boolean isCompleted();
     void cancel();
 }
diff --git a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
index d1805d96cad89..1bc026cd3b197 100644
--- a/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
+++ b/services/autofill/java/com/android/server/autofill/AutofillManagerServiceImpl.java
@@ -815,26 +815,27 @@ final class AutofillManagerServiceImpl
         }
     }
 
-    void logAugmentedAutofillSelected(int sessionId, @Nullable String suggestionId) {
+    void logAugmentedAutofillSelected(int sessionId, @Nullable String suggestionId,
+            @Nullable Bundle clientState) {
         synchronized (mLock) {
             if (mAugmentedAutofillEventHistory == null
                     || mAugmentedAutofillEventHistory.getSessionId() != sessionId) {
                 return;
             }
             mAugmentedAutofillEventHistory.addEvent(
-                    new Event(Event.TYPE_DATASET_SELECTED, suggestionId, null, null, null,
+                    new Event(Event.TYPE_DATASET_SELECTED, suggestionId, clientState, null, null,
                             null, null, null, null, null, null));
         }
     }
 
-    void logAugmentedAutofillShown(int sessionId) {
+    void logAugmentedAutofillShown(int sessionId, @Nullable Bundle clientState) {
         synchronized (mLock) {
             if (mAugmentedAutofillEventHistory == null
                     || mAugmentedAutofillEventHistory.getSessionId() != sessionId) {
                 return;
             }
             mAugmentedAutofillEventHistory.addEvent(
-                    new Event(Event.TYPE_DATASETS_SHOWN, null, null, null, null, null,
+                    new Event(Event.TYPE_DATASETS_SHOWN, null, clientState, null, null, null,
                             null, null, null, null, null));
 
         }
@@ -1185,15 +1186,16 @@ final class AutofillManagerServiceImpl
                         }
 
                         @Override
-                        public void logAugmentedAutofillShown(int sessionId) {
-                            AutofillManagerServiceImpl.this.logAugmentedAutofillShown(sessionId);
+                        public void logAugmentedAutofillShown(int sessionId, Bundle clientState) {
+                            AutofillManagerServiceImpl.this.logAugmentedAutofillShown(sessionId,
+                                    clientState);
                         }
 
                         @Override
                         public void logAugmentedAutofillSelected(int sessionId,
-                                String suggestionId) {
+                                String suggestionId, Bundle clientState) {
                             AutofillManagerServiceImpl.this.logAugmentedAutofillSelected(sessionId,
-                                    suggestionId);
+                                    suggestionId, clientState);
                         }
 
                         @Override
diff --git a/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java b/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
index b6bc7c5646a53..6cec8d82f9d40 100644
--- a/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
+++ b/services/autofill/java/com/android/server/autofill/RemoteAugmentedAutofillService.java
@@ -168,12 +168,12 @@ final class RemoteAugmentedAutofillService
                             focusedId, focusedValue, requestTime, inlineSuggestionsRequest,
                             new IFillCallback.Stub() {
                                 @Override
-                                public void onSuccess(
-                                        @Nullable List<Dataset> inlineSuggestionsData) {
+                                public void onSuccess(@Nullable List<Dataset> inlineSuggestionsData,
+                                        @Nullable Bundle clientState) {
                                     mCallbacks.resetLastResponse();
                                     maybeRequestShowInlineSuggestions(sessionId,
                                             inlineSuggestionsRequest, inlineSuggestionsData,
-                                            focusedId, inlineSuggestionsCallback,
+                                            clientState, focusedId, inlineSuggestionsCallback,
                                             client, onErrorCallback, remoteRenderService);
                                     requestAutofill.complete(null);
                                 }
@@ -238,7 +238,8 @@ final class RemoteAugmentedAutofillService
 
     private void maybeRequestShowInlineSuggestions(int sessionId,
             @Nullable InlineSuggestionsRequest request,
-            @Nullable List<Dataset> inlineSuggestionsData, @NonNull AutofillId focusedId,
+            @Nullable List<Dataset> inlineSuggestionsData, @Nullable Bundle clientState,
+            @NonNull AutofillId focusedId,
             @Nullable Function<InlineSuggestionsResponse, Boolean> inlineSuggestionsCallback,
             @NonNull IAutoFillManagerClient client, @NonNull Runnable onErrorCallback,
             @Nullable RemoteInlineSuggestionRenderService remoteRenderService) {
@@ -256,7 +257,7 @@ final class RemoteAugmentedAutofillService
                             @Override
                             public void autofill(Dataset dataset) {
                                 mCallbacks.logAugmentedAutofillSelected(sessionId,
-                                        dataset.getId());
+                                        dataset.getId(), clientState);
                                 try {
                                     final ArrayList<AutofillId> fieldIds = dataset.getFieldIds();
                                     final int size = fieldIds.size();
@@ -287,7 +288,7 @@ final class RemoteAugmentedAutofillService
             return;
         }
         if (inlineSuggestionsCallback.apply(inlineSuggestionsResponse)) {
-            mCallbacks.logAugmentedAutofillShown(sessionId);
+            mCallbacks.logAugmentedAutofillShown(sessionId, clientState);
         }
     }
 
@@ -310,8 +311,9 @@ final class RemoteAugmentedAutofillService
 
         void setLastResponse(int sessionId);
 
-        void logAugmentedAutofillShown(int sessionId);
+        void logAugmentedAutofillShown(int sessionId, @Nullable Bundle clientState);
 
-        void logAugmentedAutofillSelected(int sessionId, @Nullable String suggestionId);
+        void logAugmentedAutofillSelected(int sessionId, @Nullable String suggestionId,
+                @Nullable Bundle clientState);
     }
 }