Add shade sections for NAS classification

Test: BundleCoordinatorTest
Fixes: 350541684
Flag: android.service.notification.notification_classification
Change-Id: I0353411b3b6877cf034333dfbdca82491d1f7b39

11 files changed, 455 insertions, 16 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinatorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinatorTest.kt
new file mode 100644
index 000000000000..8a9720ea3cb0
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinatorTest.kt
@@ -0,0 +1,104 @@
+/*
+ * Copyright (C) 2024 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.collection.coordinator
+
+import android.app.NotificationChannel
+import android.app.NotificationChannel.NEWS_ID
+import android.app.NotificationChannel.PROMOTIONS_ID
+import android.app.NotificationChannel.RECS_ID
+import android.app.NotificationChannel.SOCIAL_MEDIA_ID
+import android.testing.TestableLooper
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.notification.collection.NotificationEntry
+import com.android.systemui.statusbar.notification.collection.NotificationEntryBuilder
+import com.android.systemui.statusbar.notification.collection.render.NodeController
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@TestableLooper.RunWithLooper
+class BundleCoordinatorTest : SysuiTestCase() {
+    @Mock private lateinit var newsController: NodeController
+    @Mock private lateinit var socialController: NodeController
+    @Mock private lateinit var recsController: NodeController
+    @Mock private lateinit var promoController: NodeController
+
+    private lateinit var coordinator: BundleCoordinator
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        coordinator =
+            BundleCoordinator(
+                newsController,
+                socialController,
+                recsController,
+                promoController
+            )
+    }
+
+    @Test
+    fun newsSectioner() {
+        assertThat(coordinator.newsSectioner.isInSection(makeEntryOfChannelType(NEWS_ID))).isTrue()
+        assertThat(coordinator.newsSectioner.isInSection(makeEntryOfChannelType("news"))).isFalse()
+    }
+
+    @Test
+    fun socialSectioner() {
+        assertThat(coordinator.socialSectioner.isInSection(makeEntryOfChannelType(SOCIAL_MEDIA_ID)))
+            .isTrue()
+        assertThat(coordinator.socialSectioner.isInSection(makeEntryOfChannelType("social")))
+            .isFalse()
+    }
+
+    @Test
+    fun recsSectioner() {
+        assertThat(coordinator.recsSectioner.isInSection(makeEntryOfChannelType(RECS_ID))).isTrue()
+        assertThat(coordinator.recsSectioner.isInSection(makeEntryOfChannelType("recommendations")))
+            .isFalse()
+    }
+
+    @Test
+    fun promoSectioner() {
+        assertThat(coordinator.promoSectioner.isInSection(makeEntryOfChannelType(PROMOTIONS_ID)))
+            .isTrue()
+        assertThat(coordinator.promoSectioner.isInSection(makeEntryOfChannelType("promo"))).
+        isFalse()
+    }
+
+    private fun makeEntryOfChannelType(
+        type: String,
+        buildBlock: NotificationEntryBuilder.() -> Unit = {}
+    ): NotificationEntry {
+        val channel: NotificationChannel = NotificationChannel(type, type, 2)
+        val entry =
+            NotificationEntryBuilder()
+                .updateRanking {
+                    it.setChannel(channel)
+                }
+                .also(buildBlock)
+                .build()
+        return entry
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationSectionsFeatureManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationSectionsFeatureManager.kt
index f62b24aa96a3..3dcaff3de35a 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationSectionsFeatureManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/NotificationSectionsFeatureManager.kt
@@ -21,6 +21,7 @@ import android.provider.DeviceConfig
 import com.android.internal.annotations.VisibleForTesting
 import com.android.internal.config.sysui.SystemUiDeviceConfigFlags.NOTIFICATIONS_USE_PEOPLE_FILTERING
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.statusbar.notification.collection.NotificationClassificationFlag
 import com.android.systemui.statusbar.notification.shared.NotificationMinimalismPrototype
 import com.android.systemui.statusbar.notification.shared.PriorityPeopleSection
 import com.android.systemui.statusbar.notification.stack.BUCKET_ALERTING
@@ -51,7 +52,8 @@ constructor(val proxy: DeviceConfigProxy, val context: Context) {
     }
 
     fun getNotificationBuckets(): IntArray {
-        if (PriorityPeopleSection.isEnabled || NotificationMinimalismPrototype.V2.isEnabled) {
+        if (PriorityPeopleSection.isEnabled || NotificationMinimalismPrototype.V2.isEnabled
+            || NotificationClassificationFlag.isEnabled) {
             // We don't need this list to be adaptive, it can be the superset of all features.
             return PriorityBucket.getAllInOrder()
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationClassificationFlag.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationClassificationFlag.kt
new file mode 100644
index 000000000000..139347c86d54
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/NotificationClassificationFlag.kt
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2024 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.collection
+
+import android.service.notification.Flags
+import com.android.systemui.flags.FlagToken
+import com.android.systemui.flags.RefactorFlagUtils
+
+/**
+ * Helper for android.service.notification.Flags.FLAG_NOTIFICATION_CLASSIFICATION
+ */
+@Suppress("NOTHING_TO_INLINE")
+object NotificationClassificationFlag {
+    const val FLAG_NAME = Flags.FLAG_NOTIFICATION_CLASSIFICATION
+
+    /** A token used for dependency declaration */
+    val token: FlagToken
+        get() = FlagToken(FLAG_NAME, isEnabled)
+
+    /** Are sections sorted by time? */
+    @JvmStatic
+    inline val isEnabled
+        get() = Flags.notificationClassification()
+
+    /**
+     * Called to ensure code is only run when the flag is enabled. This protects users from the
+     * unintended behaviors caused by accidentally running new logic, while also crashing on an eng
+     * build to ensure that the refactor author catches issues in testing.
+     */
+    @JvmStatic
+    inline fun isUnexpectedlyInLegacyMode() =
+            RefactorFlagUtils.isUnexpectedlyInLegacyMode(isEnabled, FLAG_NAME)
+
+    /**
+     * Called to ensure code is only run when the flag is disabled. This will throw an exception if
+     * the flag is enabled to ensure that the refactor author catches issues in testing.
+     */
+    @JvmStatic
+    inline fun assertInLegacyMode() = RefactorFlagUtils.assertInLegacyMode(isEnabled, FLAG_NAME)
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinator.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinator.kt
new file mode 100644
index 000000000000..244c5946c21e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/BundleCoordinator.kt
@@ -0,0 +1,95 @@
+/*
+ * Copyright (C) 2024 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.collection.coordinator
+
+import android.app.NotificationChannel.NEWS_ID
+import android.app.NotificationChannel.PROMOTIONS_ID
+import android.app.NotificationChannel.RECS_ID
+import android.app.NotificationChannel.SOCIAL_MEDIA_ID
+import com.android.systemui.statusbar.notification.collection.ListEntry
+import com.android.systemui.statusbar.notification.collection.NotifPipeline
+import com.android.systemui.statusbar.notification.collection.coordinator.dagger.CoordinatorScope
+import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner
+import com.android.systemui.statusbar.notification.collection.render.NodeController
+import com.android.systemui.statusbar.notification.dagger.NewsHeader
+import com.android.systemui.statusbar.notification.dagger.PromoHeader
+import com.android.systemui.statusbar.notification.dagger.RecsHeader
+import com.android.systemui.statusbar.notification.dagger.SocialHeader
+import com.android.systemui.statusbar.notification.stack.BUCKET_NEWS
+import com.android.systemui.statusbar.notification.stack.BUCKET_PROMO
+import com.android.systemui.statusbar.notification.stack.BUCKET_RECS
+import com.android.systemui.statusbar.notification.stack.BUCKET_SOCIAL
+import javax.inject.Inject
+
+/**
+ * Coordinator for sections derived from NotificationAssistantService classification.
+ */
+@CoordinatorScope
+class BundleCoordinator @Inject constructor(
+    @NewsHeader private val newsHeaderController: NodeController,
+    @SocialHeader private val socialHeaderController: NodeController,
+    @RecsHeader private val recsHeaderController: NodeController,
+    @PromoHeader private val promoHeaderController: NodeController,
+) : Coordinator {
+
+    val newsSectioner =
+            object : NotifSectioner("News", BUCKET_NEWS) {
+                override fun isInSection(entry: ListEntry): Boolean {
+                    return entry.representativeEntry?.channel?.id == NEWS_ID
+                }
+
+                override fun getHeaderNodeController(): NodeController? {
+                    return newsHeaderController
+                }
+            }
+
+    val socialSectioner =
+        object : NotifSectioner("Social", BUCKET_SOCIAL) {
+            override fun isInSection(entry: ListEntry): Boolean {
+                return entry.representativeEntry?.channel?.id == SOCIAL_MEDIA_ID
+            }
+
+            override fun getHeaderNodeController(): NodeController? {
+                return socialHeaderController
+            }
+        }
+
+    val recsSectioner =
+        object : NotifSectioner("Recommendations", BUCKET_RECS) {
+            override fun isInSection(entry: ListEntry): Boolean {
+                return entry.representativeEntry?.channel?.id == RECS_ID
+            }
+
+            override fun getHeaderNodeController(): NodeController? {
+                return recsHeaderController
+            }
+        }
+
+    val promoSectioner =
+        object : NotifSectioner("Promotions", BUCKET_PROMO) {
+            override fun isInSection(entry: ListEntry): Boolean {
+                return entry.representativeEntry?.channel?.id == PROMOTIONS_ID
+            }
+
+            override fun getHeaderNodeController(): NodeController? {
+                return promoHeaderController
+            }
+        }
+
+    override fun attach(pipeline: NotifPipeline) {
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
index e41352254bac..e0389820aedf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/collection/coordinator/NotifCoordinators.kt
@@ -17,10 +17,7 @@ package com.android.systemui.statusbar.notification.collection.coordinator
 
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags.LOCKSCREEN_WALLPAPER_DREAM_ENABLED
-import com.android.systemui.statusbar.notification.collection.NotifPipeline
-import com.android.systemui.statusbar.notification.collection.PipelineDumpable
-import com.android.systemui.statusbar.notification.collection.PipelineDumper
-import com.android.systemui.statusbar.notification.collection.SortBySectionTimeFlag
+import com.android.systemui.statusbar.notification.collection.*
 import com.android.systemui.statusbar.notification.collection.coordinator.dagger.CoordinatorScope
 import com.android.systemui.statusbar.notification.collection.listbuilder.pluggable.NotifSectioner
 import com.android.systemui.statusbar.notification.collection.provider.SectionStyleProvider
@@ -69,6 +66,7 @@ constructor(
     dismissibilityCoordinator: DismissibilityCoordinator,
     dreamCoordinator: DreamCoordinator,
     statsLoggerCoordinator: NotificationStatsLoggerCoordinator,
+    bundleCoordinator: BundleCoordinator,
 ) : NotifCoordinators {
 
     private val mCoreCoordinators: MutableList<CoreCoordinator> = ArrayList()
@@ -132,6 +130,12 @@ constructor(
             mOrderedSections.add(conversationCoordinator.peopleSilentSectioner) // People Silent
         }
         mOrderedSections.add(rankingCoordinator.alertingSectioner) // Alerting
+        if (NotificationClassificationFlag.isEnabled) {
+            mOrderedSections.add(bundleCoordinator.newsSectioner);
+            mOrderedSections.add(bundleCoordinator.socialSectioner);
+            mOrderedSections.add(bundleCoordinator.recsSectioner);
+            mOrderedSections.add(bundleCoordinator.promoSectioner);
+        }
         mOrderedSections.add(rankingCoordinator.silentSectioner) // Silent
         mOrderedSections.add(rankingCoordinator.minimizedSectioner) // Minimized
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationSectionHeadersModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationSectionHeadersModule.kt
index ca43591e0776..e661090011c6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationSectionHeadersModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/dagger/NotificationSectionHeadersModule.kt
@@ -80,6 +80,50 @@ object NotificationSectionHeadersModule {
             .build()
 
     @Provides
+    @NewsHeader
+    @SysUISingleton
+    @JvmStatic fun providesNewsHeaderSubcomponent(
+        builder: Provider<SectionHeaderControllerSubcomponent.Builder>
+    ) = builder.get()
+        .nodeLabel("news header")
+        .headerText(com.android.internal.R.string.news_notification_channel_label)
+        .clickIntentAction(Settings.ACTION_NOTIFICATION_SETTINGS)
+        .build()
+
+    @Provides
+    @SocialHeader
+    @SysUISingleton
+    @JvmStatic fun providesSocialHeaderSubcomponent(
+        builder: Provider<SectionHeaderControllerSubcomponent.Builder>
+    ) = builder.get()
+        .nodeLabel("social header")
+        .headerText(com.android.internal.R.string.social_notification_channel_label)
+        .clickIntentAction(Settings.ACTION_NOTIFICATION_SETTINGS)
+        .build()
+
+    @Provides
+    @RecsHeader
+    @SysUISingleton
+    @JvmStatic fun providesRecsHeaderSubcomponent(
+        builder: Provider<SectionHeaderControllerSubcomponent.Builder>
+    ) = builder.get()
+        .nodeLabel("recs header")
+        .headerText(com.android.internal.R.string.recs_notification_channel_label)
+        .clickIntentAction(Settings.ACTION_NOTIFICATION_SETTINGS)
+        .build()
+
+    @Provides
+    @PromoHeader
+    @SysUISingleton
+    @JvmStatic fun providesPromoHeaderSubcomponent(
+        builder: Provider<SectionHeaderControllerSubcomponent.Builder>
+    ) = builder.get()
+        .nodeLabel("promo header")
+        .headerText(com.android.internal.R.string.promotional_notification_channel_label)
+        .clickIntentAction(Settings.ACTION_NOTIFICATION_SETTINGS)
+        .build()
+
+    @Provides
     @SilentHeader
     @JvmStatic fun providesSilentHeaderNodeController(
         @SilentHeader subcomponent: SectionHeaderControllerSubcomponent
@@ -126,6 +170,54 @@ object NotificationSectionHeadersModule {
     @JvmStatic fun providesIncomingHeaderController(
         @IncomingHeader subcomponent: SectionHeaderControllerSubcomponent
     ) = subcomponent.headerController
+
+    @Provides
+    @NewsHeader
+    @JvmStatic fun providesNewsHeaderNodeController(
+        @NewsHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.nodeController
+
+    @Provides
+    @NewsHeader
+    @JvmStatic fun providesNewsHeaderController(
+        @NewsHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.headerController
+
+    @Provides
+    @SocialHeader
+    @JvmStatic fun providesSocialHeaderNodeController(
+        @SocialHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.nodeController
+
+    @Provides
+    @SocialHeader
+    @JvmStatic fun providesSocialHeaderController(
+        @SocialHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.headerController
+
+    @Provides
+    @RecsHeader
+    @JvmStatic fun providesRecsHeaderNodeController(
+        @RecsHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.nodeController
+
+    @Provides
+    @RecsHeader
+    @JvmStatic fun providesRecsHeaderController(
+        @RecsHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.headerController
+
+    @Provides
+    @PromoHeader
+    @JvmStatic fun providesPromoHeaderNodeController(
+        @PromoHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.nodeController
+
+    @Provides
+    @PromoHeader
+    @JvmStatic fun providesPromoHeaderController(
+        @PromoHeader subcomponent: SectionHeaderControllerSubcomponent
+    ) = subcomponent.headerController
 }
 
 @Subcomponent(modules = [ SectionHeaderBindingModule::class ])
@@ -183,3 +275,19 @@ annotation class HeaderClickAction
 @Scope
 @Retention(AnnotationRetention.BINARY)
 annotation class SectionHeaderScope
+
+@Qualifier
+@Retention(AnnotationRetention.BINARY)
+annotation class NewsHeader
+
+@Qualifier
+@Retention(AnnotationRetention.BINARY)
+annotation class SocialHeader
+
+@Qualifier
+@Retention(AnnotationRetention.BINARY)
+annotation class RecsHeader
+
+@Qualifier
+@Retention(AnnotationRetention.BINARY)
+annotation class PromoHeader
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
index 9e0dd8fc4d92..175512336b8e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/NotificationPanelLogger.java
@@ -20,9 +20,13 @@ import static com.android.systemui.statusbar.notification.stack.NotificationPrio
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_FOREGROUND_SERVICE;
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_HEADS_UP;
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_MEDIA_CONTROLS;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_NEWS;
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PEOPLE;
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PRIORITY_PEOPLE;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_PROMO;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_RECS;
 import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SILENT;
+import static com.android.systemui.statusbar.notification.stack.NotificationPriorityBucketKt.BUCKET_SOCIAL;
 
 import android.annotation.Nullable;
 import android.service.notification.StatusBarNotification;
@@ -135,6 +139,10 @@ public interface NotificationPanelLogger {
                 return Notifications.Notification.SECTION_PEOPLE;
             case BUCKET_ALERTING: return Notifications.Notification.SECTION_ALERTING;
             case BUCKET_SILENT: return Notifications.Notification.SECTION_SILENT;
+            case BUCKET_NEWS: return Notifications.Notification.SECTION_NEWS;
+            case BUCKET_SOCIAL: return Notifications.Notification.SECTION_SOCIAL;
+            case BUCKET_RECS: return Notifications.Notification.SECTION_RECS;
+            case BUCKET_PROMO: return Notifications.Notification.SECTION_PROMO;
         }
         return Notifications.Notification.SECTION_UNKNOWN;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
index c2ab2758dd74..ce4356ae6c2b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/logging/Notifications.proto
@@ -43,6 +43,13 @@ message Notification {
         SECTION_ALERTING = 4;
         SECTION_SILENT = 5;
         SECTION_FOREGROUND_SERVICE = 6;
+        SECTION_PRIORITY_PEOPLE = 7;
+        SECTION_TOP_ONGOING = 8;
+        SECTION_TOP_UNSEEN = 9;
+        SECTION_NEWS = 10;
+        SECTION_SOCIAL = 11;
+        SECTION_RECS = 12;
+        SECTION_PROMO = 13;
     }
     optional NotificationSection section = 6;
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt
index fabb696d9182..f4a452784267 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationPriorityBucket.kt
@@ -20,6 +20,10 @@ import android.annotation.IntDef
             BUCKET_PRIORITY_PEOPLE,
             BUCKET_PEOPLE,
             BUCKET_ALERTING,
+            BUCKET_NEWS,
+            BUCKET_SOCIAL,
+            BUCKET_RECS,
+            BUCKET_PROMO,
             BUCKET_SILENT
         ]
 )
@@ -35,6 +39,10 @@ annotation class PriorityBucket {
                 BUCKET_PRIORITY_PEOPLE,
                 BUCKET_PEOPLE,
                 BUCKET_ALERTING,
+                BUCKET_NEWS,
+                BUCKET_SOCIAL,
+                BUCKET_RECS,
+                BUCKET_PROMO,
                 BUCKET_SILENT,
             )
     }
@@ -49,4 +57,9 @@ const val BUCKET_FOREGROUND_SERVICE = 3
 const val BUCKET_PRIORITY_PEOPLE = 7
 const val BUCKET_PEOPLE = 4
 const val BUCKET_ALERTING = 5
+const val BUCKET_NEWS = 10
+const val BUCKET_SOCIAL = 11
+const val BUCKET_RECS = 12
+const val BUCKET_PROMO = 13
 const val BUCKET_SILENT = 6
+
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
index 3400ad107133..7441c7058d7b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManager.kt
@@ -22,12 +22,10 @@ import com.android.internal.annotations.VisibleForTesting
 import com.android.systemui.media.controls.ui.controller.KeyguardMediaController
 import com.android.systemui.statusbar.notification.NotificationSectionsFeatureManager
 import com.android.systemui.statusbar.notification.SourceType
+import com.android.systemui.statusbar.notification.collection.NotificationClassificationFlag
 import com.android.systemui.statusbar.notification.collection.render.MediaContainerController
 import com.android.systemui.statusbar.notification.collection.render.SectionHeaderController
-import com.android.systemui.statusbar.notification.dagger.AlertingHeader
-import com.android.systemui.statusbar.notification.dagger.IncomingHeader
-import com.android.systemui.statusbar.notification.dagger.PeopleHeader
-import com.android.systemui.statusbar.notification.dagger.SilentHeader
+import com.android.systemui.statusbar.notification.dagger.*
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.statusbar.notification.stack.StackScrollAlgorithm.SectionProvider
@@ -51,7 +49,11 @@ internal constructor(
     @IncomingHeader private val incomingHeaderController: SectionHeaderController,
     @PeopleHeader private val peopleHeaderController: SectionHeaderController,
     @AlertingHeader private val alertingHeaderController: SectionHeaderController,
-    @SilentHeader private val silentHeaderController: SectionHeaderController
+    @SilentHeader private val silentHeaderController: SectionHeaderController,
+    @NewsHeader private val newsHeaderController: SectionHeaderController,
+    @SocialHeader private val socialHeaderController: SectionHeaderController,
+    @RecsHeader private val recsHeaderController: SectionHeaderController,
+    @PromoHeader private val promoHeaderController: SectionHeaderController
 ) : SectionProvider {
 
     private val configurationListener =
@@ -84,6 +86,22 @@ internal constructor(
     val mediaControlsView: MediaContainerView?
         get() = mediaContainerController.mediaContainerView
 
+    @VisibleForTesting
+    val newsHeaderView: SectionHeaderView?
+        get() = newsHeaderController.headerView
+
+    @VisibleForTesting
+    val socialHeaderView: SectionHeaderView?
+        get() = socialHeaderController.headerView
+
+    @VisibleForTesting
+    val recsHeaderView: SectionHeaderView?
+        get() = recsHeaderController.headerView
+
+    @VisibleForTesting
+    val promoHeaderView: SectionHeaderView?
+        get() = promoHeaderController.headerView
+
     /** Must be called before use. */
     fun initialize(parent: NotificationStackScrollLayout) {
         check(!initialized) { "NotificationSectionsManager already initialized" }
@@ -107,15 +125,24 @@ internal constructor(
         incomingHeaderController.reinflateView(parent)
         mediaContainerController.reinflateView(parent)
         keyguardMediaController.attachSinglePaneContainer(mediaControlsView)
+        if (NotificationClassificationFlag.isEnabled) {
+            newsHeaderController.reinflateView(parent)
+            socialHeaderController.reinflateView(parent)
+            recsHeaderController.reinflateView(parent)
+            promoHeaderController.reinflateView(parent)
+        }
     }
 
     override fun beginsSection(view: View, previous: View?): Boolean =
         view === silentHeaderView ||
-            view === mediaControlsView ||
-            view === peopleHeaderView ||
-            view === alertingHeaderView ||
-            view === incomingHeaderView ||
-            getBucket(view) != getBucket(previous)
+                view === mediaControlsView ||
+                view === peopleHeaderView ||
+                view === alertingHeaderView ||
+                view === incomingHeaderView ||
+                (NotificationClassificationFlag.isEnabled && (view === newsHeaderView
+                        || view === socialHeaderView || view === recsHeaderView
+                        || view === promoHeaderView)) ||
+                getBucket(view) != getBucket(previous)
 
     private fun getBucket(view: View?): Int? =
         when {
@@ -124,6 +151,10 @@ internal constructor(
             view === mediaControlsView -> BUCKET_MEDIA_CONTROLS
             view === peopleHeaderView -> BUCKET_PEOPLE
             view === alertingHeaderView -> BUCKET_ALERTING
+            view === newsHeaderView -> BUCKET_NEWS
+            view === socialHeaderView -> BUCKET_SOCIAL
+            view === recsHeaderView -> BUCKET_RECS
+            view === promoHeaderView -> BUCKET_PROMO
             view is ExpandableNotificationRow -> view.entry.bucket
             else -> null
         }
@@ -255,6 +286,12 @@ internal constructor(
         peopleHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
         silentHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
         alertingHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
+        if (NotificationClassificationFlag.isEnabled) {
+            newsHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
+            socialHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
+            recsHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
+            promoHeaderView?.setForegroundColors(onSurface, onSurfaceVariant)
+        }
     }
 
     companion object {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
index 5c45b2e53047..3669e3d8fed1 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSectionsManagerTest.java
@@ -64,6 +64,10 @@ public class NotificationSectionsManagerTest extends SysuiTestCase {
     @Mock private SectionHeaderController mPeopleHeaderController;
     @Mock private SectionHeaderController mAlertingHeaderController;
     @Mock private SectionHeaderController mSilentHeaderController;
+    @Mock private SectionHeaderController mNewsHeaderController;
+    @Mock private SectionHeaderController mSocialHeaderController;
+    @Mock private SectionHeaderController mRecsHeaderController;
+    @Mock private SectionHeaderController mPromoHeaderController;
 
     private NotificationSectionsManager mSectionsManager;
 
@@ -94,7 +98,11 @@ public class NotificationSectionsManagerTest extends SysuiTestCase {
                         mIncomingHeaderController,
                         mPeopleHeaderController,
                         mAlertingHeaderController,
-                        mSilentHeaderController
+                        mSilentHeaderController,
+                        mNewsHeaderController,
+                        mSocialHeaderController,
+                        mRecsHeaderController,
+                        mPromoHeaderController
                 );
         // Required in order for the header inflation to work properly
         when(mNssl.generateLayoutParams(any(AttributeSet.class)))