Fix size and color of priority modes icons in status bar

* Allow StatusBarIcons to specify whether they want a "fixed" size instead of the variable width they normally have. This also affects their scaling calculations. This is mostly for app-provided icons which may have nonstandard sizes and proportions.
* Create a new drawable from constant state so that theming isn't shared with the other usages.

Additionally: Add a flow to ZenModeInteractor with the main active mode name and icon exclusively. While the QS tile cares about which other modes are active, neither status bar nor Smartspace do, which means we can use a flow that doesn't emit when more modes are activated (and thus avoid wasted work).

Bug: 360399800
Test: coming
Flag: android.app.modes_ui
Change-Id: I6a17d5a49f76169f15efc89432e1ce0f993d040e

14 files changed, 290 insertions, 49 deletions

diff --git a/core/java/com/android/internal/statusbar/StatusBarIcon.java b/core/java/com/android/internal/statusbar/StatusBarIcon.java
index f8a143693c83..1938cdb0ba84 100644
--- a/core/java/com/android/internal/statusbar/StatusBarIcon.java
+++ b/core/java/com/android/internal/statusbar/StatusBarIcon.java
@@ -51,6 +51,19 @@ public class StatusBarIcon implements Parcelable {
         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 class StatusBarIcon implements Parcelable {
     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 @@ public class StatusBarIcon implements Parcelable {
     @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 @@ public class StatusBarIcon implements Parcelable {
         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 @@ public class StatusBarIcon implements Parcelable {
     @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 @@ public class StatusBarIcon implements Parcelable {
         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 @@ public class StatusBarIcon implements Parcelable {
         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 b183ecb50591..149e132a0df4 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 static com.google.common.truth.Truth.assertThat;
 
 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 @@ public class StatusBarIconTest {
         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 @@ public class StatusBarIconTest {
         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 cb743bc732f0..639d34d5e74d 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
 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 @@ class ZenModeInteractorTest : SysuiTestCase() {
             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 6eadd2627399..2b44c2f9ea7f 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 androidx.core.graphics.ColorUtils;
 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 @@ public class StatusBarIconView extends AnimatedImageView implements StatusIconDi
     /** 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 @@ public class StatusBarIconView extends AnimatedImageView implements StatusIconDi
         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 @@ public class StatusBarIconView extends AnimatedImageView implements StatusIconDi
     @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 a9b886fbb073..ba39c3bb4124 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.text.format.DateFormat;
 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.SensorPrivacyController;
 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 @@ public class PhoneStatusBarPolicy
         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 @@ public class PhoneStatusBarPolicy
                 () -> 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 8871dae3c620..6c303303c8f8 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 @@ public class DarkIconManager extends IconManager {
     }
 
     @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 5ad737684ca1..91ead614ffa4 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_BIND
 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.os.Bundle;
 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 @@ public class IconManager implements DemoModeCommandReceiver {
         };
     }
 
-    @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 @@ public class IconManager implements DemoModeCommandReceiver {
             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 @@ public class IconManager implements DemoModeCommandReceiver {
 
     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 @@ public class IconManager implements DemoModeCommandReceiver {
             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 @@ public class IconManager implements DemoModeCommandReceiver {
                 );
     }
 
-    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 @@ public class IconManager implements DemoModeCommandReceiver {
     /** 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 ee528e915079..0459b9749e0a 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 @@ public interface StatusBarIconController {
      * @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 ad3a9e350c4b..9b6d32bd179d 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 @@ public class StatusBarIconControllerImpl implements Tunable,
                 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 @@ public class StatusBarIconControllerImpl implements Tunable,
                 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 @@ public class StatusBarIconControllerImpl implements Tunable,
         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 @@ public class StatusBarIconControllerImpl implements Tunable,
             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 @@ public class StatusBarIconControllerImpl implements Tunable,
         }
     }
 
-    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 520f3f45478e..93c631f65df7 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 @@ constructor(
             .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 a0b0572179bb..2ed34735db1f 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
 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 @@ class PhoneStatusBarPolicyTest : SysuiTestCase() {
                     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 @@ class PhoneStatusBarPolicyTest : SysuiTestCase() {
                     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 @@ class PhoneStatusBarPolicyTest : SysuiTestCase() {
 
         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 @@ class PhoneStatusBarPolicyTest : SysuiTestCase() {
             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 000000000000..90732d0183d2
--- /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 26a57e4c1ca9..50a13b93ea15 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 @@ class StatusBarIconControllerImplTest : SysuiTestCase() {
     @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 @@ class StatusBarIconControllerImplTest : SysuiTestCase() {
         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 2dbac670b298..0089199cfb88 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 @@ public class FakeStatusBarIconController extends BaseLeakChecker<IconManager>
 
     @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