Merge "Fix issue where turning off fixed landscape goes to default grid" into main
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index 9fa2f50..ed370ec 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -601,3 +601,10 @@
     purpose: PURPOSE_BUGFIX
   }
 }
+
+flag {
+  name: "enable_strict_mode"
+  namespace: "launcher"
+  description: "Enable Strict Mode for the Launcher app"
+  bug: "394651876"
+}
diff --git a/aconfig/launcher_accessibility.aconfig b/aconfig/launcher_accessibility.aconfig
index afee8fe..13e1127 100644
--- a/aconfig/launcher_accessibility.aconfig
+++ b/aconfig/launcher_accessibility.aconfig
@@ -1,5 +1,5 @@
 package: "com.android.launcher3"
-container: "system"
+container: "system_ext"
 
 flag {
     name: "remove_exclude_from_screen_magnification_flag_usage"
@@ -9,4 +9,4 @@
     metadata {
         purpose: PURPOSE_BUGFIX
     }
-}
\ No newline at end of file
+}
diff --git a/quickstep/res/drawable/bg_overview_add_desktop_button.xml b/quickstep/res/drawable/bg_overview_add_desktop_button.xml
new file mode 100644
index 0000000..12581bf
--- /dev/null
+++ b/quickstep/res/drawable/bg_overview_add_desktop_button.xml
@@ -0,0 +1,27 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2025 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.
+-->
+<ripple android:color="?android:attr/colorControlHighlight"
+    xmlns:android="http://schemas.android.com/apk/res/android">
+    <item>
+        <shape
+            android:shape="rectangle"
+            android:tint="?colorButtonNormal">
+            <corners android:radius="@dimen/add_desktop_button_size" />
+            <solid android:color="@color/materialColorSurfaceBright"/>
+        </shape>
+    </item>
+</ripple>
\ No newline at end of file
diff --git a/quickstep/res/drawable/bg_overview_clear_all_button.xml b/quickstep/res/drawable/bg_overview_clear_all_button.xml
index 7f58cf8..2f28689 100644
--- a/quickstep/res/drawable/bg_overview_clear_all_button.xml
+++ b/quickstep/res/drawable/bg_overview_clear_all_button.xml
@@ -15,8 +15,7 @@
      limitations under the License.
 -->
 <ripple android:color="?android:attr/colorControlHighlight"
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:androidprv="http://schemas.android.com/apk/prv/res/android">
+    xmlns:android="http://schemas.android.com/apk/res/android">
     <item>
         <shape android:shape="rectangle"
             android:tint="?colorButtonNormal">
diff --git a/quickstep/res/layout/overview_add_desktop_button.xml b/quickstep/res/layout/overview_add_desktop_button.xml
index e36cf72..a1c64f3 100644
--- a/quickstep/res/layout/overview_add_desktop_button.xml
+++ b/quickstep/res/layout/overview_add_desktop_button.xml
@@ -16,9 +16,11 @@
 -->
 <com.android.quickstep.views.AddDesktopButton
     xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:launcher="http://schemas.android.com/apgk/res-auto"
+    xmlns:launcher="http://schemas.android.com/apk/res-auto"
     android:id="@+id/add_desktop_button"
     android:layout_width="@dimen/add_desktop_button_size"
     android:layout_height="@dimen/add_desktop_button_size"
     android:src="@drawable/ic_desktop_add"
-    android:padding="10dp" />
\ No newline at end of file
+    android:background="@drawable/bg_overview_add_desktop_button"
+    launcher:focusBorderColor="@color/materialColorOutline"
+    android:padding="10dp" />
diff --git a/quickstep/res/layout/overview_clear_all_button.xml b/quickstep/res/layout/overview_clear_all_button.xml
index 18a6240..034c3c2 100644
--- a/quickstep/res/layout/overview_clear_all_button.xml
+++ b/quickstep/res/layout/overview_clear_all_button.xml
@@ -16,7 +16,6 @@
 -->
 <com.android.quickstep.views.ClearAllButton
     xmlns:android="http://schemas.android.com/apk/res/android"
-    xmlns:androidprv="http://schemas.android.com/apk/prv/res/android"
     xmlns:launcher="http://schemas.android.com/apk/res-auto"
     style="@style/OverviewClearAllButton"
     android:id="@+id/clear_all"
@@ -25,4 +24,4 @@
     android:text="@string/recents_clear_all"
     android:textColor="@color/materialColorOnSurface"
     launcher:focusBorderColor="@color/materialColorOutline"
-    android:textSize="14sp" />
\ No newline at end of file
+    android:textSize="14sp" />
diff --git a/quickstep/res/values/attrs.xml b/quickstep/res/values/attrs.xml
index 7fd6b5c..28c0d5c 100644
--- a/quickstep/res/values/attrs.xml
+++ b/quickstep/res/values/attrs.xml
@@ -36,6 +36,11 @@
         <attr name="focusBorderColor" />
     </declare-styleable>
 
+    <declare-styleable name="AddDesktopButton">
+        <!-- focus border color for overview add desktop button views -->
+        <attr name="focusBorderColor" />
+    </declare-styleable>
+
     <!--
          Gesture nav edu specific attributes. These attributes are used to customize Gesture nav edu
          view lottie animation colors in XML files.
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index b253343..f2f1ebd 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -108,6 +108,7 @@
 
     <!-- Recents add desktop button -->
     <dimen name="add_desktop_button_size">56dp</dimen>
+    <dimen name="add_desktop_button_outline_padding">2dp</dimen>
 
     <!-- The speed in dp/s at which the user needs to be scrolling in recents such that we start
              loading full resolution screenshots. -->
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
index 3c4bc91..03f5d96 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.kt
@@ -18,6 +18,7 @@
 import android.content.Context
 import android.os.Debug
 import android.util.Log
+import android.util.SparseArray
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
 import com.android.launcher3.LauncherState
 import com.android.launcher3.dagger.ApplicationContext
@@ -52,6 +53,29 @@
     systemUiProxy: SystemUiProxy,
     lifecycleTracker: DaggerSingletonTracker,
 ) {
+    /**
+     * Tracks the desks configurations on each display.
+     *
+     * (Used only when multiple desks are enabled).
+     *
+     * @property displayId The ID of the display this object represents.
+     * @property canCreateDesks true if it's possible to create new desks on the display represented
+     *   by this object.
+     * @property activeDeskId The ID of the active desk on the associated display (if any). It has a
+     *   value of `-1` if there are no active desks. Note that there can only be at most one active
+     *   desk on each display.
+     * @property deskIds a set containing the IDs of the desks on the associated display.
+     */
+    private data class DisplayDeskConfig(
+        val displayId: Int,
+        var canCreateDesks: Boolean,
+        var activeDeskId: Int = -1,
+        val deskIds: MutableSet<Int>,
+    )
+
+    /** Maps each display by its ID to its desks configuration. */
+    private val displaysDesksConfigsMap = SparseArray<DisplayDeskConfig>()
+
     private val desktopVisibilityListeners: MutableSet<DesktopVisibilityListener> = HashSet()
     private val taskbarDesktopModeListeners: MutableSet<TaskbarDesktopModeListener> = HashSet()
 
@@ -313,6 +337,81 @@
         }
     }
 
