Add onSuggestedReplySent in NotificationAssistantService

This is added to report clicks on smart reply buttons to NAS.

Also refactored the code a bit by having SmartReplyView to use
SmartReplies and SmartActions object, rather than passing a long list
of arguments.

Test: atest SystemUITests
Test: Manual. Tapped on the reply and observed the log.
BUG: 111437455

Change-Id: I897fb46a304f4f7b80b2a6bc4db0ac39f6dc6e8f
diff --git a/api/system-current.txt b/api/system-current.txt
index f963c10..76e76b0 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -5083,8 +5083,11 @@
     method public void onNotificationRemoved(android.service.notification.StatusBarNotification, android.service.notification.NotificationListenerService.RankingMap, android.service.notification.NotificationStats, int);
     method public abstract void onNotificationSnoozedUntilContext(android.service.notification.StatusBarNotification, java.lang.String);
     method public void onNotificationsSeen(java.util.List<java.lang.String>);
+    method public void onSuggestedReplySent(java.lang.String, java.lang.CharSequence, int);
     method public final void unsnoozeNotification(java.lang.String);
     field public static final java.lang.String SERVICE_INTERFACE = "android.service.notification.NotificationAssistantService";
+    field public static final int SOURCE_FROM_APP = 0; // 0x0
+    field public static final int SOURCE_FROM_ASSISTANT = 1; // 0x1
   }
 
   public final class NotificationStats implements android.os.Parcelable {
diff --git a/api/test-current.txt b/api/test-current.txt
index 0fa83f1..738caec 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -1162,8 +1162,11 @@
     method public void onNotificationExpansionChanged(java.lang.String, boolean, boolean);
     method public abstract void onNotificationSnoozedUntilContext(android.service.notification.StatusBarNotification, java.lang.String);
     method public void onNotificationsSeen(java.util.List<java.lang.String>);
+    method public void onSuggestedReplySent(java.lang.String, java.lang.CharSequence, int);
     method public final void unsnoozeNotification(java.lang.String);
     field public static final java.lang.String SERVICE_INTERFACE = "android.service.notification.NotificationAssistantService";
+    field public static final int SOURCE_FROM_APP = 0; // 0x0
+    field public static final int SOURCE_FROM_ASSISTANT = 1; // 0x1
   }
 
   public abstract class NotificationListenerService extends android.app.Service {
diff --git a/core/java/android/service/notification/INotificationListener.aidl b/core/java/android/service/notification/INotificationListener.aidl
index 0988510..ab94f43 100644
--- a/core/java/android/service/notification/INotificationListener.aidl
+++ b/core/java/android/service/notification/INotificationListener.aidl
@@ -49,4 +49,5 @@
     void onNotificationsSeen(in List<String> keys);
     void onNotificationExpansionChanged(String key, boolean userAction, boolean expanded);
     void onNotificationDirectReply(String key);
+    void onSuggestedReplySent(String key, in CharSequence reply, int source);
 }
diff --git a/core/java/android/service/notification/NotificationAssistantService.java b/core/java/android/service/notification/NotificationAssistantService.java
index 90f4792..68da83f 100644
--- a/core/java/android/service/notification/NotificationAssistantService.java
+++ b/core/java/android/service/notification/NotificationAssistantService.java
@@ -16,6 +16,9 @@
 
 package android.service.notification;
 
+import static java.lang.annotation.RetentionPolicy.SOURCE;
+
+import android.annotation.IntDef;
 import android.annotation.SdkConstant;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
@@ -33,6 +36,7 @@
 
 import com.android.internal.os.SomeArgs;
 
+import java.lang.annotation.Retention;
 import java.util.List;
 
 /**
@@ -63,6 +67,13 @@
 public abstract class NotificationAssistantService extends NotificationListenerService {
     private static final String TAG = "NotificationAssistants";
 
+    /** @hide */
+    @Retention(SOURCE)
+    @IntDef({SOURCE_FROM_APP, SOURCE_FROM_ASSISTANT})
+    public @interface Source {}
+    public static final int SOURCE_FROM_APP = 0;
+    public static final int SOURCE_FROM_ASSISTANT = 1;
+
     /**
      * The {@link Intent} that must be declared as handled by the service.
      */
