Add onActionClicked in NotificationAssistantService

This is added to report clicks on actions buttons to NAS.

BUG: 119010281
Test: atest SystemUITests
Test: atest RemoteViewsTest
Test: atest NotificationManagerServiceTest
Test: Manual. Tapped on the action (both normal and contextual) and
      observed the log.

Change-Id: I381994737d8c3185d3fabf9b6c481fd01a89a634
diff --git a/api/system-current.txt b/api/system-current.txt
index a53ab3e..21ef675 100644
--- a/api/system-current.txt
+++ b/api/system-current.txt
@@ -5074,6 +5074,7 @@
     ctor public NotificationAssistantService();
     method public final void adjustNotification(android.service.notification.Adjustment);
     method public final void adjustNotifications(java.util.List<android.service.notification.Adjustment>);
+    method public void onActionClicked(java.lang.String, android.app.Notification.Action, int);
     method public final android.os.IBinder onBind(android.content.Intent);
     method public void onNotificationDirectReply(java.lang.String);
     method public android.service.notification.Adjustment onNotificationEnqueued(android.service.notification.StatusBarNotification);
diff --git a/api/test-current.txt b/api/test-current.txt
index 738caec..9695b78 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -1155,6 +1155,7 @@
     ctor public NotificationAssistantService();
     method public final void adjustNotification(android.service.notification.Adjustment);
     method public final void adjustNotifications(java.util.List<android.service.notification.Adjustment>);
+    method public void onActionClicked(java.lang.String, android.app.Notification.Action, int);
     method public final android.os.IBinder onBind(android.content.Intent);
     method public void onNotificationDirectReply(java.lang.String);
     method public android.service.notification.Adjustment onNotificationEnqueued(android.service.notification.StatusBarNotification);
diff --git a/core/java/android/app/Notification.aidl b/core/java/android/app/Notification.aidl
index 9d8129c..8a7156e 100644
--- a/core/java/android/app/Notification.aidl
+++ b/core/java/android/app/Notification.aidl
@@ -17,3 +17,4 @@
 package android.app;
 
 parcelable Notification;
+parcelable Notification.Action;
\ No newline at end of file
diff --git a/core/java/android/app/Notification.java b/core/java/android/app/Notification.java
index 75b56f3..cf43f59 100644
--- a/core/java/android/app/Notification.java
+++ b/core/java/android/app/Notification.java
@@ -5434,6 +5434,8 @@
                             ambient ? resolveAmbientColor() : resolveContrastColor());
                 }
             }
+            button.setIntTag(R.id.action0, R.id.notification_action_index_tag,
+                    mActions.indexOf(action));
             return button;
         }
 
diff --git a/core/java/android/service/notification/INotificationListener.aidl b/core/java/android/service/notification/INotificationListener.aidl
index ab94f43..1ddc099e 100644
--- a/core/java/android/service/notification/INotificationListener.aidl
+++ b/core/java/android/service/notification/INotificationListener.aidl
@@ -16,6 +16,7 @@
 
 package android.service.notification;
 
+import android.app.Notification;
 import android.app.NotificationChannel;
 import android.app.NotificationChannelGroup;
 import android.content.pm.ParceledListSlice;
@@ -50,4 +51,5 @@
     void onNotificationExpansionChanged(String key, boolean userAction, boolean expanded);
     void onNotificationDirectReply(String key);
     void onSuggestedReplySent(String key, in CharSequence reply, int source);
+    void onActionClicked(String key, in Notification.Action action, int source);
 }
diff --git a/core/java/android/service/notification/NotificationAssistantService.java b/core/java/android/service/notification/NotificationAssistantService.java
index 68da83f..c850a4e 100644
--- a/core/java/android/service/notification/NotificationAssistantService.java
+++ b/core/java/android/service/notification/NotificationAssistantService.java
@@ -19,9 +19,11 @@
 import static java.lang.annotation.RetentionPolicy.SOURCE;
 
 import android.annotation.IntDef;
+import android.annotation.Nullable;
 import android.annotation.SdkConstant;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
+import android.app.Notification;
 import android.app.NotificationChannel;
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
@@ -189,11 +191,20 @@
      * 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
+     * @param source the source that provided the reply, e.g. SOURCE_FROM_APP
      */
     public void onSuggestedReplySent(String key, CharSequence reply, @Source int source) {}
 
     /**
+     * Implement this to know when an action is clicked.
+     * @param key the notification key
+     * @param action the action that is just clicked
+     * @param source the source that provided the action, e.g. SOURCE_FROM_APP
+     */
+    public void onActionClicked(String key, @Nullable Notification.Action action, 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.
@@ -317,6 +328,15 @@
             args.argi2 = source;
             mHandler.obtainMessage(MyHandler.MSG_ON_SUGGESTED_REPLY_SENT, args).sendToTarget();
         }
+
+        @Override
+        public void onActionClicked(String key, Notification.Action action, int source) {
+            SomeArgs args = SomeArgs.obtain();
+            args.arg1 = key;
+            args.arg2 = action;
+            args.argi2 = source;
+            mHandler.obtainMessage(MyHandler.MSG_ON_ACTION_CLICKED, args).sendToTarget();
+        }
     }
 
     private final class MyHandler extends Handler {
@@ -326,6 +346,7 @@
         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 static final int MSG_ON_ACTION_CLICKED = 7;
 
         public MyHandler(Looper looper) {
             super(looper, null, false);
@@ -395,6 +416,15 @@
                     onSuggestedReplySent(key, reply, source);
                     break;
                 }
+                case MSG_ON_ACTION_CLICKED: {
+                    SomeArgs args = (SomeArgs) msg.obj;
+                    String key = (String) args.arg1;
+                    Notification.Action action = (Notification.Action) args.arg2;
+                    int source = args.argi2;
+                    args.recycle();
+                    onActionClicked(key, action, source);
+                    break;
+                }
             }
         }
     }
diff --git a/core/java/android/service/notification/NotificationListenerService.java b/core/java/android/service/notification/NotificationListenerService.java
index 756a7c6..1fe97b7 100644
--- a/core/java/android/service/notification/NotificationListenerService.java
+++ b/core/java/android/service/notification/NotificationListenerService.java
@@ -1382,6 +1382,11 @@
         }
 
         @Override
+        public void onActionClicked(String key, Notification.Action action, 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/android/widget/RemoteViews.java b/core/java/android/widget/RemoteViews.java
index c0979fe..7b39efe 100644
--- a/core/java/android/widget/RemoteViews.java
+++ b/core/java/android/widget/RemoteViews.java
@@ -161,6 +161,7 @@
     private static final int LAYOUT_PARAM_ACTION_TAG = 19;
     private static final int OVERRIDE_TEXT_COLORS_TAG = 20;
     private static final int SET_RIPPLE_DRAWABLE_COLOR_TAG = 21;
+    private static final int SET_INT_TAG_TAG = 22;
 
     /**
      * Application that hosts the remote views.
@@ -274,6 +275,15 @@
     }
 
     /**
+     * Sets an integer tag to the view.
+     *
+     * @hide
+     */
+    public void setIntTag(int viewId, int key, int tag) {
+        addAction(new SetIntTagAction(viewId, key, tag));
+    }
+
+    /**
      * Set that it is disallowed to reapply another remoteview with the same layout as this view.
      * This should be done if an action is destroying the view tree of the base layout.
      *
@@ -2122,6 +2132,43 @@
         }
     }
 
+    private class SetIntTagAction extends Action {
+        private final int mViewId;
+        private final int mKey;
+        private final int mTag;
+
+        SetIntTagAction(int viewId, int key, int tag) {
+            mViewId = viewId;
+            mKey = key;
+            mTag = tag;
+        }
+
+        SetIntTagAction(Parcel parcel) {
+            mViewId = parcel.readInt();
+            mKey = parcel.readInt();
+            mTag = parcel.readInt();
+        }
+
+        public void writeToParcel(Parcel dest, int flags) {
+            dest.writeInt(mViewId);
+            dest.writeInt(mKey);
+            dest.writeInt(mTag);
+        }
+
+        @Override
+        public void apply(View root, ViewGroup rootParent, OnClickHandler handler) {
+            final View target = root.findViewById(mViewId);
+            if (target == null) return;
+
+            target.setTagInternal(mKey, mTag);
+        }
+
+        @Override
+        public int getActionTag() {
+            return SET_INT_TAG_TAG;
+        }
+    }
+
     /**
      * Create a new RemoteViews object that will display the views contained
      * in the specified layout file.
@@ -2326,6 +2373,8 @@
                 return new OverrideTextColorsAction(parcel);
             case SET_RIPPLE_DRAWABLE_COLOR_TAG:
                 return new SetRippleDrawableColor(parcel);
+            case SET_INT_TAG_TAG:
+                return new SetIntTagAction(parcel);
             default:
                 throw new ActionException("Tag " + tag + " not found");
         }
diff --git a/core/java/com/android/internal/statusbar/IStatusBarService.aidl b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
index 69ba070..b7ffb57 100644
--- a/core/java/com/android/internal/statusbar/IStatusBarService.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBarService.aidl
@@ -16,6 +16,7 @@
 
 package com.android.internal.statusbar;
 
+import android.app.Notification;
 import android.content.ComponentName;
 import android.graphics.Rect;
 import android.os.Bundle;
@@ -55,7 +56,7 @@
     // Mark current notifications as "seen" and stop ringing, vibrating, blinking.
     void clearNotificationEffects();
     void onNotificationClick(String key, in NotificationVisibility nv);
-    void onNotificationActionClick(String key, int actionIndex, in NotificationVisibility nv);
+    void onNotificationActionClick(String key, int actionIndex, in Notification.Action action, in NotificationVisibility nv, boolean generatedByAssistant);
     void onNotificationError(String pkg, String tag, int id,
             int uid, int initialPid, String message, int userId);
     void onClearAllNotifications(int userId);
diff --git a/core/res/res/values/ids.xml b/core/res/res/values/ids.xml
index 64e5bc0..bbe3ff9 100644
--- a/core/res/res/values/ids.xml
+++ b/core/res/res/values/ids.xml
@@ -190,4 +190,7 @@
 
   <!-- A tag used to save the view added to a transition overlay -->
   <item type="id" name="transition_overlay_view_tag" />
+
+  <!-- A tag used to save the notification action object -->
+  <item type="id" name="notification_action_index_tag" />
 </resources>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index 9264f90..e251e27 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -2683,6 +2683,7 @@
   <java-symbol type="id" name="smart_reply_container" />
   <java-symbol type="id" name="remote_input_tag" />
   <java-symbol type="id" name="pending_intent_tag" />
+  <java-symbol type="id" name="notification_action_index_tag" />
 
   <java-symbol type="attr" name="seekBarDialogPreferenceStyle" />
   <java-symbol type="string" name="ext_media_status_removed" />
diff --git a/core/tests/coretests/src/android/widget/RemoteViewsTest.java b/core/tests/coretests/src/android/widget/RemoteViewsTest.java
index 4456122..36792bb 100644
--- a/core/tests/coretests/src/android/widget/RemoteViewsTest.java
+++ b/core/tests/coretests/src/android/widget/RemoteViewsTest.java
@@ -467,6 +467,21 @@
         assertArrayEquals(container.mSharedViewNames, new String[] {"e0", "e1", "e2"});
     }
 
+    @Test
+    public void setIntTag() {
+        RemoteViews views = new RemoteViews(mPackage, R.layout.remote_views_test);
+        int index = 10;
+        views.setIntTag(
+                R.id.layout, com.android.internal.R.id.notification_action_index_tag, index);
+
+        RemoteViews recovered = parcelAndRecreate(views);
+        RemoteViews cloned = new RemoteViews(recovered);
+        View inflated = cloned.apply(mContext, mContainer);
+
+        assertEquals(
+                index, inflated.getTag(com.android.internal.R.id.notification_action_index_tag));
+    }
+
     private class WidgetContainer extends AppWidgetHostView {
         int[] mSharedViewIds;
         String[] mSharedViewNames;
diff --git a/packages/ExtServices/src/android/ext/services/notification/Assistant.java b/packages/ExtServices/src/android/ext/services/notification/Assistant.java
index 22d0e3b..4b212c2 100644
--- a/packages/ExtServices/src/android/ext/services/notification/Assistant.java
+++ b/packages/ExtServices/src/android/ext/services/notification/Assistant.java
@@ -378,6 +378,14 @@
     }
 
     @Override
+    public void onActionClicked(String key, Notification.Action action, int source) {
+        if (DEBUG) {
+            Log.d(TAG, "onActionClicked() called with: key = [" + key + "], action = [" + action.title
+                    + "], source = [" + source + "]");
+        }
+    }
+
+    @Override
     public void onListenerConnected() {
         if (DEBUG) Log.i(TAG, "CONNECTED");
         try {
diff --git a/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java b/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
index f492208..4891e50 100644
--- a/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
+++ b/packages/SystemUI/plugin/src/com/android/systemui/plugins/ActivityStarter.java
@@ -30,6 +30,12 @@
     int VERSION = 1;
 
     void startPendingIntentDismissingKeyguard(PendingIntent intent);
+
+    /**
+     * Similar to {@link #startPendingIntentDismissingKeyguard(PendingIntent, Runnable)}, but
+     * allow you to specify the callback that is executed after the intent is sent.
+     */
+    void startPendingIntentDismissingKeyguard(PendingIntent intent, Runnable intentSentCallback);
     void startActivity(Intent intent, boolean dismissShade);
     void startActivity(Intent intent, boolean onlyProvisioned, boolean dismissShade);
     void startActivity(Intent intent, boolean dismissShade, Callback callback);
diff --git a/packages/SystemUI/src/com/android/systemui/ActivityStarterDelegate.java b/packages/SystemUI/src/com/android/systemui/ActivityStarterDelegate.java
index e1b8dc8..9e7c5ba 100644
--- a/packages/SystemUI/src/com/android/systemui/ActivityStarterDelegate.java
+++ b/packages/SystemUI/src/com/android/systemui/ActivityStarterDelegate.java
@@ -36,6 +36,15 @@
     }
 
     @Override
+    public void startPendingIntentDismissingKeyguard(PendingIntent intent,
+            Runnable intentSentCallback) {
+        if (mActualStarter == null) {
+            return;
+        }
+        mActualStarter.startPendingIntentDismissingKeyguard(intent, intentSentCallback);
+    }
+
+    @Override
     public void startActivity(Intent intent, boolean dismissShade) {
         if (mActualStarter == null) {
             return;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index 2ee5443..7be5461 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
@@ -150,34 +150,42 @@
         }
 
         private void logActionClick(View view) {
+            Integer actionIndex = (Integer)
+                    view.getTag(com.android.internal.R.id.notification_action_index_tag);
+            if (actionIndex == null) {
+                Log.e(TAG, "Couldn't retrieve the actionIndex from the clicked button");
+                return;
+            }
             ViewParent parent = view.getParent();
-            String key = getNotificationKeyForParent(parent);
-            if (key == null) {
+            StatusBarNotification statusBarNotification = getNotificationForParent(parent);
+            if (statusBarNotification == null) {
                 Log.w(TAG, "Couldn't determine notification for click.");
                 return;
             }
-            int index = -1;
+            String key = statusBarNotification.getKey();
+            int buttonIndex = -1;
             // If this is a default template, determine the index of the button.
             if (view.getId() == com.android.internal.R.id.action0 &&
                     parent != null && parent instanceof ViewGroup) {
                 ViewGroup actionGroup = (ViewGroup) parent;
-                index = actionGroup.indexOfChild(view);
+                buttonIndex = actionGroup.indexOfChild(view);
             }
             final int count = mEntryManager.getNotificationData().getActiveNotifications().size();
             final int rank = mEntryManager.getNotificationData().getRank(key);
+            final Notification.Action action =
+                    statusBarNotification.getNotification().actions[actionIndex];
             final NotificationVisibility nv = NotificationVisibility.obtain(key, rank, count, true);
             try {
-                mBarService.onNotificationActionClick(key, index, nv);
+                mBarService.onNotificationActionClick(key, buttonIndex, action, nv, false);
             } catch (RemoteException e) {
                 // Ignore
             }
         }
 
-        private String getNotificationKeyForParent(ViewParent parent) {
+        private StatusBarNotification getNotificationForParent(ViewParent parent) {
             while (parent != null) {
                 if (parent instanceof ExpandableNotificationRow) {
-                    return ((ExpandableNotificationRow) parent)
-                            .getStatusBarNotification().getKey();
+                    return ((ExpandableNotificationRow) parent).getStatusBarNotification();
                 }
                 parent = parent.getParent();
             }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
index 37bdc1c..f5d6904 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/SmartReplyController.java
@@ -15,12 +15,15 @@
  */
 package com.android.systemui.statusbar;
 
+import android.app.Notification;
 import android.os.RemoteException;
 import android.util.ArraySet;
 
 import com.android.internal.statusbar.IStatusBarService;
+import com.android.internal.statusbar.NotificationVisibility;
 import com.android.systemui.Dependency;
 import com.android.systemui.statusbar.notification.NotificationData;
+import com.android.systemui.statusbar.notification.NotificationEntryManager;
 
 import java.util.Set;
 
@@ -32,6 +35,9 @@
     private IStatusBarService mBarService;
     private Set<String> mSendingKeys = new ArraySet<>();
     private Callback mCallback;
+    private final NotificationEntryManager mEntryManager =
+            Dependency.get(NotificationEntryManager.class);
+
 
     public SmartReplyController() {
         mBarService = Dependency.get(IStatusBarService.class);
@@ -57,6 +63,24 @@
     }
 
     /**
+     * Notifies StatusBarService a smart action is clicked.
+     */
+    public void smartActionClicked(
+            NotificationData.Entry entry, int actionIndex, Notification.Action action,
+            boolean generatedByAssistant) {
+        final int count = mEntryManager.getNotificationData().getActiveNotifications().size();
+        final int rank = mEntryManager.getNotificationData().getRank(entry.key);
+        final NotificationVisibility nv =
+                NotificationVisibility.obtain(entry.key, rank, count, true);
+        try {
+            mBarService.onNotificationActionClick(
+                    entry.key, actionIndex, action, nv, generatedByAssistant);
+        } catch (RemoteException e) {
+            // Nothing to do, system going down
+        }
+    }
+
+    /**
      * Have we posted an intent to an app about sending a smart reply from the
      * notification with this key.
      */
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 a4fdc08..92d1b45 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
@@ -57,7 +57,6 @@
 
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
-import java.util.Collections;
 import java.util.List;
 
 /**
@@ -1520,7 +1519,8 @@
                         smartRepliesAndActions.smartReplies, mSmartReplyController, entry);
             }
             if (smartRepliesAndActions.smartActions != null) {
-                smartReplyView.addSmartActions(smartRepliesAndActions.smartActions);
+                smartReplyView.addSmartActions(
+                        smartRepliesAndActions.smartActions, mSmartReplyController, entry);
             }
             smartReplyContainer.setVisibility(View.VISIBLE);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index 05f8f18..0a7923c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -195,7 +195,6 @@
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
 import com.android.systemui.statusbar.notification.stack.NotificationListContainer;
-import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout;
 import com.android.systemui.statusbar.phone.UnlockMethodCache.OnUnlockMethodChangedListener;
 import com.android.systemui.statusbar.policy.BatteryController;
 import com.android.systemui.statusbar.policy.BatteryController.BatteryStateChangeCallback;
@@ -4321,7 +4320,14 @@
         }, afterKeyguardGone);
     }
 
+    @Override
     public void startPendingIntentDismissingKeyguard(final PendingIntent intent) {
+        startPendingIntentDismissingKeyguard(intent, null);
+    }
+
+    @Override
+    public void startPendingIntentDismissingKeyguard(
+            final PendingIntent intent, @Nullable final Runnable intentSentCallback) {
         final boolean afterKeyguardGone = intent.isActivity()
                 && PreviewInflater.wouldLaunchResolverActivity(mContext, intent.getIntent(),
                 mLockscreenUserManager.getCurrentUserId());
@@ -4340,6 +4346,9 @@
             if (intent.isActivity()) {
                 mAssistManager.hideAssist();
             }
+            if (intentSentCallback != null) {
+                intentSentCallback.run();
+            }
         }, afterKeyguardGone);
     }
 
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 88ff078..f36066c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SmartReplyView.java
@@ -208,12 +208,14 @@
      * Add smart actions to be shown next to smart replies. Only the actions that fit into the
      * notification are shown.
      */
-    public void addSmartActions(SmartActions smartActions) {
+    public void addSmartActions(SmartActions smartActions,
+            SmartReplyController smartReplyController, NotificationData.Entry entry) {
         int numSmartActions = smartActions.actions.size();
         for (int n = 0; n < numSmartActions; n++) {
             Notification.Action action = smartActions.actions.get(n);
             if (action.actionIntent != null) {
-                Button actionButton = inflateActionButton(getContext(), this, action);
+                Button actionButton = inflateActionButton(
+                        getContext(), this, n, smartActions, smartReplyController, entry);
                 addView(actionButton);
             }
         }
@@ -270,7 +272,10 @@
     }
 
     @VisibleForTesting
-    Button inflateActionButton(Context context, ViewGroup root, Notification.Action action) {
+    Button inflateActionButton(Context context, ViewGroup root, int actionIndex,
+            SmartActions smartActions, SmartReplyController smartReplyController,
+            NotificationData.Entry entry) {
+        Notification.Action action = smartActions.actions.get(actionIndex);
         Button button = (Button) LayoutInflater.from(context).inflate(
                 R.layout.smart_action_button, root, false);
         button.setText(action.title);
@@ -283,7 +288,10 @@
         button.setCompoundDrawables(iconDrawable, null, null, null);
 
         button.setOnClickListener(view ->
-                getActivityStarter().startPendingIntentDismissingKeyguard(action.actionIntent));
+                getActivityStarter().startPendingIntentDismissingKeyguard(
+                        action.actionIntent,
+                        () -> smartReplyController.smartActionClicked(
+                                entry, actionIndex, action, smartActions.fromAssistant)));
 
         // TODO(b/119010281): handle accessibility
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
index b3b45eb..f94ba95 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/NotificationDataTest.java
@@ -121,9 +121,9 @@
         when(mEnvironment.isDeviceProvisioned()).thenReturn(true);
         when(mEnvironment.isNotificationForCurrentProfiles(any())).thenReturn(true);
         mNotificationData = new TestableNotificationData();
-        Dependency.get(InitController.class).executePostInitTasks();
         mNotificationData.updateRanking(mock(NotificationListenerService.RankingMap.class));
         mRow = new NotificationTestHelper(getContext()).createRow();
+        Dependency.get(InitController.class).executePostInitTasks();
     }
 
     @Test
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 7fee0ee..f0fa788 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
@@ -18,7 +18,6 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.hamcrest.core.IsEqual.equalTo;
 import static org.junit.Assert.assertFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyFloat;
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 b6e3fc1..df7aeab 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
@@ -62,6 +62,7 @@
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicReference;
 
@@ -430,12 +431,18 @@
 
     private void setSmartActions(String[] actionTitles) {
         mView.resetSmartSuggestions(mContainer);
-        mView.addSmartActions(new SmartReplyView.SmartActions(createActions(actionTitles), false));
+        mView.addSmartActions(
+                new SmartReplyView.SmartActions(createActions(actionTitles), false),
+                mLogger,
+                mEntry);
     }
 
     private void setSmartRepliesAndActions(CharSequence[] choices, String[] actionTitles) {
         setSmartReplies(choices);
-        mView.addSmartActions(new SmartReplyView.SmartActions(createActions(actionTitles), false));
+        mView.addSmartActions(
+                new SmartReplyView.SmartActions(createActions(actionTitles), false),
+                mLogger,
+                mEntry);
     }
 
     private ViewGroup buildExpectedView(CharSequence[] choices, int lineCount) {
@@ -553,7 +560,7 @@
 
         mView.getChildAt(2).performClick();
 
-        verify(mActivityStarter, times(1)).startPendingIntentDismissingKeyguard(any());
+        verify(mActivityStarter, times(1)).startPendingIntentDismissingKeyguard(any(), any());
     }
 
     @Test
@@ -738,7 +745,9 @@
     }
 
     private Button inflateActionButton(Notification.Action action) {
-        return mView.inflateActionButton(getContext(), mView, action);
+        return mView.inflateActionButton(getContext(), mView, 0,
+                new SmartReplyView.SmartActions(Collections.singletonList(action), false),
+                mLogger, mEntry);
     }
 
     @Test
diff --git a/services/core/java/com/android/server/notification/NotificationDelegate.java b/services/core/java/com/android/server/notification/NotificationDelegate.java
index 9222740..84bb13e 100644
--- a/services/core/java/com/android/server/notification/NotificationDelegate.java
+++ b/services/core/java/com/android/server/notification/NotificationDelegate.java
@@ -16,6 +16,7 @@
 
 package com.android.server.notification;
 
+import android.app.Notification;
 import android.service.notification.NotificationStats;
 
 import com.android.internal.statusbar.NotificationVisibility;
@@ -26,7 +27,7 @@
     void onNotificationClick(int callingUid, int callingPid, String key,
             NotificationVisibility nv);
     void onNotificationActionClick(int callingUid, int callingPid, String key, int actionIndex,
-            NotificationVisibility nv);
+            Notification.Action action, NotificationVisibility nv, boolean generatedByAssistant);
     void onNotificationClear(int callingUid, int callingPid,
             String pkg, String tag, int id, int userId, String key,
             @NotificationStats.DismissalSurface int dismissalSurface,
diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 6005872..ae27d0c 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -752,7 +752,8 @@
 
         @Override
         public void onNotificationActionClick(int callingUid, int callingPid, String key,
-                int actionIndex, NotificationVisibility nv) {
+                int actionIndex, Notification.Action action, NotificationVisibility nv,
+                boolean generatedByAssistant) {
             exitIdle();
             synchronized (mNotificationLock) {
                 NotificationRecord r = mNotificationsByKey.get(key);
@@ -772,6 +773,8 @@
                         nv.rank, nv.count);
                 nv.recycle();
                 reportUserInteraction(r);
+                mAssistants.notifyAssistantActionClicked(
+                        r.sbn, actionIndex, action, generatedByAssistant);
             }
         }
 
@@ -6923,6 +6926,27 @@
                     });
         }
 
+        @GuardedBy("mNotificationLock")
+        void notifyAssistantActionClicked(
+                final StatusBarNotification sbn, int actionIndex, Notification.Action action,
+                boolean generatedByAssistant) {
+            final String key = sbn.getKey();
+            notifyAssistantLocked(
+                    sbn,
+                    false /* sameUserOnly */,
+                    (assistant, sbnHolder) -> {
+                        try {
+                            assistant.onActionClicked(
+                                    key,
+                                    action,
+                                    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 361622f..0d66a2c8 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
@@ -19,6 +19,7 @@
 import static android.app.StatusBarManager.DISABLE2_GLOBAL_ACTIONS;
 
 import android.app.ActivityThread;
+import android.app.Notification;
 import android.app.StatusBarManager;
 import android.content.ComponentName;
 import android.content.Context;
@@ -1080,14 +1081,16 @@
     }
 
     @Override
-    public void onNotificationActionClick(String key, int actionIndex, NotificationVisibility nv) {
+    public void onNotificationActionClick(
+            String key, int actionIndex, Notification.Action action, NotificationVisibility nv,
+            boolean generatedByAssistant) {
         enforceStatusBarService();
         final int callingUid = Binder.getCallingUid();
         final int callingPid = Binder.getCallingPid();
         long identity = Binder.clearCallingIdentity();
         try {
             mNotificationDelegate.onNotificationActionClick(callingUid, callingPid, key,
-                    actionIndex, nv);
+                    actionIndex, action, nv, 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 fcd29e1..d950360 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -3711,4 +3711,23 @@
         verify(mAssistants).notifyAssistantSuggestedReplySent(
                 eq(r.sbn), eq(reply), eq(generatedByAssistant));
     }
+
+    @Test
+    public void testOnNotificationActionClick() {
+        final int actionIndex = 2;
+        final Notification.Action action =
+                new Notification.Action.Builder(null, "text", null).build();
+        final boolean generatedByAssistant = false;
+
+        NotificationRecord r = generateNotificationRecord(mTestNotificationChannel);
+        mService.addNotification(r);
+
+        NotificationVisibility notificationVisibility =
+                NotificationVisibility.obtain(r.getKey(), 1, 2, true);
+        mService.mNotificationDelegate.onNotificationActionClick(
+                10, 10, r.getKey(), actionIndex, action, notificationVisibility,
+                generatedByAssistant);
+        verify(mAssistants).notifyAssistantActionClicked(
+                eq(r.sbn), eq(actionIndex), eq(action), eq(generatedByAssistant));
+    }
 }