Track taskviews in its own repository

This allows multiple shell components to have unified
tracking of taskview state since we now have multiple
systems using taskviews simultaneously.

Bug: 384976265
Test: TaskViewTests
Flag: com.android.wm.shell.task_view_repository
Change-Id: I7b51a520677c178f7400c4d2e09e5fa806a958e0

7 files changed, 309 insertions, 58 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
index eab7f6c055bd..1408b6efc7f9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellBaseModule.java
@@ -124,6 +124,7 @@ import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.sysui.ShellInterface;
 import com.android.wm.shell.taskview.TaskViewFactory;
 import com.android.wm.shell.taskview.TaskViewFactoryController;
+import com.android.wm.shell.taskview.TaskViewRepository;
 import com.android.wm.shell.taskview.TaskViewTransitions;
 import com.android.wm.shell.transition.FocusTransitionObserver;
 import com.android.wm.shell.transition.HomeTransitionObserver;
@@ -772,8 +773,15 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
-    static TaskViewTransitions provideTaskViewTransitions(Transitions transitions) {
-        return new TaskViewTransitions(transitions);
+    static TaskViewTransitions provideTaskViewTransitions(Transitions transitions,
+            TaskViewRepository repository) {
+        return new TaskViewTransitions(transitions, repository);
+    }
+
+    @WMSingleton
+    @Provides
+    static TaskViewRepository provideTaskViewRepository() {
+        return new TaskViewRepository();
     }
 
     // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewRepository.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewRepository.java
new file mode 100644
index 000000000000..8cb39a0b84de
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewRepository.java
@@ -0,0 +1,122 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.taskview;
+
+import android.annotation.Nullable;
+import android.graphics.Rect;
+import android.window.WindowContainerToken;
+
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+
+/**
+ * Keeps track of all TaskViews known by Shell. This is separated into its own object so that
+ * different TaskView managers can share state.
+ */
+public class TaskViewRepository {
+    /**
+     * The latest visibility and bounds state that has been requested for
+     * a {@link TaskViewTaskController}.
+     */
+    public static class TaskViewState {
+        final WeakReference<TaskViewTaskController> mTaskView;
+        public boolean mVisible;
+        public Rect mBounds = new Rect();
+
+        TaskViewState(TaskViewTaskController taskView) {
+            mTaskView = new WeakReference<>(taskView);
+        }
+
+        @Nullable public TaskViewTaskController getTaskView() {
+            return mTaskView.get();
+        }
+    }
+
+    /**
+     * List of tracked TaskViews
+     */
+    private final ArrayList<TaskViewState> mTaskViews = new ArrayList<>();
+
+    private int findAndPrune(TaskViewTaskController tv) {
+        for (int i = mTaskViews.size() - 1; i >= 0; --i) {
+            final TaskViewTaskController key = mTaskViews.get(i).mTaskView.get();
+            if (key == null) {
+                mTaskViews.remove(i);
+                continue;
+            }
+            if (key != tv) continue;
+            return i;
+        }
+        return -1;
+    }
+
+    /** @return if the repository is tracking {@param tv}. */
+    public boolean contains(TaskViewTaskController tv) {
+        return findAndPrune(tv) >= 0;
+    }
+
+    /** Start tracking {@param tv}. */
+    public void add(TaskViewTaskController tv) {
+        if (contains(tv)) return;
+        mTaskViews.add(new TaskViewState(tv));
+    }
+
+    /** Remove {@param tv} from tracking. */
+    public void remove(TaskViewTaskController tv) {
+        int idx = findAndPrune(tv);
+        if (idx < 0) return;
+        mTaskViews.remove(idx);
+    }
+
+    /** @return whether there are any TaskViews */
+    public boolean isEmpty() {
+        if (mTaskViews.isEmpty()) return true;
+        for (int i = mTaskViews.size() - 1; i >= 0; --i) {
+            if (mTaskViews.get(i).mTaskView.get() != null) continue;
+            mTaskViews.remove(i);
+        }
+        return mTaskViews.isEmpty();
+    }
+
+    /** @return the state of {@param tv} if tracked, {@code null} otherwise. */
+    @Nullable
+    public TaskViewState byTaskView(TaskViewTaskController tv) {
+        int idx = findAndPrune(tv);
+        if (idx < 0) return null;
+        return mTaskViews.get(idx);
+    }
+
+    /**
+     * @return the state of the taskview containing {@param token} if tracked,
+     *         {@code null} otherwise.
+     */
+    @Nullable
+    public TaskViewState byToken(WindowContainerToken token) {
+        for (int i = mTaskViews.size() - 1; i >= 0; --i) {
+            final TaskViewTaskController key = mTaskViews.get(i).mTaskView.get();
+            if (key == null) {
+                mTaskViews.remove(i);
+                continue;
+            }
+            if (key.getTaskInfo() == null) continue;
+            if (key.getTaskInfo().token.equals(token)) {
+                return mTaskViews.get(i);
+            }
+        }
+        return null;
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
index eaeedba8ef9c..8dd14983e122 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/taskview/TaskViewTaskController.java
@@ -40,6 +40,7 @@ import android.view.WindowInsets;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SyncTransactionQueue;
@@ -619,6 +620,11 @@ public class TaskViewTaskController implements ShellTaskOrganizer.TaskListener {
         return mTaskInfo;
     }
 
+    @VisibleForTesting
+    ActivityManager.RunningTaskInfo getPendingInfo() {
+        return mPendingInfo;
+    }
+
     /**
      * Indicates that the task was not found in the start animation for the transition.
      * In this case we should clean up the task if we have the pending info. If we don't
@@ -681,6 +687,10 @@ public class TaskViewTaskController implements ShellTaskOrganizer.TaskListener {
                 wct);
     }
 
+    private TaskViewRepository.TaskViewState getState() {
+        return mTaskViewTransitions.getRepository().byTaskView(this);
+    }
+
     private void prepareOpenAnimationInternal(final boolean newTask,
             SurfaceControl.Transaction startTransaction,
             SurfaceControl.Transaction finishTransaction,
@@ -703,8 +713,16 @@ public class TaskViewTaskController implements ShellTaskOrganizer.TaskListener {
                         // TODO: maybe once b/280900002 is fixed this will be unnecessary
                         .setWindowCrop(mTaskLeash, boundsOnScreen.width(), boundsOnScreen.height());
             }
-            mTaskViewTransitions.updateBoundsState(this, boundsOnScreen);
-            mTaskViewTransitions.updateVisibilityState(this, true /* visible */);
+            if (TaskViewTransitions.useRepo()) {
+                final TaskViewRepository.TaskViewState state = getState();
+                if (state != null) {
+                    state.mBounds.set(boundsOnScreen);
+                    state.mVisible = true;
+                }
+            } else {
+                mTaskViewTransitions.updateBoundsState(this, boundsOnScreen);
+                mTaskViewTransitions.updateVisibilityState(this, true /* visible */);
+            }
             wct.setBounds(mTaskToken, boundsOnScreen);
             applyCaptionInsetsIfNeeded();
         } else {
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 ae75cb589de0..0cbb7bd14e6f 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
@@ -60,7 +60,8 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
      * Only keep a weak reference to the controller instance here to allow for it to be cleaned
      * up when its TaskView is no longer used.
      */
-    private final Map<TaskViewTaskController, TaskViewRequestedState> mTaskViews;
+    private final Map<TaskViewTaskController, TaskViewRepository.TaskViewState> mTaskViews;
+    private final TaskViewRepository mTaskViewRepo;
     private final ArrayList<PendingTransition> mPending = new ArrayList<>();
     private final Transitions mTransitions;
     private final boolean[] mRegistered = new boolean[]{false};
@@ -95,26 +96,29 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
         }
     }
 