@@ -175,6 +186,14 @@
     public void onNotificationDirectReply(String key) {}
 
     /**
+     * Implement this to know when a suggested reply is sent.
+     * @param key the notification key
+     * @param reply the reply that is just sent
+     * @param source the source of the reply, e.g. SOURCE_FROM_APP
+     */
+    public void onSuggestedReplySent(String key, CharSequence reply, @Source int source) {}
+
+    /**
      * Updates a notification.  N.B. this won’t cause
      * an existing notification to alert, but might allow a future update to
      * this notification to alert.
@@ -289,6 +308,15 @@
             mHandler.obtainMessage(MyHandler.MSG_ON_NOTIFICATION_DIRECT_REPLY_SENT, args)
                     .sendToTarget();
         }
+
+        @Override
+        public void onSuggestedReplySent(String key, CharSequence reply, int source) {
+            SomeArgs args = SomeArgs.obtain();
+            args.arg1 = key;
+            args.arg2 = reply;
+            args.argi2 = source;
+            mHandler.obtainMessage(MyHandler.MSG_ON_SUGGESTED_REPLY_SENT, args).sendToTarget();
+        }
     }
 
     private final class MyHandler extends Handler {
@@ -297,6 +325,7 @@
         public static final int MSG_ON_NOTIFICATIONS_SEEN = 3;
         public static final int MSG_ON_NOTIFICATION_EXPANSION_CHANGED = 4;
         public static final int MSG_ON_NOTIFICATION_DIRECT_REPLY_SENT = 5;
+        public static final int MSG_ON_SUGGESTED_REPLY_SENT = 6;
 
         public MyHandler(Looper looper) {
             super(looper, null, false);
@@ -357,6 +386,15 @@
                     onNotificationDirectReply(key);
                     break;
                 }
+                case MSG_ON_SUGGESTED_REPLY_SENT: {
+                    SomeArgs args = (SomeArgs) msg.obj;
+                    String key = (String) args.arg1;
+                    CharSequence reply = (CharSequence) args.arg2;
+                    int source = args.argi2;
+                    args.recycle();
+                    onSuggestedReplySent(key, reply, source);
+                    break;
+                }
             }
         }
     }
diff --git a/core/java/android/service/notification/NotificationListenerService.java b/core/java/android/service/notification/NotificationListenerService.java
index a4db451..756a7c6 100644
--- a/core/java/android/service/notification/NotificationListenerService.java
+++ b/core/java/android/service/notification/NotificationListenerService.java
@@ -1377,6 +1377,11 @@
         }
 
         @Override
+        public void onSuggestedReplySent(String key, CharSequence reply, int source) {
+            // no-op in the listener
+        }
+
+        @Override
         public void onNotificationChannelModification(String pkgName, UserHandle user,
                 NotificationChannel channel,
                 @ChannelOrGroupModificationTypes int modificationType) {
diff --git a/core/java/com/android/internal/statusbar/IStatusBarService.aidl b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
index 9a7094a..69ba070 100644
--- a/core/java/com/android/internal/statusbar/IStatusBarService.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
@@ -66,7 +66,7 @@
     void onNotificationExpansionChanged(in String key, in boolean userAction, in boolean expanded);
     void onNotificationDirectReplied(String key);
     void onNotificationSmartRepliesAdded(in String key, in int replyCount);
-    void onNotificationSmartReplySent(in String key, in int replyIndex);
+    void onNotificationSmartReplySent(in String key, in int replyIndex, in CharSequence reply, boolean generatedByAssistant);
     void onNotificationSettingsViewed(String key);
     void setSystemUiVisibility(int vis, int mask, String cause);
 
diff --git a/packages/ExtServices/src/android/ext/services/notification/Assistant.java b/packages/ExtServices/src/android/ext/services/notification/Assistant.java
index 0e4900a..22d0e3b 100644
--- a/packages/ExtServices/src/android/ext/services/notification/Assistant.java
+++ b/packages/ExtServices/src/android/ext/services/notification/Assistant.java
@@ -370,6 +370,14 @@
     }
 
     @Override
+    public void onSuggestedReplySent(String key, CharSequence reply, int source) {
+        if (DEBUG) {
+            Log.d(TAG, "onSuggestedReplySent() called with: key = [" + key + "], reply = [" + reply
+                    + "], source = [" + source + "]");
+        }
+    }
+
+    @Override
     public void onListenerConnected() {
         if (DEBUG) Log.i(TAG, "CONNECTED");
         try {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
index 758c33a..37bdc1c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
@@ -41,12 +41,16 @@
         mCallback = callback;
     }
 
-    public void smartReplySent(NotificationData.Entry entry, int replyIndex, CharSequence reply) {
+    /**
+     * Notifies StatusBarService a smart reply is sent.
+     */
+    public void smartReplySent(NotificationData.Entry entry, int replyIndex, CharSequence reply,
+            boolean generatedByAssistant) {
         mCallback.onSmartReplySent(entry, reply);
         mSendingKeys.add(entry.key);
         try {
-            mBarService.onNotificationSmartReplySent(entry.notification.getKey(),
-                    replyIndex);
+            mBarService.onNotificationSmartReplySent(
+                    entry.notification.getKey(), replyIndex, reply, generatedByAssistant);
         } catch (RemoteException e) {
             // Nothing to do, system going down
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
index 88edc0d..aa5ab71 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationContentView.java
@@ -1290,10 +1290,10 @@
             return;
         }
 
-        SmartRepliesAndActions smartRepliesAndActions = chooseSmartRepliesAndActions(
-                mSmartReplyConstants, entry);
+        SmartRepliesAndActions smartRepliesAndActions =
+                chooseSmartRepliesAndActions(mSmartReplyConstants, entry);
 
-        applyRemoteInput(entry, smartRepliesAndActions.freeformRemoteInputActionPair != null);
+        applyRemoteInput(entry, smartRepliesAndActions.hasFreeformRemoteInput);
         applySmartReplyView(smartRepliesAndActions, entry);
     }
 
