Merge bubble transition with unfold

This change merges the bubble task view bounds change with the unfold
transition in order to avoid clipped expanded view.

For bubble bar, the task view will switch surfaces after unfolding
so nothing is merged in that case.

Bug: 390050545
Flag: EXEMPT bug fix
Test: atest TaskViewTransitionsTest
Test: atest UnfoldTransitionHandlerTest
Test: manual
       - disable bubble bar flag
       - create a floating bubble in folded state
       - pull up the IME
       - unfold and observe expanded view bounds
Test: repeat steps above with bubble bar flag enabled
Test: repeat steps above with task view repo flag disabled
Change-Id: I22787311159f2940d37b5aaa071c9b05aae15d17

8 files changed, 209 insertions, 12 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
index 81eff6f7399a..63c4c6eb857a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleController.java
@@ -81,6 +81,7 @@ import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.window.ScreenCapture;
 import android.window.ScreenCapture.SynchronousScreenCaptureListener;
+import android.window.TransitionInfo;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -155,7 +156,7 @@ import java.util.function.IntConsumer;
  */
 public class BubbleController implements ConfigurationChangeListener,
         RemoteCallable<BubbleController>, Bubbles.SysuiProxy.Provider,
-        BubbleBarDragListener {
+        BubbleBarDragListener, BubbleTaskUnfoldTransitionMerger {
 
     private static final String TAG = TAG_WITH_CLASS_NAME ? "BubbleController" : TAG_BUBBLES;
 
@@ -2174,6 +2175,32 @@ public class BubbleController implements ConfigurationChangeListener,
         });
     }
 
