diff --git a/core/java/android/app/AppOpsManager.java b/core/java/android/app/AppOpsManager.java
index cb64173..9344d96 100644
--- a/core/java/android/app/AppOpsManager.java
+++ b/core/java/android/app/AppOpsManager.java
@@ -9134,8 +9134,9 @@
      */
     public int startProxyOpNoThrow(int op, @NonNull AttributionSource attributionSource,
             @Nullable String message, boolean skipProxyOperation) {
-        return startProxyOpNoThrow(op, attributionSource, message, skipProxyOperation,
-                ATTRIBUTION_FLAGS_NONE, ATTRIBUTION_FLAGS_NONE, ATTRIBUTION_CHAIN_ID_NONE);
+        return startProxyOpNoThrow(attributionSource.getToken(), op, attributionSource, message,
+                skipProxyOperation, ATTRIBUTION_FLAGS_NONE, ATTRIBUTION_FLAGS_NONE,
+                ATTRIBUTION_CHAIN_ID_NONE);
     }
 
     /**
@@ -9147,7 +9148,8 @@
      *
      * @hide
      */
-    public int startProxyOpNoThrow(int op, @NonNull AttributionSource attributionSource,
+    public int startProxyOpNoThrow(@NonNull IBinder clientId, int op,
+            @NonNull AttributionSource attributionSource,
             @Nullable String message, boolean skipProxyOperation, @AttributionFlags
             int proxyAttributionFlags, @AttributionFlags int proxiedAttributionFlags,
             int attributionChainId) {
@@ -9165,7 +9167,7 @@
                 }
             }
 
-            SyncNotedAppOp syncOp = mService.startProxyOperation(op,
+            SyncNotedAppOp syncOp = mService.startProxyOperation(clientId, op,
                     attributionSource, false, collectionMode == COLLECT_ASYNC, message,
                     shouldCollectMessage, skipProxyOperation, proxyAttributionFlags,
                     proxiedAttributionFlags, attributionChainId);
@@ -9263,9 +9265,10 @@
      */
     public void finishProxyOp(@NonNull String op, int proxiedUid,
             @NonNull String proxiedPackageName, @Nullable String proxiedAttributionTag) {
-        finishProxyOp(op, new AttributionSource(mContext.getAttributionSource(),
+        IBinder token = mContext.getAttributionSource().getToken();
+        finishProxyOp(token, op, new AttributionSource(mContext.getAttributionSource(),
                 new AttributionSource(proxiedUid, proxiedPackageName,  proxiedAttributionTag,
-                        mContext.getAttributionSource().getToken())), /*skipProxyOperation*/ false);
+                        token)), /*skipProxyOperation*/ false);
     }
 
     /**
@@ -9280,10 +9283,11 @@
      *
      * @hide
      */
-    public void finishProxyOp(@NonNull String op, @NonNull AttributionSource attributionSource,
-            boolean skipProxyOperation) {
+    public void finishProxyOp(@NonNull IBinder clientId, @NonNull String op,
+            @NonNull AttributionSource attributionSource, boolean skipProxyOperation) {
         try {
-            mService.finishProxyOperation(strOpToOp(op), attributionSource, skipProxyOperation);
+            mService.finishProxyOperation(clientId, strOpToOp(op), attributionSource,
+                    skipProxyOperation);
         } catch (RemoteException e) {
             throw e.rethrowFromSystemServer();
         }
diff --git a/core/java/android/app/AppOpsManagerInternal.java b/core/java/android/app/AppOpsManagerInternal.java
index 4d6e4ae..43023fe 100644
--- a/core/java/android/app/AppOpsManagerInternal.java
+++ b/core/java/android/app/AppOpsManagerInternal.java
@@ -26,13 +26,11 @@
 import android.util.SparseIntArray;
 
 import com.android.internal.app.IAppOpsCallback;
-import com.android.internal.util.function.DecFunction;
 import com.android.internal.util.function.HeptFunction;
 import com.android.internal.util.function.HexFunction;
 import com.android.internal.util.function.QuadFunction;
 import com.android.internal.util.function.QuintConsumer;
 import com.android.internal.util.function.QuintFunction;
-import com.android.internal.util.function.TriFunction;
 import com.android.internal.util.function.UndecFunction;
 
 /**
@@ -135,6 +133,7 @@
         /**
          * Allows overriding start proxy operation behavior.
          *
+         * @param clientId The client calling start, represented by an IBinder
          * @param code The op code to start.
          * @param attributionSource The permission identity of the caller.
          * @param startIfModeDefault Whether to start the op of the mode is default.
@@ -148,11 +147,12 @@
          * @param superImpl The super implementation.
          * @return The app op note result.
          */
-        SyncNotedAppOp startProxyOperation(int code, @NonNull AttributionSource attributionSource,
-                boolean startIfModeDefault, boolean shouldCollectAsyncNotedOp, String message,
-                boolean shouldCollectMessage, boolean skipProxyOperation, @AttributionFlags
-                int proxyAttributionFlags, @AttributionFlags int proxiedAttributionFlags,
-                int attributionChainId, @NonNull DecFunction<Integer, AttributionSource, Boolean,
+        SyncNotedAppOp startProxyOperation(@NonNull IBinder clientId, int code,
+                @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
+                boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
+                boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
+                @AttributionFlags int proxiedAttributionFlags, int attributionChainId,
+                @NonNull UndecFunction<IBinder, Integer, AttributionSource, Boolean,
                         Boolean, String, Boolean, Boolean, Integer, Integer, Integer,
                         SyncNotedAppOp> superImpl);
 
@@ -176,10 +176,15 @@
          *
          * @param code The op code to finish.
          * @param attributionSource The permission identity of the caller.
+         * @param skipProxyOperation Whether to skip the proxy in the proxy/proxied operation
+         * @param clientId The client calling finishProxyOperation
+         * @param superImpl The "standard" implementation to potentially call
          */
-        void finishProxyOperation(int code, @NonNull AttributionSource attributionSource,
+        void finishProxyOperation(@NonNull IBinder clientId, int code,
+                @NonNull AttributionSource attributionSource,
                 boolean skipProxyOperation,
-                @NonNull TriFunction<Integer, AttributionSource, Boolean, Void> superImpl);
+                @NonNull QuadFunction<IBinder, Integer, AttributionSource, Boolean,
+                        Void> superImpl);
     }
 
     /**
diff --git a/core/java/com/android/internal/app/IAppOpsService.aidl b/core/java/com/android/internal/app/IAppOpsService.aidl
index 30da4b4..88447da 100644
--- a/core/java/com/android/internal/app/IAppOpsService.aidl
+++ b/core/java/com/android/internal/app/IAppOpsService.aidl
@@ -58,11 +58,12 @@
     SyncNotedAppOp noteProxyOperation(int code, in AttributionSource attributionSource,
             boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
             boolean skipProxyOperation);
-    SyncNotedAppOp startProxyOperation(int code, in AttributionSource attributionSource,
-            boolean startIfModeDefault, boolean shouldCollectAsyncNotedOp, String message,
-            boolean shouldCollectMessage, boolean skipProxyOperation, int proxyAttributionFlags,
-            int proxiedAttributionFlags, int attributionChainId);
-    void finishProxyOperation(int code, in AttributionSource attributionSource,
+    SyncNotedAppOp startProxyOperation(IBinder clientId, int code,
+            in AttributionSource attributionSource, boolean startIfModeDefault,
+            boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
+            boolean skipProxyOperation, int proxyAttributionFlags, int proxiedAttributionFlags,
+            int attributionChainId);
+    void finishProxyOperation(IBinder clientId, int code, in AttributionSource attributionSource,
             boolean skipProxyOperation);
 
     // Remaining methods are only used in Java.
diff --git a/core/java/com/android/internal/statusbar/IStatusBar.aidl b/core/java/com/android/internal/statusbar/IStatusBar.aidl
index 2cf41bb..017bf3f 100644
--- a/core/java/com/android/internal/statusbar/IStatusBar.aidl
+++ b/core/java/com/android/internal/statusbar/IStatusBar.aidl
@@ -331,4 +331,11 @@
 
     /** Called when requested to go to fullscreen from the active split app. */
     void goToFullscreenFromSplit();
+
+    /**
+     * Enters stage split from a current running app.
+     *
+     * @param leftOrTop indicates where the stage split is.
+     */
+    void enterStageSplitFromRunningApp(boolean leftOrTop);
 }
diff --git a/core/res/res/values/strings.xml b/core/res/res/values/strings.xml
index 5763345..9410e06 100644
--- a/core/res/res/values/strings.xml
+++ b/core/res/res/values/strings.xml
@@ -58,6 +58,9 @@
     <!-- Displayed when the user dialed an MMI code whose function
          could not be performed. This will be displayed in a toast. -->
     <string name="mmiError">Connection problem or invalid MMI code.</string>
+    <!-- Displayed when the user dialed an MMI code whose function could not be performed because
+         the feature is not supported on the current mobile network. -->
+    <string name="mmiErrorNotSupported">Feature not supported.</string>
     <!-- Displayed when the user dialed an MMI code whose function
          could not be performed because FDN is enabled. This will be displayed in a toast. -->
     <string name="mmiFdnError">Operation is restricted to fixed dialing numbers only.</string>
diff --git a/core/res/res/values/symbols.xml b/core/res/res/values/symbols.xml
index 694040a..9722e79 100644
--- a/core/res/res/values/symbols.xml
+++ b/core/res/res/values/symbols.xml
@@ -837,6 +837,7 @@
   <java-symbol type="string" name="mismatchPin" />
   <java-symbol type="string" name="mmiComplete" />
   <java-symbol type="string" name="mmiError" />
+  <java-symbol type="string" name="mmiErrorNotSupported" />
   <java-symbol type="string" name="mmiFdnError" />
   <java-symbol type="string" name="mmiErrorWhileRoaming" />
   <java-symbol type="string" name="month_day_year" />
diff --git a/libs/WindowManager/Shell/res/drawable/size_compat_restart_button.xml b/libs/WindowManager/Shell/res/drawable/size_compat_restart_button.xml
index e6ae282..2994593 100644
--- a/libs/WindowManager/Shell/res/drawable/size_compat_restart_button.xml
+++ b/libs/WindowManager/Shell/res/drawable/size_compat_restart_button.xml
@@ -1,19 +1,19 @@
 <?xml version="1.0" encoding="utf-8"?>
 <!--
-     Copyright (C) 2019 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.
--->
+  ~ Copyright (C) 2022 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.
+  -->
 <vector xmlns:android="http://schemas.android.com/apk/res/android"
         android:width="48dp"
         android:height="48dp"
@@ -25,13 +25,12 @@
         android:fillAlpha="0.8"
         android:pathData="M0,24 a24,24 0 1,0 48,0 a24,24 0 1,0 -48,0"/>
     <group
-        android:translateX="12"
-        android:translateY="12">
+        android:scaleX="0.8"
+        android:scaleY="0.8"
+        android:translateX="10"
+        android:translateY="10">
         <path
-            android:fillColor="@color/compat_controls_text"
-            android:pathData="M6,13c0,-1.65 0.67,-3.15 1.76,-4.24L6.34,7.34C4.9,8.79 4,10.79 4,13c0,4.08 3.05,7.44 7,7.93v-2.02C8.17,18.43 6,15.97 6,13z"/>
-        <path
-            android:fillColor="@color/compat_controls_text"
-            android:pathData="M20,13c0,-4.42 -3.58,-8 -8,-8c-0.06,0 -0.12,0.01 -0.18,0.01v0l1.09,-1.09L11.5,2.5L8,6l3.5,3.5l1.41,-1.41l-1.08,-1.08C11.89,7.01 11.95,7 12,7c3.31,0 6,2.69 6,6c0,2.97 -2.17,5.43 -5,5.91v2.02C16.95,20.44 20,17.08 20,13z"/>
+            android:pathData="M0,36V24.5H3V30.85L10.4,23.45L12.55,25.6L5.15,33H11.5V36H0ZM24.5,36V33H30.85L23.5,25.65L25.65,23.5L33,30.85V24.5H36V36H24.5ZM10.35,12.5L3,5.15V11.5H0V0H11.5V3H5.15L12.5,10.35L10.35,12.5ZM25.65,12.5L23.5,10.35L30.85,3H24.5V0H36V11.5H33V5.15L25.65,12.5Z"
+            android:fillColor="@color/compat_controls_text"/>
     </group>
 </vector>
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
index 6e116b9..c836b95 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitDecorManager.java
@@ -51,8 +51,6 @@
 import com.android.wm.shell.common.ScreenshotUtils;
 import com.android.wm.shell.common.SurfaceUtils;
 
-import java.util.function.Consumer;
-
 /**
  * Handles split decor like showing resizing hint for a specific split.
  */
@@ -72,17 +70,18 @@
     private SurfaceControl mIconLeash;
     private SurfaceControl mBackgroundLeash;
     private SurfaceControl mGapBackgroundLeash;
+    private SurfaceControl mScreenshot;
 
     private boolean mShown;
     private boolean mIsResizing;
     private final Rect mBounds = new Rect();
-    private final Rect mResizingBounds = new Rect();
     private final Rect mTempRect = new Rect();
     private ValueAnimator mFadeAnimator;
 
     private int mIconSize;
     private int mOffsetX;
     private int mOffsetY;
+    private int mRunningAnimationCount = 0;
 
     public SplitDecorManager(Configuration configuration, IconProvider iconProvider,
             SurfaceSession surfaceSession) {
@@ -173,7 +172,6 @@
             mIsResizing = true;
             mBounds.set(newBounds);
         }
-        mResizingBounds.set(newBounds);
         mOffsetX = offsetX;
         mOffsetY = offsetY;
 
@@ -227,33 +225,41 @@
                 t.setVisibility(mBackgroundLeash, show);
                 t.setVisibility(mIconLeash, show);
             } else {
-                startFadeAnimation(show, null /* finishedConsumer */);
+                startFadeAnimation(show, false, null);
             }
             mShown = show;
         }
     }
 
     /** Stops showing resizing hint. */
-    public void onResized(SurfaceControl.Transaction t) {
-        if (!mShown && mIsResizing) {
-            mTempRect.set(mResizingBounds);
-            mTempRect.offsetTo(-mOffsetX, -mOffsetY);
-            final SurfaceControl screenshot = ScreenshotUtils.takeScreenshot(t,
-                    mHostLeash, mTempRect, Integer.MAX_VALUE - 1);
+    public void onResized(SurfaceControl.Transaction t, Runnable animFinishedCallback) {
+        if (mScreenshot != null) {
+            t.setPosition(mScreenshot, mOffsetX, mOffsetY);
 
             final SurfaceControl.Transaction animT = new SurfaceControl.Transaction();
             final ValueAnimator va = ValueAnimator.ofFloat(1, 0);
             va.addUpdateListener(valueAnimator -> {
                 final float progress = (float) valueAnimator.getAnimatedValue();
-                animT.setAlpha(screenshot, progress);
+                animT.setAlpha(mScreenshot, progress);
                 animT.apply();
             });
             va.addListener(new AnimatorListenerAdapter() {
                 @Override
+                public void onAnimationStart(Animator animation) {
+                    mRunningAnimationCount++;
+                }
+
+                @Override
                 public void onAnimationEnd(@androidx.annotation.NonNull Animator animation) {
-                    animT.remove(screenshot);
+                    mRunningAnimationCount--;
+                    animT.remove(mScreenshot);
                     animT.apply();
                     animT.close();
+                    mScreenshot = null;
+
+                    if (mRunningAnimationCount == 0 && animFinishedCallback != null) {
+                        animFinishedCallback.run();
+                    }
                 }
             });
             va.start();
@@ -285,10 +291,34 @@
             mFadeAnimator.cancel();
         }
         if (mShown) {
-            fadeOutDecor(null /* finishedCallback */);
+            fadeOutDecor(animFinishedCallback);
         } else {
             // Decor surface is hidden so release it directly.
             releaseDecor(t);
+            if (mRunningAnimationCount == 0 && animFinishedCallback != null) {
+                animFinishedCallback.run();
+            }
+        }
+    }
+
+    /** Screenshot host leash and attach on it if meet some conditions */
+    public void screenshotIfNeeded(SurfaceControl.Transaction t) {
+        if (!mShown && mIsResizing) {
+            mTempRect.set(mBounds);
+            mTempRect.offsetTo(0, 0);
+            mScreenshot = ScreenshotUtils.takeScreenshot(t, mHostLeash, mTempRect,
+                    Integer.MAX_VALUE - 1);
+        }
+    }
+
+    /** Set screenshot and attach on host leash it if meet some conditions */
+    public void setScreenshotIfNeeded(SurfaceControl screenshot, SurfaceControl.Transaction t) {
+        if (screenshot == null || !screenshot.isValid()) return;
+
+        if (!mShown && mIsResizing) {
+            mScreenshot = screenshot;
+            t.reparent(screenshot, mHostLeash);
+            t.setLayer(screenshot, Integer.MAX_VALUE - 1);
         }
     }
 
@@ -296,18 +326,15 @@
      * directly. */
     public void fadeOutDecor(Runnable finishedCallback) {
         if (mShown) {
-            startFadeAnimation(false /* show */, transaction -> {
-                releaseDecor(transaction);
-                if (finishedCallback != null) finishedCallback.run();
-            });
+            startFadeAnimation(false /* show */, true, finishedCallback);
             mShown = false;
         } else {
             if (finishedCallback != null) finishedCallback.run();
         }
     }
 
-    private void startFadeAnimation(boolean show,
-            Consumer<SurfaceControl.Transaction> finishedConsumer) {
+    private void startFadeAnimation(boolean show, boolean releaseSurface,
+            Runnable finishedCallback) {
         final SurfaceControl.Transaction animT = new SurfaceControl.Transaction();
         mFadeAnimator = ValueAnimator.ofFloat(0f, 1f);
         mFadeAnimator.setDuration(FADE_DURATION);
@@ -324,6 +351,7 @@
         mFadeAnimator.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationStart(@NonNull Animator animation) {
+                mRunningAnimationCount++;
                 if (show) {
                     animT.show(mBackgroundLeash).show(mIconLeash);
                 }
@@ -335,6 +363,7 @@
 
             @Override
             public void onAnimationEnd(@NonNull Animator animation) {
+                mRunningAnimationCount--;
                 if (!show) {
                     if (mBackgroundLeash != null) {
                         animT.hide(mBackgroundLeash);
@@ -343,11 +372,15 @@
                         animT.hide(mIconLeash);
                     }
                 }
-                if (finishedConsumer != null) {
-                    finishedConsumer.accept(animT);
+                if (releaseSurface) {
+                    releaseDecor(animT);
                 }
                 animT.apply();
                 animT.close();
+
+                if (mRunningAnimationCount == 0 && finishedCallback != null) {
+                    finishedCallback.run();
+                }
             }
         });
         mFadeAnimator.start();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/OWNERS b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/OWNERS
new file mode 100644
index 0000000..926cfb3
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/OWNERS
@@ -0,0 +1,2 @@
+# WM shell sub-module desktop owners
+madym@google.com
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/OWNERS b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/OWNERS
new file mode 100644
index 0000000..0c2d5c4
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/OWNERS
@@ -0,0 +1,2 @@
+# WM shell sub-module freeform owners
+madym@google.com
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 400039b..9329d02 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -529,10 +529,24 @@
             @SplitPosition int splitPosition, float splitRatio, RemoteAnimationAdapter adapter,
             InstanceId instanceId) {
         Intent fillInIntent = null;
-        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)
-                && supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
-            fillInIntent = new Intent();
-            fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)) {
+            if (supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
+                fillInIntent = new Intent();
+                fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                try {
+                    adapter.getRunner().onAnimationCancelled(false /* isKeyguardOccluded */);
+                    ActivityTaskManager.getService().startActivityFromRecents(taskId, options2);
+                } catch (RemoteException e) {
+                    Slog.e(TAG, "Error starting remote animation", e);
+                }
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+                return;
+            }
         }
         mStageCoordinator.startIntentAndTaskWithLegacyTransition(pendingIntent, fillInIntent,
                 options1, taskId, options2, splitPosition, splitRatio, adapter, instanceId);
@@ -542,10 +556,17 @@
             int taskId, @Nullable Bundle options2, @SplitPosition int splitPosition,
             float splitRatio, @Nullable RemoteTransition remoteTransition, InstanceId instanceId) {
         Intent fillInIntent = null;
-        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)
-                && supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
-            fillInIntent = new Intent();
-            fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent, splitPosition, taskId)) {
+            if (supportMultiInstancesSplit(pendingIntent.getIntent().getComponent())) {
+                fillInIntent = new Intent();
+                fillInIntent.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+            }
         }
         mStageCoordinator.startIntentAndTask(pendingIntent, fillInIntent, options1, taskId,
                 options2, splitPosition, splitRatio, remoteTransition, instanceId);
@@ -557,12 +578,26 @@
             float splitRatio, RemoteAnimationAdapter adapter, InstanceId instanceId) {
         Intent fillInIntent1 = null;
         Intent fillInIntent2 = null;
-        if (launchSameComponentAdjacently(pendingIntent1, pendingIntent2)
-                && supportMultiInstancesSplit(pendingIntent1.getIntent().getComponent())) {
-            fillInIntent1 = new Intent();
-            fillInIntent1.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
-            fillInIntent2 = new Intent();
-            fillInIntent2.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+        if (launchSameComponentAdjacently(pendingIntent1, pendingIntent2)) {
+            if (supportMultiInstancesSplit(pendingIntent1.getIntent().getComponent())) {
+                fillInIntent1 = new Intent();
+                fillInIntent1.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                fillInIntent2 = new Intent();
+                fillInIntent2.addFlags(FLAG_ACTIVITY_MULTIPLE_TASK);
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN, "Adding MULTIPLE_TASK");
+            } else {
+                try {
+                    adapter.getRunner().onAnimationCancelled(false /* isKeyguardOccluded */);
+                    pendingIntent1.send();
+                } catch (RemoteException | PendingIntent.CanceledException e) {
+                    Slog.e(TAG, "Error starting remote animation", e);
+                }
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
+                Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
+                        Toast.LENGTH_SHORT).show();
+                return;
+            }
         }
         mStageCoordinator.startIntentsWithLegacyTransition(pendingIntent1, fillInIntent1, options1,
                 pendingIntent2, fillInIntent2, options2, splitPosition, splitRatio, adapter,
@@ -601,6 +636,8 @@
                 mStageCoordinator.switchSplitPosition("startIntent");
                 return;
             } else {
+                ProtoLog.v(ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN,
+                        "Cancel entering split as not supporting multi-instances");
                 Toast.makeText(mContext, R.string.dock_multi_instances_not_supported_text,
                         Toast.LENGTH_SHORT).show();
                 return;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
index 21a1310..1cf3a89 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
@@ -47,6 +47,7 @@
 
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.common.split.SplitDecorManager;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.transition.OneShotRemoteHandler;
 import com.android.wm.shell.transition.Transitions;
@@ -64,6 +65,7 @@
     DismissTransition mPendingDismiss = null;
     TransitSession mPendingEnter = null;
     TransitSession mPendingRecent = null;
+    TransitSession mPendingResize = null;
 
     private IBinder mAnimatingTransition = null;
     OneShotRemoteHandler mPendingRemoteHandler = null;
@@ -177,6 +179,43 @@
         onFinish(null /* wct */, null /* wctCB */);
     }
 
+    void applyResizeTransition(@NonNull IBinder transition, @NonNull TransitionInfo info,
+            @NonNull SurfaceControl.Transaction startTransaction,
+            @NonNull SurfaceControl.Transaction finishTransaction,
+            @NonNull Transitions.TransitionFinishCallback finishCallback,
+            @NonNull WindowContainerToken mainRoot, @NonNull WindowContainerToken sideRoot,
+            @NonNull SplitDecorManager mainDecor, @NonNull SplitDecorManager sideDecor) {
+        mFinishCallback = finishCallback;
+        mAnimatingTransition = transition;
+        mFinishTransaction = finishTransaction;
+
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            final TransitionInfo.Change change = info.getChanges().get(i);
+            if (mainRoot.equals(change.getContainer()) || sideRoot.equals(change.getContainer())) {
+                final SurfaceControl leash = change.getLeash();
+                startTransaction.setPosition(leash, change.getEndAbsBounds().left,
+                        change.getEndAbsBounds().top);
+                startTransaction.setWindowCrop(leash, change.getEndAbsBounds().width(),
+                        change.getEndAbsBounds().height());
+
+                SplitDecorManager decor = mainRoot.equals(change.getContainer())
+                        ? mainDecor : sideDecor;
+                ValueAnimator va = new ValueAnimator();
+                mAnimations.add(va);
+                decor.setScreenshotIfNeeded(change.getSnapshot(), startTransaction);
+                decor.onResized(startTransaction, () -> {
+                    mTransitions.getMainExecutor().execute(() -> {
+                        mAnimations.remove(va);
+                        onFinish(null /* wct */, null /* wctCB */);
+                    });
+                });
+            }
+        }
+
+        startTransaction.apply();
+        onFinish(null /* wct */, null /* wctCB */);
+    }
+
     boolean isPendingTransition(IBinder transition) {
         return getPendingTransition(transition) != null;
     }
@@ -193,6 +232,10 @@
         return mPendingDismiss != null && mPendingDismiss.mTransition == transition;
     }
 
+    boolean isPendingResize(IBinder transition) {
+        return mPendingResize != null && mPendingResize.mTransition == transition;
+    }
+
     @Nullable
     private TransitSession getPendingTransition(IBinder transition) {
         if (isPendingEnter(transition)) {
@@ -201,11 +244,14 @@
             return mPendingRecent;
         } else if (isPendingDismiss(transition)) {
             return mPendingDismiss;
+        } else if (isPendingResize(transition)) {
+            return mPendingResize;
         }
 
         return null;
     }
 
+
     /** Starts a transition to enter split with a remote transition animator. */
     IBinder startEnterTransition(
             @WindowManager.TransitionType int transitType,
@@ -258,6 +304,21 @@
                 exitReasonToString(reason), stageTypeToString(dismissTop));
     }
 
+    IBinder startResizeTransition(WindowContainerTransaction wct,
+            Transitions.TransitionHandler handler,
+            @Nullable TransitionFinishedCallback finishCallback) {
+        IBinder transition = mTransitions.startTransition(TRANSIT_CHANGE, wct, handler);
+        setResizeTransition(transition, finishCallback);
+        return transition;
+    }
+
+    void setResizeTransition(@NonNull IBinder transition,
+            @Nullable TransitionFinishedCallback finishCallback) {
+        mPendingResize = new TransitSession(transition, null /* consumedCb */, finishCallback);
+        ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, "  splitTransition "
+                + " deduced Resize split screen");
+    }
+
     void setRecentTransition(@NonNull IBinder transition,
             @Nullable RemoteTransition remoteTransition,
             @Nullable TransitionFinishedCallback finishCallback) {
@@ -324,6 +385,9 @@
             mPendingRecent.onConsumed(aborted);
             mPendingRecent = null;
             mPendingRemoteHandler = null;
+        } else if (isPendingResize(transition)) {
+            mPendingResize.onConsumed(aborted);
+            mPendingResize = null;
         }
     }
 
@@ -340,6 +404,9 @@
         } else if (isPendingDismiss(mAnimatingTransition)) {
             mPendingDismiss.onFinished(wct, mFinishTransaction);
             mPendingDismiss = null;
+        } else if (isPendingResize(mAnimatingTransition)) {
+            mPendingResize.onFinished(wct, mFinishTransaction);
+            mPendingResize = null;
         }
 
         mPendingRemoteHandler = null;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
index 2dc4a04..1016e1b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitscreenEventLogger.java
@@ -23,6 +23,7 @@
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__APP_FINISHED;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__DEVICE_FOLDED;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__DRAG_DIVIDER;
+import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__FULLSCREEN_SHORTCUT;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__RETURN_HOME;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__ROOT_TASK_VANISHED;
 import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__SCREEN_LOCKED;
@@ -38,6 +39,7 @@
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_APP_FINISHED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_DEVICE_FOLDED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_DRAG_DIVIDER;
+import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_FULLSCREEN_SHORTCUT;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_RETURN_HOME;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_ROOT_TASK_VANISHED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_SCREEN_LOCKED;
@@ -180,6 +182,8 @@
                 return SPLITSCREEN_UICHANGED__EXIT_REASON__SCREEN_LOCKED;
             case EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP:
                 return SPLITSCREEN_UICHANGED__EXIT_REASON__SCREEN_LOCKED_SHOW_ON_TOP;
+            case EXIT_REASON_FULLSCREEN_SHORTCUT:
+                return SPLITSCREEN_UICHANGED__EXIT_REASON__FULLSCREEN_SHORTCUT;
             case EXIT_REASON_UNKNOWN:
                 // Fall through
             default:
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index aa0512b..da8dc87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -456,8 +456,6 @@
     void startIntentLegacy(PendingIntent intent, Intent fillInIntent, @SplitPosition int position,
             @Nullable Bundle options) {
         final boolean isEnteringSplit = !isSplitActive();
-        final WindowContainerTransaction evictWct = new WindowContainerTransaction();
-        prepareEvictChildTasks(position, evictWct);
 
         LegacyTransitions.ILegacyTransition transition = new LegacyTransitions.ILegacyTransition() {
             @Override
@@ -465,22 +463,21 @@
                     RemoteAnimationTarget[] wallpapers, RemoteAnimationTarget[] nonApps,
                     IRemoteAnimationFinishedCallback finishedCallback,
                     SurfaceControl.Transaction t) {
-                if (isEnteringSplit) {
-                    boolean openingToSide = false;
-                    if (apps != null) {
-                        for (int i = 0; i < apps.length; ++i) {
-                            if (apps[i].mode == MODE_OPENING
-                                    && mSideStage.containsTask(apps[i].taskId)) {
-                                openingToSide = true;
-                                break;
-                            }
+                boolean openingToSide = false;
+                if (apps != null) {
+                    for (int i = 0; i < apps.length; ++i) {
+                        if (apps[i].mode == MODE_OPENING
+                                && mSideStage.containsTask(apps[i].taskId)) {
+                            openingToSide = true;
+                            break;
                         }
                     }
-                    if (!openingToSide) {
-                        mMainExecutor.execute(() -> exitSplitScreen(
-                                mSideStage.getChildCount() == 0 ? mMainStage : mSideStage,
-                                EXIT_REASON_UNKNOWN));
-                    }
+                }
+
+                if (isEnteringSplit && !openingToSide) {
+                    mMainExecutor.execute(() -> exitSplitScreen(
+                            mSideStage.getChildCount() == 0 ? mMainStage : mSideStage,
+                            EXIT_REASON_UNKNOWN));
                 }
 
                 if (apps != null) {
@@ -500,7 +497,12 @@
                     }
                 }
 
-                mSyncQueue.queue(evictWct);
+
+                if (!isEnteringSplit && openingToSide) {
+                    final WindowContainerTransaction evictWct = new WindowContainerTransaction();
+                    prepareEvictNonOpeningChildTasks(position, apps, evictWct);
+                    mSyncQueue.queue(evictWct);
+                }
             }
         };
 
@@ -1667,15 +1669,29 @@
     public void onLayoutSizeChanged(SplitLayout layout) {
         // Reset this flag every time onLayoutSizeChanged.
         mShowDecorImmediately = false;
+
+        if (!ENABLE_SHELL_TRANSITIONS) {
+            // Only need screenshot for legacy case because shell transition should screenshot
+            // itself during transition.
+            final SurfaceControl.Transaction startT = mTransactionPool.acquire();
+            mMainStage.screenshotIfNeeded(startT);
+            mSideStage.screenshotIfNeeded(startT);
+            mTransactionPool.release(startT);
+        }
+
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         updateWindowBounds(layout, wct);
         sendOnBoundsChanged();
-        mSyncQueue.queue(wct);
-        mSyncQueue.runInSync(t -> {
-            updateSurfaceBounds(layout, t, false /* applyResizingOffset */);
-            mMainStage.onResized(t);
-            mSideStage.onResized(t);
-        });
+        if (ENABLE_SHELL_TRANSITIONS) {
+            mSplitTransitions.startResizeTransition(wct, this, null /* callback */);
+        } else {
+            mSyncQueue.queue(wct);
+            mSyncQueue.runInSync(t -> {
+                updateSurfaceBounds(layout, t, false /* applyResizingOffset */);
+                mMainStage.onResized(t);
+                mSideStage.onResized(t);
+            });
+        }
         mLogger.logResize(mSplitLayout.getDividerPositionAsFraction());
     }
 
@@ -2029,6 +2045,12 @@
         } else if (mSplitTransitions.isPendingDismiss(transition)) {
             shouldAnimate = startPendingDismissAnimation(
                     mSplitTransitions.mPendingDismiss, info, startTransaction, finishTransaction);
+        } else if (mSplitTransitions.isPendingResize(transition)) {
+            mSplitTransitions.applyResizeTransition(transition, info, startTransaction,
+                    finishTransaction, finishCallback, mMainStage.mRootTaskInfo.token,
+                    mSideStage.mRootTaskInfo.token, mMainStage.getSplitDecorManager(),
+                    mSideStage.getSplitDecorManager());
+            return true;
         }
         if (!shouldAnimate) return false;
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 358f712..8a52c87 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -292,7 +292,13 @@
 
     void onResized(SurfaceControl.Transaction t) {
         if (mSplitDecorManager != null) {
-            mSplitDecorManager.onResized(t);
+            mSplitDecorManager.onResized(t, null);
+        }
+    }
+
+    void screenshotIfNeeded(SurfaceControl.Transaction t) {
+        if (mSplitDecorManager != null) {
+            mSplitDecorManager.screenshotIfNeeded(t);
         }
     }
 
@@ -304,6 +310,10 @@
         }
     }
 
+    SplitDecorManager getSplitDecorManager() {
+        return mSplitDecorManager;
+    }
+
     void addTask(ActivityManager.RunningTaskInfo task, WindowContainerTransaction wct) {
         // Clear overridden bounds and windowing mode to make sure the child task can inherit
         // windowing mode and bounds from split root.
diff --git a/packages/SystemUI/AndroidManifest.xml b/packages/SystemUI/AndroidManifest.xml
index 844e88a..acaa008 100644
--- a/packages/SystemUI/AndroidManifest.xml
+++ b/packages/SystemUI/AndroidManifest.xml
@@ -904,6 +904,29 @@
         <service android:name=".controls.controller.AuxiliaryPersistenceWrapper$DeletionJobService"
                  android:permission="android.permission.BIND_JOB_SERVICE"/>
 
+        <!-- region Note Task -->
+        <activity
+            android:name=".notetask.shortcut.CreateNoteTaskShortcutActivity"
+            android:enabled="false"
+            android:exported="true"
+            android:excludeFromRecents="true"
+            android:theme="@android:style/Theme.NoDisplay"
+            android:label="@string/note_task_button_label"
+            android:icon="@drawable/ic_note_task_button">
+
+            <intent-filter>
+                <action android:name="android.intent.action.CREATE_SHORTCUT" />
+                <category android:name="android.intent.category.DEFAULT" />
+            </intent-filter>
+        </activity>
+
+        <activity
+            android:name=".notetask.shortcut.LaunchNoteTaskActivity"
+            android:exported="true"
+            android:excludeFromRecents="true"
+            android:theme="@android:style/Theme.NoDisplay" />
+        <!-- endregion -->
+
         <!-- started from ControlsRequestReceiver -->
         <activity
             android:name=".controls.management.ControlsRequestDialog"
diff --git a/packages/SystemUI/animation/Android.bp b/packages/SystemUI/animation/Android.bp
index 17ad55f..8acc2f8 100644
--- a/packages/SystemUI/animation/Android.bp
+++ b/packages/SystemUI/animation/Android.bp
@@ -44,23 +44,3 @@
     manifest: "AndroidManifest.xml",
     kotlincflags: ["-Xjvm-default=all"],
 }
-
-android_test {
-    name: "SystemUIAnimationLibTests",
-
-    static_libs: [
-        "SystemUIAnimationLib",
-        "androidx.test.ext.junit",
-        "androidx.test.rules",
-        "testables",
-    ],
-    libs: [
-        "android.test.base",
-    ],
-    srcs: [
-        "**/*.java",
-        "**/*.kt",
-    ],
-    kotlincflags: ["-Xjvm-default=all"],
-    test_suites: ["general-tests"],
-}
diff --git a/packages/SystemUI/animation/TEST_MAPPING b/packages/SystemUI/animation/TEST_MAPPING
deleted file mode 100644
index 3dc8510..0000000
--- a/packages/SystemUI/animation/TEST_MAPPING
+++ /dev/null
@@ -1,7 +0,0 @@
-{
-  "presubmit": [
-    {
-      "name": "SystemUIAnimationLibTests"
-    }
-  ]
-}
diff --git a/packages/SystemUI/res/drawable/ic_note_task_button.xml b/packages/SystemUI/res/drawable/ic_note_task_button.xml
new file mode 100644
index 0000000..bb5e224
--- /dev/null
+++ b/packages/SystemUI/res/drawable/ic_note_task_button.xml
@@ -0,0 +1,28 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2022 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.
+  -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportHeight="24"
+    android:viewportWidth="24">
+    <path
+        android:fillColor="#636C6F"
+        android:pathData="M17.6258,4.96L19.0358,6.37L7.4058,18.01L5.9958,16.6L17.6258,4.96ZM16.1358,3.62L4.1258,15.63L3.0158,19.83C2.9058,20.45 3.3858,21 3.9958,21C4.0558,21 4.1058,21 4.1658,20.99L8.3658,19.88L20.3758,7.86C20.7758,7.46 20.9958,6.93 20.9958,6.37C20.9958,5.81 20.7758,5.28 20.3758,4.88L19.1058,3.61C18.7158,3.22 18.1858,3 17.6258,3C17.0658,3 16.5358,3.22 16.1358,3.62Z" />
+    <path
+        android:fillColor="#636C6F"
+        android:fillType="evenOdd"
+        android:pathData="M20.1936,15.3369C20.3748,16.3837 19.9151,17.5414 18.8846,18.7597C19.1546,18.872 19.4576,18.9452 19.7724,18.9867C20.0839,19.0278 20.3683,19.0325 20.5749,19.0266C20.6772,19.0236 20.7578,19.0181 20.8101,19.0138C20.8362,19.0116 20.855,19.0097 20.8657,19.0085L20.8754,19.0074L20.875,19.0075C21.4217,18.9385 21.9214,19.325 21.9918,19.8718C22.0624,20.4195 21.6756,20.9208 21.1279,20.9914L21,19.9996C21.1279,20.9914 21.1265,20.9916 21.1265,20.9916L21.1249,20.9918L21.1211,20.9923L21.1107,20.9935L21.0795,20.997C21.0542,20.9998 21.0199,21.0032 20.9775,21.0067C20.8929,21.0138 20.7753,21.0216 20.6323,21.0257C20.3481,21.0339 19.9533,21.0279 19.5109,20.9695C18.873,20.8854 18.0393,20.6793 17.3106,20.1662C16.9605,20.3559 16.5876,20.4952 16.2299,20.6003C15.5742,20.7927 14.8754,20.8968 14.2534,20.9534C13.6801,21.0055 13.4553,21.0037 13.1015,21.0008C13.0689,21.0005 13.0352,21.0002 13,21H12.8594C12.8214,21.0002 12.785,21.0006 12.7504,21.0009C12.6524,21.0019 12.5683,21.0027 12.5,21H12.0562C12.0277,21.0003 12.0054,21.0006 11.9926,21.001L11.9751,21H9L11,19H11.9795C11.9929,18.9997 12.0064,18.9997 12.0199,19H12.4117C12.4534,18.9996 12.4864,18.9995 12.5,19H12.9675C12.977,18.9999 12.9878,18.9999 13,19C13.0446,19.0003 13.0859,19.0007 13.1249,19.0011C13.4259,19.0038 13.591,19.0054 14.0723,18.9616C14.6201,18.9118 15.1795,18.8242 15.6665,18.6813C15.753,18.6559 15.8346,18.6295 15.9114,18.6022C15.0315,17.2981 14.7125,16.1044 15.015,15.0829C15.4095,13.7511 16.6784,13.2418 17.7026,13.2864C18.7262,13.3309 19.954,13.9529 20.1936,15.3369ZM16.9327,15.6508C16.873,15.8523 16.8651,16.3878 17.4697,17.334C18.2007,16.4284 18.2585,15.8839 18.2229,15.6781C18.1939,15.5108 18.0297,15.3025 17.6157,15.2845C17.2025,15.2665 16.9885,15.4626 16.9327,15.6508Z" />
+</vector>
diff --git a/packages/SystemUI/res/values/flags.xml b/packages/SystemUI/res/values/flags.xml
index 49dd574..fd2e324 100644
--- a/packages/SystemUI/res/values/flags.xml
+++ b/packages/SystemUI/res/values/flags.xml
@@ -36,4 +36,8 @@
       avatar will no longer show on the lockscreen -->
     <bool name="flag_user_switcher_chip">false</bool>
 
+    <!-- Whether the battery icon is allowed to display a shield when battery life is being
+         protected. -->
+    <bool name="flag_battery_shield_icon">false</bool>
+
 </resources>
diff --git a/packages/SystemUI/res/values/strings.xml b/packages/SystemUI/res/values/strings.xml
index 30788cf..52a931b 100644
--- a/packages/SystemUI/res/values/strings.xml
+++ b/packages/SystemUI/res/values/strings.xml
@@ -2669,6 +2669,10 @@
         <xliff:g id="weather_condition" example="Partly cloudy">%1$s</xliff:g>, <xliff:g id="temperature" example="7°C">%2$s</xliff:g>
     </string>
 
+    <!-- TODO(b/259369672): Replace with final resource. -->
+    <!-- [CHAR LIMIT=30] Label used to open Note Task -->
+    <string name="note_task_button_label">Notetaking</string>
+
     <!-- [CHAR LIMIT=NONE] Le audio broadcast dialog, media app is broadcasting -->
     <string name="broadcasting_description_is_broadcasting">Broadcasting</string>
     <!-- [CHAR LIMIT=NONE] Le audio broadcast dialog, title -->
diff --git a/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java b/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
index 02a6d7b..e6f559b 100644
--- a/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
+++ b/packages/SystemUI/src/com/android/systemui/ScreenDecorations.java
@@ -210,8 +210,10 @@
                 (FaceScanningOverlay) getOverlayView(mFaceScanningViewId);
         if (faceScanningOverlay != null) {
             faceScanningOverlay.setHideOverlayRunnable(() -> {
+                Trace.beginSection("ScreenDecorations#hideOverlayRunnable");
                 updateOverlayWindowVisibilityIfViewExists(
                         faceScanningOverlay.findViewById(mFaceScanningViewId));
+                Trace.endSection();
             });
             faceScanningOverlay.enableShowProtection(false);
         }
@@ -273,16 +275,18 @@
             if (mOverlays == null || !shouldOptimizeVisibility()) {
                 return;
             }
-
+            Trace.beginSection("ScreenDecorations#updateOverlayWindowVisibilityIfViewExists");
             for (final OverlayWindow overlay : mOverlays) {
                 if (overlay == null) {
                     continue;
                 }
                 if (overlay.getView(view.getId()) != null) {
                     overlay.getRootView().setVisibility(getWindowVisibility(overlay, true));
+                    Trace.endSection();
                     return;
                 }
             }
+            Trace.endSection();
         });
     }
 
@@ -370,6 +374,7 @@
     }
 
     private void startOnScreenDecorationsThread() {
+        Trace.beginSection("ScreenDecorations#startOnScreenDecorationsThread");
         mWindowManager = mContext.getSystemService(WindowManager.class);
         mDisplayManager = mContext.getSystemService(DisplayManager.class);
         mContext.getDisplay().getDisplayInfo(mDisplayInfo);
@@ -472,6 +477,7 @@
 
         mDisplayManager.registerDisplayListener(mDisplayListener, mHandler);
         updateConfiguration();
+        Trace.endSection();
     }
 
     @VisibleForTesting
@@ -521,6 +527,12 @@
     }
 
     private void setupDecorations() {
+        Trace.beginSection("ScreenDecorations#setupDecorations");
+        setupDecorationsInner();
+        Trace.endSection();
+    }
+
+    private void setupDecorationsInner() {
         if (hasRoundedCorners() || shouldDrawCutout() || isPrivacyDotEnabled()
                 || mFaceScanningFactory.getHasProviders()) {
 
@@ -573,7 +585,11 @@
                 return;
             }
 
-            mMainExecutor.execute(() -> mTunerService.addTunable(this, SIZE));
+            mMainExecutor.execute(() -> {
+                Trace.beginSection("ScreenDecorations#addTunable");
+                mTunerService.addTunable(this, SIZE);
+                Trace.endSection();
+            });
 
             // Watch color inversion and invert the overlay as needed.
             if (mColorInversionSetting == null) {
@@ -593,7 +609,11 @@
             mUserTracker.addCallback(mUserChangedCallback, mExecutor);
             mIsRegistered = true;
         } else {
-            mMainExecutor.execute(() -> mTunerService.removeTunable(this));
+            mMainExecutor.execute(() -> {
+                Trace.beginSection("ScreenDecorations#removeTunable");
+                mTunerService.removeTunable(this);
+                Trace.endSection();
+            });
 
             if (mColorInversionSetting != null) {
                 mColorInversionSetting.setListening(false);
@@ -939,6 +959,7 @@
         }
 
         mExecutor.execute(() -> {
+            Trace.beginSection("ScreenDecorations#onConfigurationChanged");
             int oldRotation = mRotation;
             mPendingConfigChange = false;
             updateConfiguration();
@@ -951,6 +972,7 @@
                 // the updated rotation).
                 updateLayoutParams();
             }
+            Trace.endSection();
         });
     }
 
@@ -1119,6 +1141,7 @@
             if (mOverlays == null || !SIZE.equals(key)) {
                 return;
             }
+            Trace.beginSection("ScreenDecorations#onTuningChanged");
             try {
                 final int sizeFactor = Integer.parseInt(newValue);
                 mRoundedCornerResDelegate.setTuningSizeFactor(sizeFactor);
@@ -1132,6 +1155,7 @@
                     R.id.rounded_corner_bottom_right
             });
             updateHwLayerRoundedCornerExistAndSize();
+            Trace.endSection();
         });
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
index 7c05356..389639d 100644
--- a/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
+++ b/packages/SystemUI/src/com/android/systemui/flags/Flags.kt
@@ -83,8 +83,7 @@
     val SEMI_STABLE_SORT = unreleasedFlag(115, "semi_stable_sort", teamfood = true)
 
     @JvmField
-    val NOTIFICATION_GROUP_CORNER =
-        unreleasedFlag(116, "notification_group_corner", teamfood = true)
+    val USE_ROUNDNESS_SOURCETYPES = unreleasedFlag(116, "use_roundness_sourcetype", teamfood = true)
 
     // TODO(b/259217907)
     @JvmField
@@ -229,7 +228,9 @@
     val NEW_STATUS_BAR_WIFI_ICON_BACKEND = unreleasedFlag(609, "new_status_bar_wifi_icon_backend")
 
     // TODO(b/256623670): Tracking Bug
-    @JvmField val BATTERY_SHIELD_ICON = unreleasedFlag(610, "battery_shield_icon")
+    @JvmField
+    val BATTERY_SHIELD_ICON =
+        resourceBooleanFlag(610, R.bool.flag_battery_shield_icon, "battery_shield_icon")
 
     // TODO(b/260881289): Tracking Bug
     val NEW_STATUS_BAR_ICONS_DEBUG_COLORING =
diff --git a/packages/SystemUI/src/com/android/systemui/media/controls/ui/MediaControlPanel.java b/packages/SystemUI/src/com/android/systemui/media/controls/ui/MediaControlPanel.java
index 827ac78..df8fb91 100644
--- a/packages/SystemUI/src/com/android/systemui/media/controls/ui/MediaControlPanel.java
+++ b/packages/SystemUI/src/com/android/systemui/media/controls/ui/MediaControlPanel.java
@@ -940,19 +940,9 @@
         if (mIsSeekBarEnabled) {
             return ConstraintSet.VISIBLE;
         }
-        // If disabled and "neighbours" are visible, set progress bar to INVISIBLE instead of GONE
-        // so layout weights still work.
-        return areAnyExpandedBottomActionsVisible() ? ConstraintSet.INVISIBLE : ConstraintSet.GONE;
-    }
-
-    private boolean areAnyExpandedBottomActionsVisible() {
-        ConstraintSet expandedSet = mMediaViewController.getExpandedLayout();
-        for (int id : MediaViewHolder.Companion.getExpandedBottomActionIds()) {
-            if (expandedSet.getVisibility(id) == ConstraintSet.VISIBLE) {
-                return true;
-            }
-        }
-        return false;
+        // Set progress bar to INVISIBLE to keep the positions of text and buttons similar to the
+        // original positions when seekbar is enabled.
+        return ConstraintSet.INVISIBLE;
     }
 
     private void setGenericButton(
diff --git a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
index 647beb9..b10abb5 100644
--- a/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/media/taptotransfer/MediaTttCommandLineHelper.kt
@@ -48,52 +48,66 @@
     /** All commands for the sender device. */
     inner class SenderCommand : Command {
         override fun execute(pw: PrintWriter, args: List<String>) {
-            val commandName = args[1]
+            if (args.size < 2) {
+                help(pw)
+                return
+            }
+
+            val senderArgs = processArgs(args)
+
             @StatusBarManager.MediaTransferSenderState
             val displayState: Int?
             try {
-                displayState = ChipStateSender.getSenderStateIdFromName(commandName)
+                displayState = ChipStateSender.getSenderStateIdFromName(senderArgs.commandName)
             } catch (ex: IllegalArgumentException) {
-                pw.println("Invalid command name $commandName")
+                pw.println("Invalid command name ${senderArgs.commandName}")
                 return
             }
 
             @SuppressLint("WrongConstant") // sysui allowed to call STATUS_BAR_SERVICE
             val statusBarManager = context.getSystemService(Context.STATUS_BAR_SERVICE)
                     as StatusBarManager
-            val routeInfo = MediaRoute2Info.Builder(if (args.size >= 4) args[3] else "id", args[0])
+            val routeInfo = MediaRoute2Info.Builder(senderArgs.id, senderArgs.deviceName)
                     .addFeature("feature")
-            val useAppIcon = !(args.size >= 3 && args[2] == "useAppIcon=false")
-            if (useAppIcon) {
+            if (senderArgs.useAppIcon) {
                 routeInfo.setClientPackageName(TEST_PACKAGE_NAME)
             }
 
+            var undoExecutor: Executor? = null
+            var undoRunnable: Runnable? = null
+            if (isSucceededState(displayState) && senderArgs.showUndo) {
+                undoExecutor = mainExecutor
+                undoRunnable = Runnable { Log.i(CLI_TAG, "Undo triggered for $displayState") }
+            }
+
             statusBarManager.updateMediaTapToTransferSenderDisplay(
                 displayState,
                 routeInfo.build(),
-                getUndoExecutor(displayState),
-                getUndoCallback(displayState)
+                undoExecutor,
+                undoRunnable,
             )
         }
 
-        private fun getUndoExecutor(
-            @StatusBarManager.MediaTransferSenderState displayState: Int
-        ): Executor? {
-            return if (isSucceededState(displayState)) {
-                mainExecutor
-            } else {
-                null
-            }
-        }
+        private fun processArgs(args: List<String>): SenderArgs {
+            val senderArgs = SenderArgs(
+                deviceName = args[0],
+                commandName = args[1],
+            )
 
-        private fun getUndoCallback(
-            @StatusBarManager.MediaTransferSenderState displayState: Int
-        ): Runnable? {
-            return if (isSucceededState(displayState)) {
-                Runnable { Log.i(CLI_TAG, "Undo triggered for $displayState") }
-            } else {
-                null
+            if (args.size == 2) {
+                return senderArgs
             }
+
+            // Process any optional arguments
+            args.subList(2, args.size).forEach {
+                when {
+                    it == "useAppIcon=false" -> senderArgs.useAppIcon = false
+                    it == "showUndo=false" -> senderArgs.showUndo = false
+                    it.substring(0, 3) == "id=" -> senderArgs.id = it.substring(3)
+                }
+            }
+
+            return senderArgs
         }
 
         private fun isSucceededState(
@@ -106,14 +120,31 @@
         }
 
         override fun help(pw: PrintWriter) {
-            pw.println("Usage: adb shell cmd statusbar $SENDER_COMMAND " +
-                    "<deviceName> <chipState> useAppIcon=[true|false] <id>")
+            pw.println(
+                "Usage: adb shell cmd statusbar $SENDER_COMMAND " +
+                "<deviceName> <chipState> " +
+                "useAppIcon=[true|false] id=<id> showUndo=[true|false]"
+            )
+            pw.println("Note: useAppIcon, id, and showUndo are optional additional commands.")
         }
     }
 
+    private data class SenderArgs(
+        val deviceName: String,
+        val commandName: String,
+        var id: String = "id",
+        var useAppIcon: Boolean = true,
+        var showUndo: Boolean = true,
+    )
+
     /** All commands for the receiver device. */
     inner class ReceiverCommand : Command {
         override fun execute(pw: PrintWriter, args: List<String>) {
+            if (args.isEmpty()) {
+                help(pw)
+                return
+            }
+
             val commandName = args[0]
             @StatusBarManager.MediaTransferReceiverState
             val displayState: Int?
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
index 3fd1aa7..e2f55f0 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/NavigationBarController.java
@@ -145,7 +145,7 @@
         boolean willApplyConfig = mConfigChanges.applyNewConfig(mContext.getResources());
         boolean largeScreenChanged = mIsTablet != isOldConfigTablet;
         // TODO(b/243765256): Disable this logging once b/243765256 is fixed.
-        Log.d(DEBUG_MISSING_GESTURE_TAG, "NavbarController: newConfig=" + newConfig
+        Log.i(DEBUG_MISSING_GESTURE_TAG, "NavbarController: newConfig=" + newConfig
                 + " mTaskbarDelegate initialized=" + mTaskbarDelegate.isInitialized()
                 + " willApplyConfigToNavbars=" + willApplyConfig
                 + " navBarCount=" + mNavigationBars.size());
diff --git a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
index cb0f3e2..d03ac3b 100644
--- a/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/navigationbar/gestural/EdgeBackGestureHandler.java
@@ -958,7 +958,7 @@
         }
 
         // TODO(b/243765256): Disable this logging once b/243765256 is fixed.
-        Log.d(DEBUG_MISSING_GESTURE_TAG, "Config changed: newConfig=" + newConfig
+        Log.i(DEBUG_MISSING_GESTURE_TAG, "Config changed: newConfig=" + newConfig
                 + " lastReportedConfig=" + mLastReportedConfig);
         mLastReportedConfig.updateFrom(newConfig);
         updateDisplaySize();
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
index b964b76..6dd60d0 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskController.kt
@@ -17,10 +17,12 @@
 package com.android.systemui.notetask
 
 import android.app.KeyguardManager
+import android.content.ComponentName
 import android.content.Context
+import android.content.pm.PackageManager
 import android.os.UserManager
-import android.view.KeyEvent
 import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.notetask.shortcut.CreateNoteTaskShortcutActivity
 import com.android.systemui.util.kotlin.getOrNull
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
@@ -45,15 +47,22 @@
     @NoteTaskEnabledKey private val isEnabled: Boolean,
 ) {
 
-    fun handleSystemKey(keyCode: Int) {
+    /**
+     * Shows a note task. How the task is shown will depend on when the method is invoked.
+     *
+     * If in multi-window mode, notes will open as a full screen experience. That is particularly
+     * important for Large screen devices. These devices may support a taskbar that let users to
+     * drag and drop a shortcut into multi-window mode, and notes should comply with this behaviour.
+     *
+     * If the keyguard is locked, notes will open as a full screen experience. A locked device has
+     * no contextual information which let us use the whole screen space available.
+     *
+     * If no in multi-window or the keyguard is unlocked, notes will open as a floating experience.
+     * That will let users open other apps in full screen, and take contextual notes.
+     */
+    fun showNoteTask(isInMultiWindowMode: Boolean = false) {
         if (!isEnabled) return
 
-        if (keyCode == KeyEvent.KEYCODE_VIDEO_APP_1) {
-            showNoteTask()
-        }
-    }
-
-    private fun showNoteTask() {
         val bubbles = optionalBubbles.getOrNull() ?: return
         val keyguardManager = optionalKeyguardManager.getOrNull() ?: return
         val userManager = optionalUserManager.getOrNull() ?: return
@@ -62,11 +71,35 @@
         // TODO(b/249954038): We should handle direct boot (isUserUnlocked). For now, we do nothing.
         if (!userManager.isUserUnlocked) return
 
-        if (keyguardManager.isKeyguardLocked) {
+        if (isInMultiWindowMode || keyguardManager.isKeyguardLocked) {
             context.startActivity(intent)
         } else {
             // TODO(b/254606432): Should include Intent.EXTRA_FLOATING_WINDOW_MODE parameter.
             bubbles.showAppBubble(intent)
         }
     }
+
+    /**
+     * Set `android:enabled` property in the `AndroidManifest` associated with the Shortcut
+     * component to [value].
+     *
+     * If the shortcut entry `android:enabled` is set to `true`, the shortcut will be visible in the
+     * Widget Picker to all users.
+     */
+    fun setNoteTaskShortcutEnabled(value: Boolean) {
+        val componentName = ComponentName(context, CreateNoteTaskShortcutActivity::class.java)
+
+        val enabledState =
+            if (value) {
+                PackageManager.COMPONENT_ENABLED_STATE_ENABLED
+            } else {
+                PackageManager.COMPONENT_ENABLED_STATE_DISABLED
+            }
+
+        context.packageManager.setComponentEnabledSetting(
+            componentName,
+            enabledState,
+            PackageManager.DONT_KILL_APP,
+        )
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
index 0a5b600..d14b7a7 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskInitializer.kt
@@ -16,9 +16,10 @@
 
 package com.android.systemui.notetask
 
+import android.view.KeyEvent
+import androidx.annotation.VisibleForTesting
 import com.android.systemui.statusbar.CommandQueue
 import com.android.wm.shell.bubbles.Bubbles
-import dagger.Lazy
 import java.util.Optional
 import javax.inject.Inject
 
@@ -27,15 +28,18 @@
 @Inject
 constructor(
     private val optionalBubbles: Optional<Bubbles>,
-    private val lazyNoteTaskController: Lazy<NoteTaskController>,
+    private val noteTaskController: NoteTaskController,
     private val commandQueue: CommandQueue,
     @NoteTaskEnabledKey private val isEnabled: Boolean,
 ) {
 
-    private val callbacks =
+    @VisibleForTesting
+    val callbacks =
         object : CommandQueue.Callbacks {
             override fun handleSystemKey(keyCode: Int) {
-                lazyNoteTaskController.get().handleSystemKey(keyCode)
+                if (keyCode == KeyEvent.KEYCODE_VIDEO_APP_1) {
+                    noteTaskController.showNoteTask()
+                }
             }
         }
 
@@ -43,5 +47,6 @@
         if (isEnabled && optionalBubbles.isPresent) {
             commandQueue.addCallback(callbacks)
         }
+        noteTaskController.setNoteTaskShortcutEnabled(isEnabled)
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
index 035396a..8bdf319 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/NoteTaskModule.kt
@@ -16,32 +16,47 @@
 
 package com.android.systemui.notetask
 
+import android.app.Activity
 import android.app.KeyguardManager
 import android.content.Context
 import android.os.UserManager
 import androidx.core.content.getSystemService
 import com.android.systemui.flags.FeatureFlags
 import com.android.systemui.flags.Flags
+import com.android.systemui.notetask.shortcut.CreateNoteTaskShortcutActivity
+import com.android.systemui.notetask.shortcut.LaunchNoteTaskActivity
+import dagger.Binds
 import dagger.Module
 import dagger.Provides
-import java.util.*
+import dagger.multibindings.ClassKey
+import dagger.multibindings.IntoMap
+import java.util.Optional
 
 /** Compose all dependencies required by Note Task feature. */
 @Module
-internal class NoteTaskModule {
+internal interface NoteTaskModule {
 
-    @[Provides NoteTaskEnabledKey]
-    fun provideIsNoteTaskEnabled(featureFlags: FeatureFlags): Boolean {
-        return featureFlags.isEnabled(Flags.NOTE_TASKS)
-    }
+    @[Binds IntoMap ClassKey(LaunchNoteTaskActivity::class)]
+    fun bindNoteTaskLauncherActivity(activity: LaunchNoteTaskActivity): Activity?
 
-    @Provides
-    fun provideOptionalKeyguardManager(context: Context): Optional<KeyguardManager> {
-        return Optional.ofNullable(context.getSystemService())
-    }
+    @[Binds IntoMap ClassKey(CreateNoteTaskShortcutActivity::class)]
+    fun bindNoteTaskShortcutActivity(activity: CreateNoteTaskShortcutActivity): Activity?
 
-    @Provides
-    fun provideOptionalUserManager(context: Context): Optional<UserManager> {
-        return Optional.ofNullable(context.getSystemService())
+    companion object {
+
+        @[Provides NoteTaskEnabledKey]
+        fun provideIsNoteTaskEnabled(featureFlags: FeatureFlags): Boolean {
+            return featureFlags.isEnabled(Flags.NOTE_TASKS)
+        }
+
+        @Provides
+        fun provideOptionalKeyguardManager(context: Context): Optional<KeyguardManager> {
+            return Optional.ofNullable(context.getSystemService())
+        }
+
+        @Provides
+        fun provideOptionalUserManager(context: Context): Optional<UserManager> {
+            return Optional.ofNullable(context.getSystemService())
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/shortcut/CreateNoteTaskShortcutActivity.kt b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/CreateNoteTaskShortcutActivity.kt
new file mode 100644
index 0000000..f6a623e
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/CreateNoteTaskShortcutActivity.kt
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2022 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.notetask.shortcut
+
+import android.app.Activity
+import android.content.Intent
+import android.os.Bundle
+import androidx.activity.ComponentActivity
+import androidx.annotation.DrawableRes
+import androidx.core.content.pm.ShortcutInfoCompat
+import androidx.core.content.pm.ShortcutManagerCompat
+import androidx.core.graphics.drawable.IconCompat
+import com.android.systemui.R
+import javax.inject.Inject
+
+/**
+ * Activity responsible for create a shortcut for notes action. If the shortcut is enabled, a new
+ * shortcut will appear in the widget picker. If the shortcut is selected, the Activity here will be
+ * launched, creating a new shortcut for [CreateNoteTaskShortcutActivity], and will finish.
+ *
+ * @see <a
+ * href="https://developer.android.com/develop/ui/views/launch/shortcuts/creating-shortcuts#custom-pinned">Creating
+ * a custom shortcut activity</a>
+ */
+internal class CreateNoteTaskShortcutActivity @Inject constructor() : ComponentActivity() {
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+
+        val intent =
+            createShortcutIntent(
+                id = SHORTCUT_ID,
+                shortLabel = getString(R.string.note_task_button_label),
+                intent = LaunchNoteTaskActivity.newIntent(context = this),
+                iconResource = R.drawable.ic_note_task_button,
+            )
+        setResult(Activity.RESULT_OK, intent)
+
+        finish()
+    }
+
+    private fun createShortcutIntent(
+        id: String,
+        shortLabel: String,
+        intent: Intent,
+        @DrawableRes iconResource: Int,
+    ): Intent {
+        val shortcutInfo =
+            ShortcutInfoCompat.Builder(this, id)
+                .setIntent(intent)
+                .setShortLabel(shortLabel)
+                .setLongLived(true)
+                .setIcon(IconCompat.createWithResource(this, iconResource))
+                .build()
+
+        return ShortcutManagerCompat.createShortcutResultIntent(
+            this,
+            shortcutInfo,
+        )
+    }
+
+    private companion object {
+        private const val SHORTCUT_ID = "note-task-shortcut-id"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt
new file mode 100644
index 0000000..47fe676
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2022 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.notetask.shortcut
+
+import android.content.Context
+import android.content.Intent
+import android.os.Bundle
+import androidx.activity.ComponentActivity
+import com.android.systemui.notetask.NoteTaskController
+import com.android.systemui.notetask.NoteTaskIntentResolver
+import javax.inject.Inject
+
+/** Activity responsible for launching the note experience, and finish. */
+internal class LaunchNoteTaskActivity
+@Inject
+constructor(
+    private val noteTaskController: NoteTaskController,
+) : ComponentActivity() {
+
+    override fun onCreate(savedInstanceState: Bundle?) {
+        super.onCreate(savedInstanceState)
+
+        noteTaskController.showNoteTask(isInMultiWindowMode)
+
+        finish()
+    }
+
+    companion object {
+
+        /** Creates a new [Intent] set to start [LaunchNoteTaskActivity]. */
+        fun newIntent(context: Context): Intent {
+            return Intent(context, LaunchNoteTaskActivity::class.java).apply {
+                // Intent's action must be set in shortcuts, or an exception will be thrown.
+                // TODO(b/254606432): Use Intent.ACTION_NOTES instead.
+                action = NoteTaskIntentResolver.NOTES_ACTION
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
index 6240c10..cad296b 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/QSTileHost.java
@@ -608,7 +608,7 @@
 
         if (TextUtils.isEmpty(tileList)) {
             tileList = res.getString(R.string.quick_settings_tiles);
-            if (DEBUG) Log.d(TAG, "Loaded tile specs from config: " + tileList);
+            if (DEBUG) Log.d(TAG, "Loaded tile specs from default config: " + tileList);
         } else {
             if (DEBUG) Log.d(TAG, "Loaded tile specs from setting: " + tileList);
         }
diff --git a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
index 7130294..a6c7781 100644
--- a/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
+++ b/packages/SystemUI/src/com/android/systemui/qs/tiles/WorkModeTile.java
@@ -27,6 +27,7 @@
 import android.view.View;
 import android.widget.Switch;
 
+import androidx.annotation.MainThread;
 import androidx.annotation.Nullable;
 
 import com.android.internal.logging.MetricsLogger;
@@ -91,11 +92,13 @@
     }
 
     @Override
+    @MainThread
     public void onManagedProfileChanged() {
         refreshState(mProfileController.isWorkModeEnabled());
     }
 
     @Override
+    @MainThread
     public void onManagedProfileRemoved() {
         mHost.removeTile(getTileSpec());
         mHost.unmarkTileAsAutoAdded(getTileSpec());
diff --git a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
index 547b496..00d129a 100644
--- a/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
+++ b/packages/SystemUI/src/com/android/systemui/recents/OverviewProxyService.java
@@ -565,13 +565,25 @@
         statusBarWinController.registerCallback(mStatusBarWindowCallback);
         mScreenshotHelper = new ScreenshotHelper(context);
 
-        // Listen for tracing state changes
         commandQueue.addCallback(new CommandQueue.Callbacks() {
+
+            // Listen for tracing state changes
             @Override
             public void onTracingStateChanged(boolean enabled) {
                 mSysUiState.setFlag(SYSUI_STATE_TRACING_ENABLED, enabled)
                         .commitUpdate(mContext.getDisplayId());
             }
+
+            @Override
+            public void enterStageSplitFromRunningApp(boolean leftOrTop) {
+                if (mOverviewProxy != null) {
+                    try {
+                        mOverviewProxy.enterStageSplitFromRunningApp(leftOrTop);
+                    } catch (RemoteException e) {
+                        Log.w(TAG_OPS, "Unable to enter stage split from the current running app");
+                    }
+                }
+            }
         });
         mCommandQueue = commandQueue;
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
index d000e6e..750d004 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/CommandQueue.java
@@ -167,6 +167,7 @@
     private static final int MSG_TILE_SERVICE_REQUEST_LISTENING_STATE = 68 << MSG_SHIFT;
     private static final int MSG_SHOW_REAR_DISPLAY_DIALOG = 69 << MSG_SHIFT;
     private static final int MSG_GO_TO_FULLSCREEN_FROM_SPLIT = 70 << MSG_SHIFT;
+    private static final int MSG_ENTER_STAGE_SPLIT_FROM_RUNNING_APP = 71 << MSG_SHIFT;
 
     public static final int FLAG_EXCLUDE_NONE = 0;
     public static final int FLAG_EXCLUDE_SEARCH_PANEL = 1 << 0;
@@ -486,6 +487,11 @@
          * @see IStatusBar#goToFullscreenFromSplit
          */
         default void goToFullscreenFromSplit() {}
+
+        /**
+         * @see IStatusBar#enterStageSplitFromRunningApp
+         */
+        default void enterStageSplitFromRunningApp(boolean leftOrTop) {}
     }
 
     public CommandQueue(Context context) {
@@ -1248,6 +1254,14 @@
     }
 
     @Override
+    public void enterStageSplitFromRunningApp(boolean leftOrTop) {
+        synchronized (mLock) {
+            mHandler.obtainMessage(MSG_ENTER_STAGE_SPLIT_FROM_RUNNING_APP,
+                    leftOrTop).sendToTarget();
+        }
+    }
+
+    @Override
     public void requestAddTile(
             @NonNull ComponentName componentName,
             @NonNull CharSequence appName,
@@ -1758,6 +1772,11 @@
                         mCallbacks.get(i).goToFullscreenFromSplit();
                     }
                     break;
+                case MSG_ENTER_STAGE_SPLIT_FROM_RUNNING_APP:
+                    for (int i = 0; i < mCallbacks.size(); i++) {
+                        mCallbacks.get(i).enterStageSplitFromRunningApp((Boolean) msg.obj);
+                    }
+                    break;
             }
         }
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
index d7eddf5..56c34a0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
@@ -39,6 +39,7 @@
 import com.android.systemui.animation.Interpolators;
 import com.android.systemui.animation.ShadeInterpolation;
 import com.android.systemui.plugins.statusbar.StatusBarStateController.StateListener;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.NotificationUtils;
 import com.android.systemui.statusbar.notification.SourceType;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationView;
@@ -66,6 +67,8 @@
     // the next icon has translated out of the way, to avoid overlapping.
     private static final Interpolator ICON_ALPHA_INTERPOLATOR =
             new PathInterpolator(0.6f, 0f, 0.6f, 0f);
+    private static final SourceType BASE_VALUE = SourceType.from("BaseValue");
+    private static final SourceType SHELF_SCROLL = SourceType.from("ShelfScroll");
 
     private NotificationIconContainer mShelfIcons;
     private int[] mTmp = new int[2];
@@ -112,19 +115,24 @@
         setClipChildren(false);
         setClipToPadding(false);
         mShelfIcons.setIsStaticLayout(false);
-        requestBottomRoundness(1.0f, /* animate = */ false, SourceType.DefaultValue);
-        requestTopRoundness(1f, false, SourceType.DefaultValue);
+        requestRoundness(/* top = */ 1f, /* bottom = */ 1f, BASE_VALUE, /* animate = */ false);
 
-        // Setting this to first in section to get the clipping to the top roundness correct. This
-        // value determines the way we are clipping to the top roundness of the overall shade
-        setFirstInSection(true);
+        if (!mUseRoundnessSourceTypes) {
+            // Setting this to first in section to get the clipping to the top roundness correct.
+            // This value determines the way we are clipping to the top roundness of the overall
+            // shade
+            setFirstInSection(true);
+        }
         updateResources();
     }
 
     public void bind(AmbientState ambientState,
-            NotificationStackScrollLayoutController hostLayoutController) {
+                     NotificationStackScrollLayoutController hostLayoutController) {
         mAmbientState = ambientState;
         mHostLayoutController = hostLayoutController;
+        hostLayoutController.setOnNotificationRemovedListener((child, isTransferInProgress) -> {
+            child.requestRoundnessReset(SHELF_SCROLL);
+        });
     }
 
     private void updateResources() {
@@ -185,9 +193,11 @@
                 + " indexOfFirstViewInShelf=" + mIndexOfFirstViewInShelf + ')';
     }
 
-    /** Update the state of the shelf. */
+    /**
+     * Update the state of the shelf.
+     */
     public void updateState(StackScrollAlgorithm.StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                            AmbientState ambientState) {
         ExpandableView lastView = ambientState.getLastVisibleBackgroundChild();
         ShelfState viewState = (ShelfState) getViewState();
         if (mShowNotificationShelf && lastView != null) {
@@ -246,7 +256,7 @@
 
     /**
      * @param fractionToShade Fraction of lockscreen to shade transition
-     * @param shortestWidth Shortest width to use for lockscreen shelf
+     * @param shortestWidth   Shortest width to use for lockscreen shelf
      */
     @VisibleForTesting
     public void updateActualWidth(float fractionToShade, float shortestWidth) {
@@ -281,9 +291,9 @@
 
     /**
      * @param localX Click x from left of screen
-     * @param slop Margin of error within which we count x for valid click
-     * @param left Left of shelf, from left of screen
-     * @param right Right of shelf, from left of screen
+     * @param slop   Margin of error within which we count x for valid click
+     * @param left   Left of shelf, from left of screen
+     * @param right  Right of shelf, from left of screen
      * @return Whether click x was in view
      */
     @VisibleForTesting
@@ -293,8 +303,8 @@
 
     /**
      * @param localY Click y from top of shelf
-     * @param slop Margin of error within which we count y for valid click
-     * @param top Top of shelf
+     * @param slop   Margin of error within which we count y for valid click
+     * @param top    Top of shelf
      * @param bottom Height of shelf
      * @return Whether click y was in view
      */
@@ -306,7 +316,7 @@
     /**
      * @param localX Click x
      * @param localY Click y
-     * @param slop Margin of error for valid click
+     * @param slop   Margin of error for valid click
      * @return Whether this click was on the visible (non-clipped) part of the shelf
      */
     @Override
@@ -478,13 +488,15 @@
         }
     }
 
-    private void updateCornerRoundnessOnScroll(ActivatableNotificationView anv, float viewStart,
+    private void updateCornerRoundnessOnScroll(
+            ActivatableNotificationView anv,
+            float viewStart,
             float shelfStart) {
 
         final boolean isUnlockedHeadsUp = !mAmbientState.isOnKeyguard()
                 && !mAmbientState.isShadeExpanded()
                 && anv instanceof ExpandableNotificationRow
-                && ((ExpandableNotificationRow) anv).isHeadsUp();
+                && anv.isHeadsUp();
 
         final boolean isHunGoingToShade = mAmbientState.isShadeExpanded()
                 && anv == mAmbientState.getTrackedHeadsUpRow();
@@ -506,41 +518,40 @@
                 * mAmbientState.getExpansionFraction();
         final float cornerAnimationTop = shelfStart - cornerAnimationDistance;
 
-        if (viewEnd >= cornerAnimationTop) {
-            // Round bottom corners within animation bounds
-            final float changeFraction = MathUtils.saturate(
-                    (viewEnd - cornerAnimationTop) / cornerAnimationDistance);
-            anv.requestBottomRoundness(
-                    /* value = */ anv.isLastInSection() ? 1f : changeFraction,
-                    /* animate = */ false,
-                    SourceType.OnScroll);
-
-        } else if (viewEnd < cornerAnimationTop) {
-            // Fast scroll skips frames and leaves corners with unfinished rounding.
-            // Reset top and bottom corners outside of animation bounds.
-            anv.requestBottomRoundness(
-                    /* value = */ anv.isLastInSection() ? 1f : 0f,
-                    /* animate = */ false,
-                    SourceType.OnScroll);
+        final SourceType sourceType;
+        if (mUseRoundnessSourceTypes) {
+            sourceType = SHELF_SCROLL;
+        } else {
+            sourceType = LegacySourceType.OnScroll;
         }
 
-        if (viewStart >= cornerAnimationTop) {
+        final float topValue;
+        if (!mUseRoundnessSourceTypes && anv.isFirstInSection()) {
+            topValue = 1f;
+        } else if (viewStart >= cornerAnimationTop) {
             // Round top corners within animation bounds
-            final float changeFraction = MathUtils.saturate(
+            topValue = MathUtils.saturate(
                     (viewStart - cornerAnimationTop) / cornerAnimationDistance);
-            anv.requestTopRoundness(
-                    /* value = */ anv.isFirstInSection() ? 1f : changeFraction,
-                    /* animate = */ false,
-                    SourceType.OnScroll);
-
-        } else if (viewStart < cornerAnimationTop) {
+        } else {
             // Fast scroll skips frames and leaves corners with unfinished rounding.
             // Reset top and bottom corners outside of animation bounds.
-            anv.requestTopRoundness(
-                    /* value = */ anv.isFirstInSection() ? 1f : 0f,
-                    /* animate = */ false,
-                    SourceType.OnScroll);
+            topValue = 0f;
         }
+        anv.requestTopRoundness(topValue, sourceType, /* animate = */ false);
+
+        final float bottomValue;
+        if (!mUseRoundnessSourceTypes && anv.isLastInSection()) {
+            bottomValue = 1f;
+        } else if (viewEnd >= cornerAnimationTop) {
+            // Round bottom corners within animation bounds
+            bottomValue = MathUtils.saturate(
+                    (viewEnd - cornerAnimationTop) / cornerAnimationDistance);
+        } else {
+            // Fast scroll skips frames and leaves corners with unfinished rounding.
+            // Reset top and bottom corners outside of animation bounds.
+            bottomValue = 0f;
+        }
+        anv.requestBottomRoundness(bottomValue, sourceType, /* animate = */ false);
     }
 
     /**
@@ -626,10 +637,11 @@
 
     /**
      * Update the clipping of this view.
+     *
      * @return the amount that our own top should be clipped
      */
     private int updateNotificationClipHeight(ExpandableView view,
-            float notificationClipEnd, int childIndex) {
+                                             float notificationClipEnd, int childIndex) {
         float viewEnd = view.getTranslationY() + view.getActualHeight();
         boolean isPinned = (view.isPinned() || view.isHeadsUpAnimatingAway())
                 && !mAmbientState.isDozingAndNotPulsing(view);
@@ -657,7 +669,7 @@
 
     @Override
     public void setFakeShadowIntensity(float shadowIntensity, float outlineAlpha, int shadowYEnd,
-            int outlineTranslation) {
+                                       int outlineTranslation) {
         if (!mHasItemsInStableShelf) {
             shadowIntensity = 0.0f;
         }
@@ -665,18 +677,24 @@
     }
 
     /**
-     * @param i Index of the view in the host layout.
-     * @param view The current ExpandableView.
-     * @param scrollingFast Whether we are scrolling fast.
+     * @param i                 Index of the view in the host layout.
+     * @param view              The current ExpandableView.
+     * @param scrollingFast     Whether we are scrolling fast.
      * @param expandingAnimated Whether we are expanding a notification.
-     * @param isLastChild Whether this is the last view.
-     * @param shelfClipStart The point at which notifications start getting clipped by the shelf.
+     * @param isLastChild       Whether this is the last view.
+     * @param shelfClipStart    The point at which notifications start getting clipped by the shelf.
      * @return The amount how much this notification is in the shelf.
-     *         0f is not in shelf. 1f is completely in shelf.
+     * 0f is not in shelf. 1f is completely in shelf.
      */
     @VisibleForTesting
-    public float getAmountInShelf(int i, ExpandableView view, boolean scrollingFast,
-            boolean expandingAnimated, boolean isLastChild, float shelfClipStart) {
+    public float getAmountInShelf(
+            int i,
+            ExpandableView view,
+            boolean scrollingFast,
+            boolean expandingAnimated,
+            boolean isLastChild,
+            float shelfClipStart
+    ) {
 
         // Let's calculate how much the view is in the shelf
         float viewStart = view.getTranslationY();
@@ -755,8 +773,13 @@
         return start;
     }
 
-    private void updateIconPositioning(ExpandableView view, float iconTransitionAmount,
-            boolean scrollingFast, boolean expandingAnimated, boolean isLastChild) {
+    private void updateIconPositioning(
+            ExpandableView view,
+            float iconTransitionAmount,
+            boolean scrollingFast,
+            boolean expandingAnimated,
+            boolean isLastChild
+    ) {
         StatusBarIconView icon = view.getShelfIcon();
         NotificationIconContainer.IconState iconState = getIconState(icon);
         if (iconState == null) {
@@ -817,7 +840,7 @@
                 || row.showingPulsing()
                 || row.getTranslationZ() > mAmbientState.getBaseZHeight();
 
-        iconState.iconAppearAmount = iconState.hidden? 0f : transitionAmount;
+        iconState.iconAppearAmount = iconState.hidden ? 0f : transitionAmount;
 
         // Fade in icons at shelf start
         // This is important for conversation icons, which are badged and need x reset
@@ -847,7 +870,7 @@
     }
 
     private float getFullyClosedTranslation() {
-        return - (getIntrinsicHeight() - mStatusBarHeight) / 2;
+        return -(getIntrinsicHeight() - mStatusBarHeight) / 2;
     }
 
     @Override
@@ -904,7 +927,7 @@
 
     /**
      * @return whether the shelf has any icons in it when a potential animation has finished, i.e
-     *         if the current state would be applied right now
+     * if the current state would be applied right now
      */
     public boolean hasItemsInStableShelf() {
         return mHasItemsInStableShelf;
@@ -962,7 +985,7 @@
 
     @Override
     public void onLayoutChange(View v, int left, int top, int right, int bottom, int oldLeft,
-            int oldTop, int oldRight, int oldBottom) {
+                               int oldTop, int oldRight, int oldBottom) {
         updateRelativeOffset();
     }
 
@@ -981,12 +1004,11 @@
 
     /**
      * This method resets the OnScroll roundness of a view to 0f
-     *
+     * <p>
      * Note: This should be the only class that handles roundness {@code SourceType.OnScroll}
      */
-    public static void resetOnScrollRoundness(ExpandableView expandableView) {
-        expandableView.requestTopRoundness(0f, false, SourceType.OnScroll);
-        expandableView.requestBottomRoundness(0f, false, SourceType.OnScroll);
+    public static void resetLegacyOnScrollRoundness(ExpandableView expandableView) {
+        expandableView.requestRoundnessReset(LegacySourceType.OnScroll);
     }
 
     public class ShelfState extends ExpandableViewState {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.java
index 3b1fa17..bb84c75 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelfController.java
@@ -18,6 +18,8 @@
 
 import android.view.View;
 
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationView;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationViewController;
 import com.android.systemui.statusbar.notification.row.dagger.NotificationRowScope;
@@ -42,14 +44,17 @@
     private AmbientState mAmbientState;
 
     @Inject
-    public NotificationShelfController(NotificationShelf notificationShelf,
+    public NotificationShelfController(
+            NotificationShelf notificationShelf,
             ActivatableNotificationViewController activatableNotificationViewController,
             KeyguardBypassController keyguardBypassController,
-            SysuiStatusBarStateController statusBarStateController) {
+            SysuiStatusBarStateController statusBarStateController,
+            FeatureFlags featureFlags) {
         mView = notificationShelf;
         mActivatableNotificationViewController = activatableNotificationViewController;
         mKeyguardBypassController = keyguardBypassController;
         mStatusBarStateController = statusBarStateController;
+        mView.useRoundnessSourceTypes(featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES));
         mOnAttachStateChangeListener = new View.OnAttachStateChangeListener() {
             @Override
             public void onViewAttachedToWindow(View v) {
@@ -88,7 +93,7 @@
 
     public @View.Visibility int getVisibility() {
         return mView.getVisibility();
-    };
+    }
 
     public void setCollapsedIcons(NotificationIconContainer notificationIcons) {
         mView.setCollapsedIcons(notificationIcons);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/connectivity/NetworkControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/connectivity/NetworkControllerImpl.java
index 97a47b5..362764d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/connectivity/NetworkControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/connectivity/NetworkControllerImpl.java
@@ -78,6 +78,7 @@
 import com.android.systemui.plugins.log.LogLevel;
 import com.android.systemui.qs.tiles.dialog.InternetDialogFactory;
 import com.android.systemui.settings.UserTracker;
+import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags;
 import com.android.systemui.statusbar.policy.ConfigurationController;
 import com.android.systemui.statusbar.policy.DataSaverController;
 import com.android.systemui.statusbar.policy.DataSaverControllerImpl;
@@ -193,6 +194,7 @@
     private final Executor mBgExecutor;
     // Handler that all callbacks are made on.
     private final CallbackHandler mCallbackHandler;
+    private final StatusBarPipelineFlags mStatusBarPipelineFlags;
 
     private int mEmergencySource;
     private boolean mIsEmergency;
@@ -243,6 +245,7 @@
             TelephonyListenerManager telephonyListenerManager,
             @Nullable WifiManager wifiManager,
             AccessPointControllerImpl accessPointController,
+            StatusBarPipelineFlags statusBarPipelineFlags,
             DemoModeController demoModeController,
             CarrierConfigTracker carrierConfigTracker,
             WifiStatusTrackerFactory trackerFactory,
@@ -261,6 +264,7 @@
                 bgExecutor,
                 callbackHandler,
                 accessPointController,
+                statusBarPipelineFlags,
                 new DataUsageController(context),
                 new SubscriptionDefaults(),
                 deviceProvisionedController,
@@ -288,6 +292,7 @@
             Executor bgExecutor,
             CallbackHandler callbackHandler,
             AccessPointControllerImpl accessPointController,
+            StatusBarPipelineFlags statusBarPipelineFlags,
             DataUsageController dataUsageController,
             SubscriptionDefaults defaultsHandler,
             DeviceProvisionedController deviceProvisionedController,
@@ -309,6 +314,7 @@
         mBgLooper = bgLooper;
         mBgExecutor = bgExecutor;
         mCallbackHandler = callbackHandler;
+        mStatusBarPipelineFlags = statusBarPipelineFlags;
         mDataSaverController = new DataSaverControllerImpl(context);
         mBroadcastDispatcher = broadcastDispatcher;
         mMobileFactory = mobileFactory;
@@ -1334,7 +1340,7 @@
             mWifiSignalController.notifyListeners();
         }
         String sims = args.getString("sims");
-        if (sims != null) {
+        if (sims != null && !mStatusBarPipelineFlags.useNewMobileIcons()) {
             int num = MathUtils.constrain(Integer.parseInt(sims), 1, 8);
             List<SubscriptionInfo> subs = new ArrayList<>();
             if (num != mMobileSignalControllers.size()) {
@@ -1357,7 +1363,7 @@
             mCallbackHandler.setNoSims(mHasNoSubs, mSimDetected);
         }
         String mobile = args.getString("mobile");
-        if (mobile != null) {
+        if (mobile != null && !mStatusBarPipelineFlags.useNewMobileIcons()) {
             boolean show = mobile.equals("show");
             String datatype = args.getString("datatype");
             String slotString = args.getString("slot");
@@ -1442,7 +1448,7 @@
             controller.notifyListeners();
         }
         String carrierNetworkChange = args.getString("carriernetworkchange");
-        if (carrierNetworkChange != null) {
+        if (carrierNetworkChange != null && !mStatusBarPipelineFlags.useNewMobileIcons()) {
             boolean show = carrierNetworkChange.equals("show");
             for (int i = 0; i < mMobileSignalControllers.size(); i++) {
                 MobileSignalController controller = mMobileSignalControllers.valueAt(i);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
index ed7f648..0eb0000 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/Roundable.kt
@@ -74,8 +74,8 @@
     @JvmDefault
     fun requestTopRoundness(
         @FloatRange(from = 0.0, to = 1.0) value: Float,
-        animate: Boolean,
         sourceType: SourceType,
+        animate: Boolean,
     ): Boolean {
         val roundnessMap = roundableState.topRoundnessMap
         val lastValue = roundnessMap.values.maxOrNull() ?: 0f
@@ -105,6 +105,30 @@
     }
 
     /**
+     * Request the top roundness [value] for a specific [sourceType]. Animate the roundness if the
+     * view is shown.
+     *
+     * The top roundness of a [Roundable] can be defined by different [sourceType]. In case more
+     * origins require different roundness, for the same property, the maximum value will always be
+     * chosen.
+     *
+     * @param value a value between 0f and 1f.
+     * @param sourceType the source from which the request for roundness comes.
+     * @return Whether the roundness was changed.
+     */
+    @JvmDefault
+    fun requestTopRoundness(
+        @FloatRange(from = 0.0, to = 1.0) value: Float,
+        sourceType: SourceType,
+    ): Boolean {
+        return requestTopRoundness(
+            value = value,
+            sourceType = sourceType,
+            animate = roundableState.targetView.isShown
+        )
+    }
+
+    /**
      * Request the bottom roundness [value] for a specific [sourceType].
      *
      * The bottom roundness of a [Roundable] can be defined by different [sourceType]. In case more
@@ -119,8 +143,8 @@
     @JvmDefault
     fun requestBottomRoundness(
         @FloatRange(from = 0.0, to = 1.0) value: Float,
-        animate: Boolean,
         sourceType: SourceType,
+        animate: Boolean,
     ): Boolean {
         val roundnessMap = roundableState.bottomRoundnessMap
         val lastValue = roundnessMap.values.maxOrNull() ?: 0f
@@ -149,9 +173,101 @@
         return false
     }
 
+    /**
+     * Request the bottom roundness [value] for a specific [sourceType]. Animate the roundness if
+     * the view is shown.
+     *
+     * The bottom roundness of a [Roundable] can be defined by different [sourceType]. In case more
+     * origins require different roundness, for the same property, the maximum value will always be
+     * chosen.
+     *
+     * @param value value between 0f and 1f.
+     * @param sourceType the source from which the request for roundness comes.
+     * @return Whether the roundness was changed.
+     */
+    @JvmDefault
+    fun requestBottomRoundness(
+        @FloatRange(from = 0.0, to = 1.0) value: Float,
+        sourceType: SourceType,
+    ): Boolean {
+        return requestBottomRoundness(
+            value = value,
+            sourceType = sourceType,
+            animate = roundableState.targetView.isShown
+        )
+    }
+
+    /**
+     * Request the roundness [value] for a specific [sourceType].
+     *
+     * The top/bottom roundness of a [Roundable] can be defined by different [sourceType]. In case
+     * more origins require different roundness, for the same property, the maximum value will
+     * always be chosen.
+     *
+     * @param top top value between 0f and 1f.
+     * @param bottom bottom value between 0f and 1f.
+     * @param sourceType the source from which the request for roundness comes.
+     * @param animate true if it should animate to that value.
+     * @return Whether the roundness was changed.
+     */
+    @JvmDefault
+    fun requestRoundness(
+        @FloatRange(from = 0.0, to = 1.0) top: Float,
+        @FloatRange(from = 0.0, to = 1.0) bottom: Float,
+        sourceType: SourceType,
+        animate: Boolean,
+    ): Boolean {
+        val hasTopChanged =
+            requestTopRoundness(value = top, sourceType = sourceType, animate = animate)
+        val hasBottomChanged =
+            requestBottomRoundness(value = bottom, sourceType = sourceType, animate = animate)
+        return hasTopChanged || hasBottomChanged
+    }
+
+    /**
+     * Request the roundness [value] for a specific [sourceType]. Animate the roundness if the view
+     * is shown.
+     *
+     * The top/bottom roundness of a [Roundable] can be defined by different [sourceType]. In case
+     * more origins require different roundness, for the same property, the maximum value will
+     * always be chosen.
+     *
+     * @param top top value between 0f and 1f.
+     * @param bottom bottom value between 0f and 1f.
+     * @param sourceType the source from which the request for roundness comes.
+     * @return Whether the roundness was changed.
+     */
+    @JvmDefault
+    fun requestRoundness(
+        @FloatRange(from = 0.0, to = 1.0) top: Float,
+        @FloatRange(from = 0.0, to = 1.0) bottom: Float,
+        sourceType: SourceType,
+    ): Boolean {
+        return requestRoundness(
+            top = top,
+            bottom = bottom,
+            sourceType = sourceType,
+            animate = roundableState.targetView.isShown,
+        )
+    }
+
+    /**
+     * Request the roundness 0f for a [SourceType]. Animate the roundness if the view is shown.
+     *
+     * The top/bottom roundness of a [Roundable] can be defined by different [sourceType]. In case
+     * more origins require different roundness, for the same property, the maximum value will
+     * always be chosen.
+     *
+     * @param sourceType the source from which the request for roundness comes.
+     */
+    @JvmDefault
+    fun requestRoundnessReset(sourceType: SourceType) {
+        requestRoundness(top = 0f, bottom = 0f, sourceType = sourceType)
+    }
+
     /** Apply the roundness changes, usually means invalidate the [RoundableState.targetView]. */
     @JvmDefault
-    fun applyRoundness() {
+    fun applyRoundnessAndInvalidate() {
         roundableState.targetView.invalidate()
     }
 
@@ -227,7 +343,7 @@
     /** Set the current top roundness */
     internal fun setTopRoundness(
         value: Float,
-        animated: Boolean = targetView.isShown,
+        animated: Boolean,
     ) {
         PropertyAnimator.setProperty(targetView, topAnimatable, value, DURATION, animated)
     }
@@ -235,11 +351,19 @@
     /** Set the current bottom roundness */
     internal fun setBottomRoundness(
         value: Float,
-        animated: Boolean = targetView.isShown,
+        animated: Boolean,
     ) {
         PropertyAnimator.setProperty(targetView, bottomAnimatable, value, DURATION, animated)
     }
 
+    fun debugString() = buildString {
+        append("TargetView: ${targetView.hashCode()} ")
+        append("Top: $topRoundness ")
+        append(topRoundnessMap.map { "${it.key} ${it.value}" })
+        append(" Bottom: $bottomRoundness ")
+        append(bottomRoundnessMap.map { "${it.key} ${it.value}" })
+    }
+
     companion object {
         private val DURATION: AnimationProperties =
             AnimationProperties()
@@ -252,7 +376,7 @@
 
                     override fun setValue(view: View, value: Float) {
                         roundable.roundableState.topRoundness = value
-                        roundable.applyRoundness()
+                        roundable.applyRoundnessAndInvalidate()
                     }
                 },
                 R.id.top_roundess_animator_tag,
@@ -267,7 +391,7 @@
 
                     override fun setValue(view: View, value: Float) {
                         roundable.roundableState.bottomRoundness = value
-                        roundable.applyRoundness()
+                        roundable.applyRoundnessAndInvalidate()
                     }
                 },
                 R.id.bottom_roundess_animator_tag,
@@ -277,7 +401,31 @@
     }
 }
 
-enum class SourceType {
+/**
+ * Interface used to define the owner of a roundness. Usually the [SourceType] is defined as a
+ * private property of a class.
+ */
+interface SourceType {
+    companion object {
+        /**
+         * This is the most convenient way to define a new [SourceType].
+         *
+         * For example:
+         *
+         * ```kotlin
+         *     private val SECTION = SourceType.from("Section")
+         * ```
+         */
+        @JvmStatic
+        fun from(name: String) =
+            object : SourceType {
+                override fun toString() = name
+            }
+    }
+}
+
+@Deprecated("Use SourceType.from() instead", ReplaceWith("SourceType.from()"))
+enum class LegacySourceType : SourceType {
     DefaultValue,
     OnDismissAnimation,
     OnScroll,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
index d29298a..fbe88df 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ActivatableNotificationView.java
@@ -39,9 +39,13 @@
 import com.android.systemui.statusbar.NotificationShelf;
 import com.android.systemui.statusbar.notification.FakeShadowView;
 import com.android.systemui.statusbar.notification.NotificationUtils;
+import com.android.systemui.statusbar.notification.SourceType;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayout;
 import com.android.systemui.statusbar.notification.stack.StackStateAnimator;
 
+import java.util.HashSet;
+import java.util.Set;
+
 /**
  * Base class for both {@link ExpandableNotificationRow} and {@link NotificationShelf}
  * to implement dimming/activating on Keyguard for the double-tap gesture
@@ -91,6 +95,7 @@
             = new PathInterpolator(0.6f, 0, 0.5f, 1);
     private static final Interpolator ACTIVATE_INVERSE_ALPHA_INTERPOLATOR
             = new PathInterpolator(0, 0, 0.5f, 1);
+    private final Set<SourceType> mOnDetachResetRoundness = new HashSet<>();
     private int mTintedRippleColor;
     private int mNormalRippleColor;
     private Gefingerpoken mTouchHandler;
@@ -134,6 +139,7 @@
     private boolean mDismissed;
     private boolean mRefocusOnDismiss;
     private AccessibilityManager mAccessibilityManager;
+    protected boolean mUseRoundnessSourceTypes;
 
     public ActivatableNotificationView(Context context, AttributeSet attrs) {
         super(context, attrs);
@@ -613,9 +619,9 @@
     protected void resetAllContentAlphas() {}
 
     @Override
-    public void applyRoundness() {
-        super.applyRoundness();
+    public void applyRoundnessAndInvalidate() {
         applyBackgroundRoundness(getTopCornerRadius(), getBottomCornerRadius());
+        super.applyRoundnessAndInvalidate();
     }
 
     @Override
@@ -775,6 +781,33 @@
         mAccessibilityManager = accessibilityManager;
     }
 
+    /**
+     * Enable the support for rounded corner based on the SourceType
+     * @param enabled true if is supported
+     */
+    public void useRoundnessSourceTypes(boolean enabled) {
+        mUseRoundnessSourceTypes = enabled;
+    }
+
+    @Override
+    protected void onDetachedFromWindow() {
+        super.onDetachedFromWindow();
+        if (mUseRoundnessSourceTypes && !mOnDetachResetRoundness.isEmpty()) {
+            for (SourceType sourceType : mOnDetachResetRoundness) {
+                requestRoundnessReset(sourceType);
+            }
+            mOnDetachResetRoundness.clear();
+        }
+    }
+
+    /**
+     * SourceType which should be reset when this View is detached
+     * @param sourceType will be reset on View detached
+     */
+    public void addOnDetachResetRoundness(SourceType sourceType) {
+        mOnDetachResetRoundness.add(sourceType);
+    }
+
     public interface OnActivatedListener {
         void onActivated(ActivatableNotificationView view);
         void onActivationReset(ActivatableNotificationView view);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
index d7d5ac9..c7c1634 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRow.java
@@ -91,6 +91,7 @@
 import com.android.systemui.statusbar.notification.AboveShelfChangedListener;
 import com.android.systemui.statusbar.notification.FeedbackIcon;
 import com.android.systemui.statusbar.notification.LaunchAnimationParameters;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.NotificationFadeAware;
 import com.android.systemui.statusbar.notification.NotificationLaunchAnimatorController;
 import com.android.systemui.statusbar.notification.NotificationUtils;
@@ -143,6 +144,9 @@
     private static final int MENU_VIEW_INDEX = 0;
     public static final float DEFAULT_HEADER_VISIBLE_AMOUNT = 1.0f;
     private static final long RECENTLY_ALERTED_THRESHOLD_MS = TimeUnit.SECONDS.toMillis(30);
+    private static final SourceType BASE_VALUE = SourceType.from("BaseValue");
+    private static final SourceType FROM_PARENT = SourceType.from("FromParent(ENR)");
+    private static final SourceType PINNED = SourceType.from("Pinned");
 
     // We don't correctly track dark mode until the content views are inflated, so always update
     // the background on first content update just in case it happens to be during a theme change.
@@ -150,6 +154,7 @@
     private boolean mNotificationTranslationFinished = false;
     private boolean mIsSnoozed;
     private boolean mIsFaded;
+    private boolean mAnimatePinnedRoundness = false;
 
     /**
      * Listener for when {@link ExpandableNotificationRow} is laid out.
@@ -376,7 +381,7 @@
 
     private float mTopRoundnessDuringLaunchAnimation;
     private float mBottomRoundnessDuringLaunchAnimation;
-    private boolean mIsNotificationGroupCornerEnabled;
+    private float mSmallRoundness;
 
     /**
      * Returns whether the given {@code statusBarNotification} is a system notification.
@@ -844,7 +849,9 @@
         }
         onAttachedChildrenCountChanged();
         row.setIsChildInGroup(false, null);
-        row.requestBottomRoundness(0.0f, /* animate = */ false, SourceType.DefaultValue);
+        if (!mUseRoundnessSourceTypes) {
+            row.requestBottomRoundness(0.0f, LegacySourceType.DefaultValue, /* animate = */ false);
+        }
     }
 
     /**
@@ -860,7 +867,10 @@
             if (child.keepInParentForDismissAnimation()) {
                 mChildrenContainer.removeNotification(child);
                 child.setIsChildInGroup(false, null);
-                child.requestBottomRoundness(0.0f, /* animate = */ false, SourceType.DefaultValue);
+                if (!mUseRoundnessSourceTypes) {
+                    LegacySourceType sourceType = LegacySourceType.DefaultValue;
+                    child.requestBottomRoundness(0f, sourceType, /* animate = */ false);
+                }
                 child.setKeepInParentForDismissAnimation(false);
                 logKeepInParentChildDetached(child);
                 childCountChanged = true;
@@ -915,6 +925,9 @@
             mNotificationParent.updateBackgroundForGroupState();
         }
         updateBackgroundClipping();
+        if (mUseRoundnessSourceTypes) {
+            updateBaseRoundness();
+        }
     }
 
     @Override
@@ -1033,6 +1046,16 @@
         if (isAboveShelf() != wasAboveShelf) {
             mAboveShelfChangedListener.onAboveShelfStateChanged(!wasAboveShelf);
         }
+        if (mUseRoundnessSourceTypes) {
+            if (pinned) {
+                // Should be animated if someone explicitly set it to 0 and the row is shown.
+                boolean animated = mAnimatePinnedRoundness && isShown();
+                requestRoundness(/* top = */ 1f, /* bottom = */ 1f, PINNED, animated);
+            } else {
+                requestRoundnessReset(PINNED);
+                mAnimatePinnedRoundness = true;
+            }
+        }
     }
 
     @Override
@@ -1607,6 +1630,8 @@
         super(context, attrs);
         mImageResolver = new NotificationInlineImageResolver(context,
                 new NotificationInlineImageCache());
+        float radius = getResources().getDimension(R.dimen.notification_corner_radius_small);
+        mSmallRoundness = radius / getMaxRadius();
         initDimens();
     }
 
@@ -1839,7 +1864,7 @@
             mChildrenContainer.setIsLowPriority(mIsLowPriority);
             mChildrenContainer.setContainingNotification(ExpandableNotificationRow.this);
             mChildrenContainer.onNotificationUpdated();
-            mChildrenContainer.enableNotificationGroupCorner(mIsNotificationGroupCornerEnabled);
+            mChildrenContainer.useRoundnessSourceTypes(mUseRoundnessSourceTypes);
 
             mTranslateableViews.add(mChildrenContainer);
         });
@@ -2271,7 +2296,7 @@
 
     @Override
     public float getTopRoundness() {
-        if (mExpandAnimationRunning) {
+        if (!mUseRoundnessSourceTypes && mExpandAnimationRunning) {
             return mTopRoundnessDuringLaunchAnimation;
         }
 
@@ -2280,7 +2305,7 @@
 
     @Override
     public float getBottomRoundness() {
-        if (mExpandAnimationRunning) {
+        if (!mUseRoundnessSourceTypes && mExpandAnimationRunning) {
             return mBottomRoundnessDuringLaunchAnimation;
         }
 
@@ -3436,17 +3461,24 @@
     }
 
     @Override
-    public void applyRoundness() {
-        super.applyRoundness();
+    public void applyRoundnessAndInvalidate() {
         applyChildrenRoundness();
+        super.applyRoundnessAndInvalidate();
     }
 
     private void applyChildrenRoundness() {
         if (mIsSummaryWithChildren) {
-            mChildrenContainer.requestBottomRoundness(
-                    getBottomRoundness(),
-                    /* animate = */ false,
-                    SourceType.DefaultValue);
+            if (mUseRoundnessSourceTypes) {
+                mChildrenContainer.requestRoundness(
+                        /* top = */ getTopRoundness(),
+                        /* bottom = */ getBottomRoundness(),
+                        FROM_PARENT);
+            } else {
+                mChildrenContainer.requestBottomRoundness(
+                        getBottomRoundness(),
+                        LegacySourceType.DefaultValue,
+                        /* animate = */ false);
+            }
         }
     }
 
@@ -3605,6 +3637,7 @@
             } else {
                 pw.println("no viewState!!!");
             }
+            pw.println("Roundness: " + getRoundableState().debugString());
 
             if (mIsSummaryWithChildren) {
                 pw.println();
@@ -3649,14 +3682,38 @@
         return mTargetPoint;
     }
 
+    /** Update the minimum roundness based on current state */
+    private void updateBaseRoundness() {
+        if (isChildInGroup()) {
+            requestRoundnessReset(BASE_VALUE);
+        } else {
+            requestRoundness(mSmallRoundness, mSmallRoundness, BASE_VALUE);
+        }
+    }
+
     /**
-     * Enable the support for rounded corner in notification group
+     * Enable the support for rounded corner based on the SourceType
      * @param enabled true if is supported
      */
-    public void enableNotificationGroupCorner(boolean enabled) {
-        mIsNotificationGroupCornerEnabled = enabled;
+    @Override
+    public void useRoundnessSourceTypes(boolean enabled) {
+        super.useRoundnessSourceTypes(enabled);
         if (mChildrenContainer != null) {
-            mChildrenContainer.enableNotificationGroupCorner(mIsNotificationGroupCornerEnabled);
+            mChildrenContainer.useRoundnessSourceTypes(mUseRoundnessSourceTypes);
+        }
+    }
+
+    @Override
+    public String toString() {
+        String roundableStateDebug = "RoundableState = " + getRoundableState().debugString();
+        return "ExpandableNotificationRow:" + hashCode() + " { " + roundableStateDebug + " }";
+    }
+
+    @Override
+    protected void onAttachedToWindow() {
+        super.onAttachedToWindow();
+        if (mUseRoundnessSourceTypes) {
+            updateBaseRoundness();
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowController.java
index 8a400d5..d113860 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableNotificationRowController.java
@@ -255,8 +255,8 @@
                 mStatusBarStateController.removeCallback(mStatusBarStateListener);
             }
         });
-        mView.enableNotificationGroupCorner(
-                mFeatureFlags.isEnabled(Flags.NOTIFICATION_GROUP_CORNER));
+        mView.useRoundnessSourceTypes(
+                mFeatureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES));
     }
 
     private final StatusBarStateController.StateListener mStatusBarStateListener =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
index 2324627..0213b96 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/ExpandableOutlineView.java
@@ -219,9 +219,9 @@
     }
 
     @Override
-    public void applyRoundness() {
+    public void applyRoundnessAndInvalidate() {
         invalidateOutline();
-        super.applyRoundness();
+        super.applyRoundnessAndInvalidate();
     }
 
     protected void setBackgroundTop(int backgroundTop) {
@@ -233,7 +233,7 @@
 
     public void onDensityOrFontScaleChanged() {
         initDimens();
-        applyRoundness();
+        applyRoundnessAndInvalidate();
     }
 
     @Override
@@ -241,7 +241,7 @@
         int previousHeight = getActualHeight();
         super.setActualHeight(actualHeight, notifyListeners);
         if (previousHeight != actualHeight) {
-            applyRoundness();
+            applyRoundnessAndInvalidate();
         }
     }
 
@@ -250,7 +250,7 @@
         int previousAmount = getClipTopAmount();
         super.setClipTopAmount(clipTopAmount);
         if (previousAmount != clipTopAmount) {
-            applyRoundness();
+            applyRoundnessAndInvalidate();
         }
     }
 
@@ -259,14 +259,14 @@
         int previousAmount = getClipBottomAmount();
         super.setClipBottomAmount(clipBottomAmount);
         if (previousAmount != clipBottomAmount) {
-            applyRoundness();
+            applyRoundnessAndInvalidate();
         }
     }
 
     protected void setOutlineAlpha(float alpha) {
         if (alpha != mOutlineAlpha) {
             mOutlineAlpha = alpha;
-            applyRoundness();
+            applyRoundnessAndInvalidate();
         }
     }
 
@@ -280,7 +280,7 @@
             setOutlineRect(rect.left, rect.top, rect.right, rect.bottom);
         } else {
             mCustomOutline = false;
-            applyRoundness();
+            applyRoundnessAndInvalidate();
         }
     }
 
@@ -340,7 +340,7 @@
         // Outlines need to be at least 1 dp
         mOutlineRect.bottom = (int) Math.max(top, mOutlineRect.bottom);
         mOutlineRect.right = (int) Math.max(left, mOutlineRect.right);
-        applyRoundness();
+        applyRoundnessAndInvalidate();
     }
 
     public Path getCustomClipPath(View child) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationHeaderViewWrapper.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationHeaderViewWrapper.java
index f13e48d..1f664cb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationHeaderViewWrapper.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/row/wrapper/NotificationHeaderViewWrapper.java
@@ -71,6 +71,8 @@
     private View mFeedbackIcon;
     private boolean mIsLowPriority;
     private boolean mTransformLowPriorityTitle;
+    private boolean mUseRoundnessSourceTypes;
+    private RoundnessChangedListener mRoundnessChangedListener;
 
     protected NotificationHeaderViewWrapper(Context ctx, View view, ExpandableNotificationRow row) {
         super(ctx, view, row);
@@ -117,6 +119,20 @@
         return mRoundableState;
     }
 
+    @Override
+    public void applyRoundnessAndInvalidate() {
+        if (mUseRoundnessSourceTypes && mRoundnessChangedListener != null) {
+            // We cannot apply the rounded corner to this View, so our parents (in drawChild()) will
+            // clip our canvas. So we should invalidate our parent.
+            mRoundnessChangedListener.applyRoundnessAndInvalidate();
+        }
+        Roundable.super.applyRoundnessAndInvalidate();
+    }
+
+    public void setOnRoundnessChangedListener(RoundnessChangedListener listener) {
+        mRoundnessChangedListener = listener;
+    }
+
     protected void resolveHeaderViews() {
         mIcon = mView.findViewById(com.android.internal.R.id.icon);
         mHeaderText = mView.findViewById(com.android.internal.R.id.header_text);
@@ -343,4 +359,23 @@
             }
         }
     }
+
+    /**
+     * Enable the support for rounded corner based on the SourceType
+     *
+     * @param enabled true if is supported
+     */
+    public void useRoundnessSourceTypes(boolean enabled) {
+        mUseRoundnessSourceTypes = enabled;
+    }
+
+    /**
+     * Interface that handle the Roundness changes
+     */
+    public interface RoundnessChangedListener {
+        /**
+         * This method will be called when this class call applyRoundnessAndInvalidate()
+         */
+        void applyRoundnessAndInvalidate();
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
index d43ca823..4a8e2db 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainer.java
@@ -47,6 +47,7 @@
 import com.android.systemui.statusbar.NotificationGroupingUtil;
 import com.android.systemui.statusbar.NotificationShelf;
 import com.android.systemui.statusbar.notification.FeedbackIcon;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.NotificationFadeAware;
 import com.android.systemui.statusbar.notification.NotificationUtils;
 import com.android.systemui.statusbar.notification.Roundable;
@@ -83,6 +84,7 @@
             return mAnimationFilter;
         }
     }.setDuration(200);
+    private static final SourceType FROM_PARENT = SourceType.from("FromParent(NCC)");
 
     private final List<View> mDividers = new ArrayList<>();
     private final List<ExpandableNotificationRow> mAttachedChildren = new ArrayList<>();
@@ -131,7 +133,7 @@
     private int mUntruncatedChildCount;
     private boolean mContainingNotificationIsFaded = false;
     private RoundableState mRoundableState;
-    private boolean mIsNotificationGroupCornerEnabled;
+    private boolean mUseRoundnessSourceTypes;
 
     public NotificationChildrenContainer(Context context) {
         this(context, null);
@@ -313,9 +315,12 @@
         row.setContentTransformationAmount(0, false /* isLastChild */);
         row.setNotificationFaded(mContainingNotificationIsFaded);
 
-        // This is a workaround, the NotificationShelf should be the owner of `OnScroll` roundness.
-        // Here we should reset the `OnScroll` roundness only on top-level rows.
-        NotificationShelf.resetOnScrollRoundness(row);
+        if (!mUseRoundnessSourceTypes) {
+            // This is a workaround, the NotificationShelf should be the owner of `OnScroll`
+            // roundness.
+            // Here we should reset the `OnScroll` roundness only on top-level rows.
+            NotificationShelf.resetLegacyOnScrollRoundness(row);
+        }
 
         // It doesn't make sense to keep old animations around, lets cancel them!
         ExpandableViewState viewState = row.getViewState();
@@ -323,6 +328,10 @@
             viewState.cancelAnimations(row);
             row.cancelAppearDrawing();
         }
+
+        if (mUseRoundnessSourceTypes) {
+            applyRoundnessAndInvalidate();
+        }
     }
 
     private void ensureRemovedFromTransientContainer(View v) {
@@ -356,6 +365,11 @@
         if (!row.isRemoved()) {
             mGroupingUtil.restoreChildNotification(row);
         }
+
+        if (mUseRoundnessSourceTypes) {
+            row.requestRoundnessReset(FROM_PARENT);
+            applyRoundnessAndInvalidate();
+        }
     }
 
     /**
@@ -382,6 +396,10 @@
                             getContext(),
                             mNotificationHeader,
                             mContainingNotification);
+            mNotificationHeaderWrapper.useRoundnessSourceTypes(mUseRoundnessSourceTypes);
+            if (mUseRoundnessSourceTypes) {
+                mNotificationHeaderWrapper.setOnRoundnessChangedListener(this::invalidate);
+            }
             addView(mNotificationHeader, 0);
             invalidate();
         } else {
@@ -419,6 +437,12 @@
                                 getContext(),
                                 mNotificationHeaderLowPriority,
                                 mContainingNotification);
+                mNotificationHeaderWrapperLowPriority.useRoundnessSourceTypes(
+                        mUseRoundnessSourceTypes
+                );
+                if (mUseRoundnessSourceTypes) {
+                    mNotificationHeaderWrapper.setOnRoundnessChangedListener(this::invalidate);
+                }
                 addView(mNotificationHeaderLowPriority, 0);
                 invalidate();
             } else {
@@ -841,7 +865,7 @@
 
             isCanvasChanged = true;
             canvas.save();
-            if (mIsNotificationGroupCornerEnabled && translation != 0f) {
+            if (mUseRoundnessSourceTypes && translation != 0f) {
                 clipPath.offset(translation, 0f);
                 canvas.clipPath(clipPath);
                 clipPath.offset(-translation, 0f);
@@ -1392,24 +1416,28 @@
     }
 
     @Override
-    public void applyRoundness() {
-        Roundable.super.applyRoundness();
+    public void applyRoundnessAndInvalidate() {
         boolean last = true;
         for (int i = mAttachedChildren.size() - 1; i >= 0; i--) {
             ExpandableNotificationRow child = mAttachedChildren.get(i);
             if (child.getVisibility() == View.GONE) {
                 continue;
             }
-            child.requestTopRoundness(
-                    /* value = */ 0f,
-                    /* animate = */ isShown(),
-                    SourceType.DefaultValue);
-            child.requestBottomRoundness(
-                    /* value = */ last ? getBottomRoundness() : 0f,
-                    /* animate = */ isShown(),
-                    SourceType.DefaultValue);
+            if (mUseRoundnessSourceTypes) {
+                child.requestRoundness(
+                        /* top = */ 0f,
+                        /* bottom = */ last ? getBottomRoundness() : 0f,
+                        FROM_PARENT);
+            } else {
+                child.requestRoundness(
+                        /* top = */ 0f,
+                        /* bottom = */ last ? getBottomRoundness() : 0f,
+                        LegacySourceType.DefaultValue,
+                        /* animate = */ isShown());
+            }
             last = false;
         }
+        Roundable.super.applyRoundnessAndInvalidate();
     }
 
     public void setHeaderVisibleAmount(float headerVisibleAmount) {
@@ -1467,10 +1495,17 @@
     }
 
     /**
-     * Enable the support for rounded corner in notification group
+     * Enable the support for rounded corner based on the SourceType
+     *
      * @param enabled true if is supported
      */
-    public void enableNotificationGroupCorner(boolean enabled) {
-        mIsNotificationGroupCornerEnabled = enabled;
+    public void useRoundnessSourceTypes(boolean enabled) {
+        mUseRoundnessSourceTypes = enabled;
+    }
+
+    @Override
+    public String toString() {
+        String roundableStateDebug = "RoundableState = " + getRoundableState().debugString();
+        return "NotificationChildrenContainer:" + hashCode() + " { " + roundableStateDebug + " }";
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
index 6810055..fde8c4d 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManager.java
@@ -25,6 +25,9 @@
 import com.android.systemui.R;
 import com.android.systemui.dagger.SysUISingleton;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.NotificationSectionsFeatureManager;
 import com.android.systemui.statusbar.notification.Roundable;
 import com.android.systemui.statusbar.notification.SourceType;
@@ -45,6 +48,7 @@
 public class NotificationRoundnessManager implements Dumpable {
 
     private static final String TAG = "NotificationRoundnessManager";
+    private static final SourceType DISMISS_ANIMATION = SourceType.from("DismissAnimation");
 
     private final ExpandableView[] mFirstInSectionViews;
     private final ExpandableView[] mLastInSectionViews;
@@ -63,12 +67,14 @@
     private ExpandableView mSwipedView = null;
     private Roundable mViewBeforeSwipedView = null;
     private Roundable mViewAfterSwipedView = null;
+    private boolean mUseRoundnessSourceTypes;
 
     @Inject
     NotificationRoundnessManager(
             NotificationSectionsFeatureManager sectionsFeatureManager,
             NotificationRoundnessLogger notifLogger,
-            DumpManager dumpManager) {
+            DumpManager dumpManager,
+            FeatureFlags featureFlags) {
         int numberOfSections = sectionsFeatureManager.getNumberOfBuckets();
         mFirstInSectionViews = new ExpandableView[numberOfSections];
         mLastInSectionViews = new ExpandableView[numberOfSections];
@@ -76,6 +82,7 @@
         mTmpLastInSectionViews = new ExpandableView[numberOfSections];
         mNotifLogger = notifLogger;
         mDumpManager = dumpManager;
+        mUseRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES);
 
         mDumpManager.registerDumpable(TAG, this);
     }
@@ -94,6 +101,7 @@
     }
 
     public void updateView(ExpandableView view, boolean animate) {
+        if (mUseRoundnessSourceTypes) return;
         boolean changed = updateViewWithoutCallback(view, animate);
         if (changed) {
             mRoundingChangedCallback.run();
@@ -110,6 +118,7 @@
     boolean updateViewWithoutCallback(
             ExpandableView view,
             boolean animate) {
+        if (mUseRoundnessSourceTypes) return false;
         if (view == null
                 || view == mViewBeforeSwipedView
                 || view == mViewAfterSwipedView) {
@@ -118,13 +127,13 @@
 
         final boolean isTopChanged = view.requestTopRoundness(
                 getRoundnessDefaultValue(view, true /* top */),
-                animate,
-                SourceType.DefaultValue);
+                LegacySourceType.DefaultValue,
+                animate);
 
         final boolean isBottomChanged = view.requestBottomRoundness(
                 getRoundnessDefaultValue(view, /* top = */ false),
-                animate,
-                SourceType.DefaultValue);
+                LegacySourceType.DefaultValue,
+                animate);
 
         final boolean isFirstInSection = isFirstInSection(view);
         final boolean isLastInSection = isLastInSection(view);
@@ -139,6 +148,7 @@
     }
 
     private boolean isFirstInSection(ExpandableView view) {
+        if (mUseRoundnessSourceTypes) return false;
         for (int i = 0; i < mFirstInSectionViews.length; i++) {
             if (view == mFirstInSectionViews[i]) {
                 return true;
@@ -148,6 +158,7 @@
     }
 
     private boolean isLastInSection(ExpandableView view) {
+        if (mUseRoundnessSourceTypes) return false;
         for (int i = mLastInSectionViews.length - 1; i >= 0; i--) {
             if (view == mLastInSectionViews[i]) {
                 return true;
@@ -160,9 +171,6 @@
             Roundable viewBefore,
             ExpandableView viewSwiped,
             Roundable viewAfter) {
-        final boolean animate = true;
-        final SourceType source = SourceType.OnDismissAnimation;
-
         // This method requires you to change the roundness of the current View targets and reset
         // the roundness of the old View targets (if any) to 0f.
         // To avoid conflicts, it generates a set of old Views and removes the current Views
@@ -172,31 +180,34 @@
         if (mSwipedView != null) oldViews.add(mSwipedView);
         if (mViewAfterSwipedView != null) oldViews.add(mViewAfterSwipedView);
 
+        final SourceType source;
+        if (mUseRoundnessSourceTypes) {
+            source = DISMISS_ANIMATION;
+        } else {
+            source = LegacySourceType.OnDismissAnimation;
+        }
+
         mViewBeforeSwipedView = viewBefore;
         if (viewBefore != null) {
             oldViews.remove(viewBefore);
-            viewBefore.requestTopRoundness(0f, animate, source);
-            viewBefore.requestBottomRoundness(1f, animate, source);
+            viewBefore.requestRoundness(/* top = */ 0f, /* bottom = */ 1f, source);
         }
 
         mSwipedView = viewSwiped;
         if (viewSwiped != null) {
             oldViews.remove(viewSwiped);
-            viewSwiped.requestTopRoundness(1f, animate, source);
-            viewSwiped.requestBottomRoundness(1f, animate, source);
+            viewSwiped.requestRoundness(/* top = */ 1f, /* bottom = */ 1f, source);
         }
 
         mViewAfterSwipedView = viewAfter;
         if (viewAfter != null) {
             oldViews.remove(viewAfter);
-            viewAfter.requestTopRoundness(1f, animate, source);
-            viewAfter.requestBottomRoundness(0f, animate, source);
+            viewAfter.requestRoundness(/* top = */ 1f, /* bottom = */ 0f, source);
         }
 
         // After setting the current Views, reset the views that are still present in the set.
         for (Roundable oldView : oldViews) {
-            oldView.requestTopRoundness(0f, animate, source);
-            oldView.requestBottomRoundness(0f, animate, source);
+            oldView.requestRoundnessReset(source);
         }
     }
 
@@ -204,7 +215,23 @@
         mIsClearAllInProgress = isClearingAll;
     }
 
+    /**
+     * Check if "Clear all" notifications is in progress.
+     */
+    public boolean isClearAllInProgress() {
+        return mIsClearAllInProgress;
+    }
+
+    /**
+     * Check if we can request the `Pulsing` roundness for notification.
+     */
+    public boolean shouldRoundNotificationPulsing() {
+        return mRoundForPulsingViews;
+    }
+
     private float getRoundnessDefaultValue(Roundable view, boolean top) {
+        if (mUseRoundnessSourceTypes) return 0f;
+
         if (view == null) {
             return 0f;
         }
@@ -250,6 +277,7 @@
     }
 
     public void setExpanded(float expandedHeight, float appearFraction) {
+        if (mUseRoundnessSourceTypes) return;
         mExpanded = expandedHeight != 0.0f;
         mAppearFraction = appearFraction;
         if (mTrackedHeadsUp != null) {
@@ -258,6 +286,7 @@
     }
 
     public void updateRoundedChildren(NotificationSection[] sections) {
+        if (mUseRoundnessSourceTypes) return;
         boolean anyChanged = false;
         for (int i = 0; i < sections.length; i++) {
             mTmpFirstInSectionViews[i] = mFirstInSectionViews[i];
@@ -280,6 +309,7 @@
             NotificationSection[] sections,
             ExpandableView[] oldViews,
             boolean first) {
+        if (mUseRoundnessSourceTypes) return false;
         boolean anyChanged = false;
         for (ExpandableView oldView : oldViews) {
             if (oldView != null) {
@@ -313,6 +343,7 @@
             NotificationSection[] sections,
             ExpandableView[] oldViews,
             boolean first) {
+        if (mUseRoundnessSourceTypes) return false;
         boolean anyChanged = false;
         for (NotificationSection section : sections) {
             ExpandableView newView =
@@ -339,6 +370,15 @@
         mAnimatedChildren = animatedChildren;
     }
 
+    /**
+     * Check if the view should be animated
+     * @param view target view
+     * @return true, if is in the AnimatedChildren set
+     */
+    public boolean isAnimatedChild(ExpandableView view) {
+        return mAnimatedChildren.contains(view);
+    }
+
     public void setOnRoundingChangedCallback(Runnable roundingChangedCallback) {
         mRoundingChangedCallback = roundingChangedCallback;
     }
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 a1b77ac..070b439 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
@@ -19,8 +19,11 @@
 import android.util.Log
 import android.view.View
 import com.android.internal.annotations.VisibleForTesting
+import com.android.systemui.flags.FeatureFlags
+import com.android.systemui.flags.Flags
 import com.android.systemui.media.controls.ui.KeyguardMediaController
 import com.android.systemui.statusbar.notification.NotificationSectionsFeatureManager
+import com.android.systemui.statusbar.notification.SourceType
 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
@@ -44,12 +47,16 @@
     private val keyguardMediaController: KeyguardMediaController,
     private val sectionsFeatureManager: NotificationSectionsFeatureManager,
     private val mediaContainerController: MediaContainerController,
+    private val notificationRoundnessManager: NotificationRoundnessManager,
     @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,
+    featureFlags: FeatureFlags
 ) : SectionProvider {
 
+    private val useRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES)
+
     private val configurationListener = object : ConfigurationController.ConfigurationListener {
         override fun onLocaleListChanged() {
             reinflateViews()
@@ -177,11 +184,49 @@
                         size = sections.size,
                         operation = SectionBounds::addNotif
                 )
+
+        // Build a set of the old first/last Views of the sections
+        val oldFirstChildren = sections.mapNotNull { it.firstVisibleChild }.toSet().toMutableSet()
+        val oldLastChildren = sections.mapNotNull { it.lastVisibleChild }.toSet().toMutableSet()
+
         // Update each section with the associated boundary, tracking if there was a change
         val changed = sections.fold(false) { changed, section ->
             val bounds = sectionBounds[section.bucket] ?: SectionBounds.None
-            bounds.updateSection(section) || changed
+            val isSectionChanged = bounds.updateSection(section)
+            isSectionChanged || changed
         }
+
+        if (useRoundnessSourceTypes) {
+            val newFirstChildren = sections.mapNotNull { it.firstVisibleChild }
+            val newLastChildren = sections.mapNotNull { it.lastVisibleChild }
+
+            // Update the roundness of Views that weren't already in the first/last position
+            newFirstChildren.forEach { firstChild ->
+                val wasFirstChild = oldFirstChildren.remove(firstChild)
+                if (!wasFirstChild) {
+                    val notAnimatedChild = !notificationRoundnessManager.isAnimatedChild(firstChild)
+                    val animated = firstChild.isShown && notAnimatedChild
+                    firstChild.requestTopRoundness(1f, SECTION, animated)
+                }
+            }
+            newLastChildren.forEach { lastChild ->
+                val wasLastChild = oldLastChildren.remove(lastChild)
+                if (!wasLastChild) {
+                    val notAnimatedChild = !notificationRoundnessManager.isAnimatedChild(lastChild)
+                    val animated = lastChild.isShown && notAnimatedChild
+                    lastChild.requestBottomRoundness(1f, SECTION, animated)
+                }
+            }
+
+            // The Views left in the set are no longer in the first/last position
+            oldFirstChildren.forEach { noMoreFirstChild ->
+                noMoreFirstChild.requestTopRoundness(0f, SECTION)
+            }
+            oldLastChildren.forEach { noMoreLastChild ->
+                noMoreLastChild.requestBottomRoundness(0f, SECTION)
+            }
+        }
+
         if (DEBUG) {
             logSections(sections)
         }
@@ -215,5 +260,6 @@
     companion object {
         private const val TAG = "NotifSectionsManager"
         private const val DEBUG = false
+        private val SECTION = SourceType.from("Section")
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
index 7c3e52c..21e2bd8 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayout.java
@@ -191,10 +191,13 @@
 
     private final boolean mDebugLines;
     private Paint mDebugPaint;
-    /** Used to track the Y positions that were already used to draw debug text labels. */
+    /**
+     * Used to track the Y positions that were already used to draw debug text labels.
+     */
     private Set<Integer> mDebugTextUsedYPositions;
     private final boolean mDebugRemoveAnimation;
     private final boolean mSimplifiedAppearFraction;
+    private final boolean mUseRoundnessSourceTypes;
 
     private int mContentHeight;
     private float mIntrinsicContentHeight;
@@ -355,15 +358,14 @@
     private final Rect mQsHeaderBound = new Rect();
     private boolean mContinuousShadowUpdate;
     private boolean mContinuousBackgroundUpdate;
-    private final ViewTreeObserver.OnPreDrawListener mShadowUpdater
-            = () -> {
-                updateViewShadows();
-                return true;
-            };
+    private final ViewTreeObserver.OnPreDrawListener mShadowUpdater = () -> {
+        updateViewShadows();
+        return true;
+    };
     private final ViewTreeObserver.OnPreDrawListener mBackgroundUpdater = () -> {
-                updateBackground();
-                return true;
-            };
+        updateBackground();
+        return true;
+    };
     private final Comparator<ExpandableView> mViewPositionComparator = (view, otherView) -> {
         float endY = view.getTranslationY() + view.getActualHeight();
         float otherEndY = otherView.getTranslationY() + otherView.getActualHeight();
@@ -419,7 +421,8 @@
      */
     private int mMaxDisplayedNotifications = -1;
     private float mKeyguardBottomPadding = -1;
-    @VisibleForTesting int mStatusBarHeight;
+    @VisibleForTesting
+    int mStatusBarHeight;
     private int mMinInteractionHeight;
     private final Rect mClipRect = new Rect();
     private boolean mIsClipped;
@@ -568,6 +571,8 @@
 
     @Nullable
     private OnClickListener mManageButtonClickListener;
+    @Nullable
+    private OnNotificationRemovedListener mOnNotificationRemovedListener;
 
     public NotificationStackScrollLayout(Context context, AttributeSet attrs) {
         super(context, attrs, 0, 0);
@@ -576,6 +581,7 @@
         mDebugLines = featureFlags.isEnabled(Flags.NSSL_DEBUG_LINES);
         mDebugRemoveAnimation = featureFlags.isEnabled(Flags.NSSL_DEBUG_REMOVE_ANIMATION);
         mSimplifiedAppearFraction = featureFlags.isEnabled(Flags.SIMPLIFIED_APPEAR_FRACTION);
+        mUseRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES);
         mSectionsManager = Dependency.get(NotificationSectionsManager.class);
         mScreenOffAnimationController =
                 Dependency.get(ScreenOffAnimationController.class);
@@ -741,7 +747,7 @@
     }
 
     private void logHunSkippedForUnexpectedState(ExpandableNotificationRow enr,
-            boolean expected, boolean actual) {
+                                                 boolean expected, boolean actual) {
         if (mLogger == null) return;
         mLogger.hunSkippedForUnexpectedState(enr.getEntry(), expected, actual);
     }
@@ -868,13 +874,13 @@
 
     /**
      * Draws round rects for each background section.
-     *
+     * <p>
      * We want to draw a round rect for each background section as defined by {@link #mSections}.
      * However, if two sections are directly adjacent with no gap between them (e.g. on the
      * lockscreen where the shelf can appear directly below the high priority section, or while
      * scrolling the shade so that the top of the shelf is right at the bottom of the high priority
      * section), we don't want to round the adjacent corners.
-     *
+     * <p>
      * Since {@link Canvas} doesn't provide a way to draw a half-rounded rect, this means that we
      * need to coalesce the backgrounds for adjacent sections and draw them as a single round rect.
      * This method tracks the top of each rect we need to draw, then iterates through the visible
@@ -883,7 +889,7 @@
      * the current section.  When we're done iterating we will always have one rect left to draw.
      */
     private void drawBackgroundRects(Canvas canvas, int left, int right, int top,
-            int animationYOffset) {
+                                     int animationYOffset) {
         int backgroundRectTop = top;
         int lastSectionBottom =
                 mSections[0].getCurrentBounds().bottom + animationYOffset;
@@ -974,7 +980,7 @@
 
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     void initView(Context context, NotificationSwipeHelper swipeHelper,
-            NotificationStackSizeCalculator notificationStackSizeCalculator) {
+                  NotificationStackSizeCalculator notificationStackSizeCalculator) {
         mScroller = new OverScroller(getContext());
         mSwipeHelper = swipeHelper;
         mNotificationStackSizeCalculator = notificationStackSizeCalculator;
@@ -1304,18 +1310,19 @@
     /**
      * @return Whether we should skip stack height updates.
      * True when
-     *      1) Unlock hint is running
-     *      2) Swiping up on lockscreen or flinging down after swipe up
+     * 1) Unlock hint is running
+     * 2) Swiping up on lockscreen or flinging down after swipe up
      */
     private boolean shouldSkipHeightUpdate() {
         return mAmbientState.isOnKeyguard()
                 && (mAmbientState.isUnlockHintRunning()
-                        || mAmbientState.isSwipingUp()
-                        || mAmbientState.isFlingingAfterSwipeUpOnLockscreen());
+                || mAmbientState.isSwipingUp()
+                || mAmbientState.isFlingingAfterSwipeUpOnLockscreen());
     }
 
     /**
      * Apply expansion fraction to the y position and height of the notifications panel.
+     *
      * @param listenerNeedsAnimation does the listener need to animate?
      */
     private void updateStackPosition(boolean listenerNeedsAnimation) {
@@ -1708,7 +1715,7 @@
      */
     @ShadeViewRefactor(RefactorComponent.COORDINATOR)
     ExpandableView getChildAtPosition(float touchX, float touchY,
-            boolean requireMinHeight, boolean ignoreDecors) {
+                                      boolean requireMinHeight, boolean ignoreDecors) {
         // find the view under the pointer, accounting for GONE views
         final int count = getChildCount();
         for (int childIdx = 0; childIdx < count; childIdx++) {
@@ -1868,9 +1875,9 @@
     public void dismissViewAnimated(
             View child, Consumer<Boolean> endRunnable, int delay, long duration) {
         if (child instanceof SectionHeaderView) {
-             ((StackScrollerDecorView) child).setContentVisible(
-                     false /* visible */, true /* animate */, endRunnable);
-             return;
+            ((StackScrollerDecorView) child).setContentVisible(
+                    false /* visible */, true /* animate */, endRunnable);
+            return;
         }
         mSwipeHelper.dismissChild(
                 child,
@@ -2032,10 +2039,11 @@
     /**
      * Scrolls by the given delta, overscrolling if needed.  If called during a fling and the delta
      * would cause us to exceed the provided maximum overscroll, springs back instead.
-     *
+     * <p>
      * This method performs the determination of whether we're exceeding the overscroll and clamps
      * the scroll amount if so.  The actual scrolling/overscrolling happens in
      * {@link #onCustomOverScrolled(int, boolean)}
+     *
      * @param deltaY         The (signed) number of pixels to scroll.
      * @param scrollY        The current scroll position (absolute scrolling only).
      * @param scrollRangeY   The maximum allowable scroll position (absolute scrolling only).
@@ -2098,7 +2106,7 @@
      */
     @ShadeViewRefactor(RefactorComponent.STATE_RESOLVER)
     public void setOverScrollAmount(float amount, boolean onTop, boolean animate,
-            boolean cancelAnimators) {
+                                    boolean cancelAnimators) {
         setOverScrollAmount(amount, onTop, animate, cancelAnimators, isRubberbanded(onTop));
     }
 
@@ -2114,7 +2122,7 @@
      */
     @ShadeViewRefactor(RefactorComponent.STATE_RESOLVER)
     public void setOverScrollAmount(float amount, boolean onTop, boolean animate,
-            boolean cancelAnimators, boolean isRubberbanded) {
+                                    boolean cancelAnimators, boolean isRubberbanded) {
         if (cancelAnimators) {
             mStateAnimator.cancelOverScrollAnimators(onTop);
         }
@@ -2123,7 +2131,7 @@
 
     @ShadeViewRefactor(RefactorComponent.STATE_RESOLVER)
     private void setOverScrollAmountInternal(float amount, boolean onTop, boolean animate,
-            boolean isRubberbanded) {
+                                             boolean isRubberbanded) {
         amount = Math.max(0, amount);
         if (animate) {
             mStateAnimator.animateOverScrollToAmount(amount, onTop, isRubberbanded);
@@ -2179,7 +2187,7 @@
      * Scrolls to the given position, overscrolling if needed.  If called during a fling and the
      * position exceeds the provided maximum overscroll, springs back instead.
      *
-     * @param scrollY The target scroll position.
+     * @param scrollY  The target scroll position.
      * @param clampedY Whether this value was clamped by the calling method, meaning we've reached
      *                 the overscroll limit.
      */
@@ -2288,8 +2296,8 @@
                 if (row.isSummaryWithChildren() && row.areChildrenExpanded()) {
                     List<ExpandableNotificationRow> notificationChildren =
                             row.getAttachedChildren();
-                    for (int childIndex = 0; childIndex < notificationChildren.size();
-                            childIndex++) {
+                    int childrenSize = notificationChildren.size();
+                    for (int childIndex = 0; childIndex < childrenSize; childIndex++) {
                         ExpandableNotificationRow rowChild = notificationChildren.get(childIndex);
                         if (rowChild.getTranslationY() + rowTranslation >= translationY) {
                             return rowChild;
@@ -2365,10 +2373,9 @@
     /**
      * Calculate the gap height between two different views
      *
-     * @param previous the previousView
-     * @param current the currentView
+     * @param previous     the previousView
+     * @param current      the currentView
      * @param visibleIndex the visible index in the list
-     *
      * @return the gap height needed before the current view
      */
     public float calculateGapHeight(
@@ -2376,7 +2383,7 @@
             ExpandableView current,
             int visibleIndex
     ) {
-       return mStackScrollAlgorithm.getGapHeightForChild(mSectionsManager, visibleIndex, current,
+        return mStackScrollAlgorithm.getGapHeightForChild(mSectionsManager, visibleIndex, current,
                 previous, mAmbientState.getFractionToShade(), mAmbientState.isOnKeyguard());
     }
 
@@ -2642,15 +2649,15 @@
         return mScrolledToTopOnFirstDown
                 && !mExpandedInThisMotion
                 && (initialVelocity > mMinimumVelocity
-                        || (topOverScroll > mMinTopOverScrollToEscape && initialVelocity > 0));
+                || (topOverScroll > mMinTopOverScrollToEscape && initialVelocity > 0));
     }
 
     /**
      * Updates the top padding of the notifications, taking {@link #getIntrinsicPadding()} into
      * account.
      *
-     * @param qsHeight               the top padding imposed by the quick settings panel
-     * @param animate                whether to animate the change
+     * @param qsHeight the top padding imposed by the quick settings panel
+     * @param animate  whether to animate the change
      */
     @ShadeViewRefactor(RefactorComponent.COORDINATOR)
     public void updateTopPadding(float qsHeight, boolean animate) {
@@ -2724,21 +2731,34 @@
     }
 
 
-
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     public void setChildTransferInProgress(boolean childTransferInProgress) {
         Assert.isMainThread();
         mChildTransferInProgress = childTransferInProgress;
     }
 
+    /**
+     * Set the remove notification listener
+     * @param listener callback for notification removed
+     */
+    public void setOnNotificationRemovedListener(OnNotificationRemovedListener listener) {
+        mOnNotificationRemovedListener = listener;
+    }
+
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     @Override
     public void onViewRemoved(View child) {
         super.onViewRemoved(child);
         // we only call our internal methods if this is actually a removal and not just a
         // notification which becomes a child notification
+        ExpandableView expandableView = (ExpandableView) child;
         if (!mChildTransferInProgress) {
-            onViewRemovedInternal((ExpandableView) child, this);
+            onViewRemovedInternal(expandableView, this);
+        }
+        if (mOnNotificationRemovedListener != null) {
+            mOnNotificationRemovedListener.onNotificationRemoved(
+                    expandableView,
+                    mChildTransferInProgress);
         }
     }
 
@@ -2998,8 +3018,10 @@
             mAnimateNextSectionBoundsChange = false;
         }
         mAmbientState.setLastVisibleBackgroundChild(lastChild);
-        // TODO: Refactor SectionManager and put the RoundnessManager there.
-        mController.getNotificationRoundnessManager().updateRoundedChildren(mSections);
+        if (!mUseRoundnessSourceTypes) {
+            // TODO: Refactor SectionManager and put the RoundnessManager there.
+            mController.getNotificationRoundnessManager().updateRoundedChildren(mSections);
+        }
         mAnimateBottomOnLayout = false;
         invalidate();
     }
@@ -3206,7 +3228,7 @@
                     // Only animate if we still have pinned heads up, otherwise we just have the
                     // regular collapse animation of the lock screen
                     || (mKeyguardBypassEnabled && onKeyguard()
-                            && mInHeadsUpPinnedMode);
+                    && mInHeadsUpPinnedMode);
             if (performDisappearAnimation && !isHeadsUp) {
                 type = row.wasJustClicked()
                         ? AnimationEvent.ANIMATION_TYPE_HEADS_UP_DISAPPEAR_CLICK
@@ -3351,7 +3373,7 @@
             AnimationEvent animEvent = duration == null
                     ? new AnimationEvent(child, AnimationEvent.ANIMATION_TYPE_CHANGE_POSITION)
                     : new AnimationEvent(
-                            child, AnimationEvent.ANIMATION_TYPE_CHANGE_POSITION, duration);
+                    child, AnimationEvent.ANIMATION_TYPE_CHANGE_POSITION, duration);
             mAnimationEvents.add(animEvent);
         }
         mChildrenChangingPositions.clear();
@@ -3442,7 +3464,9 @@
 
     @ShadeViewRefactor(RefactorComponent.LAYOUT_ALGORITHM)
     protected StackScrollAlgorithm createStackScrollAlgorithm(Context context) {
-        return new StackScrollAlgorithm(context, this);
+        StackScrollAlgorithm stackScrollAlgorithm = new StackScrollAlgorithm(context, this);
+        stackScrollAlgorithm.useRoundnessSourceTypes(mUseRoundnessSourceTypes);
+        return stackScrollAlgorithm;
     }
 
     /**
@@ -3772,7 +3796,7 @@
     }
 
     private void logEmptySpaceClick(MotionEvent ev, boolean isTouchBelowLastNotification,
-            int statusBarState, boolean touchIsClick) {
+                                    int statusBarState, boolean touchIsClick) {
         if (mLogger == null) {
             return;
         }
@@ -3957,7 +3981,9 @@
         mOnEmptySpaceClickListener = listener;
     }
 
-    /** @hide */
+    /**
+     * @hide
+     */
     @Override
     @ShadeViewRefactor(RefactorComponent.INPUT)
     public boolean performAccessibilityActionInternal(int action, Bundle arguments) {
@@ -4728,7 +4754,9 @@
         return touchY > mTopPadding + mStackTranslation;
     }
 
-    /** @hide */
+    /**
+     * @hide
+     */
     @Override
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     public void onInitializeAccessibilityEventInternal(AccessibilityEvent event) {
@@ -5353,7 +5381,9 @@
         return canChildBeCleared(row) && matchesSelection(row, selection);
     }
 
-    /** Register a {@link View.OnClickListener} to be invoked when the Manage button is clicked. */
+    /**
+     * Register a {@link View.OnClickListener} to be invoked when the Manage button is clicked.
+     */
     public void setManageButtonClickListener(@Nullable OnClickListener listener) {
         mManageButtonClickListener = listener;
         if (mFooterView != null) {
@@ -5418,6 +5448,7 @@
     /**
      * Set how far the wake up is when waking up from pulsing. This is a height and will adjust the
      * notification positions accordingly.
+     *
      * @param height the new wake up height
      * @return the overflow how much the height is further than he lowest notification
      */
@@ -5649,7 +5680,7 @@
      * Set rounded rect clipping bounds on this view.
      */
     public void setRoundedClippingBounds(int left, int top, int right, int bottom, int topRadius,
-            int bottomRadius) {
+                                         int bottomRadius) {
         if (mRoundedRectClippingLeft == left && mRoundedRectClippingRight == right
                 && mRoundedRectClippingBottom == bottom && mRoundedRectClippingTop == top
                 && mBgCornerRadii[0] == topRadius && mBgCornerRadii[5] == bottomRadius) {
@@ -5710,7 +5741,7 @@
         mLaunchingNotification = launching;
         mLaunchingNotificationNeedsToBeClipped = mLaunchAnimationParams != null
                 && (mLaunchAnimationParams.getStartRoundedTopClipping() > 0
-                        || mLaunchAnimationParams.getParentStartRoundedTopClipping() > 0);
+                || mLaunchAnimationParams.getParentStartRoundedTopClipping() > 0);
         if (!mLaunchingNotificationNeedsToBeClipped || !mLaunchingNotification) {
             mLaunchedNotificationClipPath.reset();
         }
@@ -5748,7 +5779,7 @@
                 mLaunchAnimationParams.getProgress(0,
                         NotificationLaunchAnimatorController.ANIMATION_DURATION_TOP_ROUNDING));
         int top = (int) Math.min(MathUtils.lerp(mRoundedRectClippingTop,
-                mLaunchAnimationParams.getTop(), expandProgress),
+                        mLaunchAnimationParams.getTop(), expandProgress),
                 mRoundedRectClippingTop);
         float topRadius = mLaunchAnimationParams.getTopCornerRadius();
         float bottomRadius = mLaunchAnimationParams.getBottomCornerRadius();
@@ -5872,25 +5903,25 @@
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
     public interface OnOverscrollTopChangedListener {
 
-    /**
-     * Notifies a listener that the overscroll has changed.
-     *
-     * @param amount         the amount of overscroll, in pixels
-     * @param isRubberbanded if true, this is a rubberbanded overscroll; if false, this is an
-     *                       unrubberbanded motion to directly expand overscroll view (e.g
-     *                       expand
-     *                       QS)
-     */
-    void onOverscrollTopChanged(float amount, boolean isRubberbanded);
+        /**
+         * Notifies a listener that the overscroll has changed.
+         *
+         * @param amount         the amount of overscroll, in pixels
+         * @param isRubberbanded if true, this is a rubberbanded overscroll; if false, this is an
+         *                       unrubberbanded motion to directly expand overscroll view (e.g
+         *                       expand
+         *                       QS)
+         */
+        void onOverscrollTopChanged(float amount, boolean isRubberbanded);
 
-    /**
-     * Notify a listener that the scroller wants to escape from the scrolling motion and
-     * start a fling animation to the expanded or collapsed overscroll view (e.g expand the QS)
-     *
-     * @param velocity The velocity that the Scroller had when over flinging
-     * @param open     Should the fling open or close the overscroll view.
-     */
-    void flingTopOverscroll(float velocity, boolean open);
+        /**
+         * Notify a listener that the scroller wants to escape from the scrolling motion and
+         * start a fling animation to the expanded or collapsed overscroll view (e.g expand the QS)
+         *
+         * @param velocity The velocity that the Scroller had when over flinging
+         * @param open     Should the fling open or close the overscroll view.
+         */
+        void flingTopOverscroll(float velocity, boolean open);
     }
 
     @ShadeViewRefactor(RefactorComponent.SHADE_VIEW)
@@ -6252,7 +6283,9 @@
         }
     };
 
-    public HeadsUpTouchHelper.Callback getHeadsUpCallback() { return mHeadsUpCallback; }
+    public HeadsUpTouchHelper.Callback getHeadsUpCallback() {
+        return mHeadsUpCallback;
+    }
 
     void onGroupExpandChanged(ExpandableNotificationRow changedRow, boolean expanded) {
         boolean animated = mAnimationsEnabled && (mIsExpanded || changedRow.isPinned());
@@ -6357,15 +6390,25 @@
         return mLastSentExpandedHeight;
     }
 
-    /** Enum for selecting some or all notification rows (does not included non-notif views). */
+    /**
+     * Enum for selecting some or all notification rows (does not included non-notif views).
+     */
     @Retention(SOURCE)
     @IntDef({ROWS_ALL, ROWS_HIGH_PRIORITY, ROWS_GENTLE})
-    @interface SelectedRows {}
-    /** All rows representing notifs. */
+    @interface SelectedRows {
+    }
+
+    /**
+     * All rows representing notifs.
+     */
     public static final int ROWS_ALL = 0;
-    /** Only rows where entry.isHighPriority() is true. */
+    /**
+     * Only rows where entry.isHighPriority() is true.
+     */
     public static final int ROWS_HIGH_PRIORITY = 1;
-    /** Only rows where entry.isHighPriority() is false. */
+    /**
+     * Only rows where entry.isHighPriority() is false.
+     */
     public static final int ROWS_GENTLE = 2;
 
     interface ClearAllListener {
@@ -6380,4 +6423,16 @@
         void onAnimationEnd(
                 List<ExpandableNotificationRow> viewsToRemove, @SelectedRows int selectedRows);
     }
+
+    /**
+     *
+     */
+    public interface OnNotificationRemovedListener {
+        /**
+         *
+         * @param child
+         * @param isTransferInProgress
+         */
+        void onNotificationRemoved(ExpandableView child, boolean isTransferInProgress);
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
index 2d6d0a9c..4bcc0b6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationStackScrollLayoutController.java
@@ -182,10 +182,12 @@
     private NotificationStackScrollLayout mView;
     private boolean mFadeNotificationsOnDismiss;
     private NotificationSwipeHelper mSwipeHelper;
-    @Nullable private Boolean mHistoryEnabled;
+    @Nullable
+    private Boolean mHistoryEnabled;
     private int mBarState;
     private HeadsUpAppearanceController mHeadsUpAppearanceController;
     private final FeatureFlags mFeatureFlags;
+    private final boolean mUseRoundnessSourceTypes;
     private final NotificationTargetsHelper mNotificationTargetsHelper;
 
     private View mLongPressedView;
@@ -391,7 +393,7 @@
                     if (item != null) {
                         Point origin = provider.getRevealAnimationOrigin();
                         mNotificationGutsManager.openGuts(row, origin.x, origin.y, item);
-                    } else  {
+                    } else {
                         Log.e(TAG, "Provider has shouldShowGutsOnSnapOpen, but provided no "
                                 + "menu item in menuItemtoExposeOnSnap. Skipping.");
                     }
@@ -420,7 +422,7 @@
 
                 @Override
                 public void onSnooze(StatusBarNotification sbn,
-                        NotificationSwipeActionHelper.SnoozeOption snoozeOption) {
+                                     NotificationSwipeActionHelper.SnoozeOption snoozeOption) {
                     mCentralSurfaces.setNotificationSnoozed(sbn, snoozeOption);
                 }
 
@@ -544,7 +546,7 @@
 
                 @Override
                 public boolean updateSwipeProgress(View animView, boolean dismissable,
-                        float swipeProgress) {
+                                                   float swipeProgress) {
                     // Returning true prevents alpha fading.
                     return !mFadeNotificationsOnDismiss;
                 }
@@ -584,16 +586,22 @@
 
                 @Override
                 public void onHeadsUpPinned(NotificationEntry entry) {
-                    mNotificationRoundnessManager.updateView(entry.getRow(), false /* animate */);
+                    if (!mUseRoundnessSourceTypes) {
+                        mNotificationRoundnessManager.updateView(
+                                entry.getRow(),
+                                /* animate = */ false);
+                    }
                 }
 
                 @Override
                 public void onHeadsUpUnPinned(NotificationEntry entry) {
-                    ExpandableNotificationRow row = entry.getRow();
-                    // update the roundedness posted, because we might be animating away the
-                    // headsup soon, so no need to set the roundedness to 0 and then back to 1.
-                    row.post(() -> mNotificationRoundnessManager.updateView(row,
-                            true /* animate */));
+                    if (!mUseRoundnessSourceTypes) {
+                        ExpandableNotificationRow row = entry.getRow();
+                        // update the roundedness posted, because we might be animating away the
+                        // headsup soon, so no need to set the roundedness to 0 and then back to 1.
+                        row.post(() -> mNotificationRoundnessManager.updateView(row,
+                                true /* animate */));
+                    }
                 }
 
                 @Override
@@ -603,8 +611,10 @@
                     mView.setNumHeadsUp(numEntries);
                     mView.setTopHeadsUpEntry(topEntry);
                     generateHeadsUpAnimation(entry, isHeadsUp);
-                    ExpandableNotificationRow row = entry.getRow();
-                    mNotificationRoundnessManager.updateView(row, true /* animate */);
+                    if (!mUseRoundnessSourceTypes) {
+                        ExpandableNotificationRow row = entry.getRow();
+                        mNotificationRoundnessManager.updateView(row, true /* animate */);
+                    }
                 }
             };
 
@@ -697,6 +707,7 @@
         mSeenNotificationsProvider = seenNotificationsProvider;
         mShadeController = shadeController;
         mFeatureFlags = featureFlags;
+        mUseRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES);
         mNotificationTargetsHelper = notificationTargetsHelper;
         updateResources();
     }
@@ -763,8 +774,10 @@
         mLockscreenUserManager.addUserChangedListener(mLockscreenUserChangeListener);
 
         mFadeNotificationsOnDismiss = mFeatureFlags.isEnabled(Flags.NOTIFICATION_DISMISSAL_FADE);
-        mNotificationRoundnessManager.setOnRoundingChangedCallback(mView::invalidate);
-        mView.addOnExpandedHeightChangedListener(mNotificationRoundnessManager::setExpanded);
+        if (!mUseRoundnessSourceTypes) {
+            mNotificationRoundnessManager.setOnRoundingChangedCallback(mView::invalidate);
+            mView.addOnExpandedHeightChangedListener(mNotificationRoundnessManager::setExpanded);
+        }
 
         mVisibilityLocationProviderDelegator.setDelegate(this::isInVisibleLocation);
 
@@ -997,9 +1010,11 @@
                 Log.wtf(TAG, "isHistoryEnabled failed to initialize its value");
                 return false;
             }
-            mHistoryEnabled = historyEnabled =
-                    Settings.Secure.getIntForUser(mView.getContext().getContentResolver(),
-                    Settings.Secure.NOTIFICATION_HISTORY_ENABLED, 0, UserHandle.USER_CURRENT) == 1;
+            mHistoryEnabled = historyEnabled = Settings.Secure.getIntForUser(
+                    mView.getContext().getContentResolver(),
+                    Settings.Secure.NOTIFICATION_HISTORY_ENABLED,
+                    0,
+                    UserHandle.USER_CURRENT) == 1;
         }
         return historyEnabled;
     }
@@ -1029,7 +1044,7 @@
     }
 
     public void setOverScrollAmount(float amount, boolean onTop, boolean animate,
-            boolean cancelAnimators) {
+                                    boolean cancelAnimators) {
         mView.setOverScrollAmount(amount, onTop, animate, cancelAnimators);
     }
 
@@ -1207,7 +1222,7 @@
 
     /**
      * Update whether we should show the empty shade view ("no notifications" in the shade).
-     *
+     * <p>
      * When in split mode, notifications are always visible regardless of the state of the
      * QuickSettings panel. That being the case, empty view is always shown if the other conditions
      * are true.
@@ -1233,7 +1248,7 @@
 
     /**
      * @return true if {@link StatusBarStateController} is in transition to the KEYGUARD
-     *         and false otherwise.
+     * and false otherwise.
      */
     private boolean isInTransitionToKeyguard() {
         final int currentState = mStatusBarStateController.getState();
@@ -1265,7 +1280,9 @@
         mView.setExpandedHeight(expandedHeight);
     }
 
-    /** Sets the QS header. Used to check if a touch is within its bounds. */
+    /**
+     * Sets the QS header. Used to check if a touch is within its bounds.
+     */
     public void setQsHeader(ViewGroup view) {
         mView.setQsHeader(view);
     }
@@ -1328,7 +1345,7 @@
     public RemoteInputController.Delegate createDelegate() {
         return new RemoteInputController.Delegate() {
             public void setRemoteInputActive(NotificationEntry entry,
-                    boolean remoteInputActive) {
+                                             boolean remoteInputActive) {
                 mHeadsUpManager.setRemoteInputActive(entry, remoteInputActive);
                 entry.notifyHeightChanged(true /* needsAnimation */);
                 updateFooter();
@@ -1459,7 +1476,7 @@
     }
 
     private void onAnimationEnd(List<ExpandableNotificationRow> viewsToRemove,
-            @SelectedRows int selectedRows) {
+                                @SelectedRows int selectedRows) {
         if (selectedRows == ROWS_ALL) {
             mNotifCollection.dismissAllNotifications(
                     mLockscreenUserManager.getCurrentUserId());
@@ -1502,8 +1519,8 @@
 
     /**
      * @return the inset during the full shade transition, that needs to be added to the position
-     *         of the quick settings edge. This is relevant for media, that is transitioning
-     *         from the keyguard host to the quick settings one.
+     * of the quick settings edge. This is relevant for media, that is transitioning
+     * from the keyguard host to the quick settings one.
      */
     public int getFullShadeTransitionInset() {
         MediaContainerView view = mKeyguardMediaController.getSinglePaneContainer();
@@ -1517,10 +1534,10 @@
     /**
      * @param fraction The fraction of lockscreen to shade transition.
      *                 0f for all other states.
-     *
-     * Once the lockscreen to shade transition completes and the shade is 100% open,
-     * LockscreenShadeTransitionController resets amount and fraction to 0, where they remain
-     * until the next lockscreen-to-shade transition.
+     *                 <p>
+     *                 Once the lockscreen to shade transition completes and the shade is 100% open,
+     *                 LockscreenShadeTransitionController resets amount and fraction to 0, where
+     *                 they remain until the next lockscreen-to-shade transition.
      */
     public void setTransitionToFullShadeAmount(float fraction) {
         mView.setFractionToShade(fraction);
@@ -1533,7 +1550,9 @@
         mView.setExtraTopInsetForFullShadeTransition(overScrollAmount);
     }
 
-    /** */
+    /**
+     *
+     */
     public void setWillExpand(boolean willExpand) {
         mView.setWillExpand(willExpand);
     }
@@ -1549,7 +1568,7 @@
      * Set rounded rect clipping bounds on this view.
      */
     public void setRoundedClippingBounds(int left, int top, int right, int bottom, int topRadius,
-            int bottomRadius) {
+                                         int bottomRadius) {
         mView.setRoundedClippingBounds(left, top, right, bottom, topRadius, bottomRadius);
     }
 
@@ -1569,6 +1588,15 @@
     }
 
     /**
+     * Set the remove notification listener
+     * @param listener callback for notification removed
+     */
+    public void setOnNotificationRemovedListener(
+            NotificationStackScrollLayout.OnNotificationRemovedListener listener) {
+        mView.setOnNotificationRemovedListener(listener);
+    }
+
+    /**
      * Enum for UiEvent logged from this class
      */
     enum NotificationPanelEvent implements UiEventLogger.UiEventEnum {
@@ -1578,10 +1606,13 @@
         @UiEvent(doc = "User dismissed all silent notifications from notification panel.")
         DISMISS_SILENT_NOTIFICATIONS_PANEL(314);
         private final int mId;
+
         NotificationPanelEvent(int id) {
             mId = id;
         }
-        @Override public int getId() {
+
+        @Override
+        public int getId() {
             return mId;
         }
 
@@ -1722,8 +1753,12 @@
         @Override
         public void bindRow(ExpandableNotificationRow row) {
             row.setHeadsUpAnimatingAwayListener(animatingAway -> {
-                mNotificationRoundnessManager.updateView(row, false);
-                mHeadsUpAppearanceController.updateHeader(row.getEntry());
+                if (!mUseRoundnessSourceTypes) {
+                    mNotificationRoundnessManager.updateView(row, false);
+                }
+                NotificationEntry entry = row.getEntry();
+                mHeadsUpAppearanceController.updateHeader(entry);
+                mHeadsUpAppearanceController.updateHeadsUpAndPulsingRoundness(entry);
             });
         }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelper.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelper.java
index ee57411..aaf9300 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelper.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelper.java
@@ -20,6 +20,7 @@
 import static com.android.internal.jank.InteractionJankMonitor.CUJ_NOTIFICATION_SHADE_ROW_SWIPE;
 
 import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 import android.content.res.Resources;
 import android.graphics.Rect;
@@ -34,9 +35,11 @@
 import com.android.systemui.SwipeHelper;
 import com.android.systemui.dagger.qualifiers.Main;
 import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.plugins.FalsingManager;
 import com.android.systemui.plugins.statusbar.NotificationMenuRowPlugin;
 import com.android.systemui.plugins.statusbar.NotificationSwipeActionHelper;
+import com.android.systemui.statusbar.notification.SourceType;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.ExpandableView;
 
@@ -49,6 +52,7 @@
     @VisibleForTesting
     protected static final long COVER_MENU_DELAY = 4000;
     private static final String TAG = "NotificationSwipeHelper";
+    private static final SourceType SWIPE_DISMISS = SourceType.from("SwipeDismiss");
     private final Runnable mFalsingCheck;
     private View mTranslatingParentView;
     private View mMenuExposedView;
@@ -64,13 +68,21 @@
     private WeakReference<NotificationMenuRowPlugin> mCurrMenuRowRef;
     private boolean mIsExpanded;
     private boolean mPulsing;
+    private final NotificationRoundnessManager mNotificationRoundnessManager;
+    private final boolean mUseRoundnessSourceTypes;
 
     NotificationSwipeHelper(
-            Resources resources, ViewConfiguration viewConfiguration,
-            FalsingManager falsingManager, FeatureFlags featureFlags,
-            int swipeDirection, NotificationCallback callback,
-            NotificationMenuRowPlugin.OnMenuEventListener menuListener) {
+            Resources resources,
+            ViewConfiguration viewConfiguration,
+            FalsingManager falsingManager,
+            FeatureFlags featureFlags,
+            int swipeDirection,
+            NotificationCallback callback,
+            NotificationMenuRowPlugin.OnMenuEventListener menuListener,
+            NotificationRoundnessManager notificationRoundnessManager) {
         super(swipeDirection, callback, resources, viewConfiguration, falsingManager, featureFlags);
+        mNotificationRoundnessManager = notificationRoundnessManager;
+        mUseRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES);
         mMenuListener = menuListener;
         mCallback = callback;
         mFalsingCheck = () -> resetExposedMenuView(true /* animate */, true /* force */);
@@ -304,6 +316,33 @@
         handleMenuCoveredOrDismissed();
     }
 
+    @Override
+    protected void prepareDismissAnimation(View view, Animator anim) {
+        super.prepareDismissAnimation(view, anim);
+
+        if (mUseRoundnessSourceTypes
+                && view instanceof ExpandableNotificationRow
+                && mNotificationRoundnessManager.isClearAllInProgress()) {
+            ExpandableNotificationRow row = (ExpandableNotificationRow) view;
+            anim.addListener(new AnimatorListenerAdapter() {
+                @Override
+                public void onAnimationStart(Animator animation) {
+                    row.requestRoundness(/* top = */ 1f, /* bottom = */ 1f, SWIPE_DISMISS);
+                }
+
+                @Override
+                public void onAnimationCancel(Animator animation) {
+                    row.requestRoundnessReset(SWIPE_DISMISS);
+                }
+
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    row.requestRoundnessReset(SWIPE_DISMISS);
+                }
+            });
+        }
+    }
+
     @VisibleForTesting
     protected void superDismissChild(final View view, float velocity, boolean useAccelerateInterpolator) {
         super.dismissChild(view, velocity, useAccelerateInterpolator);
@@ -521,14 +560,17 @@
         private int mSwipeDirection;
         private NotificationCallback mNotificationCallback;
         private NotificationMenuRowPlugin.OnMenuEventListener mOnMenuEventListener;
+        private NotificationRoundnessManager mNotificationRoundnessManager;
 
         @Inject
         Builder(@Main Resources resources, ViewConfiguration viewConfiguration,
-                FalsingManager falsingManager, FeatureFlags featureFlags) {
+                FalsingManager falsingManager, FeatureFlags featureFlags,
+                NotificationRoundnessManager notificationRoundnessManager) {
             mResources = resources;
             mViewConfiguration = viewConfiguration;
             mFalsingManager = falsingManager;
             mFeatureFlags = featureFlags;
+            mNotificationRoundnessManager = notificationRoundnessManager;
         }
 
         Builder setSwipeDirection(int swipeDirection) {
@@ -549,7 +591,8 @@
 
         NotificationSwipeHelper build() {
             return new NotificationSwipeHelper(mResources, mViewConfiguration, mFalsingManager,
-                    mFeatureFlags, mSwipeDirection, mNotificationCallback, mOnMenuEventListener);
+                    mFeatureFlags, mSwipeDirection, mNotificationCallback, mOnMenuEventListener,
+                    mNotificationRoundnessManager);
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelper.kt b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelper.kt
index 991a14b..548d1a1 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelper.kt
@@ -20,8 +20,7 @@
 constructor(
     featureFlags: FeatureFlags,
 ) {
-    private val isNotificationGroupCornerEnabled =
-        featureFlags.isEnabled(Flags.NOTIFICATION_GROUP_CORNER)
+    private val useRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES)
 
     /**
      * This method looks for views that can be rounded (and implement [Roundable]) during a
@@ -48,7 +47,7 @@
         if (notificationParent != null && childrenContainer != null) {
             // We are inside a notification group
 
-            if (!isNotificationGroupCornerEnabled) {
+            if (!useRoundnessSourceTypes) {
                 return RoundableTargets(null, null, null)
             }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
index d8c6878..aff7b4c 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/notification/stack/StackScrollAlgorithm.java
@@ -31,6 +31,7 @@
 import com.android.systemui.animation.ShadeInterpolation;
 import com.android.systemui.statusbar.EmptyShadeView;
 import com.android.systemui.statusbar.NotificationShelf;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.SourceType;
 import com.android.systemui.statusbar.notification.row.ActivatableNotificationView;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
@@ -51,6 +52,7 @@
 
     private static final String TAG = "StackScrollAlgorithm";
     private static final Boolean DEBUG = false;
+    private static final SourceType STACK_SCROLL_ALGO = SourceType.from("StackScrollAlgorithm");
 
     private final ViewGroup mHostView;
     private float mPaddingBetweenElements;
@@ -70,6 +72,7 @@
     private float mQuickQsOffsetHeight;
     private float mSmallCornerRadius;
     private float mLargeCornerRadius;
+    private boolean mUseRoundnessSourceTypes;
 
     public StackScrollAlgorithm(
             Context context,
@@ -129,7 +132,7 @@
     }
 
     private void updateAlphaState(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                  AmbientState ambientState) {
         for (ExpandableView view : algorithmState.visibleChildren) {
             final ViewState viewState = view.getViewState();
 
@@ -229,7 +232,7 @@
     }
 
     private void getNotificationChildrenStates(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                               AmbientState ambientState) {
         int childCount = algorithmState.visibleChildren.size();
         for (int i = 0; i < childCount; i++) {
             ExpandableView v = algorithmState.visibleChildren.get(i);
@@ -241,7 +244,7 @@
     }
 
     private void updateSpeedBumpState(StackScrollAlgorithmState algorithmState,
-            int speedBumpIndex) {
+                                      int speedBumpIndex) {
         int childCount = algorithmState.visibleChildren.size();
         int belowSpeedBump = speedBumpIndex;
         for (int i = 0; i < childCount; i++) {
@@ -268,7 +271,7 @@
     }
 
     private void updateClipping(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                AmbientState ambientState) {
         float drawStart = ambientState.isOnKeyguard() ? 0
                 : ambientState.getStackY() - ambientState.getScrollY();
         float clipStart = 0;
@@ -314,7 +317,7 @@
      * Updates the dimmed, activated and hiding sensitive states of the children.
      */
     private void updateDimmedActivatedHideSensitive(AmbientState ambientState,
-            StackScrollAlgorithmState algorithmState) {
+                                                    StackScrollAlgorithmState algorithmState) {
         boolean dimmed = ambientState.isDimmed();
         boolean hideSensitive = ambientState.isHideSensitive();
         View activatedChild = ambientState.getActivatedChild();
@@ -408,7 +411,7 @@
     }
 
     private int updateNotGoneIndex(StackScrollAlgorithmState state, int notGoneIndex,
-            ExpandableView v) {
+                                   ExpandableView v) {
         ExpandableViewState viewState = v.getViewState();
         viewState.notGoneIndex = notGoneIndex;
         state.visibleChildren.add(v);
@@ -434,7 +437,7 @@
      * @param ambientState   The current ambient state
      */
     protected void updatePositionsForState(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                           AmbientState ambientState) {
         if (!ambientState.isOnKeyguard()
                 || (ambientState.isBypassEnabled() && ambientState.isPulseExpanding())) {
             algorithmState.mCurrentYPosition += mNotificationScrimPadding;
@@ -448,7 +451,7 @@
     }
 
     private void setLocation(ExpandableViewState expandableViewState, float currentYPosition,
-            int i) {
+                             int i) {
         expandableViewState.location = ExpandableViewState.LOCATION_MAIN_AREA;
         if (currentYPosition <= 0) {
             expandableViewState.location = ExpandableViewState.LOCATION_HIDDEN_TOP;
@@ -496,9 +499,13 @@
     }
 
     @VisibleForTesting
-    void maybeUpdateHeadsUpIsVisible(ExpandableViewState viewState, boolean isShadeExpanded,
-            boolean mustStayOnScreen, boolean topVisible, float viewEnd, float hunMax) {
-
+    void maybeUpdateHeadsUpIsVisible(
+            ExpandableViewState viewState,
+            boolean isShadeExpanded,
+            boolean mustStayOnScreen,
+            boolean topVisible,
+            float viewEnd,
+            float hunMax) {
         if (isShadeExpanded && mustStayOnScreen && topVisible) {
             viewState.headsUpIsVisible = viewEnd < hunMax;
         }
@@ -676,7 +683,7 @@
     }
 
     private void updatePulsingStates(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                     AmbientState ambientState) {
         int childCount = algorithmState.visibleChildren.size();
         for (int i = 0; i < childCount; i++) {
             View child = algorithmState.visibleChildren.get(i);
@@ -693,7 +700,7 @@
     }
 
     private void updateHeadsUpStates(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                     AmbientState ambientState) {
         int childCount = algorithmState.visibleChildren.size();
 
         // Move the tracked heads up into position during the appear animation, by interpolating
@@ -777,7 +784,7 @@
      */
     @VisibleForTesting
     void clampHunToTop(float quickQsOffsetHeight, float stackTranslation, float collapsedHeight,
-            ExpandableViewState viewState) {
+                       ExpandableViewState viewState) {
 
         final float newTranslation = Math.max(quickQsOffsetHeight + stackTranslation,
                 viewState.getYTranslation());
@@ -792,7 +799,7 @@
     // Pin HUN to bottom of expanded QS
     // while the rest of notifications are scrolled offscreen.
     private void clampHunToMaxTranslation(AmbientState ambientState, ExpandableNotificationRow row,
-            ExpandableViewState childState) {
+                                          ExpandableViewState childState) {
         float maxHeadsUpTranslation = ambientState.getMaxHeadsUpTranslation();
         final float maxShelfPosition = ambientState.getInnerHeight() + ambientState.getTopPadding()
                 + ambientState.getStackTranslation();
@@ -807,14 +814,19 @@
         // Animate pinned HUN bottom corners to and from original roundness.
         final float originalCornerRadius =
                 row.isLastInSection() ? 1f : (mSmallCornerRadius / mLargeCornerRadius);
-        final float roundness = computeCornerRoundnessForPinnedHun(mHostView.getHeight(),
+        final float bottomValue = computeCornerRoundnessForPinnedHun(mHostView.getHeight(),
                 ambientState.getStackY(), getMaxAllowedChildHeight(row), originalCornerRadius);
-        row.requestBottomRoundness(roundness, /* animate = */ false, SourceType.OnScroll);
+        if (mUseRoundnessSourceTypes) {
+            row.requestBottomRoundness(bottomValue, STACK_SCROLL_ALGO);
+            row.addOnDetachResetRoundness(STACK_SCROLL_ALGO);
+        } else {
+            row.requestBottomRoundness(bottomValue, LegacySourceType.OnScroll);
+        }
     }
 
     @VisibleForTesting
     float computeCornerRoundnessForPinnedHun(float hostViewHeight, float stackY,
-            float viewMaxHeight, float originalCornerRadius) {
+                                             float viewMaxHeight, float originalCornerRadius) {
 
         // Compute y where corner roundness should be in its original unpinned state.
         // We use view max height because the pinned collapsed HUN expands to max height
@@ -844,7 +856,7 @@
      * @param ambientState   The ambient state of the algorithm
      */
     private void updateZValuesForState(StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState) {
+                                       AmbientState ambientState) {
         int childCount = algorithmState.visibleChildren.size();
         float childrenOnTop = 0.0f;
 
@@ -876,9 +888,9 @@
      *                      previous HUNs whose Z positions are greater than 0.
      */
     protected float updateChildZValue(int i, float childrenOnTop,
-            StackScrollAlgorithmState algorithmState,
-            AmbientState ambientState,
-            boolean isTopHun) {
+                                      StackScrollAlgorithmState algorithmState,
+                                      AmbientState ambientState,
+                                      boolean isTopHun) {
         ExpandableView child = algorithmState.visibleChildren.get(i);
         ExpandableViewState childViewState = child.getViewState();
         float baseZ = ambientState.getBaseZHeight();
@@ -950,6 +962,14 @@
         this.mIsExpanded = isExpanded;
     }
 
+    /**
+     * Enable the support for rounded corner based on the SourceType
+     * @param enabled true if is supported
+     */
+    public void useRoundnessSourceTypes(boolean enabled) {
+        mUseRoundnessSourceTypes = enabled;
+    }
+
     public static class StackScrollAlgorithmState {
 
         /**
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceController.java
index 484441a..c217ab3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceController.java
@@ -19,11 +19,14 @@
 import static com.android.systemui.statusbar.phone.fragment.dagger.StatusBarFragmentModule.OPERATOR_NAME_FRAME_VIEW;
 
 import android.graphics.Rect;
+import android.util.MathUtils;
 import android.view.View;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.widget.ViewClippingUtil;
 import com.android.systemui.R;
+import com.android.systemui.flags.FeatureFlags;
+import com.android.systemui.flags.Flags;
 import com.android.systemui.plugins.DarkIconDispatcher;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.shade.NotificationPanelViewController;
@@ -32,8 +35,10 @@
 import com.android.systemui.statusbar.HeadsUpStatusBarView;
 import com.android.systemui.statusbar.StatusBarState;
 import com.android.systemui.statusbar.notification.NotificationWakeUpCoordinator;
+import com.android.systemui.statusbar.notification.SourceType;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
+import com.android.systemui.statusbar.notification.stack.NotificationRoundnessManager;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController;
 import com.android.systemui.statusbar.phone.fragment.dagger.StatusBarFragmentScope;
 import com.android.systemui.statusbar.policy.Clock;
@@ -51,6 +56,7 @@
 
 /**
  * Controls the appearance of heads up notifications in the icon area and the header itself.
+ * It also controls the roundness of the heads up notifications and the pulsing notifications.
  */
 @StatusBarFragmentScope
 public class HeadsUpAppearanceController extends ViewController<HeadsUpStatusBarView>
@@ -59,12 +65,17 @@
         NotificationWakeUpCoordinator.WakeUpListener {
     public static final int CONTENT_FADE_DURATION = 110;
     public static final int CONTENT_FADE_DELAY = 100;
+
+    private static final SourceType HEADS_UP = SourceType.from("HeadsUp");
+    private static final SourceType PULSING = SourceType.from("Pulsing");
     private final NotificationIconAreaController mNotificationIconAreaController;
     private final HeadsUpManagerPhone mHeadsUpManager;
     private final NotificationStackScrollLayoutController mStackScrollerController;
 
     private final DarkIconDispatcher mDarkIconDispatcher;
     private final NotificationPanelViewController mNotificationPanelViewController;
+    private final NotificationRoundnessManager mNotificationRoundnessManager;
+    private final boolean mUseRoundnessSourceTypes;
     private final Consumer<ExpandableNotificationRow>
             mSetTrackingHeadsUp = this::setTrackingHeadsUp;
     private final BiConsumer<Float, Float> mSetExpandedHeight = this::setAppearFraction;
@@ -105,11 +116,15 @@
             CommandQueue commandQueue,
             NotificationStackScrollLayoutController stackScrollerController,
             NotificationPanelViewController notificationPanelViewController,
+            NotificationRoundnessManager notificationRoundnessManager,
+            FeatureFlags featureFlags,
             HeadsUpStatusBarView headsUpStatusBarView,
             Clock clockView,
             @Named(OPERATOR_NAME_FRAME_VIEW) Optional<View> operatorNameViewOptional) {
         super(headsUpStatusBarView);
         mNotificationIconAreaController = notificationIconAreaController;
+        mNotificationRoundnessManager = notificationRoundnessManager;
+        mUseRoundnessSourceTypes = featureFlags.isEnabled(Flags.USE_ROUNDNESS_SOURCETYPES);
         mHeadsUpManager = headsUpManager;
 
         // We may be mid-HUN-expansion when this controller is re-created (for example, if the user
@@ -179,6 +194,7 @@
     public void onHeadsUpPinned(NotificationEntry entry) {
         updateTopEntry();
         updateHeader(entry);
+        updateHeadsUpAndPulsingRoundness(entry);
     }
 
     private void updateTopEntry() {
@@ -316,6 +332,7 @@
     public void onHeadsUpUnPinned(NotificationEntry entry) {
         updateTopEntry();
         updateHeader(entry);
+        updateHeadsUpAndPulsingRoundness(entry);
     }
 
     public void setAppearFraction(float expandedHeight, float appearFraction) {
@@ -346,7 +363,9 @@
         ExpandableNotificationRow previousTracked = mTrackedChild;
         mTrackedChild = trackedChild;
         if (previousTracked != null) {
-            updateHeader(previousTracked.getEntry());
+            NotificationEntry entry = previousTracked.getEntry();
+            updateHeader(entry);
+            updateHeadsUpAndPulsingRoundness(entry);
         }
     }
 
@@ -357,6 +376,7 @@
     private void updateHeadsUpHeaders() {
         mHeadsUpManager.getAllEntries().forEach(entry -> {
             updateHeader(entry);
+            updateHeadsUpAndPulsingRoundness(entry);
         });
     }
 
@@ -370,6 +390,31 @@
         row.setHeaderVisibleAmount(headerVisibleAmount);
     }
 
+    /**
+     * Update the HeadsUp and the Pulsing roundness based on current state
+     * @param entry target notification
+     */
+    public void updateHeadsUpAndPulsingRoundness(NotificationEntry entry) {
+        if (mUseRoundnessSourceTypes) {
+            ExpandableNotificationRow row = entry.getRow();
+            boolean isTrackedChild = row == mTrackedChild;
+            if (row.isPinned() || row.isHeadsUpAnimatingAway() || isTrackedChild) {
+                float roundness = MathUtils.saturate(1f - mAppearFraction);
+                row.requestRoundness(roundness, roundness, HEADS_UP);
+            } else {
+                row.requestRoundnessReset(HEADS_UP);
+            }
+            if (mNotificationRoundnessManager.shouldRoundNotificationPulsing()) {
+                if (row.showingPulsing()) {
+                    row.requestRoundness(/* top = */ 1f, /* bottom = */ 1f, PULSING);
+                } else {
+                    row.requestRoundnessReset(PULSING);
+                }
+            }
+        }
+    }
+
+
     @Override
     public void onDarkChanged(ArrayList<Rect> areas, float darkIntensity, int tint) {
         mView.onDarkChanged(areas, darkIntensity, tint);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
index 4969a1c..6811bf6 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileController.java
@@ -14,6 +14,8 @@
 
 package com.android.systemui.statusbar.phone;
 
+import androidx.annotation.MainThread;
+
 import com.android.systemui.statusbar.phone.ManagedProfileController.Callback;
 import com.android.systemui.statusbar.policy.CallbackController;
 
@@ -25,8 +27,20 @@
 
     boolean isWorkModeEnabled();
 
-    public interface Callback {
+    /**
+     * Callback to get updates about work profile status.
+     */
+    interface Callback {
+        /**
+         * Called when managed profile change is detected. This always runs on the main thread.
+         */
+        @MainThread
         void onManagedProfileChanged();
+
+        /**
+         * Called when managed profile removal is detected. This always runs on the main thread.
+         */
+        @MainThread
         void onManagedProfileRemoved();
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
index 4beb87d..abdf827 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImpl.java
@@ -33,31 +33,28 @@
 
 import javax.inject.Inject;
 
-/**
- */
 @SysUISingleton
 public class ManagedProfileControllerImpl implements ManagedProfileController {
 
     private final List<Callback> mCallbacks = new ArrayList<>();
-
+    private final UserTrackerCallback mUserTrackerCallback = new UserTrackerCallback();
     private final Context mContext;
     private final Executor mMainExecutor;
     private final UserManager mUserManager;
     private final UserTracker mUserTracker;
     private final LinkedList<UserInfo> mProfiles;
+
     private boolean mListening;
     private int mCurrentUser;
 
-    /**
-     */
     @Inject
     public ManagedProfileControllerImpl(Context context, @Main Executor mainExecutor,
-            UserTracker userTracker) {
+            UserTracker userTracker, UserManager userManager) {
         mContext = context;
         mMainExecutor = mainExecutor;
-        mUserManager = UserManager.get(mContext);
+        mUserManager = userManager;
         mUserTracker = userTracker;
-        mProfiles = new LinkedList<UserInfo>();
+        mProfiles = new LinkedList<>();
     }
 
     @Override
@@ -100,16 +97,22 @@
                 }
             }
             if (mProfiles.size() == 0 && hadProfile && (user == mCurrentUser)) {
-                for (Callback callback : mCallbacks) {
-                    callback.onManagedProfileRemoved();
-                }
+                mMainExecutor.execute(this::notifyManagedProfileRemoved);
             }
             mCurrentUser = user;
         }
     }
 
+    private void notifyManagedProfileRemoved() {
+        for (Callback callback : mCallbacks) {
+            callback.onManagedProfileRemoved();
+        }
+    }
+
     public boolean hasActiveProfile() {
-        if (!mListening) reloadManagedProfiles();
+        if (!mListening || mUserTracker.getUserId() != mCurrentUser) {
+            reloadManagedProfiles();
+        }
         synchronized (mProfiles) {
             return mProfiles.size() > 0;
         }
@@ -134,28 +137,28 @@
         mListening = listening;
         if (listening) {
             reloadManagedProfiles();
-            mUserTracker.addCallback(mUserChangedCallback, mMainExecutor);
+            mUserTracker.addCallback(mUserTrackerCallback, mMainExecutor);
         } else {
-            mUserTracker.removeCallback(mUserChangedCallback);
+            mUserTracker.removeCallback(mUserTrackerCallback);
         }
     }
 
-    private final UserTracker.Callback mUserChangedCallback =
-            new UserTracker.Callback() {
-                @Override
-                public void onUserChanged(int newUser, @NonNull Context userContext) {
-                    reloadManagedProfiles();
-                    for (Callback callback : mCallbacks) {
-                        callback.onManagedProfileChanged();
-                    }
-                }
+    private final class UserTrackerCallback implements UserTracker.Callback {
 
-                @Override
-                public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
-                    reloadManagedProfiles();
-                    for (Callback callback : mCallbacks) {
-                        callback.onManagedProfileChanged();
-                    }
-                }
-            };
+        @Override
+        public void onUserChanged(int newUser, @NonNull Context userContext) {
+            reloadManagedProfiles();
+            for (Callback callback : mCallbacks) {
+                callback.onManagedProfileChanged();
+            }
+        }
+
+        @Override
+        public void onProfilesChanged(@NonNull List<UserInfo> profiles) {
+            reloadManagedProfiles();
+            for (Callback callback : mCallbacks) {
+                callback.onManagedProfileChanged();
+            }
+        }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
index fb67f1a..c350c78 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/dagger/StatusBarPipelineModule.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.pipeline.dagger
 
+import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.log.table.TableLogBufferFactory
@@ -24,11 +25,12 @@
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModel
 import com.android.systemui.statusbar.pipeline.airplane.ui.viewmodel.AirplaneModeViewModelImpl
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
-import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepositoryImpl
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileRepositorySwitcher
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.UserSetupRepository
 import com.android.systemui.statusbar.pipeline.mobile.data.repository.UserSetupRepositoryImpl
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractor
 import com.android.systemui.statusbar.pipeline.mobile.domain.interactor.MobileIconsInteractorImpl
+import com.android.systemui.statusbar.pipeline.mobile.ui.MobileUiAdapter
 import com.android.systemui.statusbar.pipeline.mobile.util.MobileMappingsProxy
 import com.android.systemui.statusbar.pipeline.mobile.util.MobileMappingsProxyImpl
 import com.android.systemui.statusbar.pipeline.shared.data.repository.ConnectivityRepository
@@ -40,6 +42,8 @@
 import dagger.Binds
 import dagger.Module
 import dagger.Provides
+import dagger.multibindings.ClassKey
+import dagger.multibindings.IntoMap
 
 @Module
 abstract class StatusBarPipelineModule {
@@ -52,26 +56,28 @@
     @Binds
     abstract fun connectivityRepository(impl: ConnectivityRepositoryImpl): ConnectivityRepository
 
-    @Binds
-    abstract fun wifiRepository(impl: WifiRepositoryImpl): WifiRepository
+    @Binds abstract fun wifiRepository(impl: WifiRepositoryImpl): WifiRepository
 
     @Binds
     abstract fun wifiInteractor(impl: WifiInteractorImpl): WifiInteractor
 
     @Binds
     abstract fun mobileConnectionsRepository(
-        impl: MobileConnectionsRepositoryImpl
+        impl: MobileRepositorySwitcher
     ): MobileConnectionsRepository
 
-    @Binds
-    abstract fun userSetupRepository(impl: UserSetupRepositoryImpl): UserSetupRepository
+    @Binds abstract fun userSetupRepository(impl: UserSetupRepositoryImpl): UserSetupRepository
 
-    @Binds
-    abstract fun mobileMappingsProxy(impl: MobileMappingsProxyImpl): MobileMappingsProxy
+    @Binds abstract fun mobileMappingsProxy(impl: MobileMappingsProxyImpl): MobileMappingsProxy
 
     @Binds
     abstract fun mobileIconsInteractor(impl: MobileIconsInteractorImpl): MobileIconsInteractor
 
+    @Binds
+    @IntoMap
+    @ClassKey(MobileUiAdapter::class)
+    abstract fun bindFeature(impl: MobileUiAdapter): CoreStartable
+
     @Module
     companion object {
         @JvmStatic
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/DataConnectionState.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/DataConnectionState.kt
index da87f73..5479b92 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/DataConnectionState.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/model/DataConnectionState.kt
@@ -20,6 +20,7 @@
 import android.telephony.TelephonyManager.DATA_CONNECTING
 import android.telephony.TelephonyManager.DATA_DISCONNECTED
 import android.telephony.TelephonyManager.DATA_DISCONNECTING
+import android.telephony.TelephonyManager.DATA_UNKNOWN
 import android.telephony.TelephonyManager.DataState
 
 /** Internal enum representation of the telephony data connection states */
@@ -28,6 +29,7 @@
     Connecting(DATA_CONNECTING),
     Disconnected(DATA_DISCONNECTED),
     Disconnecting(DATA_DISCONNECTING),
+    Unknown(DATA_UNKNOWN),
 }
 
 fun @receiver:DataState Int.toDataConnectionType(): DataConnectionState =
@@ -36,5 +38,6 @@
         DATA_CONNECTING -> DataConnectionState.Connecting
         DATA_DISCONNECTED -> DataConnectionState.Disconnected
         DATA_DISCONNECTING -> DataConnectionState.Disconnecting
-        else -> throw IllegalArgumentException("unknown data state received")
+        DATA_UNKNOWN -> DataConnectionState.Unknown
+        else -> throw IllegalArgumentException("unknown data state received $this")
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
index 581842b..f094563 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepository.kt
@@ -16,44 +16,13 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.data.repository
 
-import android.content.Context
-import android.database.ContentObserver
-import android.provider.Settings.Global
-import android.telephony.CellSignalStrength
-import android.telephony.CellSignalStrengthCdma
-import android.telephony.ServiceState
-import android.telephony.SignalStrength
 import android.telephony.SubscriptionInfo
 import android.telephony.SubscriptionManager
 import android.telephony.TelephonyCallback
-import android.telephony.TelephonyDisplayInfo
-import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NONE
 import android.telephony.TelephonyManager
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
-import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
-import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
-import com.android.systemui.statusbar.pipeline.mobile.data.model.toDataConnectionType
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger.Companion.logOutputChange
-import com.android.systemui.util.settings.GlobalSettings
-import java.lang.IllegalStateException
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.asExecutor
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.mapLatest
-import kotlinx.coroutines.flow.merge
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 
 /**
  * Every mobile line of service can be identified via a [SubscriptionInfo] object. We set up a
@@ -80,183 +49,3 @@
      */
     val isDefaultDataSubscription: StateFlow<Boolean>
 }
-
-@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
-@OptIn(ExperimentalCoroutinesApi::class)
-class MobileConnectionRepositoryImpl(
-    private val context: Context,
-    private val subId: Int,
-    private val telephonyManager: TelephonyManager,
-    private val globalSettings: GlobalSettings,
-    defaultDataSubId: StateFlow<Int>,
-    globalMobileDataSettingChangedEvent: Flow<Unit>,
-    bgDispatcher: CoroutineDispatcher,
-    logger: ConnectivityPipelineLogger,
-    scope: CoroutineScope,
-) : MobileConnectionRepository {
-    init {
-        if (telephonyManager.subscriptionId != subId) {
-            throw IllegalStateException(
-                "TelephonyManager should be created with subId($subId). " +
-                    "Found ${telephonyManager.subscriptionId} instead."
-            )
-        }
-    }
-
-    private val telephonyCallbackEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
-
-    override val subscriptionModelFlow: StateFlow<MobileSubscriptionModel> = run {
-        var state = MobileSubscriptionModel()
-        conflatedCallbackFlow {
-                // TODO (b/240569788): log all of these into the connectivity logger
-                val callback =
-                    object :
-                        TelephonyCallback(),
-                        TelephonyCallback.ServiceStateListener,
-                        TelephonyCallback.SignalStrengthsListener,
-                        TelephonyCallback.DataConnectionStateListener,
-                        TelephonyCallback.DataActivityListener,
-                        TelephonyCallback.CarrierNetworkListener,
-                        TelephonyCallback.DisplayInfoListener {
-                        override fun onServiceStateChanged(serviceState: ServiceState) {
-                            state = state.copy(isEmergencyOnly = serviceState.isEmergencyOnly)
-                            trySend(state)
-                        }
-
-                        override fun onSignalStrengthsChanged(signalStrength: SignalStrength) {
-                            val cdmaLevel =
-                                signalStrength
-                                    .getCellSignalStrengths(CellSignalStrengthCdma::class.java)
-                                    .let { strengths ->
-                                        if (!strengths.isEmpty()) {
-                                            strengths[0].level
-                                        } else {
-                                            CellSignalStrength.SIGNAL_STRENGTH_NONE_OR_UNKNOWN
-                                        }
-                                    }
-
-                            val primaryLevel = signalStrength.level
-
-                            state =
-                                state.copy(
-                                    cdmaLevel = cdmaLevel,
-                                    primaryLevel = primaryLevel,
-                                    isGsm = signalStrength.isGsm,
-                                )
-                            trySend(state)
-                        }
-
-                        override fun onDataConnectionStateChanged(
-                            dataState: Int,
-                            networkType: Int
-                        ) {
-                            state =
-                                state.copy(dataConnectionState = dataState.toDataConnectionType())
-                            trySend(state)
-                        }
-
-                        override fun onDataActivity(direction: Int) {
-                            state = state.copy(dataActivityDirection = direction)
-                            trySend(state)
-                        }
-
-                        override fun onCarrierNetworkChange(active: Boolean) {
-                            state = state.copy(carrierNetworkChangeActive = active)
-                            trySend(state)
-                        }
-
-                        override fun onDisplayInfoChanged(
-                            telephonyDisplayInfo: TelephonyDisplayInfo
-                        ) {
-                            val networkType =
-                                if (
-                                    telephonyDisplayInfo.overrideNetworkType ==
-                                        OVERRIDE_NETWORK_TYPE_NONE
-                                ) {
-                                    DefaultNetworkType(telephonyDisplayInfo.networkType)
-                                } else {
-                                    OverrideNetworkType(telephonyDisplayInfo.overrideNetworkType)
-                                }
-                            state = state.copy(resolvedNetworkType = networkType)
-                            trySend(state)
-                        }
-                    }
-                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
-                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
-            }
-            .onEach { telephonyCallbackEvent.tryEmit(Unit) }
-            .logOutputChange(logger, "MobileSubscriptionModel")
-            .stateIn(scope, SharingStarted.WhileSubscribed(), state)
-    }
-
-    /** Produces whenever the mobile data setting changes for this subId */
-    private val localMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
-        val observer =
-            object : ContentObserver(null) {
-                override fun onChange(selfChange: Boolean) {
-                    trySend(Unit)
-                }
-            }
-
-        globalSettings.registerContentObserver(
-            globalSettings.getUriFor("${Global.MOBILE_DATA}$subId"),
-            /* notifyForDescendants */ true,
-            observer
-        )
-
-        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
-    }
-
-    /**
-     * There are a few cases where we will need to poll [TelephonyManager] so we can update some
-     * internal state where callbacks aren't provided. Any of those events should be merged into
-     * this flow, which can be used to trigger the polling.
-     */
-    private val telephonyPollingEvent: Flow<Unit> =
-        merge(
-            telephonyCallbackEvent,
-            localMobileDataSettingChangedEvent,
-            globalMobileDataSettingChangedEvent,
-        )
-
-    override val dataEnabled: StateFlow<Boolean> =
-        telephonyPollingEvent
-            .mapLatest { dataConnectionAllowed() }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), dataConnectionAllowed())
-
-    private fun dataConnectionAllowed(): Boolean = telephonyManager.isDataConnectionAllowed
-
-    override val isDefaultDataSubscription: StateFlow<Boolean> =
-        defaultDataSubId
-            .mapLatest { it == subId }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), defaultDataSubId.value == subId)
-
-    class Factory
-    @Inject
-    constructor(
-        private val context: Context,
-        private val telephonyManager: TelephonyManager,
-        private val logger: ConnectivityPipelineLogger,
-        private val globalSettings: GlobalSettings,
-        @Background private val bgDispatcher: CoroutineDispatcher,
-        @Application private val scope: CoroutineScope,
-    ) {
-        fun build(
-            subId: Int,
-            defaultDataSubId: StateFlow<Int>,
-            globalMobileDataSettingChangedEvent: Flow<Unit>,
-        ): MobileConnectionRepository {
-            return MobileConnectionRepositoryImpl(
-                context,
-                subId,
-                telephonyManager.createForSubscriptionId(subId),
-                globalSettings,
-                defaultDataSubId,
-                globalMobileDataSettingChangedEvent,
-                bgDispatcher,
-                logger,
-                scope,
-            )
-        }
-    }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
index c3c1f14..14200f0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepository.kt
@@ -16,53 +16,14 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.data.repository
 
-import android.annotation.SuppressLint
-import android.content.Context
-import android.content.IntentFilter
-import android.database.ContentObserver
-import android.net.ConnectivityManager
-import android.net.ConnectivityManager.NetworkCallback
-import android.net.Network
-import android.net.NetworkCapabilities
-import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
-import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
 import android.provider.Settings
-import android.provider.Settings.Global.MOBILE_DATA
-import android.telephony.CarrierConfigManager
 import android.telephony.SubscriptionInfo
 import android.telephony.SubscriptionManager
-import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
-import android.telephony.TelephonyCallback
-import android.telephony.TelephonyCallback.ActiveDataSubscriptionIdListener
-import android.telephony.TelephonyManager
-import androidx.annotation.VisibleForTesting
-import com.android.internal.telephony.PhoneConstants
 import com.android.settingslib.mobile.MobileMappings
 import com.android.settingslib.mobile.MobileMappings.Config
-import com.android.systemui.broadcast.BroadcastDispatcher
-import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
-import com.android.systemui.dagger.SysUISingleton
-import com.android.systemui.dagger.qualifiers.Application
-import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
-import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
-import com.android.systemui.util.settings.GlobalSettings
-import javax.inject.Inject
-import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.asExecutor
-import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.distinctUntilChanged
-import kotlinx.coroutines.flow.mapLatest
-import kotlinx.coroutines.flow.merge
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
-import kotlinx.coroutines.withContext
 
 /**
  * Repo for monitoring the complete active subscription info list, to be consumed and filtered based
@@ -90,202 +51,3 @@
     /** Observe changes to the [Settings.Global.MOBILE_DATA] setting */
     val globalMobileDataSettingChangedEvent: Flow<Unit>
 }
-
-@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
-@OptIn(ExperimentalCoroutinesApi::class)
-@SysUISingleton
-class MobileConnectionsRepositoryImpl
-@Inject
-constructor(
-    private val connectivityManager: ConnectivityManager,
-    private val subscriptionManager: SubscriptionManager,
-    private val telephonyManager: TelephonyManager,
-    private val logger: ConnectivityPipelineLogger,
-    broadcastDispatcher: BroadcastDispatcher,
-    private val globalSettings: GlobalSettings,
-    private val context: Context,
-    @Background private val bgDispatcher: CoroutineDispatcher,
-    @Application private val scope: CoroutineScope,
-    private val mobileConnectionRepositoryFactory: MobileConnectionRepositoryImpl.Factory
-) : MobileConnectionsRepository {
-    private val subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
-
-    /**
-     * State flow that emits the set of mobile data subscriptions, each represented by its own
-     * [SubscriptionInfo]. We probably only need the [SubscriptionInfo.getSubscriptionId] of each
-     * info object, but for now we keep track of the infos themselves.
-     */
-    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : SubscriptionManager.OnSubscriptionsChangedListener() {
-                        override fun onSubscriptionsChanged() {
-                            trySend(Unit)
-                        }
-                    }
-
-                subscriptionManager.addOnSubscriptionsChangedListener(
-                    bgDispatcher.asExecutor(),
-                    callback,
-                )
-
-                awaitClose { subscriptionManager.removeOnSubscriptionsChangedListener(callback) }
-            }
-            .mapLatest { fetchSubscriptionsList() }
-            .onEach { infos -> dropUnusedReposFromCache(infos) }
-            .stateIn(scope, started = SharingStarted.WhileSubscribed(), listOf())
-
-    /** StateFlow that keeps track of the current active mobile data subscription */
-    override val activeMobileDataSubscriptionId: StateFlow<Int> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : TelephonyCallback(), ActiveDataSubscriptionIdListener {
-                        override fun onActiveDataSubscriptionIdChanged(subId: Int) {
-                            trySend(subId)
-                        }
-                    }
-
-                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
-                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
-            }
-            .stateIn(scope, started = SharingStarted.WhileSubscribed(), INVALID_SUBSCRIPTION_ID)
-
-    private val defaultDataSubIdChangeEvent: MutableSharedFlow<Unit> =
-        MutableSharedFlow(extraBufferCapacity = 1)
-
-    override val defaultDataSubId: StateFlow<Int> =
-        broadcastDispatcher
-            .broadcastFlow(
-                IntentFilter(TelephonyManager.ACTION_DEFAULT_DATA_SUBSCRIPTION_CHANGED)
-            ) { intent, _ ->
-                intent.getIntExtra(PhoneConstants.SUBSCRIPTION_KEY, INVALID_SUBSCRIPTION_ID)
-            }
-            .distinctUntilChanged()
-            .onEach { defaultDataSubIdChangeEvent.tryEmit(Unit) }
-            .stateIn(
-                scope,
-                SharingStarted.WhileSubscribed(),
-                SubscriptionManager.getDefaultDataSubscriptionId()
-            )
-
-    private val carrierConfigChangedEvent =
-        broadcastDispatcher.broadcastFlow(
-            IntentFilter(CarrierConfigManager.ACTION_CARRIER_CONFIG_CHANGED)
-        )
-
-    /**
-     * [Config] is an object that tracks relevant configuration flags for a given subscription ID.
-     * In the case of [MobileMappings], it's hard-coded to check the default data subscription's
-     * config, so this will apply to every icon that we care about.
-     *
-     * Relevant bits in the config are things like
-     * [CarrierConfigManager.KEY_SHOW_4G_FOR_LTE_DATA_ICON_BOOL]
-     *
-     * This flow will produce whenever the default data subscription or the carrier config changes.
-     */
-    override val defaultDataSubRatConfig: StateFlow<Config> =
-        merge(defaultDataSubIdChangeEvent, carrierConfigChangedEvent)
-            .mapLatest { Config.readConfig(context) }
-            .stateIn(
-                scope,
-                SharingStarted.WhileSubscribed(),
-                initialValue = Config.readConfig(context)
-            )
-
-    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
-        if (!isValidSubId(subId)) {
-            throw IllegalArgumentException(
-                "subscriptionId $subId is not in the list of valid subscriptions"
-            )
-        }
-
-        return subIdRepositoryCache[subId]
-            ?: createRepositoryForSubId(subId).also { subIdRepositoryCache[subId] = it }
-    }
-
-    /**
-     * In single-SIM devices, the [MOBILE_DATA] setting is phone-wide. For multi-SIM, the individual
-     * connection repositories also observe the URI for [MOBILE_DATA] + subId.
-     */
-    override val globalMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
-        val observer =
-            object : ContentObserver(null) {
-                override fun onChange(selfChange: Boolean) {
-                    trySend(Unit)
-                }
-            }
-
-        globalSettings.registerContentObserver(
-            globalSettings.getUriFor(MOBILE_DATA),
-            true,
-            observer
-        )
-
-        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
-    }
-
-    @SuppressLint("MissingPermission")
-    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
-        conflatedCallbackFlow {
-                val callback =
-                    object : NetworkCallback(FLAG_INCLUDE_LOCATION_INFO) {
-                        override fun onLost(network: Network) {
-                            // Send a disconnected model when lost. Maybe should create a sealed
-                            // type or null here?
-                            trySend(MobileConnectivityModel())
-                        }
-
-                        override fun onCapabilitiesChanged(
-                            network: Network,
-                            caps: NetworkCapabilities
-                        ) {
-                            trySend(
-                                MobileConnectivityModel(
-                                    isConnected = caps.hasTransport(TRANSPORT_CELLULAR),
-                                    isValidated = caps.hasCapability(NET_CAPABILITY_VALIDATED),
-                                )
-                            )
-                        }
-                    }
-
-                connectivityManager.registerDefaultNetworkCallback(callback)
-
-                awaitClose { connectivityManager.unregisterNetworkCallback(callback) }
-            }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), MobileConnectivityModel())
-
-    private fun isValidSubId(subId: Int): Boolean {
-        subscriptionsFlow.value.forEach {
-            if (it.subscriptionId == subId) {
-                return true
-            }
-        }
-
-        return false
-    }
-
-    @VisibleForTesting fun getSubIdRepoCache() = subIdRepositoryCache
-
-    private fun createRepositoryForSubId(subId: Int): MobileConnectionRepository {
-        return mobileConnectionRepositoryFactory.build(
-            subId,
-            defaultDataSubId,
-            globalMobileDataSettingChangedEvent,
-        )
-    }
-
-    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
-        // Remove any connection repository from the cache that isn't in the new set of IDs. They
-        // will get garbage collected once their subscribers go away
-        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
-
-        subIdRepositoryCache.keys.forEach {
-            if (!currentValidSubscriptionIds.contains(it)) {
-                subIdRepositoryCache.remove(it)
-            }
-        }
-    }
-
-    private suspend fun fetchSubscriptionsList(): List<SubscriptionInfo> =
-        withContext(bgDispatcher) { subscriptionManager.completeActiveSubscriptionInfoList }
-}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
new file mode 100644
index 0000000..e214005
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcher.kt
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository
+
+import android.os.Bundle
+import android.telephony.SubscriptionInfo
+import androidx.annotation.VisibleForTesting
+import com.android.settingslib.mobile.MobileMappings
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileConnectionsRepositoryImpl
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * A provider for the [MobileConnectionsRepository] interface that can choose between the Demo and
+ * Prod concrete implementations at runtime. It works by defining a base flow, [activeRepo], which
+ * switches based on the latest information from [DemoModeController], and switches every flow in
+ * the interface to point to the currently-active provider. This allows us to put the demo mode
+ * interface in its own repository, completely separate from the real version, while still using all
+ * of the prod implementations for the rest of the pipeline (interactors and onward). Looks
+ * something like this:
+ *
+ * ```
+ * RealRepository
+ *                 │
+ *                 ├──►RepositorySwitcher──►RealInteractor──►RealViewModel
+ *                 │
+ * DemoRepository
+ * ```
+ *
+ * NOTE: because the UI layer for mobile icons relies on a nested-repository structure, it is likely
+ * that we will have to drain the subscription list whenever demo mode changes. Otherwise if a real
+ * subscription list [1] is replaced with a demo subscription list [1], the view models will not see
+ * a change (due to `distinctUntilChanged`) and will not refresh their data providers to the demo
+ * implementation.
+ */
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+class MobileRepositorySwitcher
+@Inject
+constructor(
+    @Application scope: CoroutineScope,
+    val realRepository: MobileConnectionsRepositoryImpl,
+    val demoMobileConnectionsRepository: DemoMobileConnectionsRepository,
+    demoModeController: DemoModeController,
+) : MobileConnectionsRepository {
+
+    val isDemoMode: StateFlow<Boolean> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : DemoMode {
+                        override fun dispatchDemoCommand(command: String?, args: Bundle?) {
+                            // Nothing, we just care about on/off
+                        }
+
+                        override fun onDemoModeStarted() {
+                            demoMobileConnectionsRepository.startProcessingCommands()
+                            trySend(true)
+                        }
+
+                        override fun onDemoModeFinished() {
+                            demoMobileConnectionsRepository.stopProcessingCommands()
+                            trySend(false)
+                        }
+                    }
+
+                demoModeController.addCallback(callback)
+                awaitClose { demoModeController.removeCallback(callback) }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), demoModeController.isInDemoMode)
+
+    // Convenient definition flow for the currently active repo (based on demo mode or not)
+    @VisibleForTesting
+    internal val activeRepo: StateFlow<MobileConnectionsRepository> =
+        isDemoMode
+            .mapLatest { demoMode ->
+                if (demoMode) {
+                    demoMobileConnectionsRepository
+                } else {
+                    realRepository
+                }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), realRepository)
+
+    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
+        activeRepo
+            .flatMapLatest { it.subscriptionsFlow }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.subscriptionsFlow.value
+            )
+
+    override val activeMobileDataSubscriptionId: StateFlow<Int> =
+        activeRepo
+            .flatMapLatest { it.activeMobileDataSubscriptionId }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.activeMobileDataSubscriptionId.value
+            )
+
+    override val defaultDataSubRatConfig: StateFlow<MobileMappings.Config> =
+        activeRepo
+            .flatMapLatest { it.defaultDataSubRatConfig }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.defaultDataSubRatConfig.value
+            )
+
+    override val defaultDataSubId: StateFlow<Int> =
+        activeRepo
+            .flatMapLatest { it.defaultDataSubId }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), realRepository.defaultDataSubId.value)
+
+    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
+        activeRepo
+            .flatMapLatest { it.defaultMobileNetworkConnectivity }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                realRepository.defaultMobileNetworkConnectivity.value
+            )
+
+    override val globalMobileDataSettingChangedEvent: Flow<Unit> =
+        activeRepo.flatMapLatest { it.globalMobileDataSettingChangedEvent }
+
+    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
+        if (isDemoMode.value) {
+            return demoMobileConnectionsRepository.getRepoForSubId(subId)
+        }
+        return realRepository.getRepoForSubId(subId)
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
new file mode 100644
index 0000000..5f2feb2
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepository.kt
@@ -0,0 +1,277 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.demo
+
+import android.content.Context
+import android.telephony.Annotation
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
+import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NR_ADVANCED
+import android.telephony.TelephonyManager.NETWORK_TYPE_GSM
+import android.telephony.TelephonyManager.NETWORK_TYPE_LTE
+import android.telephony.TelephonyManager.NETWORK_TYPE_NR
+import android.telephony.TelephonyManager.NETWORK_TYPE_UMTS
+import android.telephony.TelephonyManager.NETWORK_TYPE_UNKNOWN
+import android.util.Log
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.MobileMappings
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.ResolvedNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.Mobile
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
+
+/** This repository vends out data based on demo mode commands */
+@OptIn(ExperimentalCoroutinesApi::class)
+class DemoMobileConnectionsRepository
+@Inject
+constructor(
+    private val dataSource: DemoModeMobileConnectionDataSource,
+    @Application private val scope: CoroutineScope,
+    context: Context,
+) : MobileConnectionsRepository {
+
+    private var demoCommandJob: Job? = null
+
+    private val connectionRepoCache = mutableMapOf<Int, DemoMobileConnectionRepository>()
+    private val subscriptionInfoCache = mutableMapOf<Int, SubscriptionInfo>()
+    val demoModeFinishedEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+
+    private val _subscriptions = MutableStateFlow<List<SubscriptionInfo>>(listOf())
+    override val subscriptionsFlow =
+        _subscriptions
+            .onEach { infos -> dropUnusedReposFromCache(infos) }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), _subscriptions.value)
+
+    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
+        // Remove any connection repository from the cache that isn't in the new set of IDs. They
+        // will get garbage collected once their subscribers go away
+        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
+
+        connectionRepoCache.keys.forEach {
+            if (!currentValidSubscriptionIds.contains(it)) {
+                connectionRepoCache.remove(it)
+            }
+        }
+    }
+
+    private fun maybeCreateSubscription(subId: Int) {
+        if (!subscriptionInfoCache.containsKey(subId)) {
+            createSubscriptionForSubId(subId, subId).also { subscriptionInfoCache[subId] = it }
+
+            _subscriptions.value = subscriptionInfoCache.values.toList()
+        }
+    }
+
+    /** Mimics the old NetworkControllerImpl for now */
+    private fun createSubscriptionForSubId(subId: Int, slotIndex: Int): SubscriptionInfo {
+        return SubscriptionInfo(
+            subId,
+            "",
+            slotIndex,
+            "",
+            "",
+            0,
+            0,
+            "",
+            0,
+            null,
+            null,
+            null,
+            "",
+            false,
+            null,
+            null,
+        )
+    }
+
+    // TODO(b/261029387): add a command for this value
+    override val activeMobileDataSubscriptionId =
+        subscriptionsFlow
+            .mapLatest { infos ->
+                // For now, active is just the first in the list
+                infos.firstOrNull()?.subscriptionId ?: INVALID_SUBSCRIPTION_ID
+            }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                subscriptionsFlow.value.firstOrNull()?.subscriptionId ?: INVALID_SUBSCRIPTION_ID
+            )
+
+    /** Demo mode doesn't currently support modifications to the mobile mappings */
+    override val defaultDataSubRatConfig =
+        MutableStateFlow(MobileMappings.Config.readConfig(context))
+
+    // TODO(b/261029387): add a command for this value
+    override val defaultDataSubId =
+        activeMobileDataSubscriptionId.stateIn(
+            scope,
+            SharingStarted.WhileSubscribed(),
+            INVALID_SUBSCRIPTION_ID
+        )
+
+    // TODO(b/261029387): not yet supported
+    override val defaultMobileNetworkConnectivity = MutableStateFlow(MobileConnectivityModel())
+
+    override fun getRepoForSubId(subId: Int): DemoMobileConnectionRepository {
+        return connectionRepoCache[subId]
+            ?: DemoMobileConnectionRepository(subId).also { connectionRepoCache[subId] = it }
+    }
+
+    override val globalMobileDataSettingChangedEvent = MutableStateFlow(Unit)
+
+    fun startProcessingCommands() {
+        demoCommandJob =
+            scope.launch {
+                dataSource.mobileEvents.filterNotNull().collect { event -> processEvent(event) }
+            }
+    }
+
+    fun stopProcessingCommands() {
+        demoCommandJob?.cancel()
+        _subscriptions.value = listOf()
+        connectionRepoCache.clear()
+        subscriptionInfoCache.clear()
+    }
+
+    private fun processEvent(event: FakeNetworkEventModel) {
+        when (event) {
+            is Mobile -> {
+                processEnabledMobileState(event)
+            }
+            is MobileDisabled -> {
+                processDisabledMobileState(event)
+            }
+        }
+    }
+
+    private fun processEnabledMobileState(state: Mobile) {
+        // get or create the connection repo, and set its values
+        val subId = state.subId ?: DEFAULT_SUB_ID
+        maybeCreateSubscription(subId)
+
+        val connection = getRepoForSubId(subId)
+        // This is always true here, because we split out disabled states at the data-source level
+        connection.dataEnabled.value = true
+        connection.isDefaultDataSubscription.value = state.dataType != null
+
+        connection.subscriptionModelFlow.value = state.toMobileSubscriptionModel()
+    }
+
+    private fun processDisabledMobileState(state: MobileDisabled) {
+        if (_subscriptions.value.isEmpty()) {
+            // Nothing to do here
+            return
+        }
+
+        val subId =
+            state.subId
+                ?: run {
+                    // For sake of usability, we can allow for no subId arg if there is only one
+                    // subscription
+                    if (_subscriptions.value.size > 1) {
+                        Log.d(
+                            TAG,
+                            "processDisabledMobileState: Unable to infer subscription to " +
+                                "disable. Specify subId using '-e slot <subId>'" +
+                                "Known subIds: [${subIdsString()}]"
+                        )
+                        return
+                    }
+
+                    // Use the only existing subscription as our arg, since there is only one
+                    _subscriptions.value[0].subscriptionId
+                }
+
+        removeSubscription(subId)
+    }
+
+    private fun removeSubscription(subId: Int) {
+        val currentSubscriptions = _subscriptions.value
+        subscriptionInfoCache.remove(subId)
+        _subscriptions.value = currentSubscriptions.filter { it.subscriptionId != subId }
+    }
+
+    private fun subIdsString(): String =
+        _subscriptions.value.joinToString(",") { it.subscriptionId.toString() }
+
+    companion object {
+        private const val TAG = "DemoMobileConnectionsRepo"
+
+        private const val DEFAULT_SUB_ID = 1
+    }
+}
+
+private fun Mobile.toMobileSubscriptionModel(): MobileSubscriptionModel {
+    return MobileSubscriptionModel(
+        isEmergencyOnly = false, // TODO(b/261029387): not yet supported
+        isGsm = false, // TODO(b/261029387): not yet supported
+        cdmaLevel = level ?: 0,
+        primaryLevel = level ?: 0,
+        dataConnectionState = DataConnectionState.Connected, // TODO(b/261029387): not yet supported
+        dataActivityDirection = activity,
+        carrierNetworkChangeActive = carrierNetworkChange,
+        // TODO(b/261185097): once mobile mappings can be mocked at this layer, we can build our
+        //  own demo map
+        resolvedNetworkType = dataType.toResolvedNetworkType()
+    )
+}
+
+@Annotation.NetworkType
+private fun SignalIcon.MobileIconGroup?.toNetworkType(): Int =
+    when (this) {
+        TelephonyIcons.THREE_G -> NETWORK_TYPE_GSM
+        TelephonyIcons.LTE -> NETWORK_TYPE_LTE
+        TelephonyIcons.FOUR_G -> NETWORK_TYPE_UMTS
+        TelephonyIcons.NR_5G -> NETWORK_TYPE_NR
+        TelephonyIcons.NR_5G_PLUS -> OVERRIDE_NETWORK_TYPE_NR_ADVANCED
+        else -> NETWORK_TYPE_UNKNOWN
+    }
+
+private fun SignalIcon.MobileIconGroup?.toResolvedNetworkType(): ResolvedNetworkType =
+    when (this) {
+        TelephonyIcons.NR_5G_PLUS -> OverrideNetworkType(toNetworkType())
+        else -> DefaultNetworkType(toNetworkType())
+    }
+
+class DemoMobileConnectionRepository(val subId: Int) : MobileConnectionRepository {
+    override val subscriptionModelFlow = MutableStateFlow(MobileSubscriptionModel())
+
+    override val dataEnabled = MutableStateFlow(true)
+
+    override val isDefaultDataSubscription = MutableStateFlow(true)
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt
new file mode 100644
index 0000000..da55787
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoModeMobileConnectionDataSource.kt
@@ -0,0 +1,141 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.demo
+
+import android.os.Bundle
+import android.telephony.Annotation.DataActivityType
+import android.telephony.TelephonyManager.DATA_ACTIVITY_IN
+import android.telephony.TelephonyManager.DATA_ACTIVITY_INOUT
+import android.telephony.TelephonyManager.DATA_ACTIVITY_NONE
+import android.telephony.TelephonyManager.DATA_ACTIVITY_OUT
+import com.android.settingslib.SignalIcon.MobileIconGroup
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoMode.COMMAND_NETWORK
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.Mobile
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.shareIn
+
+/**
+ * Data source that can map from demo mode commands to inputs into the
+ * [DemoMobileConnectionsRepository]'s flows
+ */
+@SysUISingleton
+class DemoModeMobileConnectionDataSource
+@Inject
+constructor(
+    demoModeController: DemoModeController,
+    @Application scope: CoroutineScope,
+) {
+    private val demoCommandStream: Flow<Bundle> = conflatedCallbackFlow {
+        val callback =
+            object : DemoMode {
+                override fun demoCommands(): List<String> = listOf(COMMAND_NETWORK)
+
+                override fun dispatchDemoCommand(command: String, args: Bundle) {
+                    trySend(args)
+                }
+
+                override fun onDemoModeFinished() {
+                    // Handled elsewhere
+                }
+
+                override fun onDemoModeStarted() {
+                    // Handled elsewhere
+                }
+            }
+
+        demoModeController.addCallback(callback)
+        awaitClose { demoModeController.removeCallback(callback) }
+    }
+
+    // If the args contains "mobile", then all of the args are relevant. It's just the way demo mode
+    // commands work and it's a little silly
+    private val _mobileCommands = demoCommandStream.map { args -> args.toMobileEvent() }
+    val mobileEvents = _mobileCommands.shareIn(scope, SharingStarted.WhileSubscribed())
+
+    private fun Bundle.toMobileEvent(): FakeNetworkEventModel? {
+        val mobile = getString("mobile") ?: return null
+        return if (mobile == "show") {
+            activeMobileEvent()
+        } else {
+            MobileDisabled(subId = getString("slot")?.toInt())
+        }
+    }
+
+    /** Parse a valid mobile command string into a network event */
+    private fun Bundle.activeMobileEvent(): Mobile {
+        // There are many key/value pairs supported by mobile demo mode. Bear with me here
+        val level = getString("level")?.toInt()
+        val dataType = getString("datatype")?.toDataType()
+        val slot = getString("slot")?.toInt()
+        val carrierId = getString("carrierid")?.toInt()
+        val inflateStrength = getString("inflate")?.toBoolean()
+        val activity = getString("activity")?.toActivity()
+        val carrierNetworkChange = getString("carriernetworkchange") == "show"
+
+        return Mobile(
+            level = level,
+            dataType = dataType,
+            subId = slot,
+            carrierId = carrierId,
+            inflateStrength = inflateStrength,
+            activity = activity,
+            carrierNetworkChange = carrierNetworkChange,
+        )
+    }
+}
+
+private fun String.toDataType(): MobileIconGroup =
+    when (this) {
+        "1x" -> TelephonyIcons.ONE_X
+        "3g" -> TelephonyIcons.THREE_G
+        "4g" -> TelephonyIcons.FOUR_G
+        "4g+" -> TelephonyIcons.FOUR_G_PLUS
+        "5g" -> TelephonyIcons.NR_5G
+        "5ge" -> TelephonyIcons.LTE_CA_5G_E
+        "5g+" -> TelephonyIcons.NR_5G_PLUS
+        "e" -> TelephonyIcons.E
+        "g" -> TelephonyIcons.G
+        "h" -> TelephonyIcons.H
+        "h+" -> TelephonyIcons.H_PLUS
+        "lte" -> TelephonyIcons.LTE
+        "lte+" -> TelephonyIcons.LTE_PLUS
+        "dis" -> TelephonyIcons.DATA_DISABLED
+        "not" -> TelephonyIcons.NOT_DEFAULT_DATA
+        else -> TelephonyIcons.UNKNOWN
+    }
+
+@DataActivityType
+private fun String.toActivity(): Int =
+    when (this) {
+        "inout" -> DATA_ACTIVITY_INOUT
+        "in" -> DATA_ACTIVITY_IN
+        "out" -> DATA_ACTIVITY_OUT
+        else -> DATA_ACTIVITY_NONE
+    }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt
new file mode 100644
index 0000000..3f3acaf
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/model/FakeNetworkEventModel.kt
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.demo.model
+
+import android.telephony.Annotation.DataActivityType
+import com.android.settingslib.SignalIcon
+
+/**
+ * Model for the demo commands, ported from [NetworkControllerImpl]
+ *
+ * Nullable fields represent optional command line arguments
+ */
+sealed interface FakeNetworkEventModel {
+    data class Mobile(
+        val level: Int?,
+        val dataType: SignalIcon.MobileIconGroup?,
+        // Null means the default (chosen by the repository)
+        val subId: Int?,
+        val carrierId: Int?,
+        val inflateStrength: Boolean?,
+        @DataActivityType val activity: Int?,
+        val carrierNetworkChange: Boolean,
+    ) : FakeNetworkEventModel
+
+    data class MobileDisabled(
+        // Null means the default (chosen by the repository)
+        val subId: Int?
+    ) : FakeNetworkEventModel
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt
new file mode 100644
index 0000000..4c1cf4a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryImpl.kt
@@ -0,0 +1,235 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.prod
+
+import android.content.Context
+import android.database.ContentObserver
+import android.provider.Settings.Global
+import android.telephony.CellSignalStrength
+import android.telephony.CellSignalStrengthCdma
+import android.telephony.ServiceState
+import android.telephony.SignalStrength
+import android.telephony.TelephonyCallback
+import android.telephony.TelephonyDisplayInfo
+import android.telephony.TelephonyDisplayInfo.OVERRIDE_NETWORK_TYPE_NONE
+import android.telephony.TelephonyManager
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.model.toDataConnectionType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger.Companion.logOutputChange
+import com.android.systemui.util.settings.GlobalSettings
+import java.lang.IllegalStateException
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.asExecutor
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+class MobileConnectionRepositoryImpl(
+    private val context: Context,
+    private val subId: Int,
+    private val telephonyManager: TelephonyManager,
+    private val globalSettings: GlobalSettings,
+    defaultDataSubId: StateFlow<Int>,
+    globalMobileDataSettingChangedEvent: Flow<Unit>,
+    bgDispatcher: CoroutineDispatcher,
+    logger: ConnectivityPipelineLogger,
+    scope: CoroutineScope,
+) : MobileConnectionRepository {
+    init {
+        if (telephonyManager.subscriptionId != subId) {
+            throw IllegalStateException(
+                "TelephonyManager should be created with subId($subId). " +
+                    "Found ${telephonyManager.subscriptionId} instead."
+            )
+        }
+    }
+
+    private val telephonyCallbackEvent = MutableSharedFlow<Unit>(extraBufferCapacity = 1)
+
+    override val subscriptionModelFlow: StateFlow<MobileSubscriptionModel> = run {
+        var state = MobileSubscriptionModel()
+        conflatedCallbackFlow {
+                // TODO (b/240569788): log all of these into the connectivity logger
+                val callback =
+                    object :
+                        TelephonyCallback(),
+                        TelephonyCallback.ServiceStateListener,
+                        TelephonyCallback.SignalStrengthsListener,
+                        TelephonyCallback.DataConnectionStateListener,
+                        TelephonyCallback.DataActivityListener,
+                        TelephonyCallback.CarrierNetworkListener,
+                        TelephonyCallback.DisplayInfoListener {
+                        override fun onServiceStateChanged(serviceState: ServiceState) {
+                            state = state.copy(isEmergencyOnly = serviceState.isEmergencyOnly)
+                            trySend(state)
+                        }
+
+                        override fun onSignalStrengthsChanged(signalStrength: SignalStrength) {
+                            val cdmaLevel =
+                                signalStrength
+                                    .getCellSignalStrengths(CellSignalStrengthCdma::class.java)
+                                    .let { strengths ->
+                                        if (!strengths.isEmpty()) {
+                                            strengths[0].level
+                                        } else {
+                                            CellSignalStrength.SIGNAL_STRENGTH_NONE_OR_UNKNOWN
+                                        }
+                                    }
+
+                            val primaryLevel = signalStrength.level
+
+                            state =
+                                state.copy(
+                                    cdmaLevel = cdmaLevel,
+                                    primaryLevel = primaryLevel,
+                                    isGsm = signalStrength.isGsm,
+                                )
+                            trySend(state)
+                        }
+
+                        override fun onDataConnectionStateChanged(
+                            dataState: Int,
+                            networkType: Int
+                        ) {
+                            state =
+                                state.copy(dataConnectionState = dataState.toDataConnectionType())
+                            trySend(state)
+                        }
+
+                        override fun onDataActivity(direction: Int) {
+                            state = state.copy(dataActivityDirection = direction)
+                            trySend(state)
+                        }
+
+                        override fun onCarrierNetworkChange(active: Boolean) {
+                            state = state.copy(carrierNetworkChangeActive = active)
+                            trySend(state)
+                        }
+
+                        override fun onDisplayInfoChanged(
+                            telephonyDisplayInfo: TelephonyDisplayInfo
+                        ) {
+                            val networkType =
+                                if (
+                                    telephonyDisplayInfo.overrideNetworkType ==
+                                        OVERRIDE_NETWORK_TYPE_NONE
+                                ) {
+                                    DefaultNetworkType(telephonyDisplayInfo.networkType)
+                                } else {
+                                    OverrideNetworkType(telephonyDisplayInfo.overrideNetworkType)
+                                }
+                            state = state.copy(resolvedNetworkType = networkType)
+                            trySend(state)
+                        }
+                    }
+                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
+                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
+            }
+            .onEach { telephonyCallbackEvent.tryEmit(Unit) }
+            .logOutputChange(logger, "MobileSubscriptionModel")
+            .stateIn(scope, SharingStarted.WhileSubscribed(), state)
+    }
+
+    /** Produces whenever the mobile data setting changes for this subId */
+    private val localMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
+        val observer =
+            object : ContentObserver(null) {
+                override fun onChange(selfChange: Boolean) {
+                    trySend(Unit)
+                }
+            }
+
+        globalSettings.registerContentObserver(
+            globalSettings.getUriFor("${Global.MOBILE_DATA}$subId"),
+            /* notifyForDescendants */ true,
+            observer
+        )
+
+        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
+    }
+
+    /**
+     * There are a few cases where we will need to poll [TelephonyManager] so we can update some
+     * internal state where callbacks aren't provided. Any of those events should be merged into
+     * this flow, which can be used to trigger the polling.
+     */
+    private val telephonyPollingEvent: Flow<Unit> =
+        merge(
+            telephonyCallbackEvent,
+            localMobileDataSettingChangedEvent,
+            globalMobileDataSettingChangedEvent,
+        )
+
+    override val dataEnabled: StateFlow<Boolean> =
+        telephonyPollingEvent
+            .mapLatest { dataConnectionAllowed() }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), dataConnectionAllowed())
+
+    private fun dataConnectionAllowed(): Boolean = telephonyManager.isDataConnectionAllowed
+
+    override val isDefaultDataSubscription: StateFlow<Boolean> =
+        defaultDataSubId
+            .mapLatest { it == subId }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), defaultDataSubId.value == subId)
+
+    class Factory
+    @Inject
+    constructor(
+        private val context: Context,
+        private val telephonyManager: TelephonyManager,
+        private val logger: ConnectivityPipelineLogger,
+        private val globalSettings: GlobalSettings,
+        @Background private val bgDispatcher: CoroutineDispatcher,
+        @Application private val scope: CoroutineScope,
+    ) {
+        fun build(
+            subId: Int,
+            defaultDataSubId: StateFlow<Int>,
+            globalMobileDataSettingChangedEvent: Flow<Unit>,
+        ): MobileConnectionRepository {
+            return MobileConnectionRepositoryImpl(
+                context,
+                subId,
+                telephonyManager.createForSubscriptionId(subId),
+                globalSettings,
+                defaultDataSubId,
+                globalMobileDataSettingChangedEvent,
+                bgDispatcher,
+                logger,
+                scope,
+            )
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
new file mode 100644
index 0000000..08d6010
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryImpl.kt
@@ -0,0 +1,265 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.prod
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.content.IntentFilter
+import android.database.ContentObserver
+import android.net.ConnectivityManager
+import android.net.ConnectivityManager.NetworkCallback
+import android.net.Network
+import android.net.NetworkCapabilities
+import android.net.NetworkCapabilities.NET_CAPABILITY_VALIDATED
+import android.net.NetworkCapabilities.TRANSPORT_CELLULAR
+import android.provider.Settings.Global.MOBILE_DATA
+import android.telephony.CarrierConfigManager
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import android.telephony.SubscriptionManager.INVALID_SUBSCRIPTION_ID
+import android.telephony.TelephonyCallback
+import android.telephony.TelephonyCallback.ActiveDataSubscriptionIdListener
+import android.telephony.TelephonyManager
+import androidx.annotation.VisibleForTesting
+import com.android.internal.telephony.PhoneConstants
+import com.android.settingslib.mobile.MobileMappings
+import com.android.settingslib.mobile.MobileMappings.Config
+import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.common.coroutine.ConflatedCallbackFlow.conflatedCallbackFlow
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileConnectivityModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.MobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.util.settings.GlobalSettings
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.asExecutor
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.mapLatest
+import kotlinx.coroutines.flow.merge
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.withContext
+
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+@SysUISingleton
+class MobileConnectionsRepositoryImpl
+@Inject
+constructor(
+    private val connectivityManager: ConnectivityManager,
+    private val subscriptionManager: SubscriptionManager,
+    private val telephonyManager: TelephonyManager,
+    private val logger: ConnectivityPipelineLogger,
+    broadcastDispatcher: BroadcastDispatcher,
+    private val globalSettings: GlobalSettings,
+    private val context: Context,
+    @Background private val bgDispatcher: CoroutineDispatcher,
+    @Application private val scope: CoroutineScope,
+    private val mobileConnectionRepositoryFactory: MobileConnectionRepositoryImpl.Factory
+) : MobileConnectionsRepository {
+    private val subIdRepositoryCache: MutableMap<Int, MobileConnectionRepository> = mutableMapOf()
+
+    /**
+     * State flow that emits the set of mobile data subscriptions, each represented by its own
+     * [SubscriptionInfo]. We probably only need the [SubscriptionInfo.getSubscriptionId] of each
+     * info object, but for now we keep track of the infos themselves.
+     */
+    override val subscriptionsFlow: StateFlow<List<SubscriptionInfo>> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : SubscriptionManager.OnSubscriptionsChangedListener() {
+                        override fun onSubscriptionsChanged() {
+                            trySend(Unit)
+                        }
+                    }
+
+                subscriptionManager.addOnSubscriptionsChangedListener(
+                    bgDispatcher.asExecutor(),
+                    callback,
+                )
+
+                awaitClose { subscriptionManager.removeOnSubscriptionsChangedListener(callback) }
+            }
+            .mapLatest { fetchSubscriptionsList() }
+            .onEach { infos -> dropUnusedReposFromCache(infos) }
+            .stateIn(scope, started = SharingStarted.WhileSubscribed(), listOf())
+
+    /** StateFlow that keeps track of the current active mobile data subscription */
+    override val activeMobileDataSubscriptionId: StateFlow<Int> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : TelephonyCallback(), ActiveDataSubscriptionIdListener {
+                        override fun onActiveDataSubscriptionIdChanged(subId: Int) {
+                            trySend(subId)
+                        }
+                    }
+
+                telephonyManager.registerTelephonyCallback(bgDispatcher.asExecutor(), callback)
+                awaitClose { telephonyManager.unregisterTelephonyCallback(callback) }
+            }
+            .stateIn(scope, started = SharingStarted.WhileSubscribed(), INVALID_SUBSCRIPTION_ID)
+
+    private val defaultDataSubIdChangeEvent: MutableSharedFlow<Unit> =
+        MutableSharedFlow(extraBufferCapacity = 1)
+
+    override val defaultDataSubId: StateFlow<Int> =
+        broadcastDispatcher
+            .broadcastFlow(
+                IntentFilter(TelephonyManager.ACTION_DEFAULT_DATA_SUBSCRIPTION_CHANGED)
+            ) { intent, _ ->
+                intent.getIntExtra(PhoneConstants.SUBSCRIPTION_KEY, INVALID_SUBSCRIPTION_ID)
+            }
+            .distinctUntilChanged()
+            .onEach { defaultDataSubIdChangeEvent.tryEmit(Unit) }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                SubscriptionManager.getDefaultDataSubscriptionId()
+            )
+
+    private val carrierConfigChangedEvent =
+        broadcastDispatcher.broadcastFlow(
+            IntentFilter(CarrierConfigManager.ACTION_CARRIER_CONFIG_CHANGED)
+        )
+
+    /**
+     * [Config] is an object that tracks relevant configuration flags for a given subscription ID.
+     * In the case of [MobileMappings], it's hard-coded to check the default data subscription's
+     * config, so this will apply to every icon that we care about.
+     *
+     * Relevant bits in the config are things like
+     * [CarrierConfigManager.KEY_SHOW_4G_FOR_LTE_DATA_ICON_BOOL]
+     *
+     * This flow will produce whenever the default data subscription or the carrier config changes.
+     */
+    override val defaultDataSubRatConfig: StateFlow<Config> =
+        merge(defaultDataSubIdChangeEvent, carrierConfigChangedEvent)
+            .mapLatest { Config.readConfig(context) }
+            .stateIn(
+                scope,
+                SharingStarted.WhileSubscribed(),
+                initialValue = Config.readConfig(context)
+            )
+
+    override fun getRepoForSubId(subId: Int): MobileConnectionRepository {
+        if (!isValidSubId(subId)) {
+            throw IllegalArgumentException(
+                "subscriptionId $subId is not in the list of valid subscriptions"
+            )
+        }
+
+        return subIdRepositoryCache[subId]
+            ?: createRepositoryForSubId(subId).also { subIdRepositoryCache[subId] = it }
+    }
+
+    /**
+     * In single-SIM devices, the [MOBILE_DATA] setting is phone-wide. For multi-SIM, the individual
+     * connection repositories also observe the URI for [MOBILE_DATA] + subId.
+     */
+    override val globalMobileDataSettingChangedEvent: Flow<Unit> = conflatedCallbackFlow {
+        val observer =
+            object : ContentObserver(null) {
+                override fun onChange(selfChange: Boolean) {
+                    trySend(Unit)
+                }
+            }
+
+        globalSettings.registerContentObserver(
+            globalSettings.getUriFor(MOBILE_DATA),
+            true,
+            observer
+        )
+
+        awaitClose { context.contentResolver.unregisterContentObserver(observer) }
+    }
+
+    @SuppressLint("MissingPermission")
+    override val defaultMobileNetworkConnectivity: StateFlow<MobileConnectivityModel> =
+        conflatedCallbackFlow {
+                val callback =
+                    object : NetworkCallback(FLAG_INCLUDE_LOCATION_INFO) {
+                        override fun onLost(network: Network) {
+                            // Send a disconnected model when lost. Maybe should create a sealed
+                            // type or null here?
+                            trySend(MobileConnectivityModel())
+                        }
+
+                        override fun onCapabilitiesChanged(
+                            network: Network,
+                            caps: NetworkCapabilities
+                        ) {
+                            trySend(
+                                MobileConnectivityModel(
+                                    isConnected = caps.hasTransport(TRANSPORT_CELLULAR),
+                                    isValidated = caps.hasCapability(NET_CAPABILITY_VALIDATED),
+                                )
+                            )
+                        }
+                    }
+
+                connectivityManager.registerDefaultNetworkCallback(callback)
+
+                awaitClose { connectivityManager.unregisterNetworkCallback(callback) }
+            }
+            .stateIn(scope, SharingStarted.WhileSubscribed(), MobileConnectivityModel())
+
+    private fun isValidSubId(subId: Int): Boolean {
+        subscriptionsFlow.value.forEach {
+            if (it.subscriptionId == subId) {
+                return true
+            }
+        }
+
+        return false
+    }
+
+    @VisibleForTesting fun getSubIdRepoCache() = subIdRepositoryCache
+
+    private fun createRepositoryForSubId(subId: Int): MobileConnectionRepository {
+        return mobileConnectionRepositoryFactory.build(
+            subId,
+            defaultDataSubId,
+            globalMobileDataSettingChangedEvent,
+        )
+    }
+
+    private fun dropUnusedReposFromCache(newInfos: List<SubscriptionInfo>) {
+        // Remove any connection repository from the cache that isn't in the new set of IDs. They
+        // will get garbage collected once their subscribers go away
+        val currentValidSubscriptionIds = newInfos.map { it.subscriptionId }
+
+        subIdRepositoryCache.keys.forEach {
+            if (!currentValidSubscriptionIds.contains(it)) {
+                subIdRepositoryCache.remove(it)
+            }
+        }
+    }
+
+    private suspend fun fetchSubscriptionsList(): List<SubscriptionInfo> =
+        withContext(bgDispatcher) { subscriptionManager.completeActiveSubscriptionInfoList }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileUiAdapter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileUiAdapter.kt
index c7e0ce1..d9487bf 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileUiAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/mobile/ui/MobileUiAdapter.kt
@@ -16,6 +16,7 @@
 
 package com.android.systemui.statusbar.pipeline.mobile.ui
 
+import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.statusbar.phone.StatusBarIconController
@@ -29,9 +30,10 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.flow.mapLatest
-import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
 
 /**
  * This class is intended to provide a context to collect on the
@@ -50,9 +52,9 @@
     interactor: MobileIconsInteractor,
     private val iconController: StatusBarIconController,
     private val iconsViewModelFactory: MobileIconsViewModel.Factory,
-    @Application scope: CoroutineScope,
+    @Application private val scope: CoroutineScope,
     private val statusBarPipelineFlags: StatusBarPipelineFlags,
-) {
+) : CoreStartable {
     private val mobileSubIds: Flow<List<Int>> =
         interactor.filteredSubscriptions.mapLatest { infos ->
             infos.map { subscriptionInfo -> subscriptionInfo.subscriptionId }
@@ -66,18 +68,19 @@
      * NOTE: this should go away as the view presenter learns more about this data pipeline
      */
     private val mobileSubIdsState: StateFlow<List<Int>> =
-        mobileSubIds
-            .onEach {
-                // Only notify the icon controller if we want to *render* the new icons.
-                // Note that this flow may still run if
-                // [statusBarPipelineFlags.runNewMobileIconsBackend] is true because we may want to
-                // get the logging data without rendering.
-                if (statusBarPipelineFlags.useNewMobileIcons()) {
-                    // Notify the icon controller here so that it knows to add icons
-                    iconController.setNewMobileIconSubIds(it)
-                }
+        mobileSubIds.stateIn(scope, SharingStarted.WhileSubscribed(), listOf())
+
+    override fun start() {
+        // Only notify the icon controller if we want to *render* the new icons.
+        // Note that this flow may still run if
+        // [statusBarPipelineFlags.runNewMobileIconsBackend] is true because we may want to
+        // get the logging data without rendering.
+        if (statusBarPipelineFlags.useNewMobileIcons()) {
+            scope.launch {
+                mobileSubIds.collectLatest { iconController.setNewMobileIconSubIds(it) }
             }
-            .stateIn(scope, SharingStarted.WhileSubscribed(), listOf())
+        }
+    }
 
     /**
      * Create a MobileIconsViewModel for a given [IconManager], and bind it to to the manager's
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
index b816364..5223760 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/WifiUiAdapter.kt
@@ -24,6 +24,7 @@
 import com.android.systemui.statusbar.phone.StatusBarIconController
 import com.android.systemui.statusbar.phone.StatusBarLocation
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.WifiViewModel
 import javax.inject.Inject
@@ -73,7 +74,9 @@
                         // Note that this flow may still run if
                         // [statusBarPipelineFlags.runNewWifiIconBackend] is true because we may
                         // want to get the logging data without rendering.
-                        if (wifiIcon != null && statusBarPipelineFlags.useNewWifiIcon()) {
+                        if (
+                            wifiIcon is WifiIcon.Visible && statusBarPipelineFlags.useNewWifiIcon()
+                        ) {
                             iconController.setNewWifiIcon()
                         }
                     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
index 345f8cb..f5b5950 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/binder/WifiViewBinder.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.statusbar.StatusBarIconView.STATE_DOT
 import com.android.systemui.statusbar.StatusBarIconView.STATE_HIDDEN
 import com.android.systemui.statusbar.StatusBarIconView.STATE_ICON
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.LocationBasedWifiViewModel
 import kotlinx.coroutines.InternalCoroutinesApi
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -92,8 +93,10 @@
 
                 launch {
                     viewModel.wifiIcon.collect { wifiIcon ->
-                        view.isVisible = wifiIcon != null
-                        wifiIcon?.let { IconViewBinder.bind(wifiIcon, iconView) }
+                        view.isVisible = wifiIcon is WifiIcon.Visible
+                        if (wifiIcon is WifiIcon.Visible) {
+                            IconViewBinder.bind(wifiIcon.icon, iconView)
+                        }
                     }
                 }
 
@@ -135,7 +138,7 @@
 
         return object : Binding {
             override fun getShouldIconBeVisible(): Boolean {
-                return viewModel.wifiIcon.value != null
+                return viewModel.wifiIcon.value is WifiIcon.Visible
             }
 
             override fun onVisibilityStateChanged(@StatusBarIconView.VisibleState state: Int) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt
new file mode 100644
index 0000000..e491d2b
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/model/WifiIcon.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2022 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.pipeline.wifi.ui.model
+
+import android.annotation.DrawableRes
+import com.android.systemui.common.shared.model.ContentDescription
+import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.log.table.Diffable
+import com.android.systemui.log.table.TableRowLogger
+
+/** Represents the various states of the wifi icon. */
+sealed interface WifiIcon : Diffable<WifiIcon> {
+    /** Represents a wifi icon that should be hidden (not visible). */
+    object Hidden : WifiIcon {
+        override fun toString() = "hidden"
+    }
+
+    /**
+     * Represents a visible wifi icon that uses [res] as its image and [contentDescription] as its
+     * description.
+     */
+    class Visible(
+        @DrawableRes res: Int,
+        val contentDescription: ContentDescription.Loaded,
+    ) : WifiIcon {
+        val icon = Icon.Resource(res, contentDescription)
+
+        override fun toString() = contentDescription.description.toString()
+    }
+
+    override fun logDiffs(prevVal: WifiIcon, row: TableRowLogger) {
+        if (prevVal.toString() != toString()) {
+            row.logChange(COL_ICON, toString())
+        }
+    }
+
+    override fun logFull(row: TableRowLogger) {
+        row.logChange(COL_ICON, toString())
+    }
+}
+
+private const val COL_ICON = "wifiIcon"
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
index 95ab251..a29c9b9 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/HomeWifiViewModel.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
@@ -28,7 +28,7 @@
  */
 class HomeWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
index 86535d6..1e190fb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/KeyguardWifiViewModel.kt
@@ -17,15 +17,15 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /** A view model for the wifi icon shown on keyguard (lockscreen). */
 class KeyguardWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
index 7cbdf5d..e35a8fe 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/LocationBasedWifiViewModel.kt
@@ -17,8 +17,8 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.flowOf
@@ -33,8 +33,8 @@
     statusBarPipelineFlags: StatusBarPipelineFlags,
     debugTint: Int,
 
-    /** The wifi icon that should be displayed. Null if we shouldn't display any icon. */
-    val wifiIcon: StateFlow<Icon.Resource?>,
+    /** The wifi icon that should be displayed. */
+    val wifiIcon: StateFlow<WifiIcon>,
 
     /** True if the activity in view should be visible. */
     val isActivityInViewVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
index fd54c5f..18e62b2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/QsWifiViewModel.kt
@@ -17,15 +17,15 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.graphics.Color
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 /** A view model for the wifi icon shown in quick settings (when the shade is pulled down). */
 class QsWifiViewModel(
     statusBarPipelineFlags: StatusBarPipelineFlags,
-    wifiIcon: StateFlow<Icon.Resource?>,
+    wifiIcon: StateFlow<WifiIcon>,
     isActivityInViewVisible: Flow<Boolean>,
     isActivityOutViewVisible: Flow<Boolean>,
     isActivityContainerVisible: Flow<Boolean>,
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
index 0782bbb..ec7ba65 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModel.kt
@@ -17,16 +17,18 @@
 package com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel
 
 import android.content.Context
-import androidx.annotation.DrawableRes
 import androidx.annotation.StringRes
 import androidx.annotation.VisibleForTesting
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_CONNECTION_STRENGTH
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_NO_CONNECTION
 import com.android.systemui.R
 import com.android.systemui.common.shared.model.ContentDescription
-import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.statusbar.pipeline.dagger.WifiTableLog
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
+import com.android.systemui.log.table.TableLogBuffer
+import com.android.systemui.log.table.logDiffsForTable
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_FULL_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_INTERNET_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_NETWORK
@@ -71,50 +73,39 @@
     connectivityConstants: ConnectivityConstants,
     private val context: Context,
     logger: ConnectivityPipelineLogger,
+    @WifiTableLog wifiTableLogBuffer: TableLogBuffer,
     interactor: WifiInteractor,
     @Application private val scope: CoroutineScope,
     statusBarPipelineFlags: StatusBarPipelineFlags,
     wifiConstants: WifiConstants,
 ) {
-    /**
-     * Returns the drawable resource ID to use for the wifi icon based on the given network.
-     * Null if we can't compute the icon.
-     */
-    @DrawableRes
-    private fun WifiNetworkModel.iconResId(): Int? {
+    /** Returns the icon to use based on the given network. */
+    private fun WifiNetworkModel.icon(): WifiIcon {
         return when (this) {
-            is WifiNetworkModel.CarrierMerged -> null
-            is WifiNetworkModel.Inactive -> WIFI_NO_NETWORK
-            is WifiNetworkModel.Active ->
-                when {
-                    this.level == null -> null
-                    this.isValidated -> WIFI_FULL_ICONS[this.level]
-                    else -> WIFI_NO_INTERNET_ICONS[this.level]
-                }
-        }
-    }
-
-    /**
-     * Returns the content description for the wifi icon based on the given network.
-     * Null if we can't compute the content description.
-     */
-    private fun WifiNetworkModel.contentDescription(): ContentDescription? {
-        return when (this) {
-            is WifiNetworkModel.CarrierMerged -> null
-            is WifiNetworkModel.Inactive ->
+            is WifiNetworkModel.CarrierMerged -> WifiIcon.Hidden
+            is WifiNetworkModel.Inactive -> WifiIcon.Visible(
+                res = WIFI_NO_NETWORK,
                 ContentDescription.Loaded(
                     "${context.getString(WIFI_NO_CONNECTION)},${context.getString(NO_INTERNET)}"
                 )
+            )
             is WifiNetworkModel.Active ->
                 when (this.level) {
-                    null -> null
+                    null -> WifiIcon.Hidden
                     else -> {
                         val levelDesc = context.getString(WIFI_CONNECTION_STRENGTH[this.level])
                         when {
-                            this.isValidated -> ContentDescription.Loaded(levelDesc)
+                            this.isValidated ->
+                                WifiIcon.Visible(
+                                    WIFI_FULL_ICONS[this.level],
+                                    ContentDescription.Loaded(levelDesc)
+                                )
                             else ->
-                                ContentDescription.Loaded(
-                                    "$levelDesc,${context.getString(NO_INTERNET)}"
+                                WifiIcon.Visible(
+                                    WIFI_NO_INTERNET_ICONS[this.level],
+                                    ContentDescription.Loaded(
+                                        "$levelDesc,${context.getString(NO_INTERNET)}"
+                                    )
                                 )
                         }
                     }
@@ -122,8 +113,8 @@
         }
     }
 
-    /** The wifi icon that should be displayed. Null if we shouldn't display any icon. */
-    private val wifiIcon: StateFlow<Icon.Resource?> =
+    /** The wifi icon that should be displayed. */
+    private val wifiIcon: StateFlow<WifiIcon> =
         combine(
             interactor.isEnabled,
             interactor.isDefault,
@@ -131,22 +122,29 @@
             interactor.wifiNetwork,
         ) { isEnabled, isDefault, isForceHidden, wifiNetwork ->
             if (!isEnabled || isForceHidden || wifiNetwork is WifiNetworkModel.CarrierMerged) {
-                return@combine null
+                return@combine WifiIcon.Hidden
             }
 
-            val iconResId = wifiNetwork.iconResId() ?: return@combine null
-            val icon = Icon.Resource(iconResId, wifiNetwork.contentDescription())
+            val icon = wifiNetwork.icon()
 
             return@combine when {
                 isDefault -> icon
                 wifiConstants.alwaysShowIconIfEnabled -> icon
                 !connectivityConstants.hasDataCapabilities -> icon
                 wifiNetwork is WifiNetworkModel.Active && wifiNetwork.isValidated -> icon
-                else -> null
+                else -> WifiIcon.Hidden
             }
         }
-                .logOutputChange(logger, "icon") { icon -> icon?.contentDescription.toString() }
-                .stateIn(scope, started = SharingStarted.WhileSubscribed(), initialValue = null)
+            .logDiffsForTable(
+                wifiTableLogBuffer,
+                columnPrefix = "",
+                initialValue = WifiIcon.Hidden,
+            )
+            .stateIn(
+                scope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = WifiIcon.Hidden
+            )
 
     /** The wifi activity status. Null if we shouldn't display the activity status. */
     private val activity: Flow<WifiActivityModel?> =
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
index 162c915..b2ec27c 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLightRevealOverlayAnimation.kt
@@ -40,6 +40,8 @@
 import com.android.systemui.statusbar.LightRevealEffect
 import com.android.systemui.statusbar.LightRevealScrim
 import com.android.systemui.statusbar.LinearLightRevealEffect
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.FOLD
+import com.android.systemui.unfold.UnfoldLightRevealOverlayAnimation.AddOverlayReason.UNFOLD
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider.TransitionProgressListener
 import com.android.systemui.unfold.updates.RotationChangeProvider
 import com.android.systemui.unfold.util.ScaleAwareTransitionProgressProvider.Companion.areAnimationsEnabled
@@ -125,7 +127,7 @@
         try {
             // Add the view only if we are unfolding and this is the first screen on
             if (!isFolded && !isUnfoldHandled && contentResolver.areAnimationsEnabled()) {
-                executeInBackground { addView(onOverlayReady) }
+                executeInBackground { addOverlay(onOverlayReady, reason = UNFOLD) }
                 isUnfoldHandled = true
             } else {
                 // No unfold transition, immediately report that overlay is ready
@@ -137,7 +139,7 @@
         }
     }
 
-    private fun addView(onOverlayReady: Runnable? = null) {
+    private fun addOverlay(onOverlayReady: Runnable? = null, reason: AddOverlayReason) {
         if (!::wwm.isInitialized) {
             // Surface overlay is not created yet on the first SysUI launch
             onOverlayReady?.run()
@@ -152,7 +154,10 @@
             LightRevealScrim(context, null).apply {
                 revealEffect = createLightRevealEffect()
                 isScrimOpaqueChangedListener = Consumer {}
-                revealAmount = 0f
+                revealAmount = when (reason) {
+                    FOLD -> TRANSPARENT
+                    UNFOLD -> BLACK
+                }
             }
 
         val params = getLayoutParams()
@@ -228,7 +233,7 @@
     }
 
     private fun getUnfoldedDisplayInfo(): DisplayInfo =
-        displayManager.displays
+        displayManager.getDisplays(DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED)
             .asSequence()
             .map { DisplayInfo().apply { it.getDisplayInfo(this) } }
             .filter { it.type == Display.TYPE_INTERNAL }
@@ -247,7 +252,7 @@
         override fun onTransitionStarted() {
             // Add view for folding case (when unfolding the view is added earlier)
             if (scrimView == null) {
-                executeInBackground { addView() }
+                executeInBackground { addOverlay(reason = FOLD) }
             }
             // Disable input dispatching during transition.
             InputManager.getInstance().cancelCurrentTouch()
@@ -294,11 +299,17 @@
             }
         )
 
+    private enum class AddOverlayReason { FOLD, UNFOLD }
+
     private companion object {
-        private const val ROTATION_ANIMATION_OVERLAY_Z_INDEX = Integer.MAX_VALUE
+        const val ROTATION_ANIMATION_OVERLAY_Z_INDEX = Integer.MAX_VALUE
 
         // Put the unfold overlay below the rotation animation screenshot to hide the moment
         // when it is rotated but the rotation of the other windows hasn't happen yet
-        private const val UNFOLD_OVERLAY_LAYER_Z_INDEX = ROTATION_ANIMATION_OVERLAY_Z_INDEX - 1
+        const val UNFOLD_OVERLAY_LAYER_Z_INDEX = ROTATION_ANIMATION_OVERLAY_Z_INDEX - 1
+
+        // constants for revealAmount.
+        const val TRANSPARENT = 1f
+        const val BLACK = 0f
     }
 }
diff --git a/packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt b/packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
similarity index 94%
rename from packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
index 389eed0..2c680be 100644
--- a/packages/SystemUI/animation/tests/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/animation/InterpolatorsAndroidXTest.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.animation
 
 import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
 import java.lang.reflect.Modifier
 import junit.framework.Assert.assertEquals
 import org.junit.Test
@@ -25,7 +26,7 @@
 
 @SmallTest
 @RunWith(JUnit4::class)
-class InterpolatorsAndroidXTest {
+class InterpolatorsAndroidXTest : SysuiTestCase() {
 
     @Test
     fun testInterpolatorsAndInterpolatorsAndroidXPublicMethodsAreEqual() {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/MediaControlPanelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/MediaControlPanelTest.kt
index 761773b..fdef344 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/MediaControlPanelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/media/controls/ui/MediaControlPanelTest.kt
@@ -702,7 +702,7 @@
     }
 
     @Test
-    fun bind_seekBarDisabled_noActions_seekBarVisibilityIsSetToGone() {
+    fun bind_seekBarDisabled_noActions_seekBarVisibilityIsSetToInvisible() {
         useRealConstraintSets()
 
         val state = mediaData.copy(semanticActions = MediaButton())
@@ -711,7 +711,7 @@
 
         player.bindPlayer(state, PACKAGE)
 
-        assertThat(expandedSet.getVisibility(seekBar.id)).isEqualTo(ConstraintSet.GONE)
+        assertThat(expandedSet.getVisibility(seekBar.id)).isEqualTo(ConstraintSet.INVISIBLE)
     }
 
     @Test
@@ -741,7 +741,7 @@
     }
 
     @Test
-    fun seekBarChangesToDisabledAfterBind_noActions_seekBarChangesToGone() {
+    fun seekBarChangesToDisabledAfterBind_noActions_seekBarChangesToInvisible() {
         useRealConstraintSets()
 
         val state = mediaData.copy(semanticActions = MediaButton())
@@ -752,7 +752,7 @@
 
         getEnabledChangeListener().onEnabledChanged(enabled = false)
 
-        assertThat(expandedSet.getVisibility(seekBar.id)).isEqualTo(ConstraintSet.GONE)
+        assertThat(expandedSet.getVisibility(seekBar.id)).isEqualTo(ConstraintSet.INVISIBLE)
     }
 
     @Test
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
index 9758842..4a9c750 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskControllerTest.kt
@@ -16,16 +16,21 @@
 package com.android.systemui.notetask
 
 import android.app.KeyguardManager
+import android.content.ComponentName
 import android.content.Context
 import android.content.Intent
+import android.content.pm.PackageManager
 import android.os.UserManager
 import android.test.suitebuilder.annotation.SmallTest
-import android.view.KeyEvent
 import androidx.test.runner.AndroidJUnit4
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.notetask.NoteTaskIntentResolver.Companion.NOTES_ACTION
+import com.android.systemui.notetask.shortcut.CreateNoteTaskShortcutActivity
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
 import com.android.wm.shell.bubbles.Bubbles
+import com.google.common.truth.Truth.assertThat
 import java.util.Optional
 import org.junit.Before
 import org.junit.Test
@@ -48,6 +53,7 @@
     private val notesIntent = Intent(NOTES_ACTION)
 
     @Mock lateinit var context: Context
+    @Mock lateinit var packageManager: PackageManager
     @Mock lateinit var noteTaskIntentResolver: NoteTaskIntentResolver
     @Mock lateinit var bubbles: Bubbles
     @Mock lateinit var optionalBubbles: Optional<Bubbles>
@@ -60,6 +66,7 @@
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
+        whenever(context.packageManager).thenReturn(packageManager)
         whenever(noteTaskIntentResolver.resolveIntent()).thenReturn(notesIntent)
         whenever(optionalBubbles.orElse(null)).thenReturn(bubbles)
         whenever(optionalKeyguardManager.orElse(null)).thenReturn(keyguardManager)
@@ -78,89 +85,125 @@
         )
     }
 
+    // region showNoteTask
     @Test
-    fun handleSystemKey_keyguardIsLocked_shouldStartActivity() {
+    fun showNoteTask_keyguardIsLocked_shouldStartActivity() {
         whenever(keyguardManager.isKeyguardLocked).thenReturn(true)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_keyguardIsUnlocked_shouldStartBubbles() {
+    fun showNoteTask_keyguardIsUnlocked_shouldStartBubbles() {
         whenever(keyguardManager.isKeyguardLocked).thenReturn(false)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(bubbles).showAppBubble(notesIntent)
         verify(context, never()).startActivity(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_receiveInvalidSystemKey_shouldDoNothing() {
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_UNKNOWN)
+    fun showNoteTask_isInMultiWindowMode_shouldStartActivity() {
+        whenever(keyguardManager.isKeyguardLocked).thenReturn(false)
 
-        verify(context, never()).startActivity(notesIntent)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = true)
+
+        verify(context).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_bubblesIsNull_shouldDoNothing() {
+    fun showNoteTask_bubblesIsNull_shouldDoNothing() {
         whenever(optionalBubbles.orElse(null)).thenReturn(null)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_keyguardManagerIsNull_shouldDoNothing() {
+    fun showNoteTask_keyguardManagerIsNull_shouldDoNothing() {
         whenever(optionalKeyguardManager.orElse(null)).thenReturn(null)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_userManagerIsNull_shouldDoNothing() {
+    fun showNoteTask_userManagerIsNull_shouldDoNothing() {
         whenever(optionalUserManager.orElse(null)).thenReturn(null)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_intentResolverReturnsNull_shouldDoNothing() {
+    fun showNoteTask_intentResolverReturnsNull_shouldDoNothing() {
         whenever(noteTaskIntentResolver.resolveIntent()).thenReturn(null)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_flagDisabled_shouldDoNothing() {
-        createNoteTaskController(isEnabled = false).handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+    fun showNoteTask_flagDisabled_shouldDoNothing() {
+        createNoteTaskController(isEnabled = false).showNoteTask()
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
 
     @Test
-    fun handleSystemKey_userIsLocked_shouldDoNothing() {
+    fun showNoteTask_userIsLocked_shouldDoNothing() {
         whenever(userManager.isUserUnlocked).thenReturn(false)
 
-        createNoteTaskController().handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+        createNoteTaskController().showNoteTask(isInMultiWindowMode = false)
 
         verify(context, never()).startActivity(notesIntent)
         verify(bubbles, never()).showAppBubble(notesIntent)
     }
+    // endregion
+
+    // region setNoteTaskShortcutEnabled
+    @Test
+    fun setNoteTaskShortcutEnabled_setTrue() {
+        createNoteTaskController().setNoteTaskShortcutEnabled(value = true)
+
+        val argument = argumentCaptor<ComponentName>()
+        verify(context.packageManager)
+            .setComponentEnabledSetting(
+                argument.capture(),
+                eq(PackageManager.COMPONENT_ENABLED_STATE_ENABLED),
+                eq(PackageManager.DONT_KILL_APP),
+            )
+        val expected = ComponentName(context, CreateNoteTaskShortcutActivity::class.java)
+        assertThat(argument.value.flattenToString()).isEqualTo(expected.flattenToString())
+    }
+
+    @Test
+    fun setNoteTaskShortcutEnabled_setFalse() {
+        createNoteTaskController().setNoteTaskShortcutEnabled(value = false)
+
+        val argument = argumentCaptor<ComponentName>()
+        verify(context.packageManager)
+            .setComponentEnabledSetting(
+                argument.capture(),
+                eq(PackageManager.COMPONENT_ENABLED_STATE_DISABLED),
+                eq(PackageManager.DONT_KILL_APP),
+            )
+        val expected = ComponentName(context, CreateNoteTaskShortcutActivity::class.java)
+        assertThat(argument.value.flattenToString()).isEqualTo(expected.flattenToString())
+    }
+    // endregion
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
index 334089c..538131a 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/NoteTaskInitializerTest.kt
@@ -16,10 +16,10 @@
 package com.android.systemui.notetask
 
 import android.test.suitebuilder.annotation.SmallTest
+import android.view.KeyEvent
 import androidx.test.runner.AndroidJUnit4
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.statusbar.CommandQueue
-import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.android.wm.shell.bubbles.Bubbles
 import java.util.Optional
@@ -45,6 +45,7 @@
     @Mock lateinit var commandQueue: CommandQueue
     @Mock lateinit var bubbles: Bubbles
     @Mock lateinit var optionalBubbles: Optional<Bubbles>
+    @Mock lateinit var noteTaskController: NoteTaskController
 
     @Before
     fun setUp() {
@@ -57,12 +58,13 @@
     private fun createNoteTaskInitializer(isEnabled: Boolean = true): NoteTaskInitializer {
         return NoteTaskInitializer(
             optionalBubbles = optionalBubbles,
-            lazyNoteTaskController = mock(),
+            noteTaskController = noteTaskController,
             commandQueue = commandQueue,
             isEnabled = isEnabled,
         )
     }
 
+    // region initializer
     @Test
     fun initialize_shouldAddCallbacks() {
         createNoteTaskInitializer().initialize()
@@ -85,4 +87,35 @@
 
         verify(commandQueue, never()).addCallback(any())
     }
+
+    @Test
+    fun initialize_flagEnabled_shouldEnableShortcut() {
+        createNoteTaskInitializer().initialize()
+
+        verify(noteTaskController).setNoteTaskShortcutEnabled(true)
+    }
+
+    @Test
+    fun initialize_flagDisabled_shouldDisableShortcut() {
+        createNoteTaskInitializer(isEnabled = false).initialize()
+
+        verify(noteTaskController).setNoteTaskShortcutEnabled(false)
+    }
+    // endregion
+
+    // region handleSystemKey
+    @Test
+    fun handleSystemKey_receiveValidSystemKey_shouldShowNoteTask() {
+        createNoteTaskInitializer().callbacks.handleSystemKey(KeyEvent.KEYCODE_VIDEO_APP_1)
+
+        verify(noteTaskController).showNoteTask()
+    }
+
+    @Test
+    fun handleSystemKey_receiveInvalidSystemKey_shouldDoNothing() {
+        createNoteTaskInitializer().callbacks.handleSystemKey(KeyEvent.KEYCODE_UNKNOWN)
+
+        verify(noteTaskController, never()).showNoteTask()
+    }
+    // endregion
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerBaseTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerBaseTest.java
index faf4592..5431eba 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerBaseTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerBaseTest.java
@@ -72,6 +72,7 @@
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.plugins.log.LogBuffer;
 import com.android.systemui.settings.UserTracker;
+import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags;
 import com.android.systemui.statusbar.pipeline.mobile.util.FakeMobileMappingsProxy;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController.DeviceProvisionedListener;
@@ -245,6 +246,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mMockProvisionController,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerDataTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerDataTest.java
index ca75a40..9441d49 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerDataTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerDataTest.java
@@ -49,6 +49,7 @@
 import com.android.settingslib.net.DataUsageController;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.plugins.log.LogBuffer;
+import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.util.CarrierConfigTracker;
 
@@ -150,6 +151,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mock(DeviceProvisionedController.class),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerSignalTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerSignalTest.java
index 84c242c..4c1f0a8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerSignalTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/connectivity/NetworkControllerSignalTest.java
@@ -44,6 +44,7 @@
 import com.android.systemui.R;
 import com.android.systemui.dump.DumpManager;
 import com.android.systemui.plugins.log.LogBuffer;
+import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags;
 import com.android.systemui.statusbar.policy.DeviceProvisionedController;
 import com.android.systemui.util.CarrierConfigTracker;
 
@@ -78,6 +79,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mMockProvisionController,
@@ -115,6 +117,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mMockProvisionController,
@@ -150,6 +153,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mock(DeviceProvisionedController.class),
@@ -188,6 +192,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mock(DeviceProvisionedController.class),
@@ -274,6 +279,7 @@
                 mFakeExecutor,
                 mCallbackHandler,
                 mock(AccessPointControllerImpl.class),
+                mock(StatusBarPipelineFlags.class),
                 mock(DataUsageController.class),
                 mMockSubDefaults,
                 mock(DeviceProvisionedController.class),
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/RoundableTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/RoundableTest.kt
new file mode 100644
index 0000000..89faa239
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/RoundableTest.kt
@@ -0,0 +1,164 @@
+package com.android.systemui.statusbar.notification
+
+import android.view.View
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.mock
+import org.junit.Assert.assertEquals
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mockito.atLeastOnce
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+
+@SmallTest
+@RunWith(JUnit4::class)
+class RoundableTest : SysuiTestCase() {
+    val targetView: View = mock()
+    val roundable = FakeRoundable(targetView)
+
+    @Test
+    fun defaultConfig_shouldNotHaveRoundedCorner() {
+        // the expected default value for the roundness is top = 0f, bottom = 0f
+        assertEquals(0f, roundable.roundableState.topRoundness)
+        assertEquals(0f, roundable.roundableState.bottomRoundness)
+        assertEquals(false, roundable.hasRoundedCorner())
+    }
+
+    @Test
+    fun applyRoundnessAndInvalidate_should_invalidate_targetView() {
+        roundable.applyRoundnessAndInvalidate()
+
+        verify(targetView, times(1)).invalidate()
+    }
+
+    @Test
+    fun requestTopRoundness_update_and_invalidate_targetView() {
+        roundable.requestTopRoundness(value = 1f, sourceType = SOURCE1)
+
+        assertEquals(1f, roundable.roundableState.topRoundness)
+        verify(targetView, times(1)).invalidate()
+    }
+
+    @Test
+    fun requestBottomRoundness_update_and_invalidate_targetView() {
+        roundable.requestBottomRoundness(value = 1f, sourceType = SOURCE1)
+
+        assertEquals(1f, roundable.roundableState.bottomRoundness)
+        verify(targetView, times(1)).invalidate()
+    }
+
+    @Test
+    fun requestRoundness_update_and_invalidate_targetView() {
+        roundable.requestRoundness(top = 1f, bottom = 1f, sourceType = SOURCE1)
+
+        assertEquals(1f, roundable.roundableState.topRoundness)
+        assertEquals(1f, roundable.roundableState.bottomRoundness)
+        verify(targetView, atLeastOnce()).invalidate()
+    }
+
+    @Test
+    fun requestRoundnessReset_update_and_invalidate_targetView() {
+        roundable.requestRoundness(1f, 1f, SOURCE1)
+        assertEquals(1f, roundable.roundableState.topRoundness)
+        assertEquals(1f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundnessReset(sourceType = SOURCE1)
+
+        assertEquals(0f, roundable.roundableState.topRoundness)
+        assertEquals(0f, roundable.roundableState.bottomRoundness)
+        verify(targetView, atLeastOnce()).invalidate()
+    }
+
+    @Test
+    fun hasRoundedCorner_return_true_ifRoundnessIsGreaterThenZero() {
+        roundable.requestRoundness(top = 1f, bottom = 1f, sourceType = SOURCE1)
+        assertEquals(true, roundable.hasRoundedCorner())
+
+        roundable.requestRoundness(top = 1f, bottom = 0f, sourceType = SOURCE1)
+        assertEquals(true, roundable.hasRoundedCorner())
+
+        roundable.requestRoundness(top = 0f, bottom = 1f, sourceType = SOURCE1)
+        assertEquals(true, roundable.hasRoundedCorner())
+
+        roundable.requestRoundness(top = 0f, bottom = 0f, sourceType = SOURCE1)
+        assertEquals(false, roundable.hasRoundedCorner())
+    }
+
+    @Test
+    fun roundness_take_maxValue_onMultipleSources_first_lower() {
+        roundable.requestRoundness(0.1f, 0.1f, SOURCE1)
+        assertEquals(0.1f, roundable.roundableState.topRoundness)
+        assertEquals(0.1f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.2f, 0.2f, SOURCE2)
+        // SOURCE1 has 0.1f - SOURCE2 has 0.2f
+        assertEquals(0.2f, roundable.roundableState.topRoundness)
+        assertEquals(0.2f, roundable.roundableState.bottomRoundness)
+    }
+
+    @Test
+    fun roundness_take_maxValue_onMultipleSources_first_higher() {
+        roundable.requestRoundness(0.5f, 0.5f, SOURCE1)
+        assertEquals(0.5f, roundable.roundableState.topRoundness)
+        assertEquals(0.5f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.1f, 0.1f, SOURCE2)
+        // SOURCE1 has 0.5f - SOURCE2 has 0.1f
+        assertEquals(0.5f, roundable.roundableState.topRoundness)
+        assertEquals(0.5f, roundable.roundableState.bottomRoundness)
+    }
+
+    @Test
+    fun roundness_take_maxValue_onMultipleSources_first_higher_second_step() {
+        roundable.requestRoundness(0.1f, 0.1f, SOURCE1)
+        assertEquals(0.1f, roundable.roundableState.topRoundness)
+        assertEquals(0.1f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.2f, 0.2f, SOURCE2)
+        // SOURCE1 has 0.1f - SOURCE2 has 0.2f
+        assertEquals(0.2f, roundable.roundableState.topRoundness)
+        assertEquals(0.2f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.3f, 0.3f, SOURCE1)
+        // SOURCE1 has 0.3f - SOURCE2 has 0.2f
+        assertEquals(0.3f, roundable.roundableState.topRoundness)
+        assertEquals(0.3f, roundable.roundableState.bottomRoundness)
+    }
+
+    @Test
+    fun roundness_take_maxValue_onMultipleSources_first_lower_second_step() {
+        roundable.requestRoundness(0.5f, 0.5f, SOURCE1)
+        assertEquals(0.5f, roundable.roundableState.topRoundness)
+        assertEquals(0.5f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.2f, 0.2f, SOURCE2)
+        // SOURCE1 has 0.5f - SOURCE2 has 0.2f
+        assertEquals(0.5f, roundable.roundableState.topRoundness)
+        assertEquals(0.5f, roundable.roundableState.bottomRoundness)
+
+        roundable.requestRoundness(0.1f, 0.1f, SOURCE1)
+        // SOURCE1 has 0.1f - SOURCE2 has 0.2f
+        assertEquals(0.2f, roundable.roundableState.topRoundness)
+        assertEquals(0.2f, roundable.roundableState.bottomRoundness)
+    }
+
+    class FakeRoundable(
+        targetView: View,
+        radius: Float = MAX_RADIUS,
+    ) : Roundable {
+        override val roundableState =
+            RoundableState(
+                targetView = targetView,
+                roundable = this,
+                maxRadius = radius,
+            )
+    }
+
+    companion object {
+        private const val MAX_RADIUS = 10f
+        private val SOURCE1 = SourceType.from("Source1")
+        private val SOURCE2 = SourceType.from("Source2")
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
index fb31bef..728e026 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/row/NotificationTestHelper.java
@@ -215,8 +215,7 @@
             SourceType sourceType
     ) throws Exception {
         ExpandableNotificationRow row = createRow();
-        row.requestTopRoundness(topRoundness, false, sourceType);
-        row.requestBottomRoundness(bottomRoundness, /*animate = */ false, sourceType);
+        row.requestRoundness(topRoundness, bottomRoundness, sourceType, /*animate = */ false);
         assertEquals(topRoundness, row.getTopRoundness(), /* delta = */ 0f);
         assertEquals(bottomRoundness, row.getBottomRoundness(), /* delta = */ 0f);
         return row;
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
index 438b528..fd1944e 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationChildrenContainerTest.java
@@ -25,7 +25,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
-import com.android.systemui.statusbar.notification.SourceType;
+import com.android.systemui.statusbar.notification.LegacySourceType;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
 
@@ -158,7 +158,7 @@
         ExpandableNotificationRow row = mNotificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.OnScroll);
+                /* sourceType = */ LegacySourceType.OnScroll);
 
         mChildrenContainer.addNotification(row, 0);
 
@@ -171,11 +171,11 @@
         ExpandableNotificationRow row1 = mNotificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.DefaultValue);
+                /* sourceType = */ LegacySourceType.DefaultValue);
         ExpandableNotificationRow row2 = mNotificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.OnDismissAnimation);
+                /* sourceType = */ LegacySourceType.OnDismissAnimation);
 
         mChildrenContainer.addNotification(row1, 0);
         mChildrenContainer.addNotification(row2, 0);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
index 8c8b644..bd0a556 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationRoundnessManagerTest.java
@@ -35,6 +35,7 @@
 import com.android.systemui.R;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.dump.DumpManager;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.statusbar.notification.NotificationSectionsFeatureManager;
 import com.android.systemui.statusbar.notification.collection.NotificationEntry;
 import com.android.systemui.statusbar.notification.logging.NotificationRoundnessLogger;
@@ -73,7 +74,8 @@
         mRoundnessManager = new NotificationRoundnessManager(
                 new NotificationSectionsFeatureManager(new DeviceConfigProxy(), mContext),
                 mLogger,
-                mock(DumpManager.class));
+                mock(DumpManager.class),
+                mock(FeatureFlags.class));
         allowTestableLooperAsMainThread();
         NotificationTestHelper testHelper = new NotificationTestHelper(
                 mContext,
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 ecc0224..30da08e 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
@@ -30,6 +30,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.media.controls.ui.KeyguardMediaController;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.statusbar.StatusBarState;
@@ -59,10 +60,12 @@
     @Mock private KeyguardMediaController mKeyguardMediaController;
     @Mock private NotificationSectionsFeatureManager mSectionsFeatureManager;
     @Mock private MediaContainerController mMediaContainerController;
+    @Mock private NotificationRoundnessManager mNotificationRoundnessManager;
     @Mock private SectionHeaderController mIncomingHeaderController;
     @Mock private SectionHeaderController mPeopleHeaderController;
     @Mock private SectionHeaderController mAlertingHeaderController;
     @Mock private SectionHeaderController mSilentHeaderController;
+    @Mock private FeatureFlags mFeatureFlag;
 
     private NotificationSectionsManager mSectionsManager;
 
@@ -89,10 +92,12 @@
                         mKeyguardMediaController,
                         mSectionsFeatureManager,
                         mMediaContainerController,
+                        mNotificationRoundnessManager,
                         mIncomingHeaderController,
                         mPeopleHeaderController,
                         mAlertingHeaderController,
-                        mSilentHeaderController
+                        mSilentHeaderController,
+                        mFeatureFlag
                 );
         // Required in order for the header inflation to work properly
         when(mNssl.generateLayoutParams(any(AttributeSet.class)))
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
index bda2336..9d759c4 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationShelfTest.kt
@@ -9,7 +9,7 @@
 import com.android.systemui.animation.ShadeInterpolation
 import com.android.systemui.statusbar.NotificationShelf
 import com.android.systemui.statusbar.StatusBarIconView
-import com.android.systemui.statusbar.notification.SourceType
+import com.android.systemui.statusbar.notification.LegacySourceType
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow
 import com.android.systemui.statusbar.notification.row.ExpandableView
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper
@@ -314,9 +314,9 @@
         val row: ExpandableNotificationRow = notificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.OnScroll)
+                /* sourceType = */ LegacySourceType.OnScroll)
 
-        NotificationShelf.resetOnScrollRoundness(row)
+        NotificationShelf.resetLegacyOnScrollRoundness(row)
 
         assertEquals(0f, row.topRoundness)
         assertEquals(0f, row.bottomRoundness)
@@ -327,14 +327,14 @@
         val row1: ExpandableNotificationRow = notificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.DefaultValue)
+                /* sourceType = */ LegacySourceType.DefaultValue)
         val row2: ExpandableNotificationRow = notificationTestHelper.createRowWithRoundness(
                 /* topRoundness = */ 1f,
                 /* bottomRoundness = */ 1f,
-                /* sourceType = */ SourceType.OnDismissAnimation)
+                /* sourceType = */ LegacySourceType.OnDismissAnimation)
 
-        NotificationShelf.resetOnScrollRoundness(row1)
-        NotificationShelf.resetOnScrollRoundness(row2)
+        NotificationShelf.resetLegacyOnScrollRoundness(row1)
+        NotificationShelf.resetLegacyOnScrollRoundness(row2)
 
         assertEquals(1f, row1.topRoundness)
         assertEquals(1f, row1.bottomRoundness)
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
index 4ea1c71..680a323 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationSwipeHelperTest.java
@@ -74,6 +74,7 @@
     private NotificationSwipeHelper mSwipeHelper;
     private NotificationSwipeHelper.NotificationCallback mCallback;
     private NotificationMenuRowPlugin.OnMenuEventListener mListener;
+    private NotificationRoundnessManager mNotificationRoundnessManager;
     private View mView;
     private MotionEvent mEvent;
     private NotificationMenuRowPlugin mMenuRow;
@@ -92,10 +93,17 @@
     public void setUp() throws Exception {
         mCallback = mock(NotificationSwipeHelper.NotificationCallback.class);
         mListener = mock(NotificationMenuRowPlugin.OnMenuEventListener.class);
+        mNotificationRoundnessManager = mock(NotificationRoundnessManager.class);
         mFeatureFlags = mock(FeatureFlags.class);
         mSwipeHelper = spy(new NotificationSwipeHelper(
-                mContext.getResources(), ViewConfiguration.get(mContext),
-                new FalsingManagerFake(), mFeatureFlags, SwipeHelper.X, mCallback, mListener));
+                mContext.getResources(),
+                ViewConfiguration.get(mContext),
+                new FalsingManagerFake(),
+                mFeatureFlags,
+                SwipeHelper.X,
+                mCallback,
+                mListener,
+                mNotificationRoundnessManager));
         mView = mock(View.class);
         mEvent = mock(MotionEvent.class);
         mMenuRow = mock(NotificationMenuRowPlugin.class);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
index a2e9230..81a3f12 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/notification/stack/NotificationTargetsHelperTest.kt
@@ -35,7 +35,7 @@
     ) =
         NotificationTargetsHelper(
             FakeFeatureFlags().apply {
-                set(Flags.NOTIFICATION_GROUP_CORNER, notificationGroupCorner)
+                set(Flags.USE_ROUNDNESS_SOURCETYPES, notificationGroupCorner)
             }
         )
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
index 103b7b42..9727b6c5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/HeadsUpAppearanceControllerTest.java
@@ -32,6 +32,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.flags.FeatureFlags;
 import com.android.systemui.plugins.DarkIconDispatcher;
 import com.android.systemui.plugins.statusbar.StatusBarStateController;
 import com.android.systemui.shade.NotificationPanelViewController;
@@ -40,6 +41,7 @@
 import com.android.systemui.statusbar.notification.NotificationWakeUpCoordinator;
 import com.android.systemui.statusbar.notification.row.ExpandableNotificationRow;
 import com.android.systemui.statusbar.notification.row.NotificationTestHelper;
+import com.android.systemui.statusbar.notification.stack.NotificationRoundnessManager;
 import com.android.systemui.statusbar.notification.stack.NotificationStackScrollLayoutController;
 import com.android.systemui.statusbar.policy.Clock;
 import com.android.systemui.statusbar.policy.KeyguardStateController;
@@ -71,6 +73,8 @@
     private NotificationWakeUpCoordinator mWakeUpCoordinator;
     private KeyguardStateController mKeyguardStateController;
     private CommandQueue mCommandQueue;
+    private NotificationRoundnessManager mNotificationRoundnessManager;
+    private FeatureFlags mFeatureFlag;
 
     @Before
     public void setUp() throws Exception {
@@ -89,6 +93,8 @@
         mWakeUpCoordinator = mock(NotificationWakeUpCoordinator.class);
         mKeyguardStateController = mock(KeyguardStateController.class);
         mCommandQueue = mock(CommandQueue.class);
+        mNotificationRoundnessManager = mock(NotificationRoundnessManager.class);
+        mFeatureFlag = mock(FeatureFlags.class);
         mHeadsUpAppearanceController = new HeadsUpAppearanceController(
                 mock(NotificationIconAreaController.class),
                 mHeadsUpManager,
@@ -100,6 +106,8 @@
                 mCommandQueue,
                 mStackScrollerController,
                 mPanelView,
+                mNotificationRoundnessManager,
+                mFeatureFlag,
                 mHeadsUpStatusBarView,
                 new Clock(mContext, null),
                 Optional.of(mOperatorNameView));
@@ -182,6 +190,8 @@
                 mCommandQueue,
                 mStackScrollerController,
                 mPanelView,
+                mNotificationRoundnessManager,
+                mFeatureFlag,
                 mHeadsUpStatusBarView,
                 new Clock(mContext, null),
                 Optional.empty());
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
new file mode 100644
index 0000000..7eba3b46
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ManagedProfileControllerImplTest.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2022 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
+
+import android.content.pm.UserInfo
+import android.os.UserManager
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.settings.UserTracker
+import com.android.systemui.util.concurrency.FakeExecutor
+import com.android.systemui.util.time.FakeSystemClock
+import junit.framework.Assert
+import org.junit.Before
+import org.junit.Test
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+class ManagedProfileControllerImplTest : SysuiTestCase() {
+
+    private val mainExecutor: FakeExecutor = FakeExecutor(FakeSystemClock())
+
+    private lateinit var controller: ManagedProfileControllerImpl
+
+    @Mock private lateinit var userTracker: UserTracker
+    @Mock private lateinit var userManager: UserManager
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        controller = ManagedProfileControllerImpl(context, mainExecutor, userTracker, userManager)
+    }
+
+    @Test
+    fun hasWorkingProfile_isWorkModeEnabled_returnsTrue() {
+        `when`(userTracker.userId).thenReturn(1)
+        setupWorkingProfile(1)
+
+        Assert.assertEquals(true, controller.hasActiveProfile())
+    }
+
+    @Test
+    fun noWorkingProfile_isWorkModeEnabled_returnsFalse() {
+        `when`(userTracker.userId).thenReturn(1)
+
+        Assert.assertEquals(false, controller.hasActiveProfile())
+    }
+
+    @Test
+    fun listeningUserChanges_isWorkModeEnabled_returnsTrue() {
+        `when`(userTracker.userId).thenReturn(1)
+        controller.addCallback(TestCallback)
+        `when`(userTracker.userId).thenReturn(2)
+        setupWorkingProfile(2)
+
+        Assert.assertEquals(true, controller.hasActiveProfile())
+    }
+
+    private fun setupWorkingProfile(userId: Int) {
+        `when`(userManager.getEnabledProfiles(userId))
+            .thenReturn(
+                listOf(UserInfo(userId, "test_user", "", 0, UserManager.USER_TYPE_PROFILE_MANAGED))
+            )
+    }
+
+    private object TestCallback : ManagedProfileController.Callback {
+
+        override fun onManagedProfileChanged() = Unit
+
+        override fun onManagedProfileRemoved() = Unit
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
new file mode 100644
index 0000000..96a280a
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileRepositorySwitcherTest.kt
@@ -0,0 +1,219 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository
+
+import android.net.ConnectivityManager
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import android.telephony.TelephonyManager
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.demomode.DemoMode
+import com.android.systemui.demomode.DemoModeController
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoMobileConnectionsRepository
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.DemoModeMobileConnectionDataSource
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.validMobileEvent
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.prod.MobileConnectionsRepositoryImpl
+import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.kotlinArgumentCaptor
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.settings.FakeSettings
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.runBlocking
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+import org.mockito.Mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+/**
+ * The switcher acts as a dispatcher to either the `prod` or `demo` versions of the repository
+ * interface it's switching on. These tests just need to verify that the entire interface properly
+ * switches over when the value of `demoMode` changes
+ */
+@Suppress("EXPERIMENTAL_IS_NOT_ENABLED")
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(JUnit4::class)
+class MobileRepositorySwitcherTest : SysuiTestCase() {
+    private lateinit var underTest: MobileRepositorySwitcher
+    private lateinit var realRepo: MobileConnectionsRepositoryImpl
+    private lateinit var demoRepo: DemoMobileConnectionsRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Mock private lateinit var connectivityManager: ConnectivityManager
+    @Mock private lateinit var subscriptionManager: SubscriptionManager
+    @Mock private lateinit var telephonyManager: TelephonyManager
+    @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var demoModeController: DemoModeController
+
+    private val globalSettings = FakeSettings()
+    private val fakeNetworkEventsFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private val scope = CoroutineScope(IMMEDIATE)
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+
+        // Never start in demo mode
+        whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventsFlow)
+            }
+
+        realRepo =
+            MobileConnectionsRepositoryImpl(
+                connectivityManager,
+                subscriptionManager,
+                telephonyManager,
+                logger,
+                fakeBroadcastDispatcher,
+                globalSettings,
+                context,
+                IMMEDIATE,
+                scope,
+                mock(),
+            )
+
+        demoRepo =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = scope,
+                context = context,
+            )
+
+        underTest =
+            MobileRepositorySwitcher(
+                scope = scope,
+                realRepository = realRepo,
+                demoMobileConnectionsRepository = demoRepo,
+                demoModeController = demoModeController,
+            )
+    }
+
+    @After
+    fun tearDown() {
+        scope.cancel()
+    }
+
+    @Test
+    fun `active repo matches demo mode setting`() =
+        runBlocking(IMMEDIATE) {
+            whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+            var latest: MobileConnectionsRepository? = null
+            val job = underTest.activeRepo.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEqualTo(realRepo)
+
+            startDemoMode()
+
+            assertThat(latest).isEqualTo(demoRepo)
+
+            finishDemoMode()
+
+            assertThat(latest).isEqualTo(realRepo)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `subscription list updates when demo mode changes`() =
+        runBlocking(IMMEDIATE) {
+            whenever(demoModeController.isInDemoMode).thenReturn(false)
+
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList)
+                .thenReturn(listOf(SUB_1, SUB_2))
+
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            // The real subscriptions has 2 subs
+            whenever(subscriptionManager.completeActiveSubscriptionInfoList)
+                .thenReturn(listOf(SUB_1, SUB_2))
+            getSubscriptionCallback().onSubscriptionsChanged()
+
+            assertThat(latest).isEqualTo(listOf(SUB_1, SUB_2))
+
+            // Demo mode turns on, and we should see only the demo subscriptions
+            startDemoMode()
+            fakeNetworkEventsFlow.value = validMobileEvent(subId = 3)
+
+            // Demo mobile connections repository makes arbitrarily-formed subscription info
+            // objects, so just validate the data we care about
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(3)
+
+            finishDemoMode()
+
+            assertThat(latest).isEqualTo(listOf(SUB_1, SUB_2))
+
+            job.cancel()
+        }
+
+    private fun startDemoMode() {
+        whenever(demoModeController.isInDemoMode).thenReturn(true)
+        getDemoModeCallback().onDemoModeStarted()
+    }
+
+    private fun finishDemoMode() {
+        whenever(demoModeController.isInDemoMode).thenReturn(false)
+        getDemoModeCallback().onDemoModeFinished()
+    }
+
+    private fun getSubscriptionCallback(): SubscriptionManager.OnSubscriptionsChangedListener {
+        val callbackCaptor =
+            kotlinArgumentCaptor<SubscriptionManager.OnSubscriptionsChangedListener>()
+        verify(subscriptionManager)
+            .addOnSubscriptionsChangedListener(any(), callbackCaptor.capture())
+        return callbackCaptor.value
+    }
+
+    private fun getDemoModeCallback(): DemoMode {
+        val captor = kotlinArgumentCaptor<DemoMode>()
+        verify(demoModeController).addCallback(captor.capture())
+        return captor.value
+    }
+
+    companion object {
+        private val IMMEDIATE = Dispatchers.Main.immediate
+
+        private const val SUB_1_ID = 1
+        private val SUB_1 =
+            mock<SubscriptionInfo>().also { whenever(it.subscriptionId).thenReturn(SUB_1_ID) }
+
+        private const val SUB_2_ID = 2
+        private val SUB_2 =
+            mock<SubscriptionInfo>().also { whenever(it.subscriptionId).thenReturn(SUB_2_ID) }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
new file mode 100644
index 0000000..bf5ecd8
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionParameterizedTest.kt
@@ -0,0 +1,249 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.demo
+
+import android.telephony.Annotation
+import android.telephony.TelephonyManager
+import androidx.test.filters.SmallTest
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.TelephonyIcons
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.Parameterized
+import org.junit.runners.Parameterized.Parameters
+
+/**
+ * Parameterized test for all of the common values of [FakeNetworkEventModel]. This test simply
+ * verifies that passing the given model to [DemoMobileConnectionsRepository] results in the correct
+ * flows emitting from the given connection.
+ */
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(Parameterized::class)
+internal class DemoMobileConnectionParameterizedTest(private val testCase: TestCase) :
+    SysuiTestCase() {
+    private val testDispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
+    private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private lateinit var connectionsRepo: DemoMobileConnectionsRepository
+    private lateinit var underTest: DemoMobileConnectionRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Before
+    fun setUp() {
+        // The data source only provides one API, so we can mock it with a flow here for convenience
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventFlow)
+            }
+
+        connectionsRepo =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = testScope.backgroundScope,
+                context = context,
+            )
+
+        connectionsRepo.startProcessingCommands()
+    }
+
+    @After
+    fun tearDown() {
+        testScope.cancel()
+    }
+
+    @Test
+    fun demoNetworkData() =
+        testScope.runTest {
+            val networkModel =
+                FakeNetworkEventModel.Mobile(
+                    level = testCase.level,
+                    dataType = testCase.dataType,
+                    subId = testCase.subId,
+                    carrierId = testCase.carrierId,
+                    inflateStrength = testCase.inflateStrength,
+                    activity = testCase.activity,
+                    carrierNetworkChange = testCase.carrierNetworkChange,
+                )
+
+            fakeNetworkEventFlow.value = networkModel
+            underTest = connectionsRepo.getRepoForSubId(subId)
+
+            assertConnection(underTest, networkModel)
+        }
+
+    private fun assertConnection(
+        conn: DemoMobileConnectionRepository,
+        model: FakeNetworkEventModel
+    ) {
+        when (model) {
+            is FakeNetworkEventModel.Mobile -> {
+                val subscriptionModel: MobileSubscriptionModel = conn.subscriptionModelFlow.value
+                assertThat(conn.subId).isEqualTo(model.subId)
+                assertThat(subscriptionModel.cdmaLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.primaryLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.dataActivityDirection).isEqualTo(model.activity)
+                assertThat(subscriptionModel.carrierNetworkChangeActive)
+                    .isEqualTo(model.carrierNetworkChange)
+
+                // TODO(b/261029387): check these once we start handling them
+                assertThat(subscriptionModel.isEmergencyOnly).isFalse()
+                assertThat(subscriptionModel.isGsm).isFalse()
+                assertThat(subscriptionModel.dataConnectionState)
+                    .isEqualTo(DataConnectionState.Connected)
+            }
+            // MobileDisabled isn't combinatorial in nature, and is tested in
+            // DemoMobileConnectionsRepositoryTest.kt
+            else -> {}
+        }
+    }
+
+    /** Matches [FakeNetworkEventModel] */
+    internal data class TestCase(
+        val level: Int,
+        val dataType: SignalIcon.MobileIconGroup,
+        val subId: Int,
+        val carrierId: Int,
+        val inflateStrength: Boolean,
+        @Annotation.DataActivityType val activity: Int,
+        val carrierNetworkChange: Boolean,
+    ) {
+        override fun toString(): String {
+            return "INPUT(level=$level, " +
+                "dataType=${dataType.name}, " +
+                "subId=$subId, " +
+                "carrierId=$carrierId, " +
+                "inflateStrength=$inflateStrength, " +
+                "activity=$activity, " +
+                "carrierNetworkChange=$carrierNetworkChange)"
+        }
+
+        // Convenience for iterating test data and creating new cases
+        fun modifiedBy(
+            level: Int? = null,
+            dataType: SignalIcon.MobileIconGroup? = null,
+            subId: Int? = null,
+            carrierId: Int? = null,
+            inflateStrength: Boolean? = null,
+            @Annotation.DataActivityType activity: Int? = null,
+            carrierNetworkChange: Boolean? = null,
+        ): TestCase =
+            TestCase(
+                level = level ?: this.level,
+                dataType = dataType ?: this.dataType,
+                subId = subId ?: this.subId,
+                carrierId = carrierId ?: this.carrierId,
+                inflateStrength = inflateStrength ?: this.inflateStrength,
+                activity = activity ?: this.activity,
+                carrierNetworkChange = carrierNetworkChange ?: this.carrierNetworkChange
+            )
+    }
+
+    companion object {
+        private val subId = 1
+
+        private val booleanList = listOf(true, false)
+        private val levels = listOf(0, 1, 2, 3)
+        private val dataTypes =
+            listOf(
+                TelephonyIcons.THREE_G,
+                TelephonyIcons.LTE,
+                TelephonyIcons.FOUR_G,
+                TelephonyIcons.NR_5G,
+                TelephonyIcons.NR_5G_PLUS,
+            )
+        private val carrierIds = listOf(1, 10, 100)
+        private val inflateStrength = booleanList
+        private val activity =
+            listOf(
+                TelephonyManager.DATA_ACTIVITY_NONE,
+                TelephonyManager.DATA_ACTIVITY_IN,
+                TelephonyManager.DATA_ACTIVITY_OUT,
+                TelephonyManager.DATA_ACTIVITY_INOUT
+            )
+        private val carrierNetworkChange = booleanList
+
+        @Parameters(name = "{0}") @JvmStatic fun data() = testData()
+
+        /**
+         * Generate some test data. For the sake of convenience, we'll parameterize only non-null
+         * network event data. So given the lists of test data:
+         * ```
+         *    list1 = [1, 2, 3]
+         *    list2 = [false, true]
+         *    list3 = [a, b, c]
+         * ```
+         * We'll generate test cases for:
+         *
+         * Test (1, false, a) Test (2, false, a) Test (3, false, a) Test (1, true, a) Test (1,
+         * false, b) Test (1, false, c)
+         *
+         * NOTE: this is not a combinatorial product of all of the possible sets of parameters.
+         * Since this test is built to exercise demo mode, the general approach is to define a
+         * fully-formed "base case", and from there to make sure to use every valid parameter once,
+         * by defining the rest of the test cases against the base case. Specific use-cases can be
+         * added to the non-parameterized test, or manually below the generated test cases.
+         */
+        private fun testData(): List<TestCase> {
+            val testSet = mutableSetOf<TestCase>()
+
+            val baseCase =
+                TestCase(
+                    levels.first(),
+                    dataTypes.first(),
+                    subId,
+                    carrierIds.first(),
+                    inflateStrength.first(),
+                    activity.first(),
+                    carrierNetworkChange.first()
+                )
+
+            val tail =
+                sequenceOf(
+                        levels.map { baseCase.modifiedBy(level = it) },
+                        dataTypes.map { baseCase.modifiedBy(dataType = it) },
+                        carrierIds.map { baseCase.modifiedBy(carrierId = it) },
+                        inflateStrength.map { baseCase.modifiedBy(inflateStrength = it) },
+                        activity.map { baseCase.modifiedBy(activity = it) },
+                        carrierNetworkChange.map { baseCase.modifiedBy(carrierNetworkChange = it) },
+                    )
+                    .flatten()
+
+            testSet.add(baseCase)
+            tail.toCollection(testSet)
+
+            return testSet.toList()
+        }
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
new file mode 100644
index 0000000..a8f6993
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/demo/DemoMobileConnectionsRepositoryTest.kt
@@ -0,0 +1,306 @@
+/*
+ * Copyright (C) 2022 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.pipeline.mobile.data.repository.demo
+
+import android.telephony.SubscriptionInfo
+import android.telephony.TelephonyManager.DATA_ACTIVITY_INOUT
+import android.telephony.TelephonyManager.UNKNOWN_CARRIER_ID
+import androidx.test.filters.SmallTest
+import com.android.settingslib.SignalIcon
+import com.android.settingslib.mobile.TelephonyIcons.THREE_G
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.statusbar.pipeline.mobile.data.model.DataConnectionState
+import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.demo.model.FakeNetworkEventModel.MobileDisabled
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import junit.framework.Assert
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+class DemoMobileConnectionsRepositoryTest : SysuiTestCase() {
+    private val testDispatcher = UnconfinedTestDispatcher()
+    private val testScope = TestScope(testDispatcher)
+
+    private val fakeNetworkEventFlow = MutableStateFlow<FakeNetworkEventModel?>(null)
+
+    private lateinit var underTest: DemoMobileConnectionsRepository
+    private lateinit var mockDataSource: DemoModeMobileConnectionDataSource
+
+    @Before
+    fun setUp() {
+        // The data source only provides one API, so we can mock it with a flow here for convenience
+        mockDataSource =
+            mock<DemoModeMobileConnectionDataSource>().also {
+                whenever(it.mobileEvents).thenReturn(fakeNetworkEventFlow)
+            }
+
+        underTest =
+            DemoMobileConnectionsRepository(
+                dataSource = mockDataSource,
+                scope = testScope.backgroundScope,
+                context = context,
+            )
+
+        underTest.startProcessingCommands()
+    }
+
+    @Test
+    fun `network event - create new subscription`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEmpty()
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `network event - reuses subscription when same Id`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            assertThat(latest).isEmpty()
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            // Second network event comes in with the same subId, does not create a new subscription
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 2)
+
+            assertThat(latest).hasSize(1)
+            assertThat(latest!![0].subscriptionId).isEqualTo(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `multiple subscriptions`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2)
+
+            assertThat(latest).hasSize(2)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId specified - single conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = 1)
+
+            assertThat(latest).hasSize(0)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId not specified - single conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = null)
+
+            assertThat(latest).hasSize(0)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - disables connection - subId specified - multiple conn`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = 2)
+
+            assertThat(latest).hasSize(1)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `mobile disabled event - subId not specified - multiple conn - ignores command`() =
+        testScope.runTest {
+            var latest: List<SubscriptionInfo>? = null
+            val job = underTest.subscriptionsFlow.onEach { latest = it }.launchIn(this)
+
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 1, level = 1)
+            fakeNetworkEventFlow.value = validMobileEvent(subId = 2, level = 1)
+
+            fakeNetworkEventFlow.value = MobileDisabled(subId = null)
+
+            assertThat(latest).hasSize(2)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `demo connection - single subscription`() =
+        testScope.runTest {
+            var currentEvent: FakeNetworkEventModel = validMobileEvent(subId = 1)
+            var connections: List<DemoMobileConnectionRepository>? = null
+            val job =
+                underTest.subscriptionsFlow
+                    .onEach { infos ->
+                        connections =
+                            infos.map { info -> underTest.getRepoForSubId(info.subscriptionId) }
+                    }
+                    .launchIn(this)
+
+            fakeNetworkEventFlow.value = currentEvent
+
+            assertThat(connections).hasSize(1)
+            val connection1 = connections!![0]
+
+            assertConnection(connection1, currentEvent)
+
+            // Exercise the whole api
+
+            currentEvent = validMobileEvent(subId = 1, level = 2)
+            fakeNetworkEventFlow.value = currentEvent
+            assertConnection(connection1, currentEvent)
+
+            job.cancel()
+        }
+
+    @Test
+    fun `demo connection - two connections - update second - no affect on first`() =
+        testScope.runTest {
+            var currentEvent1 = validMobileEvent(subId = 1)
+            var connection1: DemoMobileConnectionRepository? = null
+            var currentEvent2 = validMobileEvent(subId = 2)
+            var connection2: DemoMobileConnectionRepository? = null
+            var connections: List<DemoMobileConnectionRepository>? = null
+            val job =
+                underTest.subscriptionsFlow
+                    .onEach { infos ->
+                        connections =
+                            infos.map { info -> underTest.getRepoForSubId(info.subscriptionId) }
+                    }
+                    .launchIn(this)
+
+            fakeNetworkEventFlow.value = currentEvent1
+            fakeNetworkEventFlow.value = currentEvent2
+            assertThat(connections).hasSize(2)
+            connections!!.forEach {
+                if (it.subId == 1) {
+                    connection1 = it
+                } else if (it.subId == 2) {
+                    connection2 = it
+                } else {
+                    Assert.fail("Unexpected subscription")
+                }
+            }
+
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            // WHEN the event changes for connection 2, it updates, and connection 1 stays the same
+            currentEvent2 = validMobileEvent(subId = 2, activity = DATA_ACTIVITY_INOUT)
+            fakeNetworkEventFlow.value = currentEvent2
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            // and vice versa
+            currentEvent1 = validMobileEvent(subId = 1, inflateStrength = true)
+            fakeNetworkEventFlow.value = currentEvent1
+            assertConnection(connection1!!, currentEvent1)
+            assertConnection(connection2!!, currentEvent2)
+
+            job.cancel()
+        }
+
+    private fun assertConnection(
+        conn: DemoMobileConnectionRepository,
+        model: FakeNetworkEventModel
+    ) {
+        when (model) {
+            is FakeNetworkEventModel.Mobile -> {
+                val subscriptionModel: MobileSubscriptionModel = conn.subscriptionModelFlow.value
+                assertThat(conn.subId).isEqualTo(model.subId)
+                assertThat(subscriptionModel.cdmaLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.primaryLevel).isEqualTo(model.level)
+                assertThat(subscriptionModel.dataActivityDirection).isEqualTo(model.activity)
+                assertThat(subscriptionModel.carrierNetworkChangeActive)
+                    .isEqualTo(model.carrierNetworkChange)
+
+                // TODO(b/261029387) check these once we start handling them
+                assertThat(subscriptionModel.isEmergencyOnly).isFalse()
+                assertThat(subscriptionModel.isGsm).isFalse()
+                assertThat(subscriptionModel.dataConnectionState)
+                    .isEqualTo(DataConnectionState.Connected)
+            }
+            else -> {}
+        }
+    }
+}
+
+/** Convenience to create a valid fake network event with minimal params */
+fun validMobileEvent(
+    level: Int? = 1,
+    dataType: SignalIcon.MobileIconGroup? = THREE_G,
+    subId: Int? = 1,
+    carrierId: Int? = UNKNOWN_CARRIER_ID,
+    inflateStrength: Boolean? = false,
+    activity: Int? = null,
+    carrierNetworkChange: Boolean = false,
+): FakeNetworkEventModel =
+    FakeNetworkEventModel.Mobile(
+        level = level,
+        dataType = dataType,
+        subId = subId,
+        carrierId = carrierId,
+        inflateStrength = inflateStrength,
+        activity = activity,
+        carrierNetworkChange = carrierNetworkChange,
+    )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
similarity index 95%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
index 5ce51bb..c8df5ac 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionRepositoryTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.pipeline.mobile.data.repository
+package com.android.systemui.statusbar.pipeline.mobile.data.repository.prod
 
 import android.os.UserHandle
 import android.provider.Settings
@@ -31,6 +31,7 @@
 import android.telephony.TelephonyManager.DATA_CONNECTING
 import android.telephony.TelephonyManager.DATA_DISCONNECTED
 import android.telephony.TelephonyManager.DATA_DISCONNECTING
+import android.telephony.TelephonyManager.DATA_UNKNOWN
 import android.telephony.TelephonyManager.NETWORK_TYPE_LTE
 import android.telephony.TelephonyManager.NETWORK_TYPE_UNKNOWN
 import androidx.test.filters.SmallTest
@@ -39,6 +40,7 @@
 import com.android.systemui.statusbar.pipeline.mobile.data.model.DefaultNetworkType
 import com.android.systemui.statusbar.pipeline.mobile.data.model.MobileSubscriptionModel
 import com.android.systemui.statusbar.pipeline.mobile.data.model.OverrideNetworkType
+import com.android.systemui.statusbar.pipeline.mobile.data.repository.FakeMobileConnectionsRepository
 import com.android.systemui.statusbar.pipeline.shared.ConnectivityPipelineLogger
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.argumentCaptor
@@ -221,6 +223,21 @@
         }
 
     @Test
+    fun testFlowForSubId_dataConnectionState_unknown() =
+        runBlocking(IMMEDIATE) {
+            var latest: MobileSubscriptionModel? = null
+            val job = underTest.subscriptionModelFlow.onEach { latest = it }.launchIn(this)
+
+            val callback =
+                getTelephonyCallbackForType<TelephonyCallback.DataConnectionStateListener>()
+            callback.onDataConnectionStateChanged(DATA_UNKNOWN, 200 /* unused */)
+
+            assertThat(latest?.dataConnectionState).isEqualTo(DataConnectionState.Unknown)
+
+            job.cancel()
+        }
+
+    @Test
     fun testFlowForSubId_dataActivity() =
         runBlocking(IMMEDIATE) {
             var latest: MobileSubscriptionModel? = null
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
similarity index 99%
rename from packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt
rename to packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
index a953a3d..359ea18 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/MobileConnectionsRepositoryTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/mobile/data/repository/prod/MobileConnectionsRepositoryTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.systemui.statusbar.pipeline.mobile.data.repository
+package com.android.systemui.statusbar.pipeline.mobile.data.repository.prod
 
 import android.content.Intent
 import android.net.ConnectivityManager
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
index 5c16e129..3d9fd96 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/view/ModernStatusBarWifiViewTest.kt
@@ -25,6 +25,7 @@
 import com.android.systemui.R
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.lifecycle.InstantTaskExecutorRule
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.StatusBarIconView.STATE_DOT
 import com.android.systemui.statusbar.StatusBarIconView.STATE_HIDDEN
 import com.android.systemui.statusbar.StatusBarIconView.STATE_ICON
@@ -64,6 +65,7 @@
     private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock
     private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock
     private lateinit var connectivityConstants: ConnectivityConstants
     @Mock
@@ -103,6 +105,7 @@
             connectivityConstants,
             context,
             logger,
+            tableLogBuffer,
             interactor,
             scope,
             statusBarPipelineFlags,
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
index 3001b81..12b93819 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelIconParameterizedTest.kt
@@ -23,6 +23,7 @@
 import com.android.settingslib.AccessibilityContentDescriptions.WIFI_NO_CONNECTION
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.common.shared.model.ContentDescription.Companion.loadContentDescription
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_FULL_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_INTERNET_ICONS
 import com.android.systemui.statusbar.connectivity.WifiIcons.WIFI_NO_NETWORK
@@ -40,6 +41,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractor
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractorImpl
 import com.android.systemui.statusbar.pipeline.wifi.shared.WifiConstants
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.android.systemui.statusbar.pipeline.wifi.ui.viewmodel.WifiViewModel.Companion.NO_INTERNET
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
@@ -67,6 +69,7 @@
 
     @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -123,6 +126,7 @@
                     connectivityConstants,
                     context,
                     logger,
+                    tableLogBuffer,
                     interactor,
                     scope,
                     statusBarPipelineFlags,
@@ -137,15 +141,21 @@
             yield()
 
             // THEN we get the expected icon
-            assertThat(iconFlow.value?.res).isEqualTo(testCase.expected?.iconResource)
-            val expectedContentDescription =
-                if (testCase.expected == null) {
-                    null
-                } else {
-                    testCase.expected.contentDescription.invoke(context)
+            val actualIcon = iconFlow.value
+            when (testCase.expected) {
+                null -> {
+                    assertThat(actualIcon).isInstanceOf(WifiIcon.Hidden::class.java)
                 }
-            assertThat(iconFlow.value?.contentDescription?.loadContentDescription(context))
-                .isEqualTo(expectedContentDescription)
+                else -> {
+                    assertThat(actualIcon).isInstanceOf(WifiIcon.Visible::class.java)
+                    val actualIconVisible = actualIcon as WifiIcon.Visible
+                    assertThat(actualIconVisible.icon.res).isEqualTo(testCase.expected.iconResource)
+                    val expectedContentDescription =
+                        testCase.expected.contentDescription.invoke(context)
+                    assertThat(actualIconVisible.contentDescription.loadContentDescription(context))
+                        .isEqualTo(expectedContentDescription)
+                }
+            }
 
             job.cancel()
         }
@@ -174,7 +184,7 @@
         val isDefault: Boolean = false,
         val network: WifiNetworkModel,
 
-        /** The expected output. Null if we expect the output to be null. */
+        /** The expected output. Null if we expect the output to be hidden. */
         val expected: Expected?
     ) {
         override fun toString(): String {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
index 6a6b2a8..7502020 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/pipeline/wifi/ui/viewmodel/WifiViewModelTest.kt
@@ -18,7 +18,7 @@
 
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
-import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.log.table.TableLogBuffer
 import com.android.systemui.statusbar.pipeline.StatusBarPipelineFlags
 import com.android.systemui.statusbar.pipeline.airplane.data.repository.FakeAirplaneModeRepository
 import com.android.systemui.statusbar.pipeline.airplane.domain.interactor.AirplaneModeInteractor
@@ -34,6 +34,7 @@
 import com.android.systemui.statusbar.pipeline.wifi.domain.interactor.WifiInteractorImpl
 import com.android.systemui.statusbar.pipeline.wifi.shared.WifiConstants
 import com.android.systemui.statusbar.pipeline.wifi.shared.model.WifiActivityModel
+import com.android.systemui.statusbar.pipeline.wifi.ui.model.WifiIcon
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -59,6 +60,7 @@
 
     @Mock private lateinit var statusBarPipelineFlags: StatusBarPipelineFlags
     @Mock private lateinit var logger: ConnectivityPipelineLogger
+    @Mock private lateinit var tableLogBuffer: TableLogBuffer
     @Mock private lateinit var connectivityConstants: ConnectivityConstants
     @Mock private lateinit var wifiConstants: WifiConstants
     private lateinit var airplaneModeRepository: FakeAirplaneModeRepository
@@ -103,21 +105,21 @@
 
     @Test
     fun wifiIcon_allLocationViewModelsReceiveSameData() = runBlocking(IMMEDIATE) {
-        var latestHome: Icon? = null
+        var latestHome: WifiIcon? = null
         val jobHome = underTest
             .home
             .wifiIcon
             .onEach { latestHome = it }
             .launchIn(this)
 
-        var latestKeyguard: Icon? = null
+        var latestKeyguard: WifiIcon? = null
         val jobKeyguard = underTest
             .keyguard
             .wifiIcon
             .onEach { latestKeyguard = it }
             .launchIn(this)
 
-        var latestQs: Icon? = null
+        var latestQs: WifiIcon? = null
         val jobQs = underTest
             .qs
             .wifiIcon
@@ -133,7 +135,7 @@
         )
         yield()
 
-        assertThat(latestHome).isInstanceOf(Icon.Resource::class.java)
+        assertThat(latestHome).isInstanceOf(WifiIcon.Visible::class.java)
         assertThat(latestHome).isEqualTo(latestKeyguard)
         assertThat(latestKeyguard).isEqualTo(latestQs)
 
@@ -541,6 +543,7 @@
             connectivityConstants,
             context,
             logger,
+            tableLogBuffer,
             interactor,
             scope,
             statusBarPipelineFlags,
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index bc083f1..19b5cc9 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -369,12 +369,10 @@
 import com.android.internal.util.FrameworkStatsLog;
 import com.android.internal.util.MemInfoReader;
 import com.android.internal.util.Preconditions;
-import com.android.internal.util.function.DecFunction;
 import com.android.internal.util.function.HeptFunction;
 import com.android.internal.util.function.HexFunction;
 import com.android.internal.util.function.QuadFunction;
 import com.android.internal.util.function.QuintFunction;
-import com.android.internal.util.function.TriFunction;
 import com.android.internal.util.function.UndecFunction;
 import com.android.server.AlarmManagerInternal;
 import com.android.server.DeviceIdleInternal;
@@ -18329,19 +18327,20 @@
         }
 
         @Override
-        public SyncNotedAppOp startProxyOperation(int code,
+        public SyncNotedAppOp startProxyOperation(@NonNull IBinder clientId, int code,
                 @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
                 boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
                 boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
                 @AttributionFlags int proxiedAttributionFlags, int attributionChainId,
-                @NonNull DecFunction<Integer, AttributionSource, Boolean, Boolean, String, Boolean,
-                        Boolean, Integer, Integer, Integer, SyncNotedAppOp> superImpl) {
+                @NonNull UndecFunction<IBinder, Integer, AttributionSource,
+                        Boolean, Boolean, String, Boolean, Boolean, Integer, Integer, Integer,
+                        SyncNotedAppOp> superImpl) {
             if (attributionSource.getUid() == mTargetUid && isTargetOp(code)) {
                 final int shellUid = UserHandle.getUid(UserHandle.getUserId(
                         attributionSource.getUid()), Process.SHELL_UID);
                 final long identity = Binder.clearCallingIdentity();
                 try {
-                    return superImpl.apply(code, new AttributionSource(shellUid,
+                    return superImpl.apply(clientId, code, new AttributionSource(shellUid,
                             "com.android.shell", attributionSource.getAttributionTag(),
                             attributionSource.getToken(), attributionSource.getNext()),
                             startIfModeDefault, shouldCollectAsyncNotedOp, message,
@@ -18351,21 +18350,22 @@
                     Binder.restoreCallingIdentity(identity);
                 }
             }
-            return superImpl.apply(code, attributionSource, startIfModeDefault,
+            return superImpl.apply(clientId, code, attributionSource, startIfModeDefault,
                     shouldCollectAsyncNotedOp, message, shouldCollectMessage, skipProxyOperation,
                     proxyAttributionFlags, proxiedAttributionFlags, attributionChainId);
         }
 
         @Override
-        public void finishProxyOperation(int code, @NonNull AttributionSource attributionSource,
-                boolean skipProxyOperation, @NonNull TriFunction<Integer, AttributionSource,
-                        Boolean, Void> superImpl) {
+        public void finishProxyOperation(@NonNull IBinder clientId, int code,
+                @NonNull AttributionSource attributionSource, boolean skipProxyOperation,
+                @NonNull QuadFunction<IBinder, Integer, AttributionSource, Boolean,
+                        Void> superImpl) {
             if (attributionSource.getUid() == mTargetUid && isTargetOp(code)) {
                 final int shellUid = UserHandle.getUid(UserHandle.getUserId(
                         attributionSource.getUid()), Process.SHELL_UID);
                 final long identity = Binder.clearCallingIdentity();
                 try {
-                    superImpl.apply(code, new AttributionSource(shellUid,
+                    superImpl.apply(clientId, code, new AttributionSource(shellUid,
                             "com.android.shell", attributionSource.getAttributionTag(),
                             attributionSource.getToken(), attributionSource.getNext()),
                             skipProxyOperation);
@@ -18373,7 +18373,7 @@
                     Binder.restoreCallingIdentity(identity);
                 }
             }
-            superImpl.apply(code, attributionSource, skipProxyOperation);
+            superImpl.apply(clientId, code, attributionSource, skipProxyOperation);
         }
 
         private boolean isTargetOp(int code) {
diff --git a/services/core/java/com/android/server/appop/AppOpsService.java b/services/core/java/com/android/server/appop/AppOpsService.java
index e31c952..20f0c17 100644
--- a/services/core/java/com/android/server/appop/AppOpsService.java
+++ b/services/core/java/com/android/server/appop/AppOpsService.java
@@ -3865,18 +3865,18 @@
     }
 
     @Override
-    public SyncNotedAppOp startProxyOperation(int code,
+    public SyncNotedAppOp startProxyOperation(@NonNull IBinder clientId, int code,
             @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
             boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
             boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
             @AttributionFlags int proxiedAttributionFlags, int attributionChainId) {
-        return mCheckOpsDelegateDispatcher.startProxyOperation(code, attributionSource,
+        return mCheckOpsDelegateDispatcher.startProxyOperation(clientId, code, attributionSource,
                 startIfModeDefault, shouldCollectAsyncNotedOp, message, shouldCollectMessage,
                 skipProxyOperation, proxyAttributionFlags, proxiedAttributionFlags,
                 attributionChainId);
     }
 
-    private SyncNotedAppOp startProxyOperationImpl(int code,
+    private SyncNotedAppOp startProxyOperationImpl(@NonNull IBinder clientId, int code,
             @NonNull AttributionSource attributionSource,
             boolean startIfModeDefault, boolean shouldCollectAsyncNotedOp, String message,
             boolean shouldCollectMessage, boolean skipProxyOperation, @AttributionFlags
@@ -3885,11 +3885,9 @@
         final int proxyUid = attributionSource.getUid();
         final String proxyPackageName = attributionSource.getPackageName();
         final String proxyAttributionTag = attributionSource.getAttributionTag();
-        final IBinder proxyToken = attributionSource.getToken();
         final int proxiedUid = attributionSource.getNextUid();
         final String proxiedPackageName = attributionSource.getNextPackageName();
         final String proxiedAttributionTag = attributionSource.getNextAttributionTag();
-        final IBinder proxiedToken = attributionSource.getNextToken();
 
         verifyIncomingProxyUid(attributionSource);
         verifyIncomingOp(code);
@@ -3928,7 +3926,7 @@
 
         if (!skipProxyOperation) {
             // Test if the proxied operation will succeed before starting the proxy operation
-            final SyncNotedAppOp testProxiedOp = startOperationUnchecked(proxiedToken, code,
+            final SyncNotedAppOp testProxiedOp = startOperationUnchecked(clientId, code,
                     proxiedUid, resolvedProxiedPackageName, proxiedAttributionTag, proxyUid,
                     resolvedProxyPackageName, proxyAttributionTag, proxiedFlags, startIfModeDefault,
                     shouldCollectAsyncNotedOp, message, shouldCollectMessage,
@@ -3940,7 +3938,7 @@
             final int proxyFlags = isProxyTrusted ? AppOpsManager.OP_FLAG_TRUSTED_PROXY
                     : AppOpsManager.OP_FLAG_UNTRUSTED_PROXY;
 
-            final SyncNotedAppOp proxyAppOp = startOperationUnchecked(proxyToken, code, proxyUid,
+            final SyncNotedAppOp proxyAppOp = startOperationUnchecked(clientId, code, proxyUid,
                     resolvedProxyPackageName, proxyAttributionTag, Process.INVALID_UID, null, null,
                     proxyFlags, startIfModeDefault, !isProxyTrusted, "proxy " + message,
                     shouldCollectMessage, proxyAttributionFlags, attributionChainId,
@@ -3950,7 +3948,7 @@
             }
         }
 
-        return startOperationUnchecked(proxiedToken, code, proxiedUid, resolvedProxiedPackageName,
+        return startOperationUnchecked(clientId, code, proxiedUid, resolvedProxiedPackageName,
                 proxiedAttributionTag, proxyUid, resolvedProxyPackageName, proxyAttributionTag,
                 proxiedFlags, startIfModeDefault, shouldCollectAsyncNotedOp, message,
                 shouldCollectMessage, proxiedAttributionFlags, attributionChainId,
@@ -4091,22 +4089,20 @@
     }
 
     @Override
-    public void finishProxyOperation(int code, @NonNull AttributionSource attributionSource,
-            boolean skipProxyOperation) {
-        mCheckOpsDelegateDispatcher.finishProxyOperation(code, attributionSource,
+    public void finishProxyOperation(@NonNull IBinder clientId, int code,
+            @NonNull AttributionSource attributionSource, boolean skipProxyOperation) {
+        mCheckOpsDelegateDispatcher.finishProxyOperation(clientId, code, attributionSource,
                 skipProxyOperation);
     }
 
-    private Void finishProxyOperationImpl(int code, @NonNull AttributionSource attributionSource,
-            boolean skipProxyOperation) {
+    private Void finishProxyOperationImpl(IBinder clientId, int code,
+            @NonNull AttributionSource attributionSource, boolean skipProxyOperation) {
         final int proxyUid = attributionSource.getUid();
         final String proxyPackageName = attributionSource.getPackageName();
         final String proxyAttributionTag = attributionSource.getAttributionTag();
-        final IBinder proxyToken = attributionSource.getToken();
         final int proxiedUid = attributionSource.getNextUid();
         final String proxiedPackageName = attributionSource.getNextPackageName();
         final String proxiedAttributionTag = attributionSource.getNextAttributionTag();
-        final IBinder proxiedToken = attributionSource.getNextToken();
 
         skipProxyOperation = skipProxyOperation
                 && isCallerAndAttributionTrusted(attributionSource);
@@ -4123,7 +4119,7 @@
         }
 
         if (!skipProxyOperation) {
-            finishOperationUnchecked(proxyToken, code, proxyUid, resolvedProxyPackageName,
+            finishOperationUnchecked(clientId, code, proxyUid, resolvedProxyPackageName,
                     proxyAttributionTag);
         }
 
@@ -4133,7 +4129,7 @@
             return null;
         }
 
-        finishOperationUnchecked(proxiedToken, code, proxiedUid, resolvedProxiedPackageName,
+        finishOperationUnchecked(clientId, code, proxiedUid, resolvedProxiedPackageName,
                 proxiedAttributionTag);
 
         return null;
@@ -7726,42 +7722,42 @@
                     attributionFlags, attributionChainId, AppOpsService.this::startOperationImpl);
         }
 
-        public SyncNotedAppOp startProxyOperation(int code,
+        public SyncNotedAppOp startProxyOperation(@NonNull IBinder clientId, int code,
                 @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
                 boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
                 boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
                 @AttributionFlags int proxiedAttributionFlags, int attributionChainId) {
             if (mPolicy != null) {
                 if (mCheckOpsDelegate != null) {
-                    return mPolicy.startProxyOperation(code, attributionSource,
+                    return mPolicy.startProxyOperation(clientId, code, attributionSource,
                             startIfModeDefault, shouldCollectAsyncNotedOp, message,
                             shouldCollectMessage, skipProxyOperation, proxyAttributionFlags,
                             proxiedAttributionFlags, attributionChainId,
                             this::startDelegateProxyOperationImpl);
                 } else {
-                    return mPolicy.startProxyOperation(code, attributionSource,
+                    return mPolicy.startProxyOperation(clientId, code, attributionSource,
                             startIfModeDefault, shouldCollectAsyncNotedOp, message,
                             shouldCollectMessage, skipProxyOperation, proxyAttributionFlags,
                             proxiedAttributionFlags, attributionChainId,
                             AppOpsService.this::startProxyOperationImpl);
                 }
             } else if (mCheckOpsDelegate != null) {
-                return startDelegateProxyOperationImpl(code, attributionSource,
+                return startDelegateProxyOperationImpl(clientId, code, attributionSource,
                         startIfModeDefault, shouldCollectAsyncNotedOp, message,
                         shouldCollectMessage, skipProxyOperation, proxyAttributionFlags,
                         proxiedAttributionFlags, attributionChainId);
             }
-            return startProxyOperationImpl(code, attributionSource, startIfModeDefault,
+            return startProxyOperationImpl(clientId, code, attributionSource, startIfModeDefault,
                     shouldCollectAsyncNotedOp, message, shouldCollectMessage, skipProxyOperation,
                     proxyAttributionFlags, proxiedAttributionFlags, attributionChainId);
         }
 
-        private SyncNotedAppOp startDelegateProxyOperationImpl(int code,
+        private SyncNotedAppOp startDelegateProxyOperationImpl(@NonNull IBinder clientId, int code,
                 @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
                 boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
                 boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
                 @AttributionFlags int proxiedAttributionFlsgs, int attributionChainId) {
-            return mCheckOpsDelegate.startProxyOperation(code, attributionSource,
+            return mCheckOpsDelegate.startProxyOperation(clientId, code, attributionSource,
                     startIfModeDefault, shouldCollectAsyncNotedOp, message, shouldCollectMessage,
                     skipProxyOperation, proxyAttributionFlags, proxiedAttributionFlsgs,
                     attributionChainId, AppOpsService.this::startProxyOperationImpl);
@@ -7790,27 +7786,28 @@
                     AppOpsService.this::finishOperationImpl);
         }
 
-        public void finishProxyOperation(int code,
+        public void finishProxyOperation(@NonNull IBinder clientId, int code,
                 @NonNull AttributionSource attributionSource, boolean skipProxyOperation) {
             if (mPolicy != null) {
                 if (mCheckOpsDelegate != null) {
-                    mPolicy.finishProxyOperation(code, attributionSource,
+                    mPolicy.finishProxyOperation(clientId, code, attributionSource,
                             skipProxyOperation, this::finishDelegateProxyOperationImpl);
                 } else {
-                    mPolicy.finishProxyOperation(code, attributionSource,
+                    mPolicy.finishProxyOperation(clientId, code, attributionSource,
                             skipProxyOperation, AppOpsService.this::finishProxyOperationImpl);
                 }
             } else if (mCheckOpsDelegate != null) {
-                finishDelegateProxyOperationImpl(code, attributionSource, skipProxyOperation);
+                finishDelegateProxyOperationImpl(clientId, code, attributionSource,
+                        skipProxyOperation);
             } else {
-                finishProxyOperationImpl(code, attributionSource, skipProxyOperation);
+                finishProxyOperationImpl(clientId, code, attributionSource, skipProxyOperation);
             }
         }
 
-        private Void finishDelegateProxyOperationImpl(int code,
+        private Void finishDelegateProxyOperationImpl(@NonNull IBinder clientId, int code,
                 @NonNull AttributionSource attributionSource, boolean skipProxyOperation) {
-            mCheckOpsDelegate.finishProxyOperation(code, attributionSource, skipProxyOperation,
-                    AppOpsService.this::finishProxyOperationImpl);
+            mCheckOpsDelegate.finishProxyOperation(clientId, code, attributionSource,
+                    skipProxyOperation, AppOpsService.this::finishProxyOperationImpl);
             return null;
         }
     }
diff --git a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
index 37538db..d648d6f 100644
--- a/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
+++ b/services/core/java/com/android/server/pm/permission/PermissionManagerService.java
@@ -1099,7 +1099,7 @@
                     if (resolvedPackageName == null) {
                         return;
                     }
-                    appOpsManager.finishOp(accessorSource.getToken(), op,
+                    appOpsManager.finishOp(attributionSourceState.token, op,
                             accessorSource.getUid(), resolvedPackageName,
                             accessorSource.getAttributionTag());
                 } else {
@@ -1108,8 +1108,9 @@
                     if (resolvedAttributionSource.getPackageName() == null) {
                         return;
                     }
-                    appOpsManager.finishProxyOp(AppOpsManager.opToPublicName(op),
-                            resolvedAttributionSource, skipCurrentFinish);
+                    appOpsManager.finishProxyOp(attributionSourceState.token,
+                            AppOpsManager.opToPublicName(op), resolvedAttributionSource,
+                            skipCurrentFinish);
                 }
 
                 RegisteredAttribution registered =
@@ -1225,10 +1226,11 @@
                         && next.getNext() == null);
                 final boolean selfAccess = singleReceiverFromDatasource || next == null;
 
-                final int opMode = performOpTransaction(context, op, current, message,
-                        forDataDelivery, /*startDataDelivery*/ false, skipCurrentChecks,
-                        selfAccess, singleReceiverFromDatasource, AppOpsManager.OP_NONE,
-                        AppOpsManager.ATTRIBUTION_FLAGS_NONE, AppOpsManager.ATTRIBUTION_FLAGS_NONE,
+                final int opMode = performOpTransaction(context, attributionSource.getToken(), op,
+                        current, message, forDataDelivery, /*startDataDelivery*/ false,
+                        skipCurrentChecks, selfAccess, singleReceiverFromDatasource,
+                        AppOpsManager.OP_NONE, AppOpsManager.ATTRIBUTION_FLAGS_NONE,
+                        AppOpsManager.ATTRIBUTION_FLAGS_NONE,
                         AppOpsManager.ATTRIBUTION_CHAIN_ID_NONE);
 
                 switch (opMode) {
@@ -1331,10 +1333,10 @@
                         attributionSource, next, fromDatasource, startDataDelivery, selfAccess,
                         isLinkTrusted) : ATTRIBUTION_FLAGS_NONE;
 
-                final int opMode = performOpTransaction(context, op, current, message,
-                        forDataDelivery, startDataDelivery, skipCurrentChecks, selfAccess,
-                        singleReceiverFromDatasource, attributedOp, proxyAttributionFlags,
-                        proxiedAttributionFlags, attributionChainId);
+                final int opMode = performOpTransaction(context, attributionSource.getToken(), op,
+                        current, message, forDataDelivery, startDataDelivery, skipCurrentChecks,
+                        selfAccess, singleReceiverFromDatasource, attributedOp,
+                        proxyAttributionFlags, proxiedAttributionFlags, attributionChainId);
 
                 switch (opMode) {
                     case AppOpsManager.MODE_ERRORED: {
@@ -1479,8 +1481,8 @@
                         attributionSource, next, /*fromDatasource*/ false, startDataDelivery,
                         selfAccess, isLinkTrusted) : ATTRIBUTION_FLAGS_NONE;
 
-                final int opMode = performOpTransaction(context, op, current, message,
-                        forDataDelivery, startDataDelivery, skipCurrentChecks, selfAccess,
+                final int opMode = performOpTransaction(context, current.getToken(), op, current,
+                        message, forDataDelivery, startDataDelivery, skipCurrentChecks, selfAccess,
                         /*fromDatasource*/ false, AppOpsManager.OP_NONE, proxyAttributionFlags,
                         proxiedAttributionFlags, attributionChainId);
 
@@ -1502,7 +1504,8 @@
         }
 
         @SuppressWarnings("ConstantConditions")
-        private static int performOpTransaction(@NonNull Context context, int op,
+        private static int performOpTransaction(@NonNull Context context,
+                @NonNull IBinder chainStartToken, int op,
                 @NonNull AttributionSource attributionSource, @Nullable String message,
                 boolean forDataDelivery, boolean startDataDelivery, boolean skipProxyOperation,
                 boolean selfAccess, boolean singleReceiverFromDatasource, int attributedOp,
@@ -1564,7 +1567,7 @@
                 if (selfAccess) {
                     try {
                         startedOpResult = appOpsManager.startOpNoThrow(
-                                resolvedAttributionSource.getToken(), startedOp,
+                                chainStartToken, startedOp,
                                 resolvedAttributionSource.getUid(),
                                 resolvedAttributionSource.getPackageName(),
                                 /*startIfModeDefault*/ false,
@@ -1575,14 +1578,14 @@
                                 + " platform defined runtime permission "
                                 + AppOpsManager.opToPermission(op) + " while not having "
                                 + Manifest.permission.UPDATE_APP_OPS_STATS);
-                        startedOpResult = appOpsManager.startProxyOpNoThrow(attributedOp,
-                                attributionSource, message, skipProxyOperation,
+                        startedOpResult = appOpsManager.startProxyOpNoThrow(chainStartToken,
+                                attributedOp, attributionSource, message, skipProxyOperation,
                                 proxyAttributionFlags, proxiedAttributionFlags, attributionChainId);
                     }
                 } else {
                     try {
-                        startedOpResult = appOpsManager.startProxyOpNoThrow(startedOp,
-                                resolvedAttributionSource, message, skipProxyOperation,
+                        startedOpResult = appOpsManager.startProxyOpNoThrow(chainStartToken,
+                                startedOp, resolvedAttributionSource, message, skipProxyOperation,
                                 proxyAttributionFlags, proxiedAttributionFlags, attributionChainId);
                     } catch (SecurityException e) {
                         //TODO 195339480: remove
diff --git a/services/core/java/com/android/server/policy/AppOpsPolicy.java b/services/core/java/com/android/server/policy/AppOpsPolicy.java
index ebd9126..b26c1b9 100644
--- a/services/core/java/com/android/server/policy/AppOpsPolicy.java
+++ b/services/core/java/com/android/server/policy/AppOpsPolicy.java
@@ -45,13 +45,11 @@
 import android.util.SparseArray;
 
 import com.android.internal.annotations.GuardedBy;
-import com.android.internal.util.function.DecFunction;
 import com.android.internal.util.function.HeptFunction;
 import com.android.internal.util.function.HexFunction;
 import com.android.internal.util.function.QuadFunction;
 import com.android.internal.util.function.QuintConsumer;
 import com.android.internal.util.function.QuintFunction;
-import com.android.internal.util.function.TriFunction;
 import com.android.internal.util.function.UndecFunction;
 import com.android.server.LocalServices;
 
@@ -256,14 +254,14 @@
     }
 
     @Override
-    public SyncNotedAppOp startProxyOperation(int code,
+    public SyncNotedAppOp startProxyOperation(@NonNull IBinder clientId, int code,
             @NonNull AttributionSource attributionSource, boolean startIfModeDefault,
             boolean shouldCollectAsyncNotedOp, String message, boolean shouldCollectMessage,
             boolean skipProxyOperation, @AttributionFlags int proxyAttributionFlags,
             @AttributionFlags int proxiedAttributionFlags, int attributionChainId,
-            @NonNull DecFunction<Integer, AttributionSource, Boolean, Boolean, String, Boolean,
-                    Boolean, Integer, Integer, Integer, SyncNotedAppOp> superImpl) {
-        return superImpl.apply(resolveDatasourceOp(code, attributionSource.getUid(),
+            @NonNull UndecFunction<IBinder, Integer, AttributionSource, Boolean, Boolean, String,
+                    Boolean, Boolean, Integer, Integer, Integer, SyncNotedAppOp> superImpl) {
+        return superImpl.apply(clientId, resolveDatasourceOp(code, attributionSource.getUid(),
                 attributionSource.getPackageName(), attributionSource.getAttributionTag()),
                 attributionSource, startIfModeDefault, shouldCollectAsyncNotedOp, message,
                 shouldCollectMessage, skipProxyOperation, proxyAttributionFlags,
@@ -279,10 +277,10 @@
     }
 
     @Override
-    public void finishProxyOperation(int code, @NonNull AttributionSource attributionSource,
-            boolean skipProxyOperation, @NonNull TriFunction<Integer, AttributionSource,
-            Boolean, Void> superImpl) {
-        superImpl.apply(resolveDatasourceOp(code, attributionSource.getUid(),
+    public void finishProxyOperation(@NonNull IBinder clientId, int code,
+            @NonNull AttributionSource attributionSource, boolean skipProxyOperation,
+            @NonNull QuadFunction<IBinder, Integer, AttributionSource, Boolean, Void> superImpl) {
+        superImpl.apply(clientId, resolveDatasourceOp(code, attributionSource.getUid(),
                 attributionSource.getPackageName(), attributionSource.getAttributionTag()),
                 attributionSource, skipProxyOperation);
     }
diff --git a/services/core/java/com/android/server/policy/DeviceStateProviderImpl.java b/services/core/java/com/android/server/policy/DeviceStateProviderImpl.java
index 2a65ea2..ced3a45 100644
--- a/services/core/java/com/android/server/policy/DeviceStateProviderImpl.java
+++ b/services/core/java/com/android/server/policy/DeviceStateProviderImpl.java
@@ -378,13 +378,14 @@
                 try {
                     conditionSatisfied = mStateConditions.get(state).getAsBoolean();
                 } catch (IllegalStateException e) {
-                    // Failed to compute the current state based on current available data. Return
+                    // Failed to compute the current state based on current available data. Continue
                     // with the expectation that notifyDeviceStateChangedIfNeeded() will be called
-                    // when a callback with the missing data is triggered.
+                    // when a callback with the missing data is triggered. May trigger another state
+                    // change if another state is satisfied currently.
                     if (DEBUG) {
                         Slog.d(TAG, "Unable to check current state", e);
                     }
-                    return;
+                    continue;
                 }
 
                 if (conditionSatisfied) {
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index a4c9684..6d5da32 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -2873,6 +2873,18 @@
                     return key_consumed;
                 }
                 break;
+            case KeyEvent.KEYCODE_DPAD_LEFT:
+                if (down && event.isMetaPressed() && event.isCtrlPressed() && repeatCount == 0) {
+                    enterStageSplitFromRunningApp(true /* leftOrTop */);
+                    return key_consumed;
+                }
+                break;
+            case KeyEvent.KEYCODE_DPAD_RIGHT:
+                if (down && event.isMetaPressed() && event.isCtrlPressed() && repeatCount == 0) {
+                    enterStageSplitFromRunningApp(false /* leftOrTop */);
+                    return key_consumed;
+                }
+                break;
             case KeyEvent.KEYCODE_SLASH:
                 if (down && repeatCount == 0 && event.isMetaPressed() && !keyguardOn) {
                     toggleKeyboardShortcutsMenu(event.getDeviceId());
@@ -3489,6 +3501,13 @@
         }
     }
 
+    private void enterStageSplitFromRunningApp(boolean leftOrTop) {
+        StatusBarManagerInternal statusbar = getStatusBarManagerInternal();
+        if (statusbar != null) {
+            statusbar.enterStageSplitFromRunningApp(leftOrTop);
+        }
+    }
+
     void launchHomeFromHotKey(int displayId) {
         launchHomeFromHotKey(displayId, true /* awakenFromDreams */, true /*respectKeyguard*/);
     }
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java b/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
index 9957140..e7221c8 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerInternal.java
@@ -184,4 +184,11 @@
      * Called when requested to go to fullscreen from the active split app.
      */
     void goToFullscreenFromSplit();
+
+    /**
+     * Enters stage split from a current running app.
+     *
+     * @see com.android.internal.statusbar.IStatusBar#enterStageSplitFromRunningApp
+     */
+    void enterStageSplitFromRunningApp(boolean leftOrTop);
 }
diff --git a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
index 5a91dc6..45748e6 100644
--- a/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
+++ b/services/core/java/com/android/server/statusbar/StatusBarManagerService.java
@@ -705,6 +705,15 @@
                 } catch (RemoteException ex) { }
             }
         }
+
+        @Override
+        public void enterStageSplitFromRunningApp(boolean leftOrTop) {
+            if (mBar != null) {
+                try {
+                    mBar.enterStageSplitFromRunningApp(leftOrTop);
+                } catch (RemoteException ex) { }
+            }
+        }
     };
 
     private final GlobalActionsProvider mGlobalActionsProvider = new GlobalActionsProvider() {
diff --git a/services/core/java/com/android/server/wm/DisplayRotation.java b/services/core/java/com/android/server/wm/DisplayRotation.java
index eaa08fd..185e06e 100644
--- a/services/core/java/com/android/server/wm/DisplayRotation.java
+++ b/services/core/java/com/android/server/wm/DisplayRotation.java
@@ -1537,6 +1537,7 @@
         private int mHalfFoldSavedRotation = -1; // No saved rotation
         private DeviceStateController.FoldState mFoldState =
                 DeviceStateController.FoldState.UNKNOWN;
+        private boolean mInHalfFoldTransition = false;
 
         boolean overrideFrozenRotation() {
             return mFoldState == DeviceStateController.FoldState.HALF_FOLDED;
@@ -1544,6 +1545,7 @@
 
         boolean shouldRevertOverriddenRotation() {
             return mFoldState == DeviceStateController.FoldState.OPEN // When transitioning to open.
+                    && mInHalfFoldTransition
                     && mHalfFoldSavedRotation != -1 // Ignore if we've already reverted.
                     && mUserRotationMode
                     == WindowManagerPolicy.USER_ROTATION_LOCKED; // Ignore if we're unlocked.
@@ -1552,6 +1554,7 @@
         int revertOverriddenRotation() {
             int savedRotation = mHalfFoldSavedRotation;
             mHalfFoldSavedRotation = -1;
+            mInHalfFoldTransition = false;
             return savedRotation;
         }
 
@@ -1577,16 +1580,11 @@
                 mService.updateRotation(false /* alwaysSendConfiguration */,
                         false /* forceRelayout */);
             } else {
-                // Revert the rotation to our saved value if we transition from HALF_FOLDED.
-                if (mHalfFoldSavedRotation != -1) {
-                    mRotation = mHalfFoldSavedRotation;
-                }
-                // Tell the device to update its orientation (mFoldState is still HALF_FOLDED here
-                // so we will override USER_ROTATION_LOCKED and allow a rotation).
+                mInHalfFoldTransition = true;
+                mFoldState = newState;
+                // Tell the device to update its orientation.
                 mService.updateRotation(false /* alwaysSendConfiguration */,
                         false /* forceRelayout */);
-                // Once we are rotated, set mFoldstate, effectively removing the lock override.
-                mFoldState = newState;
             }
         }
     }
@@ -1683,6 +1681,7 @@
 
     private static class RotationHistory {
         private static final int MAX_SIZE = 8;
+        private static final int NO_FOLD_CONTROLLER = -2;
         private static class Record {
             final @Surface.Rotation int mFromRotation;
             final @Surface.Rotation int mToRotation;
@@ -1694,6 +1693,9 @@
             final String mLastOrientationSource;
             final @ActivityInfo.ScreenOrientation int mSourceOrientation;
             final long mTimestamp = System.currentTimeMillis();
+            final int mHalfFoldSavedRotation;
+            final boolean mInHalfFoldTransition;
+            final DeviceStateController.FoldState mFoldState;
 
             Record(DisplayRotation dr, int fromRotation, int toRotation) {
                 mFromRotation = fromRotation;
@@ -1719,6 +1721,15 @@
                     mLastOrientationSource = null;
                     mSourceOrientation = SCREEN_ORIENTATION_UNSET;
                 }
+                if (dr.mFoldController != null) {
+                    mHalfFoldSavedRotation = dr.mFoldController.mHalfFoldSavedRotation;
+                    mInHalfFoldTransition = dr.mFoldController.mInHalfFoldTransition;
+                    mFoldState = dr.mFoldController.mFoldState;
+                } else {
+                    mHalfFoldSavedRotation = NO_FOLD_CONTROLLER;
+                    mInHalfFoldTransition = false;
+                    mFoldState = DeviceStateController.FoldState.UNKNOWN;
+                }
             }
 
             void dump(String prefix, PrintWriter pw) {
@@ -1735,6 +1746,12 @@
                 if (mNonDefaultRequestingTaskDisplayArea != null) {
                     pw.println(prefix + "  requestingTda=" + mNonDefaultRequestingTaskDisplayArea);
                 }
+                if (mHalfFoldSavedRotation != NO_FOLD_CONTROLLER) {
+                    pw.println(prefix + " halfFoldSavedRotation="
+                            + mHalfFoldSavedRotation
+                            + " mInHalfFoldTransition=" + mInHalfFoldTransition
+                            + " mFoldState=" + mFoldState);
+                }
             }
         }
 
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 76a1122..51eec03 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -2183,7 +2183,7 @@
     }
 
     private boolean shouldStartChangeTransition(int prevWinMode, @NonNull Rect prevBounds) {
-        if (!isLeafTask() || !canStartChangeTransition()) {
+        if (!(isLeafTask() || mCreatedByOrganizer) || !canStartChangeTransition()) {
             return false;
         }
         final int newWinMode = getWindowingMode();
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index d3d1c16..971b619 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -38,6 +38,7 @@
 import android.util.ArrayMap;
 import android.util.Slog;
 import android.util.proto.ProtoOutputStream;
+import android.view.SurfaceControl;
 import android.view.WindowManager;
 import android.window.ITransitionMetricsReporter;
 import android.window.ITransitionPlayer;
@@ -116,6 +117,8 @@
      */
     boolean mBuildingFinishLayers = false;
 
+    private final SurfaceControl.Transaction mWakeT = new SurfaceControl.Transaction();
+
     TransitionController(ActivityTaskManagerService atm,
             TaskSnapshotController taskSnapshotController,
             TransitionTracer transitionTracer) {
@@ -619,8 +622,16 @@
     private void updateRunningRemoteAnimation(Transition transition, boolean isPlaying) {
         if (mTransitionPlayerProc == null) return;
         if (isPlaying) {
+            mWakeT.setEarlyWakeupStart();
+            mWakeT.apply();
+            // Usually transitions put quite a load onto the system already (with all the things
+            // happening in app), so pause task snapshot persisting to not increase the load.
+            mAtm.mWindowManager.mTaskSnapshotController.setPersisterPaused(true);
             mTransitionPlayerProc.setRunningRemoteAnimation(true);
         } else if (mPlayingTransitions.isEmpty()) {
+            mWakeT.setEarlyWakeupEnd();
+            mWakeT.apply();
+            mAtm.mWindowManager.mTaskSnapshotController.setPersisterPaused(false);
             mTransitionPlayerProc.setRunningRemoteAnimation(false);
             mRemotePlayer.clear();
             return;