Migrate remoteinput to entryadapter

Test: manual (flag on and off) - direct reply from HUN and non-HUN
Test: NotificationRemoteInputManagerTest
Test: NotificationEntryAdapterTest
Bug: 395857098
Flag: com.android.systemui.notification_bundle_ui

Change-Id: I65b15fa659fec5d5af6276d7b552968a093dcb05
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
index c9d910c..01046cd 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/NotificationRemoteInputManagerTest.java
@@ -20,16 +20,29 @@
 
 import static junit.framework.Assert.assertTrue;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.ActivityOptions;
 import android.app.Notification;
+import android.app.PendingIntent;
 import android.content.Context;
+import android.content.Intent;
 import android.os.SystemClock;
 import android.os.UserHandle;
+import android.platform.test.flag.junit.FlagsParameterization;
 import android.testing.TestableLooper;
+import android.util.Pair;
+import android.view.View;
+import android.widget.LinearLayout;
+import android.widget.RemoteViews;
 
-import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
@@ -42,19 +55,37 @@
 import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder;
 import com.android.systemui.statusbar.notification.collection.render.NotificationVisibilityProvider;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi;
 import com.android.systemui.statusbar.policy.RemoteInputUriController;
 import com.android.systemui.util.kotlin.JavaAdapter;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
 @SmallTest
-@RunWith(AndroidJUnit4.class)
+@RunWith(ParameterizedAndroidJunit4.class)
 @TestableLooper.RunWithLooper
 public class NotificationRemoteInputManagerTest extends SysuiTestCase {
+
+    @Parameters(name = "{0}")
+    public static List<FlagsParameterization> getParams() {
+        return FlagsParameterization.allCombinationsOf(NotificationBundleUi.FLAG_NAME);
+    }
+
     private static final String TEST_PACKAGE_NAME = "test";
     private static final int TEST_UID = 0;
 
@@ -69,14 +100,34 @@
     @Mock private NotificationClickNotifier mClickNotifier;
     @Mock private NotificationLockscreenUserManager mLockscreenUserManager;
     @Mock private PowerInteractor mPowerInteractor;
+    @Mock
+    NotificationRemoteInputManager.RemoteInputListener mRemoteInputListener;
+    private ActionClickLogger mActionClickLogger;
+    @Captor
+    ArgumentCaptor<NotificationRemoteInputManager.ClickHandler> mClickHandlerArgumentCaptor;
+    private Context mSpyContext;
 
+    private NotificationTestHelper mTestHelper;
     private TestableNotificationRemoteInputManager mRemoteInputManager;
     private NotificationEntry mEntry;
 
+   public NotificationRemoteInputManagerTest(FlagsParameterization flags) {
+       super();
+       mSetFlagsRule.setFlagsParameterization(flags);
+   }
+
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mSpyContext = spy(mContext);
+        doNothing().when(mSpyContext).startIntentSender(
+                any(), any(), anyInt(), anyInt(), anyInt(), any());
+
+
+        mTestHelper = new NotificationTestHelper(mSpyContext, mDependency);
+        mActionClickLogger = spy(new ActionClickLogger(logcatLogBuffer()));
+
         mRemoteInputManager = new TestableNotificationRemoteInputManager(mContext,
                 mock(NotifPipelineFlags.class),
                 mLockscreenUserManager,
@@ -87,9 +138,10 @@
                 mRemoteInputUriController,
                 new RemoteInputControllerLogger(logcatLogBuffer()),
                 mClickNotifier,
-                new ActionClickLogger(logcatLogBuffer()),
+                mActionClickLogger,
                 mock(JavaAdapter.class),
                 mock(ShadeInteractor.class));
+        mRemoteInputManager.setRemoteInputListener(mRemoteInputListener);
         mEntry = new NotificationEntryBuilder()
                 .setPkg(TEST_PACKAGE_NAME)
                 .setOpPkg(TEST_PACKAGE_NAME)
@@ -133,6 +185,70 @@
         assertTrue(mRemoteInputManager.shouldKeepForSmartReplyHistory(mEntry));
     }
 
+    @Test
+    public void testActionClick() throws Exception {
+        RemoteViews.RemoteResponse response = mock(RemoteViews.RemoteResponse.class);
+        when(response.getLaunchOptions(any())).thenReturn(
+                Pair.create(mock(Intent.class), mock(ActivityOptions.class)));
+        ExpandableNotificationRow row = getRowWithReplyAction();
+        View actionView = ((LinearLayout) row.getPrivateLayout().getExpandedChild().findViewById(
+                com.android.internal.R.id.actions)).getChildAt(0);
+        Notification n = getNotification(row);
+        CountDownLatch latch = new CountDownLatch(1);
+        Consumer<NotificationEntry> consumer = notificationEntry -> latch.countDown();
+        if (!NotificationBundleUi.isEnabled()) {
+            mRemoteInputManager.addActionPressListener(consumer);
+        }
+
+        mRemoteInputManager.getRemoteViewsOnClickHandler().onInteraction(
+                actionView,
+                n.actions[0].actionIntent,
+                response);
+
+        verify(mActionClickLogger).logInitialClick(row.getKey(), 0, n.actions[0].actionIntent);
+        verify(mClickNotifier).onNotificationActionClick(
+                eq(row.getKey()), eq(0), eq(n.actions[0]), any(), eq(false));
+        verify(mCallback).handleRemoteViewClick(eq(actionView), eq(n.actions[0].actionIntent),
+                eq(false), eq(0), mClickHandlerArgumentCaptor.capture());
+
+        mClickHandlerArgumentCaptor.getValue().handleClick();
+        verify(mActionClickLogger).logStartingIntentWithDefaultHandler(
+                row.getKey(), n.actions[0].actionIntent, 0);
+
+        verify(mRemoteInputListener).releaseNotificationIfKeptForRemoteInputHistory(row.getKey());
+        if (NotificationBundleUi.isEnabled()) {
+            verify(row.getEntryAdapter()).onNotificationActionClicked();
+        } else {
+            latch.await(10, TimeUnit.MILLISECONDS);
+        }
+    }
+
+    private Notification getNotification(ExpandableNotificationRow row) {
+        if (NotificationBundleUi.isEnabled()) {
+            return row.getEntryAdapter().getSbn().getNotification();
+        } else {
+            return row.getEntry().getSbn().getNotification();
+        }
+    }
+
+    private ExpandableNotificationRow getRowWithReplyAction() throws Exception {
+        PendingIntent pi = PendingIntent.getBroadcast(getContext(), 0, new Intent("Action"),
+                PendingIntent.FLAG_IMMUTABLE);
+        Notification n = new Notification.Builder(mSpyContext, "")
+                .setSmallIcon(com.android.systemui.res.R.drawable.ic_person)
+                .addAction(new Notification.Action(com.android.systemui.res.R.drawable.ic_person,
+                        "reply", pi))
+                .build();
+        ExpandableNotificationRow row = mTestHelper.createRow(n);
+        row.onNotificationUpdated();
+        row.getPrivateLayout().setExpandedChild(Notification.Builder.recoverBuilder(mSpyContext, n)
+                .createBigContentView().apply(
+                        mSpyContext,
+                        row.getPrivateLayout(),
+                        mRemoteInputManager.getRemoteViewsOnClickHandler()));
+        return row;
+    }
+
     private class TestableNotificationRemoteInputManager extends NotificationRemoteInputManager {
 
         TestableNotificationRemoteInputManager(
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapterTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapterTest.kt
index b6889af..faafa07 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapterTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapterTest.kt
@@ -29,8 +29,10 @@
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.res.R
 import com.android.systemui.statusbar.RankingBuilder
+import com.android.systemui.statusbar.notification.mockNotificationActivityStarter
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.entryAdapterFactory
+import com.android.systemui.statusbar.notification.row.mockNotificationActionClickManager
 import com.android.systemui.statusbar.notification.shared.NotificationBundleUi
 import com.android.systemui.statusbar.notification.stack.BUCKET_ALERTING
 import com.android.systemui.testKosmos
@@ -355,16 +357,27 @@
         val notification: Notification =
             Notification.Builder(mContext, "").setSmallIcon(R.drawable.ic_person).build()
 
-        val entry =
-            NotificationEntryBuilder()
-                .setNotification(notification)
-                .setImportance(NotificationManager.IMPORTANCE_MIN)
-                .build()
+        val entry = NotificationEntryBuilder().setNotification(notification).build()
 
         underTest = factory.create(entry) as NotificationEntryAdapter
 
         underTest.onNotificationBubbleIconClicked()
-        verify((factory as? EntryAdapterFactoryImpl)?.getNotificationActivityStarter())
-            ?.onNotificationBubbleIconClicked(entry)
+        verify(kosmos.mockNotificationActivityStarter).onNotificationBubbleIconClicked(entry)
+    }
+
+    @Test
+    @EnableFlags(NotificationBundleUi.FLAG_NAME)
+    fun onNotificationActionClicked() {
+        val notification: Notification =
+            Notification.Builder(mContext, "")
+                .setSmallIcon(R.drawable.ic_person)
+                .addAction(Mockito.mock(Notification.Action::class.java))
+                .build()
+
+        val entry = NotificationEntryBuilder().setNotification(notification).build()
+
+        underTest = factory.create(entry) as NotificationEntryAdapter
+        underTest.onNotificationActionClicked()
+        verify(kosmos.mockNotificationActionClickManager).onNotificationActionClicked(entry)
     }
 }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
index 3098355..44d88c3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinatorTest.kt
@@ -50,6 +50,8 @@
 import com.android.systemui.statusbar.notification.interruption.NotificationInterruptStateProviderWrapper.DecisionImpl
 import com.android.systemui.statusbar.notification.interruption.NotificationInterruptStateProviderWrapper.FullScreenIntentDecisionImpl
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionDecisionProvider
+import com.android.systemui.statusbar.notification.row.mockNotificationActionClickManager
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi
 import com.android.systemui.statusbar.phone.NotificationGroupTestHelper
 import com.android.systemui.testKosmos
 import com.android.systemui.util.concurrency.FakeExecutor
@@ -138,6 +140,7 @@
                 headsUpViewBinder,
                 visualInterruptionDecisionProvider,
                 remoteInputManager,
+                kosmos.mockNotificationActionClickManager,
                 launchFullScreenIntentProvider,
                 flags,
                 statusBarNotificationChipsInteractor,
@@ -161,8 +164,14 @@
             verify(notifPipeline).addOnBeforeFinalizeFilterListener(capture())
         }
         onHeadsUpChangedListener = withArgCaptor { verify(headsUpManager).addListener(capture()) }
-        actionPressListener = withArgCaptor {
-            verify(remoteInputManager).addActionPressListener(capture())
+        actionPressListener = if (NotificationBundleUi.isEnabled) {
+            withArgCaptor {
+                verify(kosmos.mockNotificationActionClickManager).addActionClickListener(capture())
+            }
+        } else {
+            withArgCaptor {
+                verify(remoteInputManager).addActionPressListener(capture())
+            }
         }
         given(headsUpManager.allEntries).willAnswer { huns.stream() }
         given(headsUpManager.isHeadsUpEntry(anyString())).willAnswer { invocation ->
@@ -260,7 +269,7 @@
         addHUN(entry)
 
         actionPressListener.accept(entry)
-        verify(headsUpManager, times(1)).setUserActionMayIndirectlyRemove(entry)
+        verify(headsUpManager, times(1)).setUserActionMayIndirectlyRemove(entry.key)
 
         whenever(headsUpManager.canRemoveImmediately(anyString())).thenReturn(true)
         assertFalse(notifLifetimeExtender.maybeExtendLifetime(entry, 0))
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImplTest.kt
index 9804932..8560b66 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImplTest.kt
@@ -1071,7 +1071,7 @@
 
         assertThat(underTest.canRemoveImmediately(notifEntry.key)).isFalse()
 
-        underTest.setUserActionMayIndirectlyRemove(notifEntry)
+        underTest.setUserActionMayIndirectlyRemove(notifEntry.key)
 
         assertThat(underTest.canRemoveImmediately(notifEntry.key)).isTrue()
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
index 6415f8c..7847437 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
@@ -29,6 +29,7 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -743,11 +744,12 @@
                 mock(MetricsLogger.class),
                 mock(PeopleNotificationIdentifier.class),
                 mock(NotificationIconStyleProvider.class),
-                mock(VisualStabilityCoordinator.class)
+                mock(VisualStabilityCoordinator.class),
+                mock(NotificationActionClickManager.class)
         ).create(entry);
 
         row.initialize(
-                entryAdapter,
+                spy(entryAdapter),
                 entry,
                 mock(RemoteInputViewSubcomponent.Factory.class),
                 APP_NAME,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
index 0d34bdc..041ed65 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationRemoteInputManager.java
@@ -59,11 +59,13 @@
 import com.android.systemui.statusbar.dagger.CentralSurfacesDependenciesModule;
 import com.android.systemui.statusbar.notification.NotifPipelineFlags;
 import com.android.systemui.statusbar.notification.RemoteInputControllerLogger;
+import com.android.systemui.statusbar.notification.collection.EntryAdapter;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry.EditedSuggestionInfo;
 import com.android.systemui.statusbar.notification.collection.render.NotificationVisibilityProvider;
 import com.android.systemui.statusbar.notification.logging.NotificationLogger;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi;
 import com.android.systemui.statusbar.phone.ExpandHeadsUpOnInlineReply;
 import com.android.systemui.statusbar.policy.RemoteInputUriController;
 import com.android.systemui.statusbar.policy.RemoteInputView;
@@ -134,20 +136,21 @@
                     view.getTag(com.android.internal.R.id.notification_action_index_tag);
 
             final ExpandableNotificationRow row = getNotificationRowForParent(view.getParent());
-            final NotificationEntry entry = getNotificationForParent(view.getParent());
-            mLogger.logInitialClick(
-                    row != null ? row.getLoggingKey() : null, actionIndex, pendingIntent);
+            if (row == null) {
+                return false;
+            }
+            mLogger.logInitialClick(row.getLoggingKey(), actionIndex, pendingIntent);
 
             if (handleRemoteInput(view, pendingIntent)) {
-                mLogger.logRemoteInputWasHandled(
-                        row != null ? row.getLoggingKey() : null, actionIndex);
+                mLogger.logRemoteInputWasHandled(row.getLoggingKey(), actionIndex);
                 return true;
             }
 
             if (DEBUG) {
                 Log.v(TAG, "Notification click handler invoked for intent: " + pendingIntent);
             }
-            logActionClick(view, entry, pendingIntent);
+            Notification.Action action = getActionFromView(view, row, pendingIntent);
+            logActionClick(view, row.getKey(), action);
             // The intent we are sending is for the application, which
             // won't have permission to immediately start an activity after
             // the user switches to home.  We know it is safe to do at this
@@ -156,33 +159,47 @@
                 ActivityManager.getService().resumeAppSwitches();
             } catch (RemoteException e) {
             }
-            Notification.Action action = getActionFromView(view, entry, pendingIntent);
             return mCallback.handleRemoteViewClick(view, pendingIntent,
                     action == null ? false : action.isAuthenticationRequired(), actionIndex, () -> {
                     Pair<Intent, ActivityOptions> options = response.getLaunchOptions(view);
                     mLogger.logStartingIntentWithDefaultHandler(
-                            row != null ? row.getLoggingKey() : null, pendingIntent, actionIndex);
+                             row.getLoggingKey(), pendingIntent, actionIndex);
                     boolean started = RemoteViews.startPendingIntent(view, pendingIntent, options);
-                    if (started) releaseNotificationIfKeptForRemoteInputHistory(entry);
+                    if (started) {
+                        if (NotificationBundleUi.isEnabled()) {
+                            releaseNotificationIfKeptForRemoteInputHistory(row.getEntryAdapter());
+                        } else {
+                            releaseNotificationIfKeptForRemoteInputHistory(row.getEntry());
+                        }
+                    }
                     return started;
             });
         }
 
         private @Nullable Notification.Action getActionFromView(View view,
-                NotificationEntry entry, PendingIntent actionIntent) {
+                ExpandableNotificationRow row, PendingIntent actionIntent) {
             Integer actionIndex = (Integer)
                     view.getTag(com.android.internal.R.id.notification_action_index_tag);
             if (actionIndex == null) {
                 return null;
             }
-            if (entry == null) {
+            StatusBarNotification statusBarNotification = null;
+            if (NotificationBundleUi.isEnabled()) {
+                if (row.getEntryAdapter() != null) {
+                    statusBarNotification = row.getEntryAdapter().getSbn();
+                }
+            } else {
+                if (row.getEntry() != null) {
+                    statusBarNotification = row.getEntry().getSbn();
+                }
+            }
+            if (statusBarNotification == null) {
                 Log.w(TAG, "Couldn't determine notification for click.");
                 return null;
             }
 
             // Notification may be updated before this function is executed, and thus play safe
             // here and verify that the action object is still the one that where the click happens.
-            StatusBarNotification statusBarNotification = entry.getSbn();
             Notification.Action[] actions = statusBarNotification.getNotification().actions;
             if (actions == null || actionIndex >= actions.length) {
                 Log.w(TAG, "statusBarNotification.getNotification().actions is null or invalid");
@@ -199,14 +216,12 @@
 
         private void logActionClick(
                 View view,
-                NotificationEntry entry,
-                PendingIntent actionIntent) {
-            Notification.Action action = getActionFromView(view, entry, actionIntent);
+                String key,
+                Notification.Action action) {
             if (action == null) {
                 return;
             }
             ViewParent parent = view.getParent();
-            String key = entry.getSbn().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 &&
@@ -214,20 +229,10 @@
                 ViewGroup actionGroup = (ViewGroup) parent;
                 buttonIndex = actionGroup.indexOfChild(view);
             }
-            final NotificationVisibility nv = mVisibilityProvider.obtain(entry, true);
+            final NotificationVisibility nv = mVisibilityProvider.obtain(key, true);
             mClickNotifier.onNotificationActionClick(key, buttonIndex, action, nv, false);
         }
 
-        private NotificationEntry getNotificationForParent(ViewParent parent) {
-            while (parent != null) {
-                if (parent instanceof ExpandableNotificationRow) {
-                    return ((ExpandableNotificationRow) parent).getEntry();
-                }
-                parent = parent.getParent();
-            }
-            return null;
-        }
-
         private @Nullable ExpandableNotificationRow getNotificationRowForParent(ViewParent parent) {
             while (parent != null) {
                 if (parent instanceof ExpandableNotificationRow) {
@@ -394,11 +399,21 @@
         }
     }
 
+    /**
+     * Use {@link com.android.systemui.statusbar.notification.row.NotificationActionClickManager}
+     * instead
+     */
     public void addActionPressListener(Consumer<NotificationEntry> listener) {
+        NotificationBundleUi.assertInLegacyMode();
         mActionPressListeners.addIfAbsent(listener);
     }
 
+    /**
+     * Use {@link com.android.systemui.statusbar.notification.row.NotificationActionClickManager}
+     * instead
+     */
     public void removeActionPressListener(Consumer<NotificationEntry> listener) {
+        NotificationBundleUi.assertInLegacyMode();
         mActionPressListeners.remove(listener);
     }
 
@@ -681,12 +696,30 @@
      * (after unlock, if applicable), and will then wait a short time to allow the app to update the
      * notification in response to the action.
      */
+    private void releaseNotificationIfKeptForRemoteInputHistory(EntryAdapter entryAdapter) {
+        if (entryAdapter == null) {
+            return;
+        }
+        if (mRemoteInputListener != null) {
+            mRemoteInputListener.releaseNotificationIfKeptForRemoteInputHistory(
+                    entryAdapter.getKey());
+        }
+        entryAdapter.onNotificationActionClicked();
+    }
+
+    /**
+     * Checks if the notification is being kept due to the user sending an inline reply, and if
+     * so, releases that hold.  This is called anytime an action on the notification is dispatched
+     * (after unlock, if applicable), and will then wait a short time to allow the app to update the
+     * notification in response to the action.
+     */
     private void releaseNotificationIfKeptForRemoteInputHistory(NotificationEntry entry) {
+        NotificationBundleUi.assertInLegacyMode();
         if (entry == null) {
             return;
         }
         if (mRemoteInputListener != null) {
-            mRemoteInputListener.releaseNotificationIfKeptForRemoteInputHistory(entry);
+            mRemoteInputListener.releaseNotificationIfKeptForRemoteInputHistory(entry.getKey());
         }
         for (Consumer<NotificationEntry> listener : mActionPressListeners) {
             listener.accept(entry);
@@ -866,7 +899,7 @@
         boolean isNotificationKeptForRemoteInputHistory(@NonNull String key);
 
         /** Called on user interaction to end lifetime extension for history */
-        void releaseNotificationIfKeptForRemoteInputHistory(@NonNull NotificationEntry entry);
+        void releaseNotificationIfKeptForRemoteInputHistory(@NonNull String entryKey);
 
         /** Called when the RemoteInputController is attached to the manager */
         void setRemoteInputController(@NonNull RemoteInputController remoteInputController);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntryAdapter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntryAdapter.kt
index 64db9df..26c302b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntryAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/BundleEntryAdapter.kt
@@ -20,6 +20,7 @@
 import android.content.Context
 import android.os.Build
 import android.service.notification.StatusBarNotification
+import android.util.Log
 import com.android.systemui.statusbar.notification.icon.IconPack
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import kotlinx.coroutines.flow.StateFlow
@@ -118,5 +119,13 @@
 
     override fun onNotificationBubbleIconClicked() {
         // do nothing. these cannot be a bubble
+        Log.wtf(TAG, "onNotificationBubbleIconClicked() called")
+    }
+
+    override fun onNotificationActionClicked() {
+        // do nothing. these have no actions
+        Log.wtf(TAG, "onNotificationActionClicked() called")
     }
 }
+
+private const val TAG = "BundleEntryAdapter"
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
index 0e75b60..3118ce5 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapter.java
@@ -140,5 +140,10 @@
      * Process a click on a notification bubble icon
      */
     void onNotificationBubbleIconClicked();
+
+    /**
+     * Processes a click on a notification action
+     */
+    void onNotificationActionClicked();
 }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapterFactoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapterFactoryImpl.kt
index 779c25a..a516986 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapterFactoryImpl.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/EntryAdapterFactoryImpl.kt
@@ -16,11 +16,11 @@
 
 package com.android.systemui.statusbar.notification.collection
 
-import androidx.annotation.VisibleForTesting
 import com.android.internal.logging.MetricsLogger
 import com.android.systemui.statusbar.notification.NotificationActivityStarter
 import com.android.systemui.statusbar.notification.collection.coordinator.VisualStabilityCoordinator
 import com.android.systemui.statusbar.notification.people.PeopleNotificationIdentifier
+import com.android.systemui.statusbar.notification.row.NotificationActionClickManager
 import com.android.systemui.statusbar.notification.row.icon.NotificationIconStyleProvider
 import javax.inject.Inject
 
@@ -33,6 +33,7 @@
     private val peopleNotificationIdentifier: PeopleNotificationIdentifier,
     private val iconStyleProvider: NotificationIconStyleProvider,
     private val visualStabilityCoordinator: VisualStabilityCoordinator,
+    private val notificationActionClickManager: NotificationActionClickManager,
 ) : EntryAdapterFactory {
     override fun create(entry: PipelineEntry): EntryAdapter {
         return if (entry is NotificationEntry) {
@@ -42,15 +43,11 @@
                 peopleNotificationIdentifier,
                 iconStyleProvider,
                 visualStabilityCoordinator,
+                notificationActionClickManager,
                 entry,
             )
         } else {
             BundleEntryAdapter((entry as BundleEntry))
         }
     }
-
-    @VisibleForTesting
-    fun getNotificationActivityStarter() : NotificationActivityStarter {
-        return notificationActivityStarter
-    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapter.kt
index 0ff2dd7..1168c64 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationEntryAdapter.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.statusbar.notification.icon.IconPack
 import com.android.systemui.statusbar.notification.people.PeopleNotificationIdentifier
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
+import com.android.systemui.statusbar.notification.row.NotificationActionClickManager
 import com.android.systemui.statusbar.notification.row.icon.NotificationIconStyleProvider
 import kotlinx.coroutines.flow.StateFlow
 
@@ -33,6 +34,7 @@
     private val peopleNotificationIdentifier: PeopleNotificationIdentifier,
     private val iconStyleProvider: NotificationIconStyleProvider,
     private val visualStabilityCoordinator: VisualStabilityCoordinator,
+    private val notificationActionClickManager: NotificationActionClickManager,
     private val entry: NotificationEntry,
 ) : EntryAdapter {
 
@@ -142,4 +144,8 @@
     override fun onNotificationBubbleIconClicked() {
         notificationActivityStarter.onNotificationBubbleIconClicked(entry)
     }
+
+    override fun onNotificationActionClicked() {
+        notificationActionClickManager.onNotificationActionClicked(entry)
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
index fdb8cd8..9637b7c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/HeadsUpCoordinator.kt
@@ -47,7 +47,9 @@
 import com.android.systemui.statusbar.notification.interruption.HeadsUpViewBinder
 import com.android.systemui.statusbar.notification.interruption.VisualInterruptionDecisionProvider
 import com.android.systemui.statusbar.notification.logKey
+import com.android.systemui.statusbar.notification.row.NotificationActionClickManager
 import com.android.systemui.statusbar.notification.shared.GroupHunAnimationFix
+import com.android.systemui.statusbar.notification.shared.NotificationBundleUi
 import com.android.systemui.statusbar.notification.stack.BUCKET_HEADS_UP
 import com.android.systemui.util.concurrency.DelayableExecutor
 import com.android.systemui.util.time.SystemClock
@@ -82,6 +84,7 @@
     private val mHeadsUpViewBinder: HeadsUpViewBinder,
     private val mVisualInterruptionDecisionProvider: VisualInterruptionDecisionProvider,
     private val mRemoteInputManager: NotificationRemoteInputManager,
+    private val notificationActionClickManager: NotificationActionClickManager,
     private val mLaunchFullScreenIntentProvider: LaunchFullScreenIntentProvider,
     private val mFlags: NotifPipelineFlags,
     private val statusBarNotificationChipsInteractor: StatusBarNotificationChipsInteractor,
@@ -107,7 +110,11 @@
         pipeline.addOnBeforeFinalizeFilterListener(::onBeforeFinalizeFilter)
         pipeline.addPromoter(mNotifPromoter)
         pipeline.addNotificationLifetimeExtender(mLifetimeExtender)
-        mRemoteInputManager.addActionPressListener(mActionPressListener)
+        if (NotificationBundleUi.isEnabled) {
+            notificationActionClickManager.addActionClickListener(mActionPressListener)
+        } else {
+            mRemoteInputManager.addActionPressListener(mActionPressListener)
+        }
 
         if (StatusBarNotifChips.isEnabled) {
             applicationScope.launch {
@@ -781,7 +788,7 @@
      */
     private val mActionPressListener =
         Consumer<NotificationEntry> { entry ->
-            mHeadsUpManager.setUserActionMayIndirectlyRemove(entry)
+            mHeadsUpManager.setUserActionMayIndirectlyRemove(entry.key)
             mExecutor.execute { endNotifLifetimeExtensionIfExtended(entry) }
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
index e7c767f..27c0dcc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/RemoteInputCoordinator.kt
@@ -221,20 +221,20 @@
                 mSmartReplyHistoryExtender.isExtending(key)
         } else false
 
-    override fun releaseNotificationIfKeptForRemoteInputHistory(entry: NotificationEntry) {
-        if (DEBUG) Log.d(TAG, "releaseNotificationIfKeptForRemoteInputHistory(entry=${entry.key})")
+    override fun releaseNotificationIfKeptForRemoteInputHistory(entryKey: String) {
+        if (DEBUG) Log.d(TAG, "releaseNotificationIfKeptForRemoteInputHistory(entry=${entryKey})")
         if (!lifetimeExtensionRefactor()) {
             mRemoteInputHistoryExtender.endLifetimeExtensionAfterDelay(
-                entry.key,
+                entryKey,
                 REMOTE_INPUT_EXTENDER_RELEASE_DELAY,
             )
             mSmartReplyHistoryExtender.endLifetimeExtensionAfterDelay(
-                entry.key,
+                entryKey,
                 REMOTE_INPUT_EXTENDER_RELEASE_DELAY,
             )
         }
         mRemoteInputActiveExtender.endLifetimeExtensionAfterDelay(
-            entry.key,
+            entryKey,
             REMOTE_INPUT_EXTENDER_RELEASE_DELAY,
         )
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationsModule.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationsModule.java
index ef3da94..1e5aa01 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationsModule.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationsModule.java
@@ -83,6 +83,7 @@
 import com.android.systemui.statusbar.notification.promoted.PromotedNotificationContentExtractor;
 import com.android.systemui.statusbar.notification.promoted.PromotedNotificationContentExtractorImpl;
 import com.android.systemui.statusbar.notification.promoted.shared.model.PromotedNotificationContentModel;
+import com.android.systemui.statusbar.notification.row.NotificationActionClickManager;
 import com.android.systemui.statusbar.notification.row.NotificationEntryProcessorFactory;
 import com.android.systemui.statusbar.notification.row.NotificationEntryProcessorFactoryLooperImpl;
 import com.android.systemui.statusbar.notification.row.NotificationGutsManager;
@@ -346,4 +347,5 @@
     /** Provides an instance of {@link EntryAdapterFactory} */
     @Binds
     EntryAdapterFactory provideEntryAdapterFactory(EntryAdapterFactoryImpl impl);
+
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
index 9728fcf..25ae50c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManager.kt
@@ -19,7 +19,6 @@
 import android.graphics.Region
 import com.android.systemui.Dumpable
 import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.statusbar.notification.collection.EntryAdapter
 import com.android.systemui.statusbar.notification.collection.NotificationEntry
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import dagger.Binds
@@ -155,9 +154,9 @@
     fun setAnimationStateHandler(handler: AnimationStateHandler)
 
     /**
-    * Set an entry to be expanded and therefore stick in the heads up area if it's pinned until
-    * it's collapsed again.
-    */
+     * Set an entry to be expanded and therefore stick in the heads up area if it's pinned until
+     * it's collapsed again.
+     */
     fun setExpanded(key: String, row: ExpandableNotificationRow, expanded: Boolean)
 
     /**
@@ -199,12 +198,12 @@
      * Notes that the user took an action on an entry that might indirectly cause the system or the
      * app to remove the notification.
      *
-     * @param entry the entry that might be indirectly removed by the user's action
+     * @param entry the key of the entry that might be indirectly removed by the user's action
      * @see
      *   com.android.systemui.statusbar.notification.collection.coordinator.HeadsUpCoordinator.mActionPressListener
      * @see .canRemoveImmediately
      */
-    fun setUserActionMayIndirectlyRemove(entry: NotificationEntry)
+    fun setUserActionMayIndirectlyRemove(entryKey: String)
 
     /**
      * Decides whether a click is invalid for a notification, i.e. it has not been shown long enough
@@ -332,7 +331,7 @@
 
     override fun setUser(user: Int) {}
 
-    override fun setUserActionMayIndirectlyRemove(entry: NotificationEntry) {}
+    override fun setUserActionMayIndirectlyRemove(entryKey: String) {}
 
     override fun shouldSwallowClick(key: String): Boolean = false
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
index ca94655..ca83666 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/headsup/HeadsUpManagerImpl.java
@@ -1126,8 +1126,8 @@
      * @see HeadsUpCoordinator.mActionPressListener
      * @see #canRemoveImmediately(String)
      */
-    public void setUserActionMayIndirectlyRemove(@NonNull NotificationEntry entry) {
-        HeadsUpEntry headsUpEntry = getHeadsUpEntry(entry.getKey());
+    public void setUserActionMayIndirectlyRemove(@NonNull String entryKey) {
+        HeadsUpEntry headsUpEntry = getHeadsUpEntry(entryKey);
         if (headsUpEntry != null) {
             headsUpEntry.mUserActionMayIndirectlyRemove = true;
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManager.kt
new file mode 100644
index 0000000..2b45140
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManager.kt
@@ -0,0 +1,46 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.row
+
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.util.ListenerSet
+import java.util.function.Consumer
+import javax.inject.Inject
+
+/**
+ * Pipeline components can register consumers here to be informed when a notification action is
+ * clicked
+ */
+@SysUISingleton
+class NotificationActionClickManager @Inject constructor() {
+    private val actionClickListeners = ListenerSet<Consumer<NotificationEntry>>()
+
+    fun addActionClickListener(listener: Consumer<NotificationEntry>) {
+        actionClickListeners.addIfAbsent(listener)
+    }
+
+    fun removeActionClickListener(listener: Consumer<NotificationEntry>) {
+        actionClickListeners.remove(listener)
+    }
+
+    fun onNotificationActionClicked(entry: NotificationEntry) {
+        for (listener in actionClickListeners) {
+            listener.accept(entry)
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/EntryAdapterFactoryKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/EntryAdapterFactoryKosmos.kt
index e99f61e..067e420 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/EntryAdapterFactoryKosmos.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/EntryAdapterFactoryKosmos.kt
@@ -25,12 +25,13 @@
 import com.android.systemui.statusbar.notification.row.icon.notificationIconStyleProvider
 
 val Kosmos.entryAdapterFactory by
-Kosmos.Fixture {
-    EntryAdapterFactoryImpl(
-        mockNotificationActivityStarter,
-        metricsLogger,
-        peopleNotificationIdentifier,
-        notificationIconStyleProvider,
-        visualStabilityCoordinator,
-    )
-}
\ No newline at end of file
+    Kosmos.Fixture {
+        EntryAdapterFactoryImpl(
+            mockNotificationActivityStarter,
+            metricsLogger,
+            peopleNotificationIdentifier,
+            notificationIconStyleProvider,
+            visualStabilityCoordinator,
+            mockNotificationActionClickManager,
+        )
+    }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
index 7f012ba..6a674ca 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowBuilder.kt
@@ -375,6 +375,7 @@
                     Mockito.mock(PeopleNotificationIdentifier::class.java),
                     Mockito.mock(NotificationIconStyleProvider::class.java),
                     Mockito.mock(VisualStabilityCoordinator::class.java),
+                    Mockito.mock(NotificationActionClickManager::class.java),
                 )
                 .create(entry)
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManagerKosmos.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManagerKosmos.kt
new file mode 100644
index 0000000..8e62ae8
--- /dev/null
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/statusbar/notification/row/NotificationActionClickManagerKosmos.kt
@@ -0,0 +1,23 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.statusbar.notification.row
+
+import com.android.systemui.kosmos.Kosmos
+import org.mockito.kotlin.mock
+
+var Kosmos.mockNotificationActionClickManager: NotificationActionClickManager by
+    Kosmos.Fixture { mock<NotificationActionClickManager>() }