@@ -1320,62 +1320,47 @@
         boolean appGeneratedSmartRepliesExist =
                 enableAppGeneratedSmartReplies
                         && remoteInputActionPair != null
-                        && !ArrayUtils.isEmpty(remoteInputActionPair.first.getChoices());
+                        && !ArrayUtils.isEmpty(remoteInputActionPair.first.getChoices())
+                        && remoteInputActionPair.second.actionIntent != null;
 
         List<Notification.Action> appGeneratedSmartActions = notification.getContextualActions();
         boolean appGeneratedSmartActionsExist = !appGeneratedSmartActions.isEmpty();
 
-        List<Notification.Action> sysGeneratedSmartActions =
-                notification.getAllowSystemGeneratedContextualActions()
-                        ? entry.systemGeneratedSmartActions : Collections.emptyList();
-
+        SmartReplyView.SmartReplies smartReplies = null;
+        SmartReplyView.SmartActions smartActions = null;
         if (appGeneratedSmartRepliesExist) {
-            return new SmartRepliesAndActions(remoteInputActionPair.first,
-                    remoteInputActionPair.second.actionIntent,
+            smartReplies = new SmartReplyView.SmartReplies(
                     remoteInputActionPair.first.getChoices(),
-                    appGeneratedSmartActions,
-                    freeformRemoteInputActionPair);
-        } else if (appGeneratedSmartActionsExist) {
-            return new SmartRepliesAndActions(null, null, null, appGeneratedSmartActions,
-                    freeformRemoteInputActionPair);
-        } else if (!ArrayUtils.isEmpty(entry.smartReplies)
-                && freeformRemoteInputActionPair != null
-                && freeformRemoteInputActionPair.second.getAllowGeneratedReplies()) {
-            // App didn't generate anything, use NAS-generated replies and actions
-            return new SmartRepliesAndActions(freeformRemoteInputActionPair.first,
-                    freeformRemoteInputActionPair.second.actionIntent,
-                    entry.smartReplies,
-                    sysGeneratedSmartActions,
-                    freeformRemoteInputActionPair);
+                    remoteInputActionPair.first,
+                    remoteInputActionPair.second.actionIntent,
+                    false /* fromAssistant */);
         }