-    /**
-     * Visibility and bounds state that has been requested for a {@link TaskViewTaskController}.
-     */
-    private static class TaskViewRequestedState {
-        boolean mVisible;
-        Rect mBounds = new Rect();
-    }
-
-    public TaskViewTransitions(Transitions transitions) {
+    public TaskViewTransitions(Transitions transitions, TaskViewRepository repository) {
         mTransitions = transitions;
-        if (Flags.enableTaskViewControllerCleanup()) {
+        if (useRepo()) {
+            mTaskViews = null;
+        } else if (Flags.enableTaskViewControllerCleanup()) {
             mTaskViews = new WeakHashMap<>();
         } else {
             mTaskViews = new ArrayMap<>();
         }
+        mTaskViewRepo = repository;
         // Defer registration until the first TaskView because we want this to be the "first" in
         // priority when handling requests.
         // TODO(210041388): register here once we have an explicit ordering mechanism.
     }
 
+    static boolean useRepo() {
+        return Flags.taskViewRepository() || Flags.enableBubbleAnything();
+    }
+
+    public TaskViewRepository getRepository() {
+        return mTaskViewRepo;
+    }
+
     void addTaskView(TaskViewTaskController tv) {
         synchronized (mRegistered) {
             if (!mRegistered[0]) {
@@ -122,11 +126,19 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
                 mTransitions.addHandler(this);
             }
         }
-        mTaskViews.put(tv, new TaskViewRequestedState());
+        if (useRepo()) {
+            mTaskViewRepo.add(tv);
+        } else {
+            mTaskViews.put(tv, new TaskViewRepository.TaskViewState(null));
+        }
     }
 
     void removeTaskView(TaskViewTaskController tv) {
-        mTaskViews.remove(tv);
+        if (useRepo()) {
+            mTaskViewRepo.remove(tv);
+        } else {
+            mTaskViews.remove(tv);
+        }
         // Note: Don't unregister handler since this is a singleton with lifetime bound to Shell
     }
 
@@ -223,6 +235,10 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
     }
 
     private TaskViewTaskController findTaskView(ActivityManager.RunningTaskInfo taskInfo) {
+        if (useRepo()) {
+            final TaskViewRepository.TaskViewState state = mTaskViewRepo.byToken(taskInfo.token);
+            return state != null ? state.getTaskView() : null;
+        }
         if (Flags.enableTaskViewControllerCleanup()) {
             for (TaskViewTaskController controller : mTaskViews.keySet()) {
                 if (controller.getTaskInfo() == null) continue;
@@ -231,8 +247,8 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
                 }
             }
         } else {
-            ArrayMap<TaskViewTaskController, TaskViewRequestedState> taskViews =
-                    (ArrayMap<TaskViewTaskController, TaskViewRequestedState>) mTaskViews;
+            ArrayMap<TaskViewTaskController, TaskViewRepository.TaskViewState> taskViews =
+                    (ArrayMap<TaskViewTaskController, TaskViewRepository.TaskViewState>) mTaskViews;
             for (int i = 0; i < taskViews.size(); ++i) {
                 if (taskViews.keyAt(i).getTaskInfo() == null) continue;
                 if (taskInfo.token.equals(taskViews.keyAt(i).getTaskInfo().token)) {
@@ -279,23 +295,26 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
      *
      * @param taskView the task view which the visibility is being changed for
      * @param visible  the new visibility of the task view
-     * @param reorder  whether to reorder the task or not. If this is {@code true}, the task will be
-     *                 reordered as per the given {@code visible}. For {@code visible = true}, task
-     *                 will be reordered to top. For {@code visible = false}, task will be reordered
-     *                 to the bottom
+     * @param reorder  whether to reorder the task or not. If this is {@code true}, the task will
+     *                 be reordered as per the given {@code visible}. For {@code visible = true},
+     *                 task will be reordered to top. For {@code visible = false}, task will be
+     *                 reordered to the bottom
      */
     public void setTaskViewVisible(TaskViewTaskController taskView, boolean visible,
             boolean reorder) {
-        if (mTaskViews.get(taskView) == null) return;
-        if (mTaskViews.get(taskView).mVisible == visible) return;
+        final TaskViewRepository.TaskViewState state = useRepo()
+                ? mTaskViewRepo.byTaskView(taskView)
+                : mTaskViews.get(taskView);
+        if (state == null) return;
+        if (state.mVisible == visible) return;
         if (taskView.getTaskInfo() == null) {
             // Nothing to update, task is not yet available
             return;
         }
-        mTaskViews.get(taskView).mVisible = visible;
+        state.mVisible = visible;
         final WindowContainerTransaction wct = new WindowContainerTransaction();
         wct.setHidden(taskView.getTaskInfo().token, !visible /* hidden */);
-        wct.setBounds(taskView.getTaskInfo().token, mTaskViews.get(taskView).mBounds);
+        wct.setBounds(taskView.getTaskInfo().token, state.mBounds);
         if (reorder) {
             wct.reorder(taskView.getTaskInfo().token, visible /* onTop */);
         }
@@ -308,7 +327,10 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
 
     /** Starts a new transition to reorder the given {@code taskView}'s task. */
     public void reorderTaskViewTask(TaskViewTaskController taskView, boolean onTop) {
-        if (mTaskViews.get(taskView) == null) return;
+        final TaskViewRepository.TaskViewState state = useRepo()
+                ? mTaskViewRepo.byTaskView(taskView)
+                : mTaskViews.get(taskView);
+        if (state == null) return;
         if (taskView.getTaskInfo() == null) {
             // Nothing to update, task is not yet available
             return;
@@ -323,19 +345,24 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
     }
 
     void updateBoundsState(TaskViewTaskController taskView, Rect boundsOnScreen) {
-        TaskViewRequestedState state = mTaskViews.get(taskView);
+        if (useRepo()) return;
+        final TaskViewRepository.TaskViewState state = mTaskViews.get(taskView);
         if (state == null) return;
         state.mBounds.set(boundsOnScreen);
     }
 
     void updateVisibilityState(TaskViewTaskController taskView, boolean visible) {
-        TaskViewRequestedState state = mTaskViews.get(taskView);
+        final TaskViewRepository.TaskViewState state = useRepo()
+                ? mTaskViewRepo.byTaskView(taskView)
+                : mTaskViews.get(taskView);
         if (state == null) return;
         state.mVisible = visible;
     }
 
     void setTaskBounds(TaskViewTaskController taskView, Rect boundsOnScreen) {
-        TaskViewRequestedState state = mTaskViews.get(taskView);
+        final TaskViewRepository.TaskViewState state = useRepo()
+                ? mTaskViewRepo.byTaskView(taskView)
+                : mTaskViews.get(taskView);
         if (state == null || Objects.equals(boundsOnScreen, state.mBounds)) {
             return;
         }
@@ -385,7 +412,7 @@ public class TaskViewTransitions implements Transitions.TransitionHandler {
         if (pending != null) {
             mPending.remove(pending);
         }
-        if (mTaskViews.isEmpty()) {
+        if (useRepo() ? mTaskViewRepo.isEmpty() : mTaskViews.isEmpty()) {
             if (pending != null) {
                 Slog.e(TAG, "Pending taskview transition but no task-views");
             }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
index 17fd95b69dba..66636c56bb64 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/taskview/TaskViewTest.java
@@ -20,6 +20,11 @@ import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeFalse;
 import static org.junit.Assume.assumeTrue;
@@ -45,7 +50,8 @@ import android.graphics.Insets;
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.os.Looper;
-import android.testing.AndroidTestingRunner;
+import android.platform.test.flag.junit.FlagsParameterization;
+import android.platform.test.flag.junit.SetFlagsRule;
 import android.testing.TestableLooper;
 import android.view.SurfaceControl;
 import android.view.SurfaceHolder;
@@ -55,6 +61,7 @@ import android.window.WindowContainerTransaction;
 
 import androidx.test.filters.SmallTest;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestHandler;
@@ -65,17 +72,33 @@ import com.android.wm.shell.transition.Transitions;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.invocation.InvocationOnMock;
 
+import java.util.List;
+
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4;
+import platform.test.runner.parameterized.Parameters;
+
 @SmallTest
-@RunWith(AndroidTestingRunner.class)
+@RunWith(ParameterizedAndroidJunit4.class)
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 public class TaskViewTest extends ShellTestCase {
 
+    @Parameters(name = "{0}")
+    public static List<FlagsParameterization> getParams() {
+        return FlagsParameterization.allCombinationsOf(Flags.FLAG_TASK_VIEW_REPOSITORY);
+    }
+
+    @Rule
+    public final SetFlagsRule mSetFlagsRule;
+
     @Mock
     TaskView.Listener mViewListener;
     @Mock
@@ -84,6 +107,8 @@ public class TaskViewTest extends ShellTestCase {
     WindowContainerToken mToken;
     @Mock
     ShellTaskOrganizer mOrganizer;
+    @Captor
+    ArgumentCaptor<WindowContainerTransaction> mWctCaptor;
     @Mock
     HandlerExecutor mExecutor;
     @Mock
@@ -98,9 +123,14 @@ public class TaskViewTest extends ShellTestCase {
 
     Context mContext;
     TaskView mTaskView;
+    TaskViewRepository mTaskViewRepository;
     TaskViewTransitions mTaskViewTransitions;
     TaskViewTaskController mTaskViewTaskController;
 
+    public TaskViewTest(FlagsParameterization flags) {
+        mSetFlagsRule = new SetFlagsRule(flags);
+    }
+
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
@@ -134,9 +164,10 @@ public class TaskViewTest extends ShellTestCase {
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             doReturn(true).when(mTransitions).isRegistered();
         }
-        mTaskViewTransitions = spy(new TaskViewTransitions(mTransitions));
-        mTaskViewTaskController = spy(new TaskViewTaskController(mContext, mOrganizer,
-                mTaskViewTransitions, mSyncQueue));
+        mTaskViewRepository = new TaskViewRepository();
+        mTaskViewTransitions = spy(new TaskViewTransitions(mTransitions, mTaskViewRepository));
+        mTaskViewTaskController = new TaskViewTaskController(mContext, mOrganizer,
+                mTaskViewTransitions, mSyncQueue);
         mTaskView = new TaskView(mContext, mTaskViewTaskController);
         mTaskView.setHandler(mViewHandler);
         mTaskView.setListener(mExecutor, mViewListener);
@@ -482,8 +513,14 @@ public class TaskViewTest extends ShellTestCase {
 
         // Surface created, but task not available so bounds / visibility isn't set
         mTaskView.surfaceCreated(mock(SurfaceHolder.class));
-        verify(mTaskViewTransitions, never()).updateVisibilityState(
-                eq(mTaskViewTaskController), eq(true));
+        if (TaskViewTransitions.useRepo()) {
+            assertNotNull(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController));
+            assertFalse(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mVisible);
+        } else {
+            verify(mTaskViewTransitions, never()).updateVisibilityState(
+                    eq(mTaskViewTaskController), eq(true));
+        }
 
         // Make the task available
         WindowContainerTransaction wct = mock(WindowContainerTransaction.class);
@@ -492,8 +529,16 @@ public class TaskViewTest extends ShellTestCase {
         // Bounds got set
         verify(wct).setBounds(any(WindowContainerToken.class), eq(bounds));
         // Visibility & bounds state got set
-        verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController), eq(true));
-        verify(mTaskViewTransitions).updateBoundsState(eq(mTaskViewTaskController), eq(bounds));
+        if (TaskViewTransitions.useRepo()) {
+            assertTrue(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mVisible);
+            assertEquals(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mBounds, bounds);
+        } else {
+            verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController),
+                    eq(true));
+            verify(mTaskViewTransitions).updateBoundsState(eq(mTaskViewTaskController), eq(bounds));
+        }
     }
 
     @Test
@@ -507,8 +552,15 @@ public class TaskViewTest extends ShellTestCase {
 
         // Surface created, but task not available so bounds / visibility isn't set
         mTaskView.surfaceCreated(mock(SurfaceHolder.class));
-        verify(mTaskViewTransitions, never()).updateVisibilityState(
-                eq(mTaskViewTaskController), eq(true));
+        if (TaskViewTransitions.useRepo()) {
+            assertNotNull(mTaskViewTransitions.getRepository().byTaskView(
+                    mTaskViewTaskController));
+            assertFalse(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mVisible);
+        } else {
+            verify(mTaskViewTransitions, never()).updateVisibilityState(
+                    eq(mTaskViewTaskController), eq(true));
+        }
 
         // Make the task available / start prepareOpen
         WindowContainerTransaction wct = mock(WindowContainerTransaction.class);
@@ -519,8 +571,16 @@ public class TaskViewTest extends ShellTestCase {
         // Bounds got set
         verify(wct).setBounds(any(WindowContainerToken.class), eq(bounds));
         // Visibility & bounds state got set
-        verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController), eq(true));
-        verify(mTaskViewTransitions).updateBoundsState(eq(mTaskViewTaskController), eq(bounds));
+        if (TaskViewTransitions.useRepo()) {
+            assertTrue(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mVisible);
+            assertEquals(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mBounds, bounds);
+        } else {
+            verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController),
+                    eq(true));
+            verify(mTaskViewTransitions).updateBoundsState(eq(mTaskViewTaskController), eq(bounds));
+        }
     }
 
     @Test
@@ -541,8 +601,17 @@ public class TaskViewTest extends ShellTestCase {
         // Bounds do not get set as there is no surface
         verify(wct, never()).setBounds(any(WindowContainerToken.class), any());
         // Visibility is set to false, bounds aren't set
-        verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController), eq(false));
-        verify(mTaskViewTransitions, never()).updateBoundsState(eq(mTaskViewTaskController), any());
+        if (TaskViewTransitions.useRepo()) {
+            assertFalse(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mVisible);
+            assertTrue(mTaskViewTransitions.getRepository().byTaskView(mTaskViewTaskController)
+                    .mBounds.isEmpty());
+        } else {
+            verify(mTaskViewTransitions).updateVisibilityState(eq(mTaskViewTaskController),
+                    eq(false));
+            verify(mTaskViewTransitions, never()).updateBoundsState(eq(mTaskViewTaskController),
+                    any());
+        }
     }
 
     @Test
@@ -576,7 +645,7 @@ public class TaskViewTest extends ShellTestCase {
         mTaskViewTaskController.setTaskNotFound();
         mTaskViewTaskController.onTaskAppeared(mTaskInfo, mLeash);
 
-        verify(mTaskViewTaskController).cleanUpPendingTask();
+        assertNull(mTaskViewTaskController.getTaskInfo());
         verify(mTaskViewTransitions).closeTaskView(any(), eq(mTaskViewTaskController));
     }
 
@@ -585,7 +654,8 @@ public class TaskViewTest extends ShellTestCase {
         assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
 
         mTaskViewTaskController.onTaskAppeared(mTaskInfo, mLeash);
-        verify(mTaskViewTaskController, never()).cleanUpPendingTask();
+        assertEquals(mTaskInfo, mTaskViewTaskController.getPendingInfo());
+        verify(mTaskViewTransitions, never()).closeTaskView(any(), any());
     }
 
     @Test
@@ -596,20 +666,20 @@ public class TaskViewTest extends ShellTestCase {
         mTaskView.setCaptionInsets(Insets.of(insets));
         mTaskView.onComputeInternalInsets(new ViewTreeObserver.InternalInsetsInfo());
 
-        verify(mTaskViewTaskController).applyCaptionInsetsIfNeeded();
         verify(mOrganizer, never()).applyTransaction(any());
 
         mTaskView.surfaceCreated(mock(SurfaceHolder.class));
         reset(mOrganizer);
-        reset(mTaskViewTaskController);
         WindowContainerTransaction wct = new WindowContainerTransaction();
         mTaskViewTaskController.prepareOpenAnimation(true /* newTask */,
                 new SurfaceControl.Transaction(), new SurfaceControl.Transaction(), mTaskInfo,
                 mLeash, wct);
         mTaskView.onComputeInternalInsets(new ViewTreeObserver.InternalInsetsInfo());
 
-        verify(mTaskViewTaskController).applyCaptionInsetsIfNeeded();
-        verify(mOrganizer).applyTransaction(any());
+        verify(mOrganizer).applyTransaction(mWctCaptor.capture());
+        assertTrue(mWctCaptor.getValue().getHierarchyOps().stream().anyMatch(hop ->
+                hop.getType() == WindowContainerTransaction.HierarchyOp
+                        .HIERARCHY_OP_TYPE_ADD_INSETS_FRAME_PROVIDER));
     }
 
     @Test
@@ -621,14 +691,15 @@ public class TaskViewTest extends ShellTestCase {
         mTaskViewTaskController.prepareOpenAnimation(true /* newTask */,
                 new SurfaceControl.Transaction(), new SurfaceControl.Transaction(), mTaskInfo,
                 mLeash, wct);
-        reset(mTaskViewTaskController);
         reset(mOrganizer);
 
         Rect insets = new Rect(0, 400, 0, 0);
         mTaskView.setCaptionInsets(Insets.of(insets));
         mTaskView.onComputeInternalInsets(new ViewTreeObserver.InternalInsetsInfo());
-        verify(mTaskViewTaskController).applyCaptionInsetsIfNeeded();
-        verify(mOrganizer).applyTransaction(any());
+        verify(mOrganizer).applyTransaction(mWctCaptor.capture());
+        assertTrue(mWctCaptor.getValue().getHierarchyOps().stream().anyMatch(hop ->
+                hop.getType() == WindowContainerTransaction.HierarchyOp
+                        .HIERARCHY_OP_TYPE_ADD_INSETS_FRAME_PROVIDER));
     }
 
     @Test
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 60e203009e8c..5f6f18f82fb4 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
@@ -83,6 +83,7 @@ public class TaskViewTransitionsTest extends ShellTestCase {
     @Mock
     WindowContainerToken mToken;
 
+    TaskViewRepository mTaskViewRepository;
     TaskViewTransitions mTaskViewTransitions;
 
     public TaskViewTransitionsTest(FlagsParameterization flags) {
@@ -102,7 +103,8 @@ public class TaskViewTransitionsTest extends ShellTestCase {
         mTaskInfo.taskId = 314;
         mTaskInfo.taskDescription = mock(ActivityManager.TaskDescription.class);
 
-        mTaskViewTransitions = spy(new TaskViewTransitions(mTransitions));
+        mTaskViewRepository = new TaskViewRepository();
+        mTaskViewTransitions = spy(new TaskViewTransitions(mTransitions, mTaskViewRepository));
         mTaskViewTransitions.addTaskView(mTaskViewTaskController);
         when(mTaskViewTaskController.getTaskInfo()).thenReturn(mTaskInfo);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
index af12d0119c58..7508838810ad 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -198,6 +198,7 @@ import com.android.wm.shell.sysui.ShellCommandHandler;
 import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.taskview.TaskView;
+import com.android.wm.shell.taskview.TaskViewRepository;
 import com.android.wm.shell.taskview.TaskViewTransitions;
 import com.android.wm.shell.transition.Transitions;
 
@@ -362,6 +363,7 @@ public class BubblesTest extends SysuiTestCase {
     private ShadeInteractor mShadeInteractor;
     private NotificationShadeWindowModel mNotificationShadeWindowModel;
     private ShellTaskOrganizer mShellTaskOrganizer;
+    private TaskViewRepository mTaskViewRepository;
     private TaskViewTransitions mTaskViewTransitions;
 
     private TestableBubblePositioner mPositioner;
@@ -398,7 +400,8 @@ public class BubblesTest extends SysuiTestCase {
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             doReturn(true).when(mTransitions).isRegistered();
         }
-        mTaskViewTransitions = new TaskViewTransitions(mTransitions);
+        mTaskViewRepository = new TaskViewRepository();
+        mTaskViewTransitions = new TaskViewTransitions(mTransitions, mTaskViewRepository);
 
         mTestableLooper = TestableLooper.get(this);