diff --git a/services/core/java/com/android/server/notification/NotificationManagerService.java b/services/core/java/com/android/server/notification/NotificationManagerService.java
index 56e0a89..9866f42 100644
--- a/services/core/java/com/android/server/notification/NotificationManagerService.java
+++ b/services/core/java/com/android/server/notification/NotificationManagerService.java
@@ -12009,6 +12009,10 @@
         if (record != null && (record.getSbn().getNotification().flags
                 & FLAG_LIFETIME_EXTENDED_BY_DIRECT_REPLY) > 0
                 && !record.isCanceledAfterLifetimeExtension()) {
+            // Mark that the notification is being updated due to cancelation, so it won't
+            // be updated again if the app cancels multiple times.
+            record.setCanceledAfterLifetimeExtension(true);
+
             boolean isAppForeground = pkg != null && packageImportance == IMPORTANCE_FOREGROUND;
 
             // Save the original Record's post silently value, so we can restore it after we send
@@ -12024,9 +12028,6 @@
             PostNotificationTracker tracker = mPostNotificationTrackerFactory.newTracker(null);
             tracker.addCleanupRunnable(() -> {
                 synchronized (mNotificationLock) {
-                    // Mark that the notification has been updated due to cancelation, so it won't
-                    // be updated again if the app cancels multiple times.
-                    record.setCanceledAfterLifetimeExtension(true);
                     // Set the post silently status to the record's previous value.
                     record.setPostSilently(savedPostSilentlyState);
                     // Remove FLAG_ONLY_ALERT_ONCE if the notification did not previously have it.
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 3c120e1..1349ee0 100644
--- a/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
+++ b/services/tests/uiservicestests/src/com/android/server/notification/NotificationManagerServiceTest.java
@@ -3206,7 +3206,6 @@
         // Send two cancelations.
         mBinderService.cancelNotificationWithTag(mPkg, mPkg, sbn.getTag(), sbn.getId(),
                 sbn.getUserId());
-        waitForIdle();
         mBinderService.cancelNotificationWithTag(mPkg, mPkg, sbn.getTag(), sbn.getId(),
                 sbn.getUserId());
         waitForIdle();