+    private fun onListenerConnected(displayDeskStates: Array<DisplayDeskState>) {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return
+        }
+
+        displaysDesksConfigsMap.clear()
+
+        displayDeskStates.forEach { displayDeskState ->
+            displaysDesksConfigsMap[displayDeskState.displayId] =
+                DisplayDeskConfig(
+                    displayId = displayDeskState.displayId,
+                    canCreateDesks = displayDeskState.canCreateDesk,
+                    activeDeskId = displayDeskState.activeDeskId,
+                    deskIds = displayDeskState.deskIds.toMutableSet(),
+                )
+        }
+    }
+
+    private fun getDisplayDeskConfig(displayId: Int): DisplayDeskConfig {
+        return checkNotNull(displaysDesksConfigsMap[displayId]) {
+            "Expected non-null desk config for display: $displayId"
+        }
+    }
+
+    private fun onCanCreateDesksChanged(displayId: Int, canCreateDesks: Boolean) {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return
+        }
+
+        getDisplayDeskConfig(displayId).canCreateDesks = canCreateDesks
+    }
+
+    private fun onDeskAdded(displayId: Int, deskId: Int) {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return
+        }
+
+        getDisplayDeskConfig(displayId).also {
+            check(it.deskIds.add(deskId)) {
+                "Found a duplicate desk Id: $deskId on display: $displayId"
+            }
+        }
+    }
+
+    private fun onDeskRemoved(displayId: Int, deskId: Int) {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return
+        }
+
+        getDisplayDeskConfig(displayId).also {
+            check(it.deskIds.remove(deskId)) {
+                "Removing non-existing desk Id: $deskId on display: $displayId"
+            }
+            if (it.activeDeskId == deskId) {
+                it.activeDeskId = -1
+            }
+        }
+    }
+
+    private fun onActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {
+        if (!DesktopModeStatus.enableMultipleDesktops(context)) {
+            return
+        }
+
+        getDisplayDeskConfig(displayId).also {
+            check(oldActiveDesk == it.activeDeskId) {
+                "Mismatch between the Shell's oldActiveDesk: $oldActiveDesk, and Launcher's: ${it.activeDeskId}"
+            }
+            check(it.deskIds.contains(newActiveDesk)) {
+                "newActiveDesk: $newActiveDesk was never added to display: $displayId"
+            }
+            it.activeDeskId = newActiveDesk
+        }
+    }
+
     /** TODO: b/333533253 - Remove after flag rollout */
     private fun markLauncherPaused() {
         if (ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY.isTrue) {
@@ -366,10 +465,11 @@
     ) : Stub() {
         private val controller = WeakReference(controller)
 
-        // TODO: b/392986431 - Implement the new desks APIs.
-        override fun onListenerConnected(
-            displayDeskStates: Array<DisplayDeskState>,
-        ) {}
+        override fun onListenerConnected(displayDeskStates: Array<DisplayDeskState>) {
+            Executors.MAIN_EXECUTOR.execute {
+                controller.get()?.onListenerConnected(displayDeskStates)
+            }
+        }
 
         override fun onTasksVisibilityChanged(displayId: Int, visibleTasksCount: Int) {
             if (displayId != this.displayId) return
@@ -405,14 +505,25 @@
 
         override fun onExitDesktopModeTransitionStarted(transitionDuration: Int) {}
 
-        // TODO: b/392986431 - Implement all the below new desks APIs.
-        override fun onCanCreateDesksChanged(displayId: Int, canCreateDesks: Boolean) {}
+        override fun onCanCreateDesksChanged(displayId: Int, canCreateDesks: Boolean) {
+            Executors.MAIN_EXECUTOR.execute {
+                controller.get()?.onCanCreateDesksChanged(displayId, canCreateDesks)
+            }
+        }
 
-        override fun onDeskAdded(displayId: Int, deskId: Int) {}
+        override fun onDeskAdded(displayId: Int, deskId: Int) {
+            Executors.MAIN_EXECUTOR.execute { controller.get()?.onDeskAdded(displayId, deskId) }
+        }
 
-        override fun onDeskRemoved(displayId: Int, deskId: Int) {}
+        override fun onDeskRemoved(displayId: Int, deskId: Int) {
+            Executors.MAIN_EXECUTOR.execute { controller.get()?.onDeskRemoved(displayId, deskId) }
+        }
 
-        override fun onActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {}
+        override fun onActiveDeskChanged(displayId: Int, newActiveDesk: Int, oldActiveDesk: Int) {
+            Executors.MAIN_EXECUTOR.execute {
+                controller.get()?.onActiveDeskChanged(displayId, newActiveDesk, oldActiveDesk)
+            }
+        }
     }
 
     /** A listener for Taskbar in Desktop Mode. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
index 3a83db2..f36c481 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
@@ -34,6 +34,7 @@
 import android.content.ClipDescription;
 import android.content.Intent;
 import android.content.pm.LauncherApps;
+import android.content.pm.ShortcutInfo;
 import android.content.res.Resources;
 import android.graphics.Canvas;
 import android.graphics.Point;
@@ -84,6 +85,7 @@
 import com.android.quickstep.util.MultiValueUpdateListener;
 import com.android.quickstep.util.SingleTask;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.wm.shell.shared.bubbles.BubbleAnythingFlagHelper;
 import com.android.wm.shell.shared.draganddrop.DragAndDropConstants;
 
 import java.io.PrintWriter;
@@ -416,6 +418,10 @@
                                 item.user));
                 intent.putExtra(Intent.EXTRA_PACKAGE_NAME, item.getIntent().getPackage());
                 intent.putExtra(Intent.EXTRA_SHORTCUT_ID, deepShortcutId);
+                ShortcutInfo shortcutInfo = ((WorkspaceItemInfo) item).getDeepShortcutInfo();
+                if (BubbleAnythingFlagHelper.enableCreateAnyBubble() && shortcutInfo != null) {
+                    intent.putExtra(DragAndDropConstants.EXTRA_SHORTCUT_INFO, shortcutInfo);
+                }
             } else if (item.itemType == ITEM_TYPE_SEARCH_ACTION) {
                 // TODO(b/289261756): Buggy behavior when split opposite to an existing search pane.
                 intent.putExtra(
diff --git a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
index 36a4865..b25f999 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/PredictedAppIcon.java
@@ -44,6 +44,7 @@
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.Flags;
 import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
 import com.android.launcher3.anim.AnimatedFloat;
@@ -68,8 +69,7 @@
 
     private static final float RING_SCALE_START_VALUE = 0.75f;
     private static final int RING_SHADOW_COLOR = 0x99000000;
-    private static final float RING_EFFECT_RATIO = 0.095f;
-
+    private static final float RING_EFFECT_RATIO = Flags.enableLauncherIconShapes() ? 0.1f : 0.095f;
     private static final long ICON_CHANGE_ANIM_DURATION = 360;
     private static final long ICON_CHANGE_ANIM_STAGGER = 50;
 
@@ -150,12 +150,12 @@
         int count = canvas.save();
         boolean isSlotMachineAnimRunning = mSlotMachineIcon != null;
         if (!mIsPinned) {
-            drawEffect(canvas);
+            drawRingEffect(canvas);
             if (isSlotMachineAnimRunning) {
                 // Clip to to outside of the ring during the slot machine animation.
                 canvas.clipPath(mRingPath);
             }
-            canvas.scale(1 - 2 * RING_EFFECT_RATIO, 1 - 2 * RING_EFFECT_RATIO,
+            canvas.scale(1 - 2f * RING_EFFECT_RATIO, 1 - 2f * RING_EFFECT_RATIO,
                     getWidth() * .5f, getHeight() * .5f);
             if (isSlotMachineAnimRunning) {
                 canvas.translate(0, mSlotMachineIconTranslationY);
@@ -388,7 +388,7 @@
         mRingScaleAnim.start();
     }
 
-    private void drawEffect(Canvas canvas) {
+    private void drawRingEffect(Canvas canvas) {
         // Don't draw ring effect if item is about to be dragged or if the icon is not visible.
         if (mDrawForDrag || !mIsIconVisible || mForceHideRing) {
             return;
@@ -396,12 +396,28 @@
         mIconRingPaint.setColor(RING_SHADOW_COLOR);
         mIconRingPaint.setMaskFilter(mShadowFilter);
         int count = canvas.save();
-        if (Float.compare(1, mRingScale) != 0) {
+        if (Flags.enableLauncherIconShapes()) {
+            // Scale canvas properly to for ring to be inner stroke and not exceed bounds.
+            // Since STROKE draws half on either side of Path, scale canvas down by 1x stroke ratio.
+            canvas.scale(
+                    mRingScale * (1f - RING_EFFECT_RATIO),
+                    mRingScale * (1f - RING_EFFECT_RATIO),
+                    canvas.getWidth() / 2f,
+                    canvas.getHeight() / 2f);
+        } else if (Float.compare(1, mRingScale) != 0) {
             canvas.scale(mRingScale, mRingScale, canvas.getWidth() / 2f, canvas.getHeight() / 2f);
         }
+        // Draw ring shadow around canvas.
         canvas.drawPath(mRingPath, mIconRingPaint);
         mIconRingPaint.setColor(mPlateColor.currentColor);
+        if (Flags.enableLauncherIconShapes()) {
+            mIconRingPaint.setStrokeWidth(canvas.getWidth() * RING_EFFECT_RATIO);
+            // Using FILL_AND_STROKE as there is still some gap to fill,
+            // between inner curve of ring / outer curve of icon.
+            mIconRingPaint.setStyle(Paint.Style.FILL_AND_STROKE);
+        }
         mIconRingPaint.setMaskFilter(null);
+        // Draw ring around canvas.
         canvas.drawPath(mRingPath, mIconRingPaint);
         canvas.restoreToCount(count);
     }
diff --git a/quickstep/src/com/android/quickstep/TaskIconCache.kt b/quickstep/src/com/android/quickstep/TaskIconCache.kt
index bf94d41..b82c110 100644
--- a/quickstep/src/com/android/quickstep/TaskIconCache.kt
+++ b/quickstep/src/com/android/quickstep/TaskIconCache.kt
@@ -40,6 +40,7 @@
 import com.android.launcher3.util.FlagOp
 import com.android.launcher3.util.Preconditions
 import com.android.quickstep.task.thumbnail.data.TaskIconDataSource
+import com.android.quickstep.util.IconLabelUtil.getBadgedContentDescription
 import com.android.quickstep.util.TaskKeyLruCache
 import com.android.quickstep.util.TaskVisualsChangeListener
 import com.android.systemui.shared.recents.model.Task
@@ -206,6 +207,7 @@
                 TaskCacheEntry(
                     entryIcon,
                     getBadgedContentDescription(
+                        context,
                         activityInfo,
                         task.key.userId,
                         task.taskDescription,
@@ -215,7 +217,12 @@
             else ->
                 TaskCacheEntry(
                     entryIcon,
-                    getBadgedContentDescription(activityInfo, task.key.userId, task.taskDescription),
+                    getBadgedContentDescription(
+                        context,
+                        activityInfo,
+                        task.key.userId,
+                        task.taskDescription,
+                    ),
                 )
         }.also { iconCache.put(task.key, it) }
     }
@@ -224,28 +231,6 @@
         desc.inMemoryIcon
             ?: ActivityManager.TaskDescription.loadTaskDescriptionIcon(desc.iconFilename, userId)
 
-    private fun getBadgedContentDescription(
-        info: ActivityInfo,
-        userId: Int,
-        taskDescription: ActivityManager.TaskDescription?,
-    ): String {
-        val packageManager = context.packageManager
-        var taskLabel = taskDescription?.let { Utilities.trim(it.label) }
-        if (taskLabel.isNullOrEmpty()) {
-            taskLabel = Utilities.trim(info.loadLabel(packageManager))
-        }
-
-        val applicationLabel = Utilities.trim(info.applicationInfo.loadLabel(packageManager))
-        val badgedApplicationLabel =
-            if (userId != UserHandle.myUserId())
-                packageManager
-                    .getUserBadgedLabel(applicationLabel, UserHandle.of(userId))
-                    .toString()
-            else applicationLabel
-        return if (applicationLabel == taskLabel) badgedApplicationLabel
-        else "$badgedApplicationLabel $taskLabel"
-    }
-
     @WorkerThread
     private fun getDefaultIcon(userId: Int): Drawable {
         synchronized(defaultIcons) {
diff --git a/quickstep/src/com/android/quickstep/recents/domain/usecase/OrganizeDesktopTasksUseCase.kt b/quickstep/src/com/android/quickstep/recents/domain/usecase/OrganizeDesktopTasksUseCase.kt
index 3e7d142..4ea39d8 100644
--- a/quickstep/src/com/android/quickstep/recents/domain/usecase/OrganizeDesktopTasksUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/domain/usecase/OrganizeDesktopTasksUseCase.kt
@@ -17,77 +17,295 @@
 package com.android.quickstep.recents.domain.usecase
 
 import android.graphics.Rect
-import android.util.Size
+import android.graphics.RectF
+import androidx.core.graphics.toRect
 import com.android.quickstep.recents.domain.model.DesktopTaskBoundsData
-import kotlin.math.ceil
-import kotlin.math.sqrt
 
-/**
- * This usecase is responsible for organizing desktop windows in a non-overlapping way. Note: this
- * is currently a placeholder implementation.
- */
+/** This usecase is responsible for organizing desktop windows in a non-overlapping way. */
 class OrganizeDesktopTasksUseCase {
+    /**
+     * Run to layout [taskBounds] within the screen [desktopBounds]. Layout is done in 2 stages:
+     * 1. Optimal height is determined. In this stage height is bisected to find maximum height
+     *    which still allows all the windows to fit.
+     * 2. Row widths are balanced. In this stage the available width is reduced until some windows
+     *    are no longer fitting or until the difference between the narrowest and the widest rows
+     *    starts growing. Overall this achieves the goals of maximum size for previews (or maximum
+     *    row height which is equivalent assuming fixed height), balanced rows and minimal wasted
+     *    space.
+     */
     fun run(
-        desktopSize: Size,
+        desktopBounds: Rect,
         taskBounds: List<DesktopTaskBoundsData>,
     ): List<DesktopTaskBoundsData> {
-        return getRects(desktopSize, taskBounds.size).zip(taskBounds) { rect, task ->
-            shrinkRect(rect, 0.8f)
-            DesktopTaskBoundsData(task.taskId, fitRect(task.bounds, rect))
+        if (desktopBounds.isEmpty || taskBounds.isEmpty()) {
+            return emptyList()
         }
+
+        // Filter out [taskBounds] with empty rects before calculating layout.
+        val validTaskBounds = taskBounds.filterNot { it.bounds.isEmpty }
+
+        if (validTaskBounds.isEmpty()) {
+            return emptyList()
+        }
+
+        val availableLayoutBounds = desktopBounds.getLayoutEffectiveBounds()
+        val resultRects = findOptimalHeightAndBalancedWidth(availableLayoutBounds, validTaskBounds)
+
+        centerTaskWindows(
+            availableLayoutBounds,
+            resultRects.maxOf { it.bottom }.toInt(),
+            resultRects,
+        )
+
+        val result = mutableListOf<DesktopTaskBoundsData>()
+        for (i in validTaskBounds.indices) {
+            result.add(DesktopTaskBoundsData(validTaskBounds[i].taskId, resultRects[i].toRect()))
+        }
+        return result
     }
 
-    private fun shrinkRect(bounds: Rect, fraction: Float) {
-        val xMargin = (bounds.width() * ((1.0f - fraction) / 2.0f)).toInt()
-        val yMargin = (bounds.height() * ((1.0f - fraction) / 2.0f)).toInt()
-        bounds.inset(xMargin, yMargin, xMargin, yMargin)
-    }
+    /** Calculates the effective bounds for layout by applying insets to the raw desktop bounds. */
+    private fun Rect.getLayoutEffectiveBounds() =
+        Rect(this).apply { inset(OVERVIEW_INSET_TOP_BOTTOM, OVERVIEW_INSET_LEFT_RIGHT) }
 
-    /** Generates `tasks` number of non-overlapping rects that fit into `desktopSize`. */
-    private fun getRects(desktopSize: Size, tasks: Int): List<Rect> {
-        val (xSlots, ySlots) =
-            when (tasks) {
-                2 -> Pair(2, 1)
-                3,
-                4 -> Pair(2, 2)
-                5,
-                6 -> Pair(3, 2)
-                else -> {
-                    val sides = ceil(sqrt(tasks.toDouble())).toInt()
-                    Pair(sides, sides)
+    /**
+     * Determines the optimal height for task windows and balances the row widths to minimize wasted
+     * space. Returns the bounds for each task window after layout.
+     */
+    private fun findOptimalHeightAndBalancedWidth(
+        availableLayoutBounds: Rect,
+        validTaskBounds: List<DesktopTaskBoundsData>,
+    ): List<RectF> {
+        // Right bound of the narrowest row.
+        var minRight: Int
+        // Right bound of the widest row.
+        var maxRight: Int
+
+        // Keep track of the difference between the narrowest and the widest row.
+        // Initially this is set to the worst it can ever be assuming the windows fit.
+        var widthDiff = availableLayoutBounds.width()
+
+        // Initially allow the windows to occupy all available width. Shrink this available space
+        // horizontally to find the breakdown into rows that achieves the minimal [widthDiff].
+        var rightBound = availableLayoutBounds.right
+
+        // Determine the optimal height bisecting between [lowHeight] and [highHeight]. Once this
+        // optimal height is known, [heightFixed] is set to `true` and the rows are balanced by
+        // repeatedly squeezing the widest row to cause windows to overflow to the subsequent rows.
+        var lowHeight = VERTICAL_SPACE_BETWEEN_TASKS
+        var highHeight = maxOf(lowHeight, availableLayoutBounds.height() + 1)
+        var optimalHeight = 0.5f * (lowHeight + highHeight)
+        var heightFixed = false
+
+        // Repeatedly try to fit the windows [resultRects] within [rightBound]. If a maximum
+        // [optimalHeight] is found such that all window [resultRects] fit, this fitting continues
+        // while shrinking the [rightBound] in order to balance the rows. If the windows fit the
+        // [rightBound] would have been decremented at least once so it needs to be incremented once
+        // before getting out of this loop and one additional pass made to actually fit the
+        // [resultRects]. If the [resultRects] cannot fit (e.g. there are too many windows) the
+        // bisection will still finish and we might increment the [rightBound] one pixel extra
+        // which is acceptable since there is an unused margin on the right.
+        var makeLastAdjustment = false
+        var resultRects: List<RectF>
+
+        while (true) {
+            val fitWindowResult =
+                fitWindowRectsInBounds(
+                    Rect(availableLayoutBounds).apply { right = rightBound },
+                    validTaskBounds,
+                    minOf(MAXIMUM_TASK_HEIGHT, optimalHeight.toInt()),
+                )
+            val allWindowsFit = fitWindowResult.allWindowsFit
+            resultRects = fitWindowResult.calculatedBounds
+            minRight = fitWindowResult.minRight
+            maxRight = fitWindowResult.maxRight
+
+            if (heightFixed) {
+                if (!allWindowsFit) {
+                    // Revert the previous change to [rightBound] and do one last pass.
+                    rightBound++
+                    makeLastAdjustment = true
+                    break
+                }
+                // Break if all the windows are zero-width at the current scale.
+                if (maxRight <= availableLayoutBounds.left) {
+                    break
+                }
+            } else {
+                // Find the optimal row height bisecting between [lowHeight] and [highHeight].
+                if (allWindowsFit) {
+                    lowHeight = optimalHeight.toInt()
+                } else {
+                    highHeight = optimalHeight.toInt()
+                }
+                optimalHeight = 0.5f * (lowHeight + highHeight)
+                // When height can no longer be improved, start balancing the rows.
+                if (optimalHeight.toInt() == lowHeight) {
+                    heightFixed = true
                 }
             }
 
-        // The width and height of one of the boxes.
-        val boxWidth = desktopSize.width / xSlots
-        val boxHeight = desktopSize.height / ySlots
-
-        return (0 until tasks).map {
-            val x = it % xSlots
-            val y = it / xSlots
-            Rect(x * boxWidth, y * boxHeight, (x + 1) * boxWidth, (y + 1) * boxHeight)
+            if (allWindowsFit && heightFixed) {
+                if (maxRight - minRight <= widthDiff) {
+                    // Row alignment is getting better. Try to shrink the [rightBound] in order to
+                    // squeeze the widest row.
+                    rightBound = maxRight - 1
+                    widthDiff = maxRight - minRight
+                } else {
+                    // Row alignment is getting worse.
+                    // Revert the previous change to [rightBound] and do one last pass.
+                    rightBound++
+                    makeLastAdjustment = true
+                    break
+                }
+            }
         }
+
+        // Once the windows no longer fit, the change to [rightBound] was reverted. Perform one last
+        // pass to position the [resultRects].
+        if (makeLastAdjustment) {
+            val fitWindowResult =
+                fitWindowRectsInBounds(
+                    Rect(availableLayoutBounds).apply { right = rightBound },
+                    validTaskBounds,
+                    minOf(MAXIMUM_TASK_HEIGHT, optimalHeight.toInt()),
+                )
+            resultRects = fitWindowResult.calculatedBounds
+        }
+
+        return resultRects
     }
 
-    /** Centers and fits `rect` into `bounds`, while preserving the former's aspect ratio. */
-    private fun fitRect(rect: Rect, bounds: Rect): Rect {
-        val boundsAspect = bounds.width().toFloat() / bounds.height()
-        val rectAspect = rect.width().toFloat() / rect.height()
+    /**
+     * Data structure to hold the returned result of [fitWindowRectsInBounds] function.
+     * [allWindowsFit] specifies whether all windows can be fit into the provided layout bounds.
+     * [calculatedBounds] specifies the output bounds for all provided task windows. [minRight]
+     * specifies the right bound of the narrowest row. [maxRight] specifies the right bound of the
+     * widest rows.
+     */
+    data class FitWindowResult(
+        val allWindowsFit: Boolean,
+        val calculatedBounds: List<RectF>,
+        val minRight: Int,
+        val maxRight: Int,
+    )
 
-        if (rectAspect > boundsAspect) {
-            // The width is the limiting dimension.
-            val scale = bounds.width().toFloat() / rect.width()
-            val width = bounds.width()
-            val height = (rect.height() * scale).toInt()
-            val top = (bounds.top + bounds.height() / 2.0f - height / 2.0f).toInt()
-            return Rect(bounds.left, top, bounds.left + width, top + height)
-        } else {
-            // The height is the limiting dimension.
-            val scale = bounds.height().toFloat() / rect.height()
-            val width = (rect.width() * scale).toInt()
-            val height = bounds.height()
-            val left = (bounds.left + bounds.width() / 2.0f - width / 2.0f).toInt()
-            return Rect(left, bounds.top, left + width, bounds.top + height)
+    /**
+     * Attempts to fit all [taskBounds] inside [layoutBounds]. The method ensures that the returned
+     * output bounds list has appropriate size and populates it with the values placing task windows
+     * next to each other left-to-right in rows of equal [optimalWindowHeight].
+     */
+    private fun fitWindowRectsInBounds(
+        layoutBounds: Rect,
+        taskBounds: List<DesktopTaskBoundsData>,
+        optimalWindowHeight: Int,
+    ): FitWindowResult {
+        val numTasks = taskBounds.size
+        val outRects = mutableListOf<RectF>()
+
+        // Start in the top-left corner of [layoutBounds].
+        var left = layoutBounds.left
+        var top = layoutBounds.top
+
+        // Right bound of the narrowest row.
+        var minRight = layoutBounds.right
+        // Right bound of the widest row.
+        var maxRight = layoutBounds.left
+
+        var allWindowsFit = true
+        for (i in 0 until numTasks) {
+            val taskBounds = taskBounds[i].bounds
+
+            // Use the height to calculate the width
+            val scale = optimalWindowHeight / taskBounds.height().toFloat()
+            val width = (taskBounds.width() * scale).toInt()
+            val optimalRowHeight = optimalWindowHeight + VERTICAL_SPACE_BETWEEN_TASKS
+
+            if ((left + width + HORIZONTAL_SPACE_BETWEEN_TASKS) > layoutBounds.right) {
+                // Move to the next row if possible.
+                minRight = minOf(minRight, left)
+                maxRight = maxOf(maxRight, left)
+                top += optimalRowHeight
+
+                // Check if the new row reaches the bottom or if the first item in the new
+                // row does not fit within the available width.
+                if (
+                    (top + optimalRowHeight) > layoutBounds.bottom ||
+                        layoutBounds.left + width + HORIZONTAL_SPACE_BETWEEN_TASKS >
+                            layoutBounds.right
+                ) {
+                    allWindowsFit = false
+                    break
+                }
+                left = layoutBounds.left
+            }
+
+            // Position the current rect.
+            outRects.add(
+                RectF(
+                    left.toFloat(),
+                    top.toFloat(),
+                    (left + width).toFloat(),
+                    (top + optimalWindowHeight).toFloat(),
+                )
+            )
+
+            // Increment horizontal position.
+            left += (width + HORIZONTAL_SPACE_BETWEEN_TASKS)
         }
+
+        // Update the narrowest and widest row width for the last row.
+        minRight = minOf(minRight, left)
+        maxRight = maxOf(maxRight, left)
+
+        return FitWindowResult(allWindowsFit, outRects, minRight, maxRight)
+    }
+
+    /** Centers task windows in the center of Overview. */
+    private fun centerTaskWindows(layoutBounds: Rect, maxBottom: Int, outWindowRects: List<RectF>) {
+        if (outWindowRects.isEmpty()) {
+            return
+        }
+
+        val currentRowUnionRange = RectF(outWindowRects[0])
+        var currentRowY = outWindowRects[0].top
+        var currentRowFirstItemIndex = 0
+        val offsetY = (layoutBounds.bottom - maxBottom) / 2f
+
+        // Batch process to center overview desktop task windows within the same row.
+        fun batchCenterDesktopTaskWindows(endIndex: Int) {
+            // Calculate the shift amount required to center the desktop task items.
+            val rangeCenterX = (currentRowUnionRange.left + currentRowUnionRange.right) / 2f
+            val currentDiffX = (layoutBounds.centerX() - rangeCenterX).coerceAtLeast(0f)
+            for (j in currentRowFirstItemIndex until endIndex) {
+                outWindowRects[j].offset(currentDiffX, offsetY)
+            }
+        }
+
+        outWindowRects.forEachIndexed { index, rect ->
+            if (rect.top != currentRowY) {
+                // As a new row begins processing, batch-shift the previous row's rects
+                // and reset its parameters.
+                batchCenterDesktopTaskWindows(index)
+                currentRowUnionRange.set(rect)
+                currentRowY = rect.top
+                currentRowFirstItemIndex = index
+            }
+
+            // Extend the range by adding the [rect]'s width and extra in-between items
+            // spacing.
+            currentRowUnionRange.right = rect.right
+        }
+
+        // Post-processing rects in the last row.
+        batchCenterDesktopTaskWindows(outWindowRects.size)
+    }
+
+    private companion object {
+        const val VERTICAL_SPACE_BETWEEN_TASKS = 24
+        const val HORIZONTAL_SPACE_BETWEEN_TASKS = 24
+        const val OVERVIEW_INSET_TOP_BOTTOM = 16
+        const val OVERVIEW_INSET_LEFT_RIGHT = 16
+        const val MAXIMUM_TASK_HEIGHT = 800
     }
 }
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/DesktopTaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/DesktopTaskViewModel.kt
index 0a60ee9..4de0b90 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/DesktopTaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/DesktopTaskViewModel.kt
@@ -16,17 +16,30 @@
 
 package com.android.quickstep.recents.ui.viewmodel
 
+import android.graphics.Rect
 import android.util.Size
 import com.android.quickstep.recents.domain.model.DesktopTaskBoundsData
 import com.android.quickstep.recents.domain.usecase.OrganizeDesktopTasksUseCase
 
 /** ViewModel used for [com.android.quickstep.views.DesktopTaskView]. */
 class DesktopTaskViewModel(private val organizeDesktopTasksUseCase: OrganizeDesktopTasksUseCase) {
+    /** Positions for desktop tasks as calculated by [organizeDesktopTasksUseCase] */
     var organizedDesktopTaskPositions = emptyList<DesktopTaskBoundsData>()
         private set
 
+    /**
+     * Computes new task positions using [organizeDesktopTasksUseCase]. The result is stored in
+     * [organizedDesktopTaskPositions]. This is used for the exploded desktop view where the usecase
+     * will scale and translate tasks so that they don't overlap.
+     *
+     * @param desktopSize the size available for organizing the tasks.
+     * @param defaultPositions the tasks and their bounds as they appear on a desktop.
+     */
     fun organizeDesktopTasks(desktopSize: Size, defaultPositions: List<DesktopTaskBoundsData>) {
         organizedDesktopTaskPositions =
-            organizeDesktopTasksUseCase.run(desktopSize, defaultPositions)
+            organizeDesktopTasksUseCase.run(
+                desktopBounds = Rect(0, 0, desktopSize.width, desktopSize.height),
+                taskBounds = defaultPositions,
+            )
     }
 }
diff --git a/quickstep/src/com/android/quickstep/util/IconLabelUtil.kt b/quickstep/src/com/android/quickstep/util/IconLabelUtil.kt
new file mode 100644
index 0000000..a876bca
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/IconLabelUtil.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2025 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.quickstep.util
+
+import android.app.ActivityManager
+import android.content.Context
+import android.content.pm.ActivityInfo
+import android.os.UserHandle
+import com.android.launcher3.Utilities
+
+object IconLabelUtil {
+    @JvmStatic
+    @JvmOverloads
+    fun getBadgedContentDescription(
+        context: Context,
+        info: ActivityInfo,
+        userId: Int,
+        taskDescription: ActivityManager.TaskDescription? = null,
+    ): String {
+        val packageManager = context.packageManager
+        var taskLabel = taskDescription?.let { Utilities.trim(it.label) }
+        if (taskLabel.isNullOrEmpty()) {
+            taskLabel = Utilities.trim(info.loadLabel(packageManager))
+        }
+
+        val applicationLabel = Utilities.trim(info.applicationInfo.loadLabel(packageManager))
+        val badgedApplicationLabel =
+            if (userId != UserHandle.myUserId())
+                packageManager
+                    .getUserBadgedLabel(applicationLabel, UserHandle.of(userId))
+                    .toString()
+            else applicationLabel
+        return if (applicationLabel == taskLabel) badgedApplicationLabel
+        else "$badgedApplicationLabel $taskLabel"
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java b/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java
index 498078b..ceffbe4 100644
--- a/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java
+++ b/quickstep/src/com/android/quickstep/util/TaskGridNavHelper.java
@@ -29,6 +29,7 @@
  */
 public class TaskGridNavHelper {
     public static final int CLEAR_ALL_PLACEHOLDER_ID = -1;
+    public static final int ADD_DESK_PLACEHOLDER_ID = -2;
 
     public static final int DIRECTION_UP = 0;
     public static final int DIRECTION_DOWN = 1;
@@ -41,44 +42,42 @@
     public @interface TASK_NAV_DIRECTION {}
 
     private final IntArray mOriginalTopRowIds;
-    private IntArray mTopRowIds;
-    private IntArray mBottomRowIds;
+    private final IntArray mTopRowIds = new IntArray();
+    private final IntArray mBottomRowIds = new IntArray();
 
     public TaskGridNavHelper(IntArray topIds, IntArray bottomIds,
-            List<Integer> largeTileIds) {
+            List<Integer> largeTileIds, boolean hasAddDesktopButton) {
         mOriginalTopRowIds = topIds.clone();
-        generateTaskViewIdGrid(topIds, bottomIds, largeTileIds);
+        generateTaskViewIdGrid(topIds, bottomIds, largeTileIds, hasAddDesktopButton);
     }
 
     private void generateTaskViewIdGrid(IntArray topRowIdArray, IntArray bottomRowIdArray,
-            List<Integer> largeTileIds) {
-
-        int maxSize = Math.max(topRowIdArray.size(), bottomRowIdArray.size())
-                + largeTileIds.size();
-        int minSize = Math.min(topRowIdArray.size(), bottomRowIdArray.size())
-                + largeTileIds.size();
-
-        // Add Large tile task views first at the beginning
-        for (int i = 0; i < largeTileIds.size(); i++) {
-            topRowIdArray.add(i, largeTileIds.get(i));
-            bottomRowIdArray.add(i, largeTileIds.get(i));
+            List<Integer> largeTileIds, boolean hasAddDesktopButton) {
+        // Add AddDesktopButton and lage tiles to both rows.
+        if (hasAddDesktopButton) {
+            mTopRowIds.add(ADD_DESK_PLACEHOLDER_ID);
+            mBottomRowIds.add(ADD_DESK_PLACEHOLDER_ID);
         }
+        for (Integer tileId : largeTileIds) {
+            mTopRowIds.add(tileId);
+            mBottomRowIds.add(tileId);
+        }
+
+        // Add row ids to their respective rows.
+        mTopRowIds.addAll(topRowIdArray);
+        mBottomRowIds.addAll(bottomRowIdArray);
 
         // Fill in the shorter array with the ids from the longer one.
-        for (int i = minSize; i < maxSize; i++) {
-            if (i >= topRowIdArray.size()) {
-                topRowIdArray.add(bottomRowIdArray.get(i));
-            } else {
-                bottomRowIdArray.add(topRowIdArray.get(i));
-            }
+        while (mTopRowIds.size() > mBottomRowIds.size()) {
+            mBottomRowIds.add(mTopRowIds.get(mBottomRowIds.size()));
+        }
+        while (mBottomRowIds.size() > mTopRowIds.size()) {
+            mTopRowIds.add(mBottomRowIds.get(mTopRowIds.size()));
         }
 
-        // Add the clear all button to the end of both arrays
-        topRowIdArray.add(CLEAR_ALL_PLACEHOLDER_ID);
-        bottomRowIdArray.add(CLEAR_ALL_PLACEHOLDER_ID);
-
-        mTopRowIds = topRowIdArray;
-        mBottomRowIds = bottomRowIdArray;
+        // Add the clear all button to the end of both arrays.
+        mTopRowIds.add(CLEAR_ALL_PLACEHOLDER_ID);
+        mBottomRowIds.add(CLEAR_ALL_PLACEHOLDER_ID);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/util/TransformParams.java b/quickstep/src/com/android/quickstep/util/TransformParams.java
index bb88818..1c1fbd8 100644
--- a/quickstep/src/com/android/quickstep/util/TransformParams.java
+++ b/quickstep/src/com/android/quickstep/util/TransformParams.java
@@ -22,10 +22,14 @@
 import android.view.SurfaceControl;
 import android.window.TransitionInfo;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.quickstep.RemoteAnimationTargets;
 import com.android.quickstep.util.SurfaceTransaction.SurfaceProperties;
 import com.android.window.flags.Flags;
 
+import java.util.function.Supplier;
+
 public class TransformParams {
 
     public static FloatProperty<TransformParams> PROGRESS =
@@ -60,15 +64,23 @@
     private float mCornerRadius;
     private RemoteAnimationTargets mTargetSet;
     private TransitionInfo mTransitionInfo;
+    private boolean mCornerRadiusIsOverridden;
     private SurfaceTransactionApplier mSyncTransactionApplier;
+    private Supplier<SurfaceTransaction> mSurfaceTransactionSupplier;
 
     private BuilderProxy mHomeBuilderProxy = BuilderProxy.ALWAYS_VISIBLE;
     private BuilderProxy mBaseBuilderProxy = BuilderProxy.ALWAYS_VISIBLE;
 
     public TransformParams() {
+        this(SurfaceTransaction::new);
+    }
+
+    @VisibleForTesting
+    public TransformParams(Supplier<SurfaceTransaction> surfaceTransactionSupplier) {
         mProgress = 0;
         mTargetAlpha = 1;
         mCornerRadius = -1;
+        mSurfaceTransactionSupplier = surfaceTransactionSupplier;
     }
 
     /**
@@ -115,6 +127,7 @@
      */
     public TransformParams setTransitionInfo(TransitionInfo transitionInfo) {
         mTransitionInfo = transitionInfo;
+        mCornerRadiusIsOverridden = false;
         return this;
     }
 
@@ -148,7 +161,7 @@
     /** Builds the SurfaceTransaction from the given BuilderProxy params. */
     public SurfaceTransaction createSurfaceParams(BuilderProxy proxy) {
         RemoteAnimationTargets targets = mTargetSet;
-        SurfaceTransaction transaction = new SurfaceTransaction();
+        SurfaceTransaction transaction = mSurfaceTransactionSupplier.get();
         if (targets == null) {
             return transaction;
         }
@@ -166,8 +179,13 @@
             targetProxy.onBuildTargetParams(builder, app, this);
             // Override the corner radius for {@code app} with the leash used by Shell, so that it
             // doesn't interfere with the window clip and corner radius applied here.
-            overrideChangeLeashCornerRadiusToZero(app, transaction.getTransaction());
+            // Only override the corner radius once - so that we don't accidentally override at the
+            // end of transition after WM Shell has reset the corner radius of the task.
+            if (!mCornerRadiusIsOverridden) {
+                overrideFreeformChangeLeashCornerRadiusToZero(app, transaction.getTransaction());
+            }
         }
+        mCornerRadiusIsOverridden = true;
 
         // always put wallpaper layer to bottom.
         final int wallpaperLength = targets.wallpapers != null ? targets.wallpapers.length : 0;
@@ -178,11 +196,15 @@
         return transaction;
     }
 
-    private void overrideChangeLeashCornerRadiusToZero(
+    private void overrideFreeformChangeLeashCornerRadiusToZero(
             RemoteAnimationTarget app, SurfaceControl.Transaction transaction) {
         if (!Flags.enableDesktopRecentsTransitionsCornersBugfix()) {
             return;
         }
+        if (app.taskInfo == null || !app.taskInfo.isFreeform()) {
+            return;
+        }
+
         SurfaceControl changeLeash = getChangeLeashForApp(app);
         if (changeLeash != null) {
             transaction.setCornerRadius(changeLeash, 0);
diff --git a/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt b/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
index e353160..9f3c017 100644
--- a/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
+++ b/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
@@ -17,13 +17,15 @@
 package com.android.quickstep.views
 
 import android.content.Context
-import android.graphics.drawable.ShapeDrawable
-import android.graphics.drawable.shapes.RoundRectShape
+import android.graphics.Canvas
+import android.graphics.Rect
 import android.util.AttributeSet
 import android.widget.ImageButton
 import com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_X
 import com.android.launcher3.R
 import com.android.launcher3.util.MultiPropertyFactory
+import com.android.quickstep.util.BorderAnimator
+import com.android.quickstep.util.BorderAnimator.Companion.createSimpleBorderAnimator
 
 /**
  * Button for supporting multiple desktop sessions. The button will be next to the first TaskView
@@ -55,20 +57,49 @@
             multiTranslationX[TranslationX.OFFSET.ordinal].value = value
         }
 
-    override fun onFinishInflate() {
-        super.onFinishInflate()
+    private val focusBorderAnimator: BorderAnimator =
+        createSimpleBorderAnimator(
+            context.resources.getDimensionPixelSize(R.dimen.add_desktop_button_size),
+            context.resources.getDimensionPixelSize(R.dimen.keyboard_quick_switch_border_width),
+            this::getBorderBounds,
+            this,
+            context
+                .obtainStyledAttributes(attrs, R.styleable.AddDesktopButton)
+                .getColor(
+                    R.styleable.AddDesktopButton_focusBorderColor,
+                    BorderAnimator.DEFAULT_BORDER_COLOR,
+                ),
+        )
 
-        background =
-            ShapeDrawable().apply {
-                shape =
-                    RoundRectShape(
-                        FloatArray(8) { R.dimen.add_desktop_button_size.toFloat() },
-                        null,
-                        null,
-                    )
-                setTint(
-                    resources.getColor(android.R.color.system_surface_bright_light, context.theme)
-                )
+    var borderEnabled = false
+        set(value) {
+            if (field == value) {
+                return
             }
+            field = value
+            focusBorderAnimator.setBorderVisibility(visible = field && isFocused, animated = true)
+        }
+
+    public override fun onFocusChanged(
+        gainFocus: Boolean,
+        direction: Int,
+        previouslyFocusedRect: Rect?,
+    ) {
+        super.onFocusChanged(gainFocus, direction, previouslyFocusedRect)
+        if (borderEnabled) {
+            focusBorderAnimator.setBorderVisibility(gainFocus, /* animated= */ true)
+        }
+    }
+
+    private fun getBorderBounds(bounds: Rect) {
+        bounds.set(0, 0, width, height)
+        val outlinePadding =
+            context.resources.getDimensionPixelSize(R.dimen.add_desktop_button_outline_padding)
+        bounds.inset(-outlinePadding, -outlinePadding)
+    }
+
+    override fun draw(canvas: Canvas) {
+        focusBorderAnimator.drawBorder(canvas)
+        super.draw(canvas)
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 9d3b23a..99bfa7e 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1633,6 +1633,9 @@
             taskView.setBorderEnabled(enabled);
         }
         mClearAllButton.setBorderEnabled(enabled);
+        if (mAddDesktopButton != null) {
+            mAddDesktopButton.setBorderEnabled(enabled);
+        }
     }
 
     /**
@@ -3238,6 +3241,7 @@
 
         int topRowWidth = 0;
         int bottomRowWidth = 0;
+        int largeTileRowWidth = 0;
         float topAccumulatedTranslationX = 0;
         float bottomAccumulatedTranslationX = 0;
 
@@ -3248,9 +3252,12 @@
         int focusedTaskViewShift = 0;
         int largeTaskWidthAndSpacing = 0;
         int snappedTaskRowWidth = 0;
+        int expectedCurrentTaskRowWidth = 0;
         int snappedPage = isKeyboardTaskFocusPending() ? mKeyboardTaskFocusIndex : getNextPage();
         TaskView snappedTaskView = getTaskViewAt(snappedPage);
         TaskView homeTaskView = getHomeTaskView();
+        TaskView expectedCurrentTaskView = mUtils.getExpectedCurrentTask(getFocusedTaskView(),
+                getRunningTaskView());
         TaskView nextFocusedTaskView = null;
 
         // Don't clear the top row, if the user has dismissed a task, to maintain the task order.
@@ -3289,6 +3296,7 @@
                 if (!(taskView instanceof DesktopTaskView && isSplitSelectionActive())) {
                     topRowWidth += taskWidthAndSpacing;
                     bottomRowWidth += taskWidthAndSpacing;
+                    largeTileRowWidth += taskWidthAndSpacing;
                 }
                 gridTranslation += focusedTaskViewShift;
                 gridTranslation += mIsRtl ? taskWidthAndSpacing : -taskWidthAndSpacing;
@@ -3300,8 +3308,10 @@
                 largeTaskWidthAndSpacing = taskWidthAndSpacing;
 
                 if (taskView == snappedTaskView) {
-                    // If focused task is snapped, the row width is just task width and spacing.
-                    snappedTaskRowWidth = taskWidthAndSpacing;
+                    snappedTaskRowWidth = largeTileRowWidth;
+                }
+                if (taskView == expectedCurrentTaskView) {
+                    expectedCurrentTaskRowWidth = largeTileRowWidth;
                 }
             } else {
                 if (encounteredLastLargeTaskView) {
@@ -3370,8 +3380,12 @@
                     lastBottomTaskViews.add(taskView);
                     lastTopTaskViews.clear();
                 }
+                int taskViewRowWidth = isTopRow ? topRowWidth : bottomRowWidth;
                 if (taskView == snappedTaskView) {
-                    snappedTaskRowWidth = isTopRow ? topRowWidth : bottomRowWidth;
+                    snappedTaskRowWidth = taskViewRowWidth;
+                }
+                if (taskView == expectedCurrentTaskView) {
+                    expectedCurrentTaskRowWidth = taskViewRowWidth;
                 }
             }
             gridTranslations.put(taskView, gridTranslation);
@@ -3412,17 +3426,16 @@
         float clearAllShortTotalWidthTranslation = 0;
         int longRowWidth = Math.max(topRowWidth, bottomRowWidth);
 
-        // If Recents contains only large task sizes, it should only consider 1 large size
-        // for ClearAllButton translation. The space at the left side of the large task will be
-        // empty and it should be move ClearAllButton further away as well.
-        // TODO(b/359573248): Validate the translation for ClearAllButton for grid only.
-        if (enableLargeDesktopWindowingTile() && largeTasksCount == getTaskViewCount()) {
-            longRowWidth = largeTaskWidthAndSpacing;
-        }
-
         // If first task is not in the expected position (mLastComputedTaskSize) and being too close
         // to ClearAllButton, then apply extra translation to ClearAllButton.
-        int firstTaskStart = mLastComputedGridSize.left + longRowWidth;
+        int rowWidthAfterExpectedCurrentTask = longRowWidth - expectedCurrentTaskRowWidth;
+        int expectedCurrentTaskWidthAndSpacing =
+                (expectedCurrentTaskView != null
+                        ? expectedCurrentTaskView.getLayoutParams().width
+                        : 0
+                ) + mPageSpacing;
+        int firstTaskStart = mLastComputedGridSize.left + rowWidthAfterExpectedCurrentTask
+                + expectedCurrentTaskWidthAndSpacing;
         int expectedFirstTaskStart = mLastComputedTaskSize.right;
         if (firstTaskStart < expectedFirstTaskStart) {
             mClearAllShortTotalWidthTranslation = expectedFirstTaskStart - firstTaskStart;
@@ -4561,24 +4574,32 @@
 
         // Init task grid nav helper with top/bottom id arrays.
         TaskGridNavHelper taskGridNavHelper = new TaskGridNavHelper(getTopRowIdArray(),
-                getBottomRowIdArray(), mUtils.getLargeTaskViewIds());
+                getBottomRowIdArray(), mUtils.getLargeTaskViewIds(), mAddDesktopButton != null);
 
         // Get current page's task view ID.
         TaskView currentPageTaskView = getCurrentPageTaskView();
         int currentPageTaskViewId;
+        final int clearAllButtonIndex = indexOfChild(mClearAllButton);
+        final int addDesktopButtonIndex = indexOfChild(mAddDesktopButton);
         if (currentPageTaskView != null) {
             currentPageTaskViewId = currentPageTaskView.getTaskViewId();
-        } else if (mCurrentPage == indexOfChild(mClearAllButton)) {
+        } else if (mCurrentPage == clearAllButtonIndex) {
             currentPageTaskViewId = TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID;
+        } else if (mCurrentPage == addDesktopButtonIndex) {
+            currentPageTaskViewId = TaskGridNavHelper.ADD_DESK_PLACEHOLDER_ID;
         } else {
             return INVALID_PAGE;
         }
 
-        int nextGridPage =
+        final int nextGridPage =
                 taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, cycle);
-        return nextGridPage == TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID
-                ? indexOfChild(mClearAllButton)
-                : indexOfChild(getTaskViewFromTaskViewId(nextGridPage));
+        if (nextGridPage == TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID) {
+            return clearAllButtonIndex;
+        }
+        if (nextGridPage == TaskGridNavHelper.ADD_DESK_PLACEHOLDER_ID) {
+            return addDesktopButtonIndex;
+        }
+        return indexOfChild(getTaskViewFromTaskViewId(nextGridPage));
     }
 
     private void runDismissAnimation(PendingAnimation pendingAnim) {
@@ -6116,8 +6137,10 @@
     }
 
     private int getFirstViewIndex() {
-        final TaskView firstView;
-        if (mShowAsGridLastOnLayout) {
+        final View firstView;
+        if (mAddDesktopButton != null) {
+            firstView = mAddDesktopButton;
+        } else if (mShowAsGridLastOnLayout) {
             // For grid Overview, it always start if a large tile (focused task or desktop task) if
             // they exist, otherwise it start with the first task.
             TaskView firstLargeTaskView = mUtils.getFirstLargeTaskView();
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 4b1b8dc..0465dbc 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -248,8 +248,43 @@
             )
 
     private val tempCoordinates = FloatArray(2)
-    private val focusBorderAnimator: BorderAnimator?
-    private val hoverBorderAnimator: BorderAnimator?
+    private val focusBorderAnimator: BorderAnimator? =
+        focusBorderAnimator
+            ?: createSimpleBorderAnimator(
+                TaskCornerRadius.get(context).toInt(),
+                context.resources.getDimensionPixelSize(R.dimen.keyboard_quick_switch_border_width),
+                this::getThumbnailBounds,
+                this,
+                context
+                    .obtainStyledAttributes(attrs, R.styleable.TaskView, defStyleAttr, defStyleRes)
+                    .getColor(
+                        R.styleable.TaskView_focusBorderColor,
+                        BorderAnimator.DEFAULT_BORDER_COLOR,
+                    ),
+            )
+
+    private val hoverBorderAnimator: BorderAnimator? =
+        hoverBorderAnimator
+            ?: if (enableCursorHoverStates())
+                createSimpleBorderAnimator(
+                    TaskCornerRadius.get(context).toInt(),
+                    context.resources.getDimensionPixelSize(R.dimen.task_hover_border_width),
+                    this::getThumbnailBounds,
+                    this,
+                    context
+                        .obtainStyledAttributes(
+                            attrs,
+                            R.styleable.TaskView,
+                            defStyleAttr,
+                            defStyleRes,
+                        )
+                        .getColor(
+                            R.styleable.TaskView_hoverBorderColor,
+                            BorderAnimator.DEFAULT_BORDER_COLOR,
+                        ),
+                )
+            else null
+
     private val rootViewDisplayId: Int
         get() = rootView.display?.displayId ?: Display.DEFAULT_DISPLAY
 
@@ -519,40 +554,7 @@
     init {
         setOnClickListener { _ -> onClick() }
 
-        val cursorHoverStatesEnabled = enableCursorHoverStates()
-        setWillNotDraw(!cursorHoverStatesEnabled)
-        context.obtainStyledAttributes(attrs, R.styleable.TaskView, defStyleAttr, defStyleRes).use {
-            this.focusBorderAnimator =
-                focusBorderAnimator
-                    ?: createSimpleBorderAnimator(
-                        TaskCornerRadius.get(context).toInt(),
-                        context.resources.getDimensionPixelSize(
-                            R.dimen.keyboard_quick_switch_border_width
-                        ),
-                        { bounds: Rect -> getThumbnailBounds(bounds) },
-                        this,
-                        it.getColor(
-                            R.styleable.TaskView_focusBorderColor,
-                            BorderAnimator.DEFAULT_BORDER_COLOR,
-                        ),
-                    )
-            this.hoverBorderAnimator =
-                hoverBorderAnimator
-                    ?: if (cursorHoverStatesEnabled)
-                        createSimpleBorderAnimator(
-                            TaskCornerRadius.get(context).toInt(),
-                            context.resources.getDimensionPixelSize(
-                                R.dimen.task_hover_border_width
-                            ),
-                            { bounds: Rect -> getThumbnailBounds(bounds) },
-                            this,
-                            it.getColor(
-                                R.styleable.TaskView_hoverBorderColor,
-                                BorderAnimator.DEFAULT_BORDER_COLOR,
-                            ),
-                        )
-                    else null
-        }
+        setWillNotDraw(!enableCursorHoverStates())
     }
 
     @VisibleForTesting(otherwise = VisibleForTesting.PROTECTED)
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
index 7aab75f..7066d21 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TaskGridNavHelperTest.kt
@@ -16,6 +16,7 @@
 package com.android.quickstep.util
 
 import com.android.launcher3.util.IntArray
+import com.android.quickstep.util.TaskGridNavHelper.ADD_DESK_PLACEHOLDER_ID
 import com.android.quickstep.util.TaskGridNavHelper.CLEAR_ALL_PLACEHOLDER_ID
 import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_DOWN
 import com.android.quickstep.util.TaskGridNavHelper.DIRECTION_LEFT
@@ -619,6 +620,161 @@
             .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
     }
 
+    /*
+                        5   3   1→----|
+                                      ↓
+         CLEAR_ALL                   ADD_DESKTOP
+                        6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressRightFromTop_goesToAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = 1,
+                    DIRECTION_RIGHT,
+                    delta = -1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
+    /*
+                    5   3   1
+       CLEAR_ALL                 ADD_DESKTOP
+                                  ↑
+                    6   4   2→----↑
+    */
+    @Test
+    fun withAddDesktopButton_pressRightFromBottom_goesToAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = 2,
+                    DIRECTION_RIGHT,
+                    delta = -1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
+    /*
+         ↓-------------------------------←|
+         |                                ↑
+         ↓      5   3   1                 |
+    CLEAR_ALL               ADD_DESKTOP--→
+                6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressRightFromAddDesktopButton_goesToClearAllButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
+                    DIRECTION_RIGHT,
+                    delta = -1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(CLEAR_ALL_PLACEHOLDER_ID)
+    }
+
+    /*
+           |→--------------------------------|
+           |                                 |
+           ↑                5   3   1        ↓
+           ←------CLEAR_ALL             ADD_DESKTOP
+
+                            6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressLeftFromClearAllButton_goesToAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = CLEAR_ALL_PLACEHOLDER_ID,
+                    DIRECTION_LEFT,
+                    delta = 1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
+    /*
+                        5   3   1
+                                   ←--↑
+        CLEAR_ALL                  ↓-→ADD_DESKTOP
+                        6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressUpOnAddDesktop_stayOnAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
+                    DIRECTION_UP,
+                    delta = 1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
+    /*
+                        5   3   1
+        CLEAR_ALL                  ↑--→ADD_DESKTOP
+                                   ↑←--↓
+                        6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressDownOnAddDesktop_stayOnAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
+                    DIRECTION_DOWN,
+                    delta = 1,
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
+    /*
+                        5   3   1
+           CLEAR_ALL                DESKTOP--→ADD_DESKTOP
+                        6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressRightFromDesktopTask_goesToAddDesktopButton() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = ADD_DESK_PLACEHOLDER_ID,
+                    DIRECTION_LEFT,
+                    delta = 1,
+                    largeTileIds = listOf(DESKTOP_TASK_ID),
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(DESKTOP_TASK_ID)
+    }
+
+    /*
+                        5   3   1
+           CLEAR_ALL                DESKTOP←--ADD_DESKTOP
+                        6   4   2
+    */
+    @Test
+    fun withAddDesktopButton_pressLeftFromAddDesktopButton_goesToDesktopTask() {
+        assertThat(
+                getNextGridPage(
+                    currentPageTaskViewId = DESKTOP_TASK_ID,
+                    DIRECTION_RIGHT,
+                    delta = -1,
+                    largeTileIds = listOf(DESKTOP_TASK_ID),
+                    hasAddDesktopButton = true,
+                )
+            )
+            .isEqualTo(ADD_DESK_PLACEHOLDER_ID)
+    }
+
     private fun getNextGridPage(
         currentPageTaskViewId: Int,
         direction: Int,
@@ -626,8 +782,10 @@
         topIds: IntArray = IntArray.wrap(1, 3, 5),
         bottomIds: IntArray = IntArray.wrap(2, 4, 6),
         largeTileIds: List<Int> = emptyList(),
+        hasAddDesktopButton: Boolean = false,
     ): Int {
-        val taskGridNavHelper = TaskGridNavHelper(topIds, bottomIds, largeTileIds)
+        val taskGridNavHelper =
+            TaskGridNavHelper(topIds, bottomIds, largeTileIds, hasAddDesktopButton)
         return taskGridNavHelper.getNextGridPage(currentPageTaskViewId, delta, direction, true)
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TransformParamsTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TransformParamsTest.kt
new file mode 100644
index 0000000..6dbb667
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/TransformParamsTest.kt
@@ -0,0 +1,191 @@
+/*
+ * Copyright (C) 2025 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.quickstep.util
+
+import android.app.ActivityManager.RunningTaskInfo
+import android.app.WindowConfiguration
+import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
+import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
+import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import android.view.RemoteAnimationTarget
+import android.view.SurfaceControl
+import android.view.WindowManager.TRANSIT_OPEN
+import android.window.TransitionInfo
+import android.window.TransitionInfo.Change
+import android.window.TransitionInfo.FLAG_NONE
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.quickstep.RemoteAnimationTargets
+import com.android.quickstep.util.TransformParams.BuilderProxy.NO_OP
+import com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_RECENTS_TRANSITIONS_CORNERS_BUGFIX
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.anyOrNull
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class TransformParamsTest {
+    private val surfaceTransaction = mock<SurfaceTransaction>()
+    private val transaction = mock<SurfaceControl.Transaction>()
+    private val transformParams = TransformParams(::surfaceTransaction)
+
+    private val freeformTaskInfo1 =
+        createTaskInfo(taskId = 1, windowingMode = WINDOWING_MODE_FREEFORM)
+    private val freeformTaskInfo2 =
+        createTaskInfo(taskId = 2, windowingMode = WINDOWING_MODE_FREEFORM)
+    private val fullscreenTaskInfo1 =
+        createTaskInfo(taskId = 1, windowingMode = WINDOWING_MODE_FULLSCREEN)
+
+    @get:Rule val setFlagsRule: SetFlagsRule = SetFlagsRule()
+
+    @Before
+    fun setUp() {
+        whenever(surfaceTransaction.transaction).thenReturn(transaction)
+        whenever(surfaceTransaction.forSurface(anyOrNull()))
+            .thenReturn(mock<SurfaceTransaction.SurfaceProperties>())
+        transformParams.setCornerRadius(CORNER_RADIUS)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_RECENTS_TRANSITIONS_CORNERS_BUGFIX)
+    fun createSurfaceParams_freeformTasks_overridesCornerRadius() {
+        val transitionInfo = TransitionInfo(TRANSIT_OPEN, FLAG_NONE)
+        val leash1 = mock<SurfaceControl>()
+        val leash2 = mock<SurfaceControl>()
+        transitionInfo.addChange(createChange(freeformTaskInfo1, leash = leash1))
+        transitionInfo.addChange(createChange(freeformTaskInfo2, leash = leash2))
+        transformParams.setTransitionInfo(transitionInfo)
+        transformParams.setTargetSet(createTargetSet(listOf(freeformTaskInfo1, freeformTaskInfo2)))
+
+        transformParams.createSurfaceParams(NO_OP)
+
+        verify(transaction).setCornerRadius(leash1, 0f)
+        verify(transaction).setCornerRadius(leash2, 0f)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_RECENTS_TRANSITIONS_CORNERS_BUGFIX)
+    fun createSurfaceParams_freeformTasks_overridesCornerRadiusOnlyOnce() {
+        val transitionInfo = TransitionInfo(TRANSIT_OPEN, FLAG_NONE)
+        val leash1 = mock<SurfaceControl>()
+        val leash2 = mock<SurfaceControl>()
+        transitionInfo.addChange(createChange(freeformTaskInfo1, leash = leash1))
+        transitionInfo.addChange(createChange(freeformTaskInfo2, leash = leash2))
+        transformParams.setTransitionInfo(transitionInfo)
+        transformParams.setTargetSet(createTargetSet(listOf(freeformTaskInfo1, freeformTaskInfo2)))
+        transformParams.createSurfaceParams(NO_OP)
+
+        transformParams.createSurfaceParams(NO_OP)
+
+        verify(transaction).setCornerRadius(leash1, 0f)
+        verify(transaction).setCornerRadius(leash2, 0f)
+    }
+
+    @Test
+    @DisableFlags(FLAG_ENABLE_DESKTOP_RECENTS_TRANSITIONS_CORNERS_BUGFIX)
+    fun createSurfaceParams_flagDisabled_doesntOverrideCornerRadius() {
+        val transitionInfo = TransitionInfo(TRANSIT_OPEN, FLAG_NONE)
+        val leash1 = mock<SurfaceControl>()
+        val leash2 = mock<SurfaceControl>()
+        transitionInfo.addChange(createChange(freeformTaskInfo1, leash = leash1))
+        transitionInfo.addChange(createChange(freeformTaskInfo2, leash = leash2))
+        transformParams.setTransitionInfo(transitionInfo)
+        transformParams.setTargetSet(createTargetSet(listOf(freeformTaskInfo1, freeformTaskInfo2)))
+
+        transformParams.createSurfaceParams(NO_OP)
+
+        verify(transaction, never()).setCornerRadius(leash1, 0f)
+        verify(transaction, never()).setCornerRadius(leash2, 0f)
+    }
+
+    @Test
+    @EnableFlags(FLAG_ENABLE_DESKTOP_RECENTS_TRANSITIONS_CORNERS_BUGFIX)
+    fun createSurfaceParams_fullscreenTasks_doesntOverrideCornerRadius() {
+        val transitionInfo = TransitionInfo(TRANSIT_OPEN, FLAG_NONE)
+        val leash = mock<SurfaceControl>()
+        transitionInfo.addChange(createChange(fullscreenTaskInfo1, leash = leash))
+        transformParams.setTransitionInfo(transitionInfo)
+        transformParams.setTargetSet(createTargetSet(listOf(fullscreenTaskInfo1)))
+
+        transformParams.createSurfaceParams(NO_OP)
+
+        verify(transaction, never()).setCornerRadius(leash, 0f)
+    }
+
+    private fun createTargetSet(taskInfos: List<RunningTaskInfo>): RemoteAnimationTargets {
+        val remoteAnimationTargets = mutableListOf<RemoteAnimationTarget>()
+        taskInfos.map { remoteAnimationTargets.add(createRemoteAnimationTarget(it)) }
+        return RemoteAnimationTargets(
+            remoteAnimationTargets.toTypedArray(),
+            /* wallpapers= */ null,
+            /* nonApps= */ null,
+            /* targetMode= */ TRANSIT_OPEN,
+        )
+    }
+
+    private fun createRemoteAnimationTarget(taskInfo: RunningTaskInfo): RemoteAnimationTarget {
+        val windowConfig = mock<WindowConfiguration>()
+        whenever(windowConfig.activityType).thenReturn(ACTIVITY_TYPE_STANDARD)
+        return RemoteAnimationTarget(
+            taskInfo.taskId,
+            /* mode= */ TRANSIT_OPEN,
+            /* leash= */ null,
+            /* isTranslucent= */ false,
+            /* clipRect= */ null,
+            /* contentInsets= */ null,
+            /* prefixOrderIndex= */ 0,
+            /* position= */ null,
+            /* localBounds= */ null,
+            /* screenSpaceBounds= */ null,
+            windowConfig,
+            /* isNotInRecents= */ false,
+            /* startLeash= */ null,
+            /* startBounds= */ null,
+            taskInfo,
+            /* allowEnterPip= */ false,
+        )
+    }
+
+    private fun createTaskInfo(taskId: Int, windowingMode: Int): RunningTaskInfo {
+        val taskInfo = RunningTaskInfo()
+        taskInfo.taskId = taskId
+        taskInfo.configuration.windowConfiguration.windowingMode = windowingMode
+        return taskInfo
+    }
+
+    private fun createChange(taskInfo: RunningTaskInfo, leash: SurfaceControl): Change {
+        val taskInfo = createTaskInfo(taskInfo.taskId, taskInfo.windowingMode)
+        val change = Change(taskInfo.token, mock<SurfaceControl>())
+        change.mode = TRANSIT_OPEN
+        change.taskInfo = taskInfo
+        change.leash = leash
+        return change
+    }
+
+    private companion object {
+        private const val CORNER_RADIUS = 30f
+    }
+}
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 30ef24b..58fd154 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -29,6 +29,7 @@
 import static com.android.launcher3.AbstractFloatingView.getTopOpenViewWithType;
 import static com.android.launcher3.Flags.enableAddAppWidgetViaConfigActivityV2;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.Flags.enableStrictMode;
 import static com.android.launcher3.Flags.enableWorkspaceInflation;
 import static com.android.launcher3.LauncherAnimUtils.HOTSEAT_SCALE_PROPERTY_FACTORY;
 import static com.android.launcher3.LauncherAnimUtils.SCALE_INDEX_WIDGET_TRANSITION;
@@ -459,7 +460,8 @@
         Trace.beginAsyncSection(DISPLAY_ALL_APPS_TRACE_METHOD_NAME,
                 DISPLAY_ALL_APPS_TRACE_COOKIE);
         TraceHelper.INSTANCE.beginSection(ON_CREATE_EVT);
-        if (DEBUG_STRICT_MODE) {
+        if (DEBUG_STRICT_MODE
+                || (FeatureFlags.IS_STUDIO_BUILD && enableStrictMode())) {
             StrictMode.setThreadPolicy(new StrictMode.ThreadPolicy.Builder()
                     .detectDiskReads()
                     .detectDiskWrites()