-        // App didn't generate anything, and there are no NAS-generated smart replies.
-        return new SmartRepliesAndActions(null, null, null, sysGeneratedSmartActions,
-                freeformRemoteInputActionPair);
-    }
-
-    @VisibleForTesting
-    static class SmartRepliesAndActions {
-        public final RemoteInput remoteInputWithChoices;
-        public final PendingIntent pendingIntentForSmartReplies;
-        public final CharSequence[] smartReplies;
-        public final List<Notification.Action> smartActions;
-        public final Pair<RemoteInput, Notification.Action> freeformRemoteInputActionPair;
-
-        SmartRepliesAndActions(RemoteInput remoteInput, PendingIntent pendingIntent,
-                CharSequence[] choices, List<Notification.Action> smartActions,
-                Pair<RemoteInput, Notification.Action> freeformRemoteInputActionPair) {
-            this.remoteInputWithChoices = remoteInput;
-            this.pendingIntentForSmartReplies = pendingIntent;
-            this.smartReplies = choices;
-            this.smartActions = smartActions;
-            this.freeformRemoteInputActionPair = freeformRemoteInputActionPair;
+        if (appGeneratedSmartActionsExist) {
+            smartActions = new SmartReplyView.SmartActions(appGeneratedSmartActions,
+                    false /* fromAssistant */);
         }
-
-        boolean smartRepliesExist() {
-            return remoteInputWithChoices != null
-                    && pendingIntentForSmartReplies != null
-                    && !ArrayUtils.isEmpty(smartReplies);
+        // Apps didn't provide any smart replies / actions, use those from NAS (if any).
+        if (!appGeneratedSmartRepliesExist && !appGeneratedSmartActionsExist) {
+            boolean useGeneratedReplies = !ArrayUtils.isEmpty(entry.smartReplies)
+                    && freeformRemoteInputActionPair != null
+                    && freeformRemoteInputActionPair.second.getAllowGeneratedReplies()
+                    && freeformRemoteInputActionPair.second.actionIntent != null;
+            if (useGeneratedReplies) {
+                smartReplies = new SmartReplyView.SmartReplies(
+                        entry.smartReplies,
+                        freeformRemoteInputActionPair.first,
+                        freeformRemoteInputActionPair.second.actionIntent,
+                        true /* fromAssistant */);
+            }
+            boolean useSmartActions = !ArrayUtils.isEmpty(entry.systemGeneratedSmartActions)
+                    && notification.getAllowSystemGeneratedContextualActions();
+            if (useSmartActions) {
+                smartActions = new SmartReplyView.SmartActions(
+                        entry.systemGeneratedSmartActions, true /* fromAssistant */);
+            }
         }
+        return new SmartRepliesAndActions(
+                smartReplies, smartActions, freeformRemoteInputActionPair != null);
     }
 
     private void applyRemoteInput(NotificationData.Entry entry, boolean hasFreeformRemoteInput) {
@@ -1482,12 +1467,9 @@
         if (mExpandedChild != null) {
             mExpandedSmartReplyView =
                     applySmartReplyView(mExpandedChild, smartRepliesAndActions, entry);
-            if (mExpandedSmartReplyView != null
-                    && smartRepliesAndActions.remoteInputWithChoices != null
-                    && smartRepliesAndActions.smartReplies != null
-                    && smartRepliesAndActions.smartReplies.length > 0) {
-                mSmartReplyController.smartRepliesAdded(entry,
-                        smartRepliesAndActions.smartReplies.length);
+            if (mExpandedSmartReplyView != null && smartRepliesAndActions.smartReplies != null) {
+                mSmartReplyController.smartRepliesAdded(
+                        entry, smartRepliesAndActions.smartReplies.choices.length);
             }
         }
     }
@@ -1501,8 +1483,8 @@
         }
         LinearLayout smartReplyContainer = (LinearLayout) smartReplyContainerCandidate;
         // If there are no smart replies and no smart actions - early out.
-        if (!smartRepliesAndActions.smartRepliesExist()
-                && smartRepliesAndActions.smartActions.isEmpty()) {
+        if (smartRepliesAndActions.smartReplies == null
+                && smartRepliesAndActions.smartActions == null) {
             smartReplyContainer.setVisibility(View.GONE);
             return null;
         }
@@ -1532,10 +1514,13 @@
         }
         if (smartReplyView != null) {
             smartReplyView.resetSmartSuggestions(smartReplyContainer);
-            smartReplyView.addRepliesFromRemoteInput(smartRepliesAndActions.remoteInputWithChoices,
-                    smartRepliesAndActions.pendingIntentForSmartReplies, mSmartReplyController,
-                    entry, smartRepliesAndActions.smartReplies);
-            smartReplyView.addSmartActions(smartRepliesAndActions.smartActions);
+            if (smartRepliesAndActions.smartReplies != null) {
+                smartReplyView.addRepliesFromRemoteInput(
+                        smartRepliesAndActions.smartReplies, mSmartReplyController, entry);
+            }
+            if (smartRepliesAndActions.smartActions != null) {
+                smartReplyView.addSmartActions(smartRepliesAndActions.smartActions);
+            }
             smartReplyContainer.setVisibility(View.VISIBLE);
         }
         return smartReplyView;
@@ -1954,4 +1939,22 @@
         }
         pw.println();
     }
+
+    @VisibleForTesting
+    static class SmartRepliesAndActions {
+        @Nullable
+        public final SmartReplyView.SmartReplies smartReplies;
+        @Nullable
+        public final SmartReplyView.SmartActions smartActions;
+        public final boolean hasFreeformRemoteInput;
+
+        SmartRepliesAndActions(
+                @Nullable SmartReplyView.SmartReplies smartReplies,
+                @Nullable SmartReplyView.SmartActions smartActions,
+                boolean hasFreeformRemoteInput) {
+            this.smartReplies = smartReplies;
+            this.smartActions = smartActions;
+            this.hasFreeformRemoteInput = hasFreeformRemoteInput;
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
index 0186683..88ff078 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
@@ -1,6 +1,7 @@
 package com.android.systemui.statusbar.policy;
 
 import android.annotation.ColorInt;
+import android.annotation.NonNull;
 import android.app.Notification;
 import android.app.PendingIntent;
 import android.app.RemoteInput;
@@ -189,15 +190,13 @@
      * into the notification are shown.
      */
     public void addRepliesFromRemoteInput(
-            RemoteInput remoteInput, PendingIntent pendingIntent,
-            SmartReplyController smartReplyController, NotificationData.Entry entry,
-            CharSequence[] choices) {
-        if (remoteInput != null && pendingIntent != null) {
-            if (choices != null) {
-                for (int i = 0; i < choices.length; ++i) {
+            SmartReplies smartReplies,
+            SmartReplyController smartReplyController, NotificationData.Entry entry) {
+        if (smartReplies.remoteInput != null && smartReplies.pendingIntent != null) {
+            if (smartReplies.choices != null) {
+                for (int i = 0; i < smartReplies.choices.length; ++i) {
                     Button replyButton = inflateReplyButton(
-                            getContext(), this, i, choices[i], remoteInput, pendingIntent,
-                            smartReplyController, entry);
+                            getContext(), this, i, smartReplies, smartReplyController, entry);
                     addView(replyButton);
                 }
             }
@@ -209,10 +208,10 @@
      * Add smart actions to be shown next to smart replies. Only the actions that fit into the
      * notification are shown.
      */
-    public void addSmartActions(List<Notification.Action> smartActions) {
-        int numSmartActions = smartActions.size();
+    public void addSmartActions(SmartActions smartActions) {
+        int numSmartActions = smartActions.actions.size();
         for (int n = 0; n < numSmartActions; n++) {
-            Notification.Action action = smartActions.get(n);
+            Notification.Action action = smartActions.actions.get(n);
             if (action.actionIntent != null) {
                 Button actionButton = inflateActionButton(getContext(), this, action);
                 addView(actionButton);
@@ -228,22 +227,25 @@
 
     @VisibleForTesting
     Button inflateReplyButton(Context context, ViewGroup root, int replyIndex,
-            CharSequence choice, RemoteInput remoteInput, PendingIntent pendingIntent,
-            SmartReplyController smartReplyController, NotificationData.Entry entry) {
+            SmartReplies smartReplies, SmartReplyController smartReplyController,
+            NotificationData.Entry entry) {
         Button b = (Button) LayoutInflater.from(context).inflate(
                 R.layout.smart_reply_button, root, false);
+        CharSequence choice = smartReplies.choices[replyIndex];
         b.setText(choice);
 
         OnDismissAction action = () -> {
-            smartReplyController.smartReplySent(entry, replyIndex, b.getText());
+            smartReplyController.smartReplySent(
+                    entry, replyIndex, b.getText(), smartReplies.fromAssistant);
             Bundle results = new Bundle();
-            results.putString(remoteInput.getResultKey(), choice.toString());
+            results.putString(smartReplies.remoteInput.getResultKey(), choice.toString());
             Intent intent = new Intent().addFlags(Intent.FLAG_RECEIVER_FOREGROUND);
-            RemoteInput.addResultsToIntent(new RemoteInput[]{remoteInput}, intent, results);
+            RemoteInput.addResultsToIntent(new RemoteInput[]{smartReplies.remoteInput}, intent,
+                    results);
             RemoteInput.setResultsSource(intent, RemoteInput.SOURCE_CHOICE);
             entry.setHasSentReply();
             try {
-                pendingIntent.send(context, 0, intent);
+                smartReplies.pendingIntent.send(context, 0, intent);
             } catch (PendingIntent.CanceledException e) {
                 Log.w(TAG, "Unable to send smart reply", e);
             }
@@ -741,4 +743,40 @@
             return show;
         }
     }
+
+    /**
+     * Data class for smart replies.
+     */
+    public static class SmartReplies {
+        @NonNull
+        public final RemoteInput remoteInput;
+        @NonNull
+        public final PendingIntent pendingIntent;
+        @NonNull
+        public final CharSequence[] choices;
+        public final boolean fromAssistant;
+
+        public SmartReplies(CharSequence[] choices, RemoteInput remoteInput,
+                PendingIntent pendingIntent, boolean fromAssistant) {
+            this.choices = choices;
+            this.remoteInput = remoteInput;
+            this.pendingIntent = pendingIntent;
+            this.fromAssistant = fromAssistant;
+        }
+    }
+
+
+    /**
+     * Data class for smart actions.
+     */
+    public static class SmartActions {
+        @NonNull
+        public final List<Notification.Action> actions;
+        public final boolean fromAssistant;
+
+        public SmartActions(List<Notification.Action> actions, boolean fromAssistant) {
+            this.actions = actions;
+            this.fromAssistant = fromAssistant;
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
index 1d977d8..76b992f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/SmartReplyControllerTest.java
@@ -16,18 +16,15 @@
 
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.assertTrue;
+
 import static org.mockito.ArgumentMatchers.argThat;
-import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.isNull;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager;
 import android.app.Notification;
 import android.os.RemoteException;
 import android.os.UserHandle;
-import android.service.notification.NotificationListenerService;
 import android.service.notification.StatusBarNotification;
 import android.support.test.filters.SmallTest;
 import android.testing.AndroidTestingRunner;
@@ -38,7 +35,6 @@
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.notification.NotificationData;
 import com.android.systemui.statusbar.notification.NotificationEntryManager;
-import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -93,7 +89,7 @@
 
     @Test
     public void testSendSmartReply_updatesRemoteInput() {
-        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, false);
 
         // Sending smart reply should make calls to NotificationEntryManager
         // to update the notification with reply and spinner.
@@ -103,11 +99,21 @@
 
     @Test
     public void testSendSmartReply_logsToStatusBar() throws RemoteException {
-        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, false);
 
         // Check we log the result to the status bar service.
         verify(mIStatusBarService).onNotificationSmartReplySent(mSbn.getKey(),
-                TEST_CHOICE_INDEX);
+                TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, false);
+    }
+
+
+    @Test
+    public void testSendSmartReply_logsToStatusBar_generatedByAssistant() throws RemoteException {
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, true);
+
+        // Check we log the result to the status bar service.
+        verify(mIStatusBarService).onNotificationSmartReplySent(mSbn.getKey(),
+                TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, true);
     }
 
     @Test
@@ -121,14 +127,14 @@
 
     @Test
     public void testSendSmartReply_reportsSending() {
-        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, false);
 
         assertTrue(mSmartReplyController.isSendingSmartReply(mSbn.getKey()));
     }
 
     @Test
     public void testSendingSmartReply_afterRemove_shouldReturnFalse() {
-        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT);
+        mSmartReplyController.smartReplySent(mEntry, TEST_CHOICE_INDEX, TEST_CHOICE_TEXT, false);
         mSmartReplyController.stopSending(mEntry);
 
         assertFalse(mSmartReplyController.isSendingSmartReply(mSbn.getKey()));
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.java
index 5e137a7..7fee0ee 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationContentViewTest.java
@@ -16,11 +16,10 @@
 
 package com.android.systemui.statusbar.notification.row;
 
-import static org.hamcrest.Matchers.empty;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.is;
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.hamcrest.core.IsEqual.equalTo;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyFloat;
 import static org.mockito.Mockito.doNothing;
@@ -170,8 +169,10 @@
     private void setupAppGeneratedReplies(
             CharSequence[] smartReplyTitles,
             Notification.Action freeFormRemoteInputAction) {
+        PendingIntent pendingIntent =
+                PendingIntent.getBroadcast(mContext, 0, new Intent(TEST_ACTION), 0);
         Notification.Action action =
-                new Notification.Action.Builder(null, "Test Action", null).build();
+                new Notification.Action.Builder(null, "Test Action", pendingIntent).build();
         when(mRemoteInput.getChoices()).thenReturn(smartReplyTitles);
         Pair<RemoteInput, Notification.Action> remoteInputActionPair =
                 Pair.create(mRemoteInput, action);
@@ -191,7 +192,7 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertFalse(repliesAndActions.smartRepliesExist());
+        assertThat(repliesAndActions.smartReplies).isNull();
     }
 
     @Test
@@ -203,7 +204,9 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(smartReplies));
+        assertThat(repliesAndActions.smartReplies.choices).isEqualTo(smartReplies);
+        assertThat(repliesAndActions.smartReplies.fromAssistant).isFalse();
+        assertThat(repliesAndActions.smartActions).isNull();
     }
 
     @Test
@@ -219,8 +222,10 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(smartReplies));
-        assertThat(repliesAndActions.smartActions, equalTo(smartActions));
+        assertThat(repliesAndActions.smartReplies.choices).isEqualTo(smartReplies);
+        assertThat(repliesAndActions.smartReplies.fromAssistant).isFalse();
+        assertThat(repliesAndActions.smartActions.actions).isEqualTo(smartActions);
+        assertThat(repliesAndActions.smartActions.fromAssistant).isFalse();
     }
 
     @Test
@@ -237,8 +242,9 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(mEntry.smartReplies));
-        assertThat(repliesAndActions.smartActions, is(empty()));
+        assertThat(repliesAndActions.smartReplies.choices).isEqualTo(mEntry.smartReplies);
+        assertThat(repliesAndActions.smartReplies.fromAssistant).isTrue();
+        assertThat(repliesAndActions.smartActions).isNull();
     }
 
     @Test
@@ -255,8 +261,8 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(null));
-        assertThat(repliesAndActions.smartActions, is(empty()));
+        assertThat(repliesAndActions.smartReplies).isNull();
+        assertThat(repliesAndActions.smartActions).isNull();
     }
 
     @Test
@@ -270,8 +276,10 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(null));
-        assertThat(repliesAndActions.smartActions, equalTo(mEntry.systemGeneratedSmartActions));
+        assertThat(repliesAndActions.smartReplies).isNull();
+        assertThat(repliesAndActions.smartActions.actions)
+                .isEqualTo(mEntry.systemGeneratedSmartActions);
+        assertThat(repliesAndActions.smartActions.fromAssistant).isTrue();
     }
 
     @Test
@@ -296,8 +304,10 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(appGenSmartReplies));
-        assertThat(repliesAndActions.smartActions, equalTo(appGenSmartActions));
+        assertThat(repliesAndActions.smartReplies.choices).isEqualTo(appGenSmartReplies);
+        assertThat(repliesAndActions.smartReplies.fromAssistant).isFalse();
+        assertThat(repliesAndActions.smartActions.actions).isEqualTo(appGenSmartActions);
+        assertThat(repliesAndActions.smartActions.fromAssistant).isFalse();
     }
 
     @Test
@@ -313,8 +323,8 @@
         NotificationContentView.SmartRepliesAndActions repliesAndActions =
                 NotificationContentView.chooseSmartRepliesAndActions(mSmartReplyConstants, mEntry);
 
-        assertThat(repliesAndActions.smartReplies, equalTo(null));
-        assertThat(repliesAndActions.smartActions, is(empty()));
+        assertThat(repliesAndActions.smartActions).isNull();
+        assertThat(repliesAndActions.smartReplies).isNull();
     }
 
     private Notification.Action.Builder createActionBuilder(String actionTitle) {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
index 9e659c8..b6e3fc1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SmartReplyViewTest.java
@@ -181,7 +181,16 @@
     public void testSendSmartReply_controllerCalled() {
         setSmartReplies(TEST_CHOICES);
         mView.getChildAt(2).performClick();
-        verify(mLogger).smartReplySent(mEntry, 2, TEST_CHOICES[2]);
+        verify(mLogger).smartReplySent(mEntry, 2, TEST_CHOICES[2],
+                false /* generatedByAsssitant */);
+    }
+
+    @Test
+    public void testSendSmartReply_controllerCalled_generatedByAssistant() {
+        setSmartReplies(TEST_CHOICES, true);
+        mView.getChildAt(2).performClick();
+        verify(mLogger).smartReplySent(mEntry, 2, TEST_CHOICES[2],
+                true /* generatedByAsssitant */);
     }
 
     @Test
@@ -392,11 +401,17 @@
     }
 
     private void setSmartReplies(CharSequence[] choices) {
+        setSmartReplies(choices, false);
+    }
+
+    private void setSmartReplies(CharSequence[] choices, boolean fromAssistant) {
         PendingIntent pendingIntent = PendingIntent.getBroadcast(mContext, 0,
                 new Intent(TEST_ACTION), 0);
         RemoteInput input = new RemoteInput.Builder(TEST_RESULT_KEY).setChoices(choices).build();
+        SmartReplyView.SmartReplies smartReplies =
+                new SmartReplyView.SmartReplies(choices, input, pendingIntent, fromAssistant);
         mView.resetSmartSuggestions(mContainer);
-        mView.addRepliesFromRemoteInput(input, pendingIntent, mLogger, mEntry, choices);
+        mView.addRepliesFromRemoteInput(smartReplies, mLogger, mEntry);
     }
 
     private Notification.Action createAction(String actionTitle) {
@@ -415,12 +430,12 @@
 
     private void setSmartActions(String[] actionTitles) {
         mView.resetSmartSuggestions(mContainer);
-        mView.addSmartActions(createActions(actionTitles));
+        mView.addSmartActions(new SmartReplyView.SmartActions(createActions(actionTitles), false));
     }
 
     private void setSmartRepliesAndActions(CharSequence[] choices, String[] actionTitles) {
         setSmartReplies(choices);
-        mView.addSmartActions(createActions(actionTitles));
+        mView.addSmartActions(new SmartReplyView.SmartActions(createActions(actionTitles), false));
     }
 
     private ViewGroup buildExpectedView(CharSequence[] choices, int lineCount) {
@@ -453,9 +468,11 @@
 
         // Add smart replies
         Button previous = null;
+        SmartReplyView.SmartReplies smartReplies =
+                new SmartReplyView.SmartReplies(choices, null, null, false);
         for (int i = 0; i < choices.length; ++i) {
-            Button current = mView.inflateReplyButton(mContext, mView, i, choices[i],
-                    null, null, null, null);
+            Button current = mView.inflateReplyButton(mContext, mView, i, smartReplies,
+                    null, null);
             current.setPadding(paddingHorizontal, current.getPaddingTop(), paddingHorizontal,
                     current.getPaddingBottom());
             if (previous != null) {
diff --git a/services/core/java/com/android/server/notification/NotificationDelegate.java b/services/core/java/com/android/server/notification/NotificationDelegate.java
index decdac6..9222740 100644
--- a/services/core/java/com/android/server/notification/NotificationDelegate.java
+++ b/services/core/java/com/android/server/notification/NotificationDelegate.java
@@ -45,5 +45,15 @@
     void onNotificationDirectReplied(String key);
     void onNotificationSettingsViewed(String key);
     void onNotificationSmartRepliesAdded(String key, int replyCount);
-    void onNotificationSmartReplySent(String key, int replyIndex);
+
+    /**
+     * Notifies a smart reply is sent.
+     *
+     * @param key the notification key
+     * @param clickedIndex the index of clicked reply
+     * @param reply the reply that is sent
+     * @param generatedByAssistant specifies is the reply generated by NAS
+     */
+    void onNotificationSmartReplySent(String key, int clickedIndex, CharSequence reply,
+            boolean generatedByAssistant);
 }
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 2048d5f..6005872 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -920,7 +920,8 @@
         }
 
         @Override
-        public void onNotificationSmartReplySent(String key, int replyIndex) {
+        public void onNotificationSmartReplySent(String key, int replyIndex, CharSequence reply,
+                boolean generatedByAssistant) {
             synchronized (mNotificationLock) {
                 NotificationRecord r = mNotificationsByKey.get(key);
                 if (r != null) {
@@ -930,6 +931,8 @@
                     mMetricsLogger.write(logMaker);
                     // Treat clicking on a smart reply as a user interaction.
                     reportUserInteraction(r);
+                    mAssistants.notifyAssistantSuggestedReplySent(
+                            r.sbn, reply, generatedByAssistant);
                 }
             }
         }
@@ -6899,6 +6902,27 @@
                     });
         }
 
+        @GuardedBy("mNotificationLock")
+        void notifyAssistantSuggestedReplySent(
+                final StatusBarNotification sbn, CharSequence reply, boolean generatedByAssistant) {
+            final String key = sbn.getKey();
+            notifyAssistantLocked(
+                    sbn,
+                    false /* sameUserOnly */,
+                    (assistant, sbnHolder) -> {
+                        try {
+                            assistant.onSuggestedReplySent(
+                                    key,
+                                    reply,
+                                    generatedByAssistant
+                                            ? NotificationAssistantService.SOURCE_FROM_ASSISTANT
+                                            : NotificationAssistantService.SOURCE_FROM_APP);
+                        } catch (RemoteException ex) {
+                            Log.e(TAG, "unable to notify assistant (snoozed): " + assistant, ex);
+                        }
+                    });
+        }
+
 
         /**
          * asynchronously notify the assistant that a notification has been snoozed until a
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
index 1eb44a0..361622f 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
@@ -1177,12 +1177,14 @@
     }
 
     @Override
-    public void onNotificationSmartReplySent(String key, int replyIndex)
+    public void onNotificationSmartReplySent(
+            String key, int replyIndex, CharSequence reply, boolean generatedByAssistant)
             throws RemoteException {
         enforceStatusBarService();
         long identity = Binder.clearCallingIdentity();
         try {
-            mNotificationDelegate.onNotificationSmartReplySent(key, replyIndex);
+            mNotificationDelegate.onNotificationSmartReplySent(key, replyIndex, reply,
+                    generatedByAssistant);
         } finally {
             Binder.restoreCallingIdentity(identity);
         }
diff --git a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
index f11492a..fcd29e1 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -3696,4 +3696,19 @@
                 new TestableToastCallback(), 2000, 0);
         assertEquals(1, mService.mToastQueue.size());
     }
+
+    @Test
+    public void testOnNotificationSmartReplySent() {
+        final int replyIndex = 2;
+        final String reply = "Hello";
+        final boolean generatedByAssistant = true;
+
+        NotificationRecord r = generateNotificationRecord(mTestNotificationChannel);
+        mService.addNotification(r);
+
+        mService.mNotificationDelegate.onNotificationSmartReplySent(
+                r.getKey(), replyIndex, reply, generatedByAssistant);
+        verify(mAssistants).notifyAssistantSuggestedReplySent(
+                eq(r.sbn), eq(reply), eq(generatedByAssistant));
+    }
 }