+    @Override
+    public boolean mergeTaskWithUnfold(@NonNull ActivityManager.RunningTaskInfo taskInfo,
+            @NonNull TransitionInfo.Change change,
+            @NonNull SurfaceControl.Transaction startT,
+            @NonNull SurfaceControl.Transaction finishT) {
+        if (!mBubbleTransitions.mTaskViewTransitions.isTaskViewTask(taskInfo)) {
+            // if this task isn't managed by bubble transitions just bail.
+            return false;
+        }
+        if (isShowingAsBubbleBar()) {
+            // if bubble bar is enabled, the task view will switch to a new surface on unfold, so we
+            // should not merge the transition.
+            return false;
+        }
+
+        boolean merged = mBubbleTransitions.mTaskViewTransitions.updateBoundsForUnfold(
+                change.getEndAbsBounds(), startT, finishT, change.getTaskInfo(), change.getLeash());
+        if (merged) {
+            BubbleViewProvider selectedBubble = mBubbleData.getSelectedBubble();
+            if (selectedBubble != null && selectedBubble.getExpandedView() != null) {
+                selectedBubble.getExpandedView().onContainerClipUpdate();
+            }
+        }
+        return merged;
+    }
+
     /** When bubbles are floating, this will be used to notify the floating views. */
     private final BubbleViewCallback mBubbleStackViewCallback = new BubbleViewCallback() {
         @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
index 290ef1633819..ac8393576477 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
@@ -843,7 +843,8 @@ public class BubbleExpandedView extends LinearLayout {
         onContainerClipUpdate();
     }
 
-    private void onContainerClipUpdate() {
+    /** Updates the clip bounds. */
+    public void onContainerClipUpdate() {
         if (mTopClip == 0 && mBottomClip == 0 && mRightClip == 0 && mLeftClip == 0) {
             if (mIsClipping) {
                 mIsClipping = false;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskUnfoldTransitionMerger.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskUnfoldTransitionMerger.kt
new file mode 100644
index 000000000000..13fabc8b1d91
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskUnfoldTransitionMerger.kt
@@ -0,0 +1,33 @@
+/*
+ * 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.wm.shell.bubbles
+
+import android.app.ActivityManager
+import android.view.SurfaceControl
+import android.window.TransitionInfo
+
+/** Merges a bubble task transition with the unfold transition. */
+interface BubbleTaskUnfoldTransitionMerger {
+
+    /** Attempts to merge the transition. Returns `true` if the change was merged. */
+    fun mergeTaskWithUnfold(
+        taskInfo: ActivityManager.RunningTaskInfo,
+        change: TransitionInfo.Change,
+        startT: SurfaceControl.Transaction,
+        finishT: SurfaceControl.Transaction
+    ): Boolean
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index 99d4a4066a9a..7ad891d6074e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -62,6 +62,7 @@ import com.android.wm.shell.bubbles.BubbleEducationController;
 import com.android.wm.shell.bubbles.BubbleLogger;
 import com.android.wm.shell.bubbles.BubblePositioner;
 import com.android.wm.shell.bubbles.BubbleResizabilityChecker;
+import com.android.wm.shell.bubbles.BubbleTaskUnfoldTransitionMerger;
 import com.android.wm.shell.bubbles.bar.BubbleBarDragListener;
 import com.android.wm.shell.bubbles.storage.BubblePersistentRepository;
 import com.android.wm.shell.common.DisplayController;
@@ -242,6 +243,13 @@ public abstract class WMShellModule {
                 context, logger, positioner, educationController, mainExecutor, bgExecutor);
     }
 
+    @WMSingleton
+    @Provides
+    static Optional<BubbleTaskUnfoldTransitionMerger> provideBubbleTaskUnfoldTransitionMerger(
+            Optional<BubbleController> bubbleController) {
+        return bubbleController.map(controller -> controller);
+    }
+
     // Note: Handler needed for LauncherApps.register
     @WMSingleton
     @Provides
@@ -703,7 +711,8 @@ public abstract class WMShellModule {
             Transitions transitions,
             @ShellMainThread ShellExecutor executor,
             @ShellMainThread Handler handler,
-            ShellInit shellInit) {
+            ShellInit shellInit,
+            Optional<BubbleTaskUnfoldTransitionMerger> bubbleTaskUnfoldTransitionMerger) {
         return new UnfoldTransitionHandler(
                 shellInit,
                 progressProvider.get(),
@@ -712,7 +721,8 @@ public abstract class WMShellModule {
                 transactionPool,
                 executor,
                 handler,
-                transitions);
+                transitions,
+                bubbleTaskUnfoldTransitionMerger);
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
index a6f872634ee9..22848c38bb1c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTransitions.java
@@ -728,6 +728,48 @@ public class TaskViewTransitions implements Transitions.TransitionHandler, TaskV
         taskView.notifyAppeared(newTask);
     }
 
+    /**
+     * Updates bounds for the task view during an unfold transition.
+     *
+     * @return true if the task was found and a transition for this task is pending. false
+     * otherwise.
+     */
+    public boolean updateBoundsForUnfold(Rect bounds, SurfaceControl.Transaction startTransaction,
+            SurfaceControl.Transaction finishTransaction,
+            ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
+        final TaskViewTaskController taskView = findTaskView(taskInfo);
+        if (taskView == null) {
+            return false;
+        }
+
+        final PendingTransition pendingTransition = findPending(taskView, TRANSIT_CHANGE);
+        if (pendingTransition == null) {
+            return false;
+        }
+
+        mPending.remove(pendingTransition);
+
+        // reparent the task under the task view surface and set the bounds on it
+        startTransaction.reparent(leash, taskView.getSurfaceControl())
+                .setPosition(leash, 0, 0)
+                .setWindowCrop(leash, bounds.width(), bounds.height())
+                .show(leash);
+        // the finish transaction would reparent the task back to the transition root, so reparent
+        // it again to the task view surface
+        finishTransaction.reparent(leash, taskView.getSurfaceControl())
+                .setPosition(leash, 0, 0)
+                .setWindowCrop(leash, bounds.width(), bounds.height());
+        if (useRepo()) {
+            final TaskViewRepository.TaskViewState state = mTaskViewRepo.byTaskView(taskView);
+            if (state != null) {
+                state.mBounds.set(bounds);
+            }
+        } else {
+            updateBoundsState(taskView, bounds);
+        }
+        return true;
+    }
+
     private void updateBounds(TaskViewTaskController taskView, Rect boundsOnScreen,
             SurfaceControl.Transaction startTransaction,
             SurfaceControl.Transaction finishTransaction,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
index 7fd19a7d2a88..706a366441cd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldTransitionHandler.java
@@ -38,6 +38,7 @@ import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.internal.protolog.ProtoLog;
+import com.android.wm.shell.bubbles.BubbleTaskUnfoldTransitionMerger;
 import com.android.wm.shell.shared.TransactionPool;
 import com.android.wm.shell.shared.TransitionUtil;
 import com.android.wm.shell.sysui.ShellInit;
@@ -53,6 +54,7 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.Executor;
 
 /**
@@ -80,6 +82,7 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
 
     private final ShellUnfoldProgressProvider mUnfoldProgressProvider;
     private final Transitions mTransitions;
+    private final Optional<BubbleTaskUnfoldTransitionMerger> mBubbleTaskUnfoldTransitionMerger;
     private final Executor mExecutor;
     private final TransactionPool mTransactionPool;
     private final Handler mHandler;
@@ -108,12 +111,14 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
             TransactionPool transactionPool,
             Executor executor,
             Handler handler,
-            Transitions transitions) {
+            Transitions transitions,
+            Optional<BubbleTaskUnfoldTransitionMerger> bubbleTaskUnfoldTransitionMerger) {
         mUnfoldProgressProvider = unfoldProgressProvider;
         mTransitions = transitions;
         mTransactionPool = transactionPool;
         mExecutor = executor;
         mHandler = handler;
+        mBubbleTaskUnfoldTransitionMerger = bubbleTaskUnfoldTransitionMerger;
 
         mAnimators.add(splitUnfoldTaskAnimator);
         mAnimators.add(fullscreenUnfoldAnimator);
@@ -237,14 +242,26 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
         }
         // TODO (b/286928742) unfold transition handler should be part of mixed handler to
         //  handle merges better.
+
         for (int i = 0; i < info.getChanges().size(); ++i) {
             final TransitionInfo.Change change = info.getChanges().get(i);
             final ActivityManager.RunningTaskInfo taskInfo = change.getTaskInfo();
             if (taskInfo != null
                     && taskInfo.configuration.windowConfiguration.isAlwaysOnTop()) {
-                // Tasks that are always on top (e.g. bubbles), will handle their own transition
-                // as they are on top of everything else. So skip merging transitions here.
-                return;
+                // Tasks that are always on top, excluding bubbles, will handle their own transition
+                // as they are on top of everything else. If this is a transition for a bubble task,
+                // attempt to merge it. Otherwise skip merging transitions.
+                if (mBubbleTaskUnfoldTransitionMerger.isPresent()) {
+                    boolean merged =
+                            mBubbleTaskUnfoldTransitionMerger
+                                    .get()
+                                    .mergeTaskWithUnfold(taskInfo, change, startT, finishT);
+                    if (!merged) {
+                        return;
+                    }
+                } else {
+                    return;
+                }
             }
         }
         // Apply changes happening during the unfold animation immediately
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
index 3a455ba6b5df..f11839ad4e72 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTransitionsTest.java
@@ -24,8 +24,12 @@ import static android.view.WindowManager.TRANSIT_TO_FRONT;
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assume.assumeTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyFloat;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -77,16 +81,15 @@ public class TaskViewTransitionsTest extends ShellTestCase {
     @Mock
     TaskViewTaskController mTaskViewTaskController;
     @Mock
-    ActivityManager.RunningTaskInfo mTaskInfo;
-    @Mock
     WindowContainerToken mToken;
     @Mock
     ShellTaskOrganizer mOrganizer;
     @Mock
     SyncTransactionQueue mSyncQueue;
 
-    Executor mExecutor = command -> command.run();
+    Executor mExecutor = Runnable::run;
 
+    ActivityManager.RunningTaskInfo mTaskInfo;
     TaskViewRepository mTaskViewRepository;
     TaskViewTransitions mTaskViewTransitions;
 
@@ -305,4 +308,66 @@ public class TaskViewTransitionsTest extends ShellTestCase {
 
         verify(mTaskViewTaskController).setTaskNotFound();
     }
+
+    @Test
+    public void updateBoundsForUnfold_taskNotFound_doesNothing() {
+        assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
+
+        ActivityManager.RunningTaskInfo taskInfo = new ActivityManager.RunningTaskInfo();
+        taskInfo.token = mock(WindowContainerToken.class);
+        taskInfo.taskId = 666;
+        Rect bounds = new Rect(100, 50, 200, 250);
+        SurfaceControl.Transaction startTransaction = mock(SurfaceControl.Transaction.class);
+        SurfaceControl.Transaction finishTransaction = mock(SurfaceControl.Transaction.class);
+        assertThat(
+                mTaskViewTransitions.updateBoundsForUnfold(bounds, startTransaction,
+                        finishTransaction, taskInfo, mock(SurfaceControl.class)))
+                .isFalse();
+
+        verify(startTransaction, never()).reparent(any(), any());
+    }
+
+    @Test
+    public void updateBoundsForUnfold_noPendingTransition_doesNothing() {
+        assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
+
+        Rect bounds = new Rect(100, 50, 200, 250);
+        mTaskViewTransitions.setTaskBounds(mTaskViewTaskController, bounds);
+        assertThat(mTaskViewTransitions.hasPending()).isFalse();
+
+        SurfaceControl.Transaction startTransaction = mock(SurfaceControl.Transaction.class);
+        SurfaceControl.Transaction finishTransaction = mock(SurfaceControl.Transaction.class);
+        assertThat(
+                mTaskViewTransitions.updateBoundsForUnfold(bounds, startTransaction,
+                        finishTransaction, mTaskInfo, mock(SurfaceControl.class)))
+                .isFalse();
+        verify(startTransaction, never()).reparent(any(), any());
+    }
+
+    @Test
+    public void updateBoundsForUnfold() {
+        assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
+
+        Rect bounds = new Rect(100, 50, 200, 250);
+        mTaskViewTransitions.updateVisibilityState(mTaskViewTaskController, /* visible= */ true);
+        mTaskViewTransitions.setTaskBounds(mTaskViewTaskController, bounds);
+        assertThat(mTaskViewTransitions.hasPending()).isTrue();
+
+        SurfaceControl.Transaction startTransaction = createMockTransaction();
+        SurfaceControl.Transaction finishTransaction = createMockTransaction();
+        assertThat(
+                mTaskViewTransitions.updateBoundsForUnfold(bounds, startTransaction,
+                        finishTransaction, mTaskInfo, mock(SurfaceControl.class)))
+                .isTrue();
+        assertThat(mTaskViewRepository.byTaskView(mTaskViewTaskController).mBounds)
+                .isEqualTo(bounds);
+    }
+
+    private SurfaceControl.Transaction createMockTransaction() {
+        SurfaceControl.Transaction transaction = mock(SurfaceControl.Transaction.class);
+        when(transaction.reparent(any(), any())).thenReturn(transaction);
+        when(transaction.setPosition(any(), anyFloat(), anyFloat())).thenReturn(transaction);
+        when(transaction.setWindowCrop(any(), anyInt(), anyInt())).thenReturn(transaction);
+        return transaction;
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldTransitionHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldTransitionHandlerTest.java
index aad18cba4436..e28d0acb579f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldTransitionHandlerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldTransitionHandlerTest.java
@@ -60,6 +60,7 @@ import org.mockito.InOrder;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Optional;
 import java.util.concurrent.Executor;
 
 public class UnfoldTransitionHandlerTest extends ShellTestCase {
@@ -98,7 +99,8 @@ public class UnfoldTransitionHandlerTest extends ShellTestCase {
                 mTransactionPool,
                 executor,
                 mHandler,
-                mTransitions
+                mTransitions,
+                /* bubbleTaskUnfoldTransitionMerger= */ Optional.empty()
         );
 
         shellInit.init();