diff --git a/libs/WindowManager/Shell/multivalentTests/Android.bp b/libs/WindowManager/Shell/multivalentTests/Android.bp
index ee0d5bb..41d1b5c 100644
--- a/libs/WindowManager/Shell/multivalentTests/Android.bp
+++ b/libs/WindowManager/Shell/multivalentTests/Android.bp
@@ -53,6 +53,8 @@
         "mockito-robolectric-prebuilt",
         "mockito-kotlin2",
         "truth",
+        "flag-junit-base",
+        "flag-junit",
     ],
     auto_gen_config: true,
 }
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
index 398fd55..4214e6f 100644
--- a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleTaskViewTest.kt
@@ -18,14 +18,19 @@
 
 import android.content.ComponentName
 import android.content.Context
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.wm.shell.Flags
 import com.android.wm.shell.taskview.TaskView
 
 import com.google.common.truth.Truth.assertThat
 import com.google.common.util.concurrent.MoreExecutors.directExecutor
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.kotlin.mock
@@ -36,6 +41,9 @@
 @RunWith(AndroidJUnit4::class)
 class BubbleTaskViewTest {
 
+    @get:Rule
+    val setFlagsRule = SetFlagsRule()
+
     private lateinit var bubbleTaskView: BubbleTaskView
     private val context = ApplicationProvider.getApplicationContext<Context>()
     private lateinit var taskView: TaskView
@@ -75,14 +83,33 @@
         assertThat(actualComponentName).isEqualTo(componentName)
     }
 
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
     @Test
-    fun cleanup_invalidTaskId_doesNotRemoveTask() {
+    fun cleanup_flagOff_invalidTaskId_doesNotRemoveTask() {
         bubbleTaskView.cleanup()
         verify(taskView, never()).removeTask()
     }
 
+    @EnableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
     @Test
-    fun cleanup_validTaskId_removesTask() {
+    fun cleanup_flagOn_invalidTaskId_removesTask() {
+        bubbleTaskView.cleanup()
+        verify(taskView).removeTask()
+    }
+
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
+    @Test
+    fun cleanup_flagOff_validTaskId_removesTask() {
+        val componentName = ComponentName(context, "TestClass")
+        bubbleTaskView.listener.onTaskCreated(123, componentName)
+
+        bubbleTaskView.cleanup()
+        verify(taskView).removeTask()
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_TASK_VIEW_CONTROLLER_CLEANUP)
+    @Test
+    fun cleanup_flagOn_validTaskId_removesTask() {
         val componentName = ComponentName(context, "TestClass")
         bubbleTaskView.listener.onTaskCreated(123, componentName)
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
index 65f8e48..68fc0c9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleTaskView.kt
@@ -16,14 +16,11 @@
 
 package com.android.wm.shell.bubbles
 
-import android.app.ActivityTaskManager
 import android.app.ActivityTaskManager.INVALID_TASK_ID
 import android.content.ComponentName
-import android.os.RemoteException
-import android.util.Log
 import androidx.annotation.VisibleForTesting
+import com.android.wm.shell.Flags
 import com.android.wm.shell.taskview.TaskView
-import com.android.wm.shell.transition.Transitions.ENABLE_SHELL_TRANSITIONS
 import java.util.concurrent.Executor
 
 /**
@@ -89,21 +86,8 @@
      * This should be called after all other cleanup animations have finished.
      */
     fun cleanup() {
-        if (taskId != INVALID_TASK_ID) {
-            // Ensure the task is removed from WM
-            if (ENABLE_SHELL_TRANSITIONS) {
-                taskView.removeTask()
-            } else {
-                try {
-                    ActivityTaskManager.getService().removeTask(taskId)
-                } catch (e: RemoteException) {
-                    Log.w(TAG, e.message ?: "")
-                }
-            }
+        if (Flags.enableTaskViewControllerCleanup() || taskId != INVALID_TASK_ID) {
+            taskView.removeTask()
         }
     }
-
-    private companion object {
-        const val TAG = "BubbleTaskView"
-    }
 }
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 e74342e..eaeedba 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.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SyncTransactionQueue;
 
@@ -525,8 +526,13 @@
      */
     void removeTask() {
         if (mTaskToken == null) {
-            // Call to remove task before we have one, do nothing
-            Slog.w(TAG, "Trying to remove a task that was never added? (no taskToken)");
+            if (Flags.enableTaskViewControllerCleanup()) {
+                // We don't have a task yet. Only clean up the controller
+                mTaskViewTransitions.removeTaskView(this);
+            } else {
+                // Call to remove task before we have one, do nothing
+                Slog.w(TAG, "Trying to remove a task that was never added? (no taskToken)");
+            }
             return;
         }
         mShellExecutor.execute(() -> {