Merge "Fix size and color of priority modes icons in status bar" into main
diff --git a/core/java/com/android/internal/statusbar/StatusBarIcon.java b/core/java/com/android/internal/statusbar/StatusBarIcon.java
index f8a1436..1938cdb 100644
--- a/core/java/com/android/internal/statusbar/StatusBarIcon.java
+++ b/core/java/com/android/internal/statusbar/StatusBarIcon.java
@@ -51,6 +51,19 @@
         ResourceIcon
     }
 
+    public enum Shape {
+        /**
+         * Icon view should use WRAP_CONTENT -- so that the horizontal space occupied depends on the
+         * icon's shape (skinny/fat icons take less/more). Most icons will want to use this option
+         * for a nicer-looking overall spacing in the status bar, as long as the icon is "known"
+         * (i.e. not coming from a 3P package).
+         */
+        WRAP_CONTENT,
+
+        /** Icon should always be displayed in a space as wide as the status bar is tall. */
+        FIXED_SPACE,
+    }
+
     public UserHandle user;
     public String pkg;
     public Icon icon;
@@ -59,6 +72,7 @@
     public int number;
     public CharSequence contentDescription;
     public Type type;
+    public Shape shape;
 
     /**
      * Optional {@link Drawable} corresponding to {@link #icon}. This field is not parcelable, so
@@ -68,7 +82,7 @@
     @Nullable public Drawable preloadedIcon;
 
     public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
-            CharSequence contentDescription, Type type) {
+            CharSequence contentDescription, Type type, Shape shape) {
         if (icon.getType() == Icon.TYPE_RESOURCE
                 && TextUtils.isEmpty(icon.getResPackage())) {
             // This is an odd situation where someone's managed to hand us an icon without a
@@ -83,6 +97,13 @@
         this.number = number;
         this.contentDescription = contentDescription;
         this.type = type;
+        this.shape = shape;
+    }
+
+    public StatusBarIcon(UserHandle user, String resPackage, Icon icon, int iconLevel, int number,
+            CharSequence contentDescription, Type type) {
+        this(user, resPackage, icon, iconLevel, number, contentDescription, type,
+                Shape.WRAP_CONTENT);
     }
 
     public StatusBarIcon(String iconPackage, UserHandle user,
@@ -107,7 +128,7 @@
     @Override
     public StatusBarIcon clone() {
         StatusBarIcon that = new StatusBarIcon(this.user, this.pkg, this.icon,
-                this.iconLevel, this.number, this.contentDescription, this.type);
+                this.iconLevel, this.number, this.contentDescription, this.type, this.shape);
         that.visible = this.visible;
         that.preloadedIcon = this.preloadedIcon;
         return that;
@@ -129,6 +150,7 @@
         this.number = in.readInt();
         this.contentDescription = in.readCharSequence();
         this.type = Type.valueOf(in.readString());
+        this.shape = Shape.valueOf(in.readString());
     }
 
     public void writeToParcel(Parcel out, int flags) {
@@ -140,6 +162,7 @@
         out.writeInt(this.number);
         out.writeCharSequence(this.contentDescription);
         out.writeString(this.type.name());
+        out.writeString(this.shape.name());
     }
 
     public int describeContents() {
diff --git a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
index b183ecb..149e132 100644
--- a/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
+++ b/core/tests/coretests/src/com/android/internal/statusbar/StatusBarIconTest.java
@@ -20,6 +20,7 @@
 
 import android.graphics.Color;
 import android.graphics.drawable.ColorDrawable;
+import android.graphics.drawable.Icon;
 import android.os.Parcel;
 import android.os.UserHandle;
 
@@ -69,22 +70,22 @@
         assertThat(copy.preloadedIcon).isEqualTo(original.preloadedIcon);
     }
 
-
     private static StatusBarIcon newStatusBarIcon() {
         final UserHandle dummyUserHandle = UserHandle.of(100);
         final String dummyIconPackageName = "com.android.internal.statusbar.test";
-        final int dummyIconId = 123;
+        final Icon dummyIcon = Icon.createWithResource(dummyIconPackageName, 123);
         final int dummyIconLevel = 1;
         final int dummyIconNumber = 2;
         final CharSequence dummyIconContentDescription = "dummyIcon";
         return new StatusBarIcon(
-                dummyIconPackageName,
                 dummyUserHandle,
-                dummyIconId,
+                dummyIconPackageName,
+                dummyIcon,
                 dummyIconLevel,
                 dummyIconNumber,
                 dummyIconContentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.FIXED_SPACE);
     }
 
     private static void assertSerializableFieldsEqual(StatusBarIcon copy, StatusBarIcon original) {
@@ -96,6 +97,7 @@
         assertThat(copy.number).isEqualTo(original.number);
         assertThat(copy.contentDescription).isEqualTo(original.contentDescription);
         assertThat(copy.type).isEqualTo(original.type);
+        assertThat(copy.shape).isEqualTo(original.shape);
     }
 
     private static StatusBarIcon parcelAndUnparcel(StatusBarIcon original) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
index cb743bc..639d34d 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractorTest.kt
@@ -22,8 +22,10 @@
 import android.provider.Settings.Secure.ZEN_DURATION
 import android.provider.Settings.Secure.ZEN_DURATION_FOREVER
 import android.provider.Settings.Secure.ZEN_DURATION_PROMPT
+import android.service.notification.SystemZenRules
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.R
 import com.android.settingslib.notification.data.repository.updateNotificationPolicy
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.SysuiTestCase
@@ -253,4 +255,56 @@
             assertThat(activeModes?.modeNames).hasSize(0)
             assertThat(activeModes?.mainMode).isNull()
         }
+
+    @Test
+    fun mainActiveMode_flows() =
+        testScope.runTest {
+            val mainActiveMode by collectLastValue(underTest.mainActiveMode)
+
+            zenModeRepository.addModes(
+                listOf(
+                    TestModeBuilder()
+                        .setId("Bedtime")
+                        .setName("Mode Bedtime")
+                        .setType(AutomaticZenRule.TYPE_BEDTIME)
+                        .setActive(false)
+                        .setPackage(mContext.packageName)
+                        .setIconResId(R.drawable.ic_zen_mode_type_bedtime)
+                        .build(),
+                    TestModeBuilder()
+                        .setId("Other")
+                        .setName("Mode Other")
+                        .setType(AutomaticZenRule.TYPE_OTHER)
+                        .setActive(false)
+                        .setPackage(SystemZenRules.PACKAGE_ANDROID)
+                        .setIconResId(R.drawable.ic_zen_mode_type_other)
+                        .build(),
+                )
+            )
+
+            runCurrent()
+            assertThat(mainActiveMode).isNull()
+
+            zenModeRepository.activateMode("Other")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Other")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_other)
+
+            zenModeRepository.activateMode("Bedtime")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
+
+            zenModeRepository.deactivateMode("Other")
+            runCurrent()
+            assertThat(mainActiveMode?.name).isEqualTo("Mode Bedtime")
+            assertThat(mainActiveMode?.icon?.key?.resId)
+                .isEqualTo(R.drawable.ic_zen_mode_type_bedtime)
+
+            zenModeRepository.deactivateMode("Bedtime")
+            runCurrent()
+            assertThat(mainActiveMode).isNull()
+        }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
index 6eadd26..2b44c2f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/StatusBarIconView.java
@@ -58,6 +58,7 @@
 import com.android.app.animation.Interpolators;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.internal.util.ContrastColorUtil;
 import com.android.systemui.Flags;
 import com.android.systemui.res.R;
@@ -211,16 +212,19 @@
     /** Should always be preceded by {@link #reloadDimens()} */
     @VisibleForTesting
     public void maybeUpdateIconScaleDimens() {
-        // We do not resize and scale system icons (on the right), only notification icons (on the
-        // left).
-        if (isNotification()) {
-            updateIconScaleForNotifications();
+        // We scale notification icons (on the left) plus icons on the right that explicitly
+        // want FIXED_SPACE.
+        boolean useNonSystemIconScaling = isNotification()
+                || (usesModeIcons() && mIcon != null && mIcon.shape == Shape.FIXED_SPACE);
+
+        if (useNonSystemIconScaling) {
+            updateIconScaleForNonSystemIcons();
         } else {
             updateIconScaleForSystemIcons();
         }
     }
 
-    private void updateIconScaleForNotifications() {
+    private void updateIconScaleForNonSystemIcons() {
         float iconScale;
         // we need to scale the image size to be same as the original size
         // (fit mOriginalStatusBarIconSize), then we can scale it with mScaleToFitNewIconSize
@@ -411,7 +415,9 @@
         if (!levelEquals) {
             setImageLevel(icon.iconLevel);
         }
-
+        if (usesModeIcons()) {
+            setScaleType(icon.shape == Shape.FIXED_SPACE ? ScaleType.FIT_CENTER : ScaleType.CENTER);
+        }
         if (!visibilityEquals) {
             setVisibility(icon.visible && !mBlocked ? VISIBLE : GONE);
         }
@@ -501,7 +507,12 @@
     @Nullable
     private Drawable loadDrawable(Context context, StatusBarIcon statusBarIcon) {
         if (usesModeIcons() && statusBarIcon.preloadedIcon != null) {
-            return statusBarIcon.preloadedIcon.mutate();
+            Drawable.ConstantState cached = statusBarIcon.preloadedIcon.getConstantState();
+            if (cached != null) {
+                return cached.newDrawable(mContext.getResources()).mutate();
+            } else {
+                return statusBarIcon.preloadedIcon.mutate();
+            }
         } else {
             int userId = statusBarIcon.user.getIdentifier();
             if (userId == UserHandle.USER_ALL) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
index a9b886f..ba39c3b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicy.java
@@ -42,9 +42,9 @@
 import android.util.Log;
 import android.view.View;
 
-import androidx.annotation.NonNull;
 import androidx.lifecycle.Observer;
 
+import com.android.internal.statusbar.StatusBarIcon;
 import com.android.systemui.Flags;
 import com.android.systemui.broadcast.BroadcastDispatcher;
 import com.android.systemui.dagger.qualifiers.DisplayId;
@@ -80,7 +80,6 @@
 import com.android.systemui.statusbar.policy.UserInfoController;
 import com.android.systemui.statusbar.policy.ZenModeController;
 import com.android.systemui.statusbar.policy.domain.interactor.ZenModeInteractor;
-import com.android.systemui.statusbar.policy.domain.model.ActiveZenModes;
 import com.android.systemui.statusbar.policy.domain.model.ZenModeInfo;
 import com.android.systemui.util.RingerModeTracker;
 import com.android.systemui.util.kotlin.JavaAdapter;
@@ -364,8 +363,8 @@
         if (usesModeIcons()) {
             // Note that we're not fully replacing ZenModeController with ZenModeInteractor, so
             // we listen for the extra event here but still add the ZMC callback.
-            mJavaAdapter.alwaysCollectFlow(mZenModeInteractor.getActiveModes(),
-                    this::onActiveModesChanged);
+            mJavaAdapter.alwaysCollectFlow(mZenModeInteractor.getMainActiveMode(),
+                    this::onMainActiveModeChanged);
         }
         mZenController.addCallback(mZenControllerCallback);
         if (!Flags.statusBarScreenSharingChips()) {
@@ -397,21 +396,23 @@
                 () -> mResources.getString(R.string.accessibility_managed_profile));
     }
 
-    private void onActiveModesChanged(@NonNull ActiveZenModes activeModes) {
+    private void onMainActiveModeChanged(@Nullable ZenModeInfo mainActiveMode) {
         if (!usesModeIcons()) {
-            Log.wtf(TAG, "onActiveModeChanged shouldn't be called if MODES_UI_ICONS is disabled");
+            Log.wtf(TAG, "onMainActiveModeChanged shouldn't run if MODES_UI_ICONS is disabled");
             return;
         }
 
-        ZenModeInfo mainActiveMode = activeModes.getMainMode();
         boolean visible = mainActiveMode != null;
-
         if (visible) {
+            // Shape=FIXED_SPACE because mode icons can be from 3P packages and may not be square;
+            // we don't want to allow apps to set incredibly wide icons and take up too much space
+            // in the status bar.
             mIconController.setResourceIcon(mSlotZen,
                     mainActiveMode.getIcon().key().resPackage(),
                     mainActiveMode.getIcon().key().resId(),
                     mainActiveMode.getIcon().drawable(),
-                    mainActiveMode.getName());
+                    mainActiveMode.getName(),
+                    StatusBarIcon.Shape.FIXED_SPACE);
         }
         if (visible != mZenVisible) {
             mIconController.setIconVisibility(mSlotZen, visible);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
index 8871dae..6c30330 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/DarkIconManager.java
@@ -16,7 +16,6 @@
 
 package com.android.systemui.statusbar.phone.ui;
 
-import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
 import com.android.internal.statusbar.StatusBarIcon;
@@ -64,9 +63,8 @@
     }
 
     @Override
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        LinearLayout.LayoutParams lp = new LinearLayout.LayoutParams(
-                ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(StatusBarIcon.Shape shape) {
+        LinearLayout.LayoutParams lp = super.onCreateLayoutParams(shape);
         lp.setMargins(mIconHorizontalMargin, 0, mIconHorizontalMargin, 0);
         return lp;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
index 5ad7376..91ead61 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/IconManager.java
@@ -20,6 +20,7 @@
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_ICON;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_MOBILE_NEW;
 import static com.android.systemui.statusbar.phone.StatusBarIconHolder.TYPE_WIFI_NEW;
+import static com.android.systemui.statusbar.phone.ui.StatusBarIconControllerImpl.usesModeIcons;
 
 import android.annotation.Nullable;
 import android.content.Context;
@@ -27,9 +28,8 @@
 import android.view.ViewGroup;
 import android.widget.LinearLayout;
 
-import androidx.annotation.VisibleForTesting;
-
 import com.android.internal.statusbar.StatusBarIcon;
+import com.android.internal.statusbar.StatusBarIcon.Shape;
 import com.android.systemui.demomode.DemoModeCommandReceiver;
 import com.android.systemui.statusbar.BaseStatusBarFrameLayout;
 import com.android.systemui.statusbar.StatusBarIconView;
@@ -155,12 +155,11 @@
         };
     }
 
-    @VisibleForTesting
     protected StatusBarIconView addIcon(int index, String slot, boolean blocked,
             StatusBarIcon icon) {
         StatusBarIconView view = onCreateStatusBarIconView(slot, blocked);
         view.set(icon);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(icon.shape));
         return view;
     }
 
@@ -174,7 +173,7 @@
             int index) {
         mBindableIcons.put(holder.getSlot(), holder);
         ModernStatusBarView view = holder.getInitializer().createAndBind(mContext);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
         if (mIsInDemoMode) {
             mDemoStatusIcons.addBindableIcon(holder);
         }
@@ -183,7 +182,7 @@
 
     protected StatusIconDisplayable addNewWifiIcon(int index, String slot) {
         ModernStatusBarWifiView view = onCreateModernStatusBarWifiView(slot);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             mDemoStatusIcons.addModernWifiView(mWifiViewModel);
@@ -199,7 +198,7 @@
             int subId
     ) {
         BaseStatusBarFrameLayout view = onCreateModernStatusBarMobileView(slot, subId);
-        mGroup.addView(view, index, onCreateLayoutParams());
+        mGroup.addView(view, index, onCreateLayoutParams(Shape.WRAP_CONTENT));
 
         if (mIsInDemoMode) {
             Context mobileContext = mMobileContextProvider
@@ -233,8 +232,12 @@
                 );
     }
 
-    protected LinearLayout.LayoutParams onCreateLayoutParams() {
-        return new LinearLayout.LayoutParams(ViewGroup.LayoutParams.WRAP_CONTENT, mIconSize);
+    protected LinearLayout.LayoutParams onCreateLayoutParams(Shape shape) {
+        int width = usesModeIcons() && shape == StatusBarIcon.Shape.FIXED_SPACE
+                ? mIconSize
+                : ViewGroup.LayoutParams.WRAP_CONTENT;
+
+        return new LinearLayout.LayoutParams(width, mIconSize);
     }
 
     protected void destroy() {
@@ -256,6 +259,13 @@
     /** Called once an icon has been set. */
     public void onSetIcon(int viewIndex, StatusBarIcon icon) {
         StatusBarIconView view = (StatusBarIconView) mGroup.getChildAt(viewIndex);
+        if (usesModeIcons()) {
+            ViewGroup.LayoutParams current = view.getLayoutParams();
+            ViewGroup.LayoutParams desired = onCreateLayoutParams(icon.shape);
+            if (desired.width != current.width || desired.height != current.height) {
+                view.setLayoutParams(desired);
+            }
+        }
         view.set(icon);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
index ee528e9..0459b97 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconController.java
@@ -70,7 +70,8 @@
      * @param preloadedIcon optional drawable corresponding to {@code iconResId}, if known
      */
     void setResourceIcon(String slot, @Nullable String resPackage, @DrawableRes int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription);
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape);
 
     /**
      * Sets up a wifi icon using the new data pipeline. No effect if the wifi icon has already been
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
index ad3a9e3..9b6d32b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImpl.java
@@ -234,13 +234,14 @@
                 Icon.createWithResource(mContext, resourceId),
                 /* preloadedIcon= */ null,
                 contentDescription,
-                StatusBarIcon.Type.SystemIcon);
+                StatusBarIcon.Type.SystemIcon,
+                StatusBarIcon.Shape.WRAP_CONTENT);
     }
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage,
             @DrawableRes int iconResId, @Nullable Drawable preloadedIcon,
-            CharSequence contentDescription) {
+            CharSequence contentDescription, StatusBarIcon.Shape shape) {
         if (!usesModeIcons()) {
             Log.wtf("TAG",
                     "StatusBarIconController.setResourceIcon() should not be called without "
@@ -260,12 +261,13 @@
                 icon,
                 preloadedIcon,
                 contentDescription,
-                StatusBarIcon.Type.ResourceIcon);
+                StatusBarIcon.Type.ResourceIcon,
+                shape);
     }
 
     private void setResourceIconInternal(String slot, Icon resourceIcon,
             @Nullable Drawable preloadedIcon, CharSequence contentDescription,
-            StatusBarIcon.Type type) {
+            StatusBarIcon.Type type, StatusBarIcon.Shape shape) {
         checkArgument(resourceIcon.getType() == Icon.TYPE_RESOURCE,
                 "Expected Icon of TYPE_RESOURCE, but got " + resourceIcon.getType());
         String resPackage = resourceIcon.getResPackage();
@@ -277,7 +279,7 @@
         if (holder == null) {
             StatusBarIcon icon = new StatusBarIcon(UserHandle.SYSTEM, resPackage,
                     resourceIcon, /* iconLevel= */ 0, /* number=*/ 0,
-                    contentDescription, type);
+                    contentDescription, type, shape);
             icon.preloadedIcon = preloadedIcon;
             holder = StatusBarIconHolder.fromIcon(icon);
             setIcon(slot, holder);
@@ -286,6 +288,7 @@
             holder.getIcon().icon = resourceIcon;
             holder.getIcon().contentDescription = contentDescription;
             holder.getIcon().type = type;
+            holder.getIcon().shape = shape;
             holder.getIcon().preloadedIcon = preloadedIcon;
             handleSet(slot, holder);
         }
@@ -578,7 +581,7 @@
         }
     }
 
-    private static boolean usesModeIcons() {
+    static boolean usesModeIcons() {
         return android.app.Flags.modesApi() && android.app.Flags.modesUi()
                 && android.app.Flags.modesUiIcons();
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
index 520f3f4..93c631f 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/domain/interactor/ZenModeInteractor.kt
@@ -96,6 +96,9 @@
             .flowOn(bgDispatcher)
             .distinctUntilChanged()
 
+    val mainActiveMode: Flow<ZenModeInfo?> =
+        activeModes.map { a -> a.mainMode }.distinctUntilChanged()
+
     suspend fun getModeIcon(mode: ZenMode): ZenIcon {
         return iconLoader.getIcon(context, mode).await()
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
index a0b0572..2ed3473 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/PhoneStatusBarPolicyTest.kt
@@ -34,6 +34,7 @@
 import android.testing.TestableLooper.RunWithLooper
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
 import com.android.settingslib.notification.modes.TestModeBuilder
 import com.android.systemui.Flags
 import com.android.systemui.SysuiTestCase
@@ -429,7 +430,8 @@
                     eq(mContext.packageName),
                     eq(android.R.drawable.ic_lock_lock),
                     any(), // non-null
-                    eq("Bedtime Mode")
+                    eq("Bedtime Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
                 )
 
             zenModeRepository.deactivateMode("bedtime")
@@ -441,7 +443,8 @@
                     eq(null),
                     eq(android.R.drawable.ic_media_play),
                     any(), // non-null
-                    eq("Other Mode")
+                    eq("Other Mode"),
+                    eq(StatusBarIcon.Shape.FIXED_SPACE)
                 )
 
             zenModeRepository.deactivateMode("other")
@@ -460,7 +463,8 @@
 
         verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
         verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
-        verify(iconController, never()).setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+        verify(iconController, never())
+            .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
     }
 
     @Test
@@ -476,7 +480,7 @@
             verify(iconController, never()).setIconVisibility(eq(ZEN_SLOT), any())
             verify(iconController, never()).setIcon(eq(ZEN_SLOT), anyInt(), any())
             verify(iconController, never())
-                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any())
+                .setResourceIcon(eq(ZEN_SLOT), any(), any(), any(), any(), any())
         }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
new file mode 100644
index 0000000..90732d01
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/IconManagerTest.kt
@@ -0,0 +1,122 @@
+/*
+ * 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.phone.ui
+
+import android.app.Flags
+import android.graphics.drawable.Icon
+import android.os.UserHandle
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.view.ViewGroup
+import android.widget.ImageView
+import android.widget.LinearLayout
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.internal.statusbar.StatusBarIcon
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.StatusBarIconView
+import com.android.systemui.statusbar.connectivity.ui.MobileContextProvider
+import com.android.systemui.statusbar.phone.StatusBarLocation
+import com.android.systemui.statusbar.pipeline.mobile.ui.MobileUiAdapter
+import com.android.systemui.statusbar.pipeline.wifi.ui.WifiUiAdapter
+import com.android.systemui.util.Assert
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.RETURNS_DEEP_STUBS
+import org.mockito.kotlin.mock
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class IconManagerTest : SysuiTestCase() {
+
+    private lateinit var underTest: IconManager
+    private lateinit var viewGroup: ViewGroup
+
+    @Before
+    fun setUp() {
+        Assert.setTestThread(Thread.currentThread())
+        viewGroup = LinearLayout(context)
+        underTest =
+            IconManager(
+                viewGroup,
+                StatusBarLocation.HOME,
+                mock<WifiUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileUiAdapter>(defaultAnswer = RETURNS_DEEP_STUBS),
+                mock<MobileContextProvider>(defaultAnswer = RETURNS_DEEP_STUBS),
+            )
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeWrapContent_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.WRAP_CONTENT)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_MODES_UI, Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_shapeFixedSpace_addsIconViewWithFixedWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isNotEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.layoutParams.width).isEqualTo(iconView.layoutParams.height)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.FIT_CENTER)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_MODES_UI_ICONS)
+    fun addIcon_iconsFlagOff_addsIconViewWithVariableWidth() {
+        val sbIcon = newStatusBarIcon(StatusBarIcon.Shape.FIXED_SPACE)
+
+        underTest.addIcon(0, "slot", false, sbIcon)
+
+        assertThat(viewGroup.childCount).isEqualTo(1)
+        val iconView = viewGroup.getChildAt(0) as StatusBarIconView
+        assertThat(iconView).isNotNull()
+
+        assertThat(iconView.layoutParams.width).isEqualTo(ViewGroup.LayoutParams.WRAP_CONTENT)
+        assertThat(iconView.scaleType).isEqualTo(ImageView.ScaleType.CENTER)
+    }
+
+    private fun newStatusBarIcon(shape: StatusBarIcon.Shape) =
+        StatusBarIcon(
+            UserHandle.CURRENT,
+            context.packageName,
+            Icon.createWithResource(context, android.R.drawable.ic_media_next),
+            0,
+            0,
+            "",
+            StatusBarIcon.Type.ResourceIcon,
+            shape,
+        )
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
index 26a57e4..50a13b9 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ui/StatusBarIconControllerImplTest.kt
@@ -424,7 +424,14 @@
     @EnableFlags(android.app.Flags.FLAG_MODES_UI, android.app.Flags.FLAG_MODES_UI_ICONS)
     fun setResourceIcon_setsIconAndPreloadedIconInHolder() {
         val drawable = ColorDrawable(1)
-        underTest.setResourceIcon("slot", "some.package", 123, drawable, "description")
+        underTest.setResourceIcon(
+            "slot",
+            "some.package",
+            123,
+            drawable,
+            "description",
+            StatusBarIcon.Shape.FIXED_SPACE
+        )
 
         val iconHolder = iconList.getIconHolder("slot", 0)
         assertThat(iconHolder).isNotNull()
@@ -432,6 +439,7 @@
         assertThat(iconHolder?.icon?.icon?.resId).isEqualTo(123)
         assertThat(iconHolder?.icon?.icon?.resPackage).isEqualTo("some.package")
         assertThat(iconHolder?.icon?.contentDescription).isEqualTo("description")
+        assertThat(iconHolder?.icon?.shape).isEqualTo(StatusBarIcon.Shape.FIXED_SPACE)
         assertThat(iconHolder?.icon?.preloadedIcon).isEqualTo(drawable)
     }
 
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
index 2dbac67..0089199 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/utils/leaks/FakeStatusBarIconController.java
@@ -64,7 +64,8 @@
 
     @Override
     public void setResourceIcon(String slot, @Nullable String resPackage, int iconResId,
-            @Nullable Drawable preloadedIcon, CharSequence contentDescription) {
+            @Nullable Drawable preloadedIcon, CharSequence contentDescription,
+            StatusBarIcon.Shape shape) {
     }
 
     @Override