3/ Fix inverted ShellInit dependency order (no functional changes)

- We were previously initializing components manually because some startup
  is order-dependent, which led to an increasing number of dependencies
  from ShellTaskOrganizer on things that it really doesn't need to know
  about.  The initialization order actually matches the dependency order
  pretty closely, so instead we can have these components add callbacks
  in the order they are created and initalize in sequence.
- Also move some logic out of constructors to the init callback
- The other subtle issue is that the provider for ShellInitImpl was
  implicitly also ensuring that certain shell components were created at
  startup (since they aren't otherwise referenced by any SysUI component
  directly or indirectly).  Instead, separate that out into another
  more explicit provider, with the @ShellCreateTrigger annotation, and
  allow subclasses to provide a @ShellCreateTriggerOverride annotated
  provider for product-specific initialization

Test: atest WMShellUnitTests
Test: atest SystemUITests
Bug: 238217847

Change-Id: I4a2583ceeaf469130ac07fc8c4e267c28843b7b8
Signed-off-by: Winson Chung <winsonc@google.com>

49 files changed, 1035 insertions, 453 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
index 4080b99d30cf..5cba3b4ec130 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/ShellTaskOrganizer.java
@@ -31,7 +31,6 @@ import android.annotation.Nullable;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
 import android.app.WindowConfiguration;
-import android.content.Context;
 import android.content.LocusId;
 import android.content.pm.ActivityInfo;
 import android.graphics.Rect;
@@ -56,6 +55,7 @@ import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.compatui.CompatUIController;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.startingsurface.StartingWindowController;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.unfold.UnfoldAnimationController;
 
 import java.io.PrintWriter;
@@ -186,39 +186,43 @@ public class ShellTaskOrganizer extends TaskOrganizer implements
     @Nullable
     private RunningTaskInfo mLastFocusedTaskInfo;
 
-    public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context) {
-        this(null /* taskOrganizerController */, mainExecutor, context, null /* compatUI */,
+    public ShellTaskOrganizer(ShellExecutor mainExecutor) {
+        this(null /* shellInit */, null /* taskOrganizerController */, null /* compatUI */,
                 Optional.empty() /* unfoldAnimationController */,
-                Optional.empty() /* recentTasksController */);
+                Optional.empty() /* recentTasksController */,
+                mainExecutor);
     }
 
-    public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context, @Nullable
-            CompatUIController compatUI) {
-        this(null /* taskOrganizerController */, mainExecutor, context, compatUI,
-                Optional.empty() /* unfoldAnimationController */,
-                Optional.empty() /* recentTasksController */);
-    }
-
-    public ShellTaskOrganizer(ShellExecutor mainExecutor, Context context, @Nullable
-            CompatUIController compatUI,
+    public ShellTaskOrganizer(ShellInit shellInit,
+            @Nullable CompatUIController compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<RecentTasksController> recentTasks) {
-        this(null /* taskOrganizerController */, mainExecutor, context, compatUI,
-                unfoldAnimationController, recentTasks);
+            Optional<RecentTasksController> recentTasks,
+            ShellExecutor mainExecutor) {
+        this(shellInit, null /* taskOrganizerController */, compatUI,
+                unfoldAnimationController, recentTasks, mainExecutor);
     }
 
     @VisibleForTesting
-    protected ShellTaskOrganizer(ITaskOrganizerController taskOrganizerController,
-            ShellExecutor mainExecutor, Context context, @Nullable CompatUIController compatUI,
+    protected ShellTaskOrganizer(ShellInit shellInit,
+            ITaskOrganizerController taskOrganizerController,
+            @Nullable CompatUIController compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<RecentTasksController> recentTasks) {
+            Optional<RecentTasksController> recentTasks,
+            ShellExecutor mainExecutor) {
         super(taskOrganizerController, mainExecutor);
         mCompatUI = compatUI;
         mRecentTasks = recentTasks;
         mUnfoldAnimationController = unfoldAnimationController.orElse(null);
-        if (compatUI != null) {
-            compatUI.setCompatUICallback(this);
+        if (shellInit != null) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
+    }
+
+    private void onInit() {
+        if (mCompatUI != null) {
+            mCompatUI.setCompatUICallback(this);
         }
+        registerOrganizer();
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingController.java
index 82b0270698e4..b305897b77ae 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/activityembedding/ActivityEmbeddingController.java
@@ -28,6 +28,7 @@ import android.window.WindowContainerTransaction;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 
 /**
@@ -38,13 +39,17 @@ public class ActivityEmbeddingController implements Transitions.TransitionHandle
     private final Context mContext;
     private final Transitions mTransitions;
 
-    public ActivityEmbeddingController(Context context, Transitions transitions) {
+    public ActivityEmbeddingController(Context context, ShellInit shellInit,
+            Transitions transitions) {
         mContext = context;
         mTransitions = transitions;
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
     }
 
     /** Registers to handle transitions. */
-    public void init() {
+    public void onInit() {
         mTransitions.addHandler(this);
     }
 
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 1c2f0d8ca7b2..398261600dc7 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
@@ -101,6 +101,7 @@ import com.android.wm.shell.onehanded.OneHandedTransitionCallback;
 import com.android.wm.shell.pip.PinnedStackListenerForwarder;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -227,6 +228,7 @@ public class BubbleController implements ConfigurationChangeListener {
 
   
     public BubbleController(Context context,
+            ShellInit shellInit,
             ShellController shellController,
             BubbleData data,
             @Nullable BubbleStackView.SurfaceSynchronizer synchronizer,
@@ -279,6 +281,7 @@ public class BubbleController implements ConfigurationChangeListener {
         mOneHandedOptional = oneHandedOptional;
         mDragAndDropController = dragAndDropController;
         mSyncQueue = syncQueue;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     private void registerOneHandedState(OneHandedController oneHanded) {
@@ -300,7 +303,7 @@ public class BubbleController implements ConfigurationChangeListener {
                 });
     }
 
-    public void initialize() {
+    protected void onInit() {
         mBubbleData.setListener(mBubbleDataListener);
         mBubbleData.setSuppressionChangedListener(this::onBubbleMetadataFlagChanged);
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayChangeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayChangeController.java
index 28c7367662a2..ae1f43320c8b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayChangeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayChangeController.java
@@ -28,6 +28,7 @@ import android.window.WindowContainerTransaction;
 import androidx.annotation.BinderThread;
 
 import com.android.wm.shell.common.annotations.ShellMainThread;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.concurrent.CopyOnWriteArrayList;
 
@@ -47,10 +48,15 @@ public class DisplayChangeController {
     private final CopyOnWriteArrayList<OnDisplayChangingListener> mDisplayChangeListener =
             new CopyOnWriteArrayList<>();
 
-    public DisplayChangeController(IWindowManager wmService, ShellExecutor mainExecutor) {
+    public DisplayChangeController(IWindowManager wmService, ShellInit shellInit,
+            ShellExecutor mainExecutor) {
         mMainExecutor = mainExecutor;
         mWmService = wmService;
         mControllerImpl = new DisplayChangeWindowControllerImpl();
+        shellInit.addInitCallback(this::onInit, this);
+    }
+
+    private void onInit() {
         try {
             mWmService.setDisplayChangeWindowController(mControllerImpl);
         } catch (RemoteException e) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayController.java
index 764936cceb01..f07ea751b044 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayController.java
@@ -34,6 +34,7 @@ import androidx.annotation.BinderThread;
 
 import com.android.wm.shell.common.DisplayChangeController.OnDisplayChangingListener;
 import com.android.wm.shell.common.annotations.ShellMainThread;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -57,19 +58,23 @@ public class DisplayController {
     private final SparseArray<DisplayRecord> mDisplays = new SparseArray<>();
     private final ArrayList<OnDisplaysChangedListener> mDisplayChangedListeners = new ArrayList<>();
 
-    public DisplayController(Context context, IWindowManager wmService,
+    public DisplayController(Context context, IWindowManager wmService, ShellInit shellInit,
             ShellExecutor mainExecutor) {
         mMainExecutor = mainExecutor;
         mContext = context;
         mWmService = wmService;
-        mChangeController = new DisplayChangeController(mWmService, mainExecutor);
+        // TODO: Inject this instead
+        mChangeController = new DisplayChangeController(mWmService, shellInit, mainExecutor);
         mDisplayContainerListener = new DisplayWindowListenerImpl();
+        // Note, add this after DisplaceChangeController is constructed to ensure that is
+        // initialized first
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     /**
      * Initializes the window listener.
      */
-    public void initialize() {
+    public void onInit() {
         try {
             int[] displayIds = mWmService.registerDisplayWindowListener(mDisplayContainerListener);
             for (int i = 0; i < displayIds.length; i++) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
index b3f62477077c..266cf294a950 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayImeController.java
@@ -44,6 +44,7 @@ import android.view.animation.PathInterpolator;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.internal.view.IInputMethodManager;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.ArrayList;
 import java.util.concurrent.Executor;
@@ -74,18 +75,24 @@ public class DisplayImeController implements DisplayController.OnDisplaysChanged
     private final ArrayList<ImePositionProcessor> mPositionProcessors = new ArrayList<>();
 
 
-    public DisplayImeController(IWindowManager wmService, DisplayController displayController,
+    public DisplayImeController(IWindowManager wmService,
+            ShellInit shellInit,
+            DisplayController displayController,
             DisplayInsetsController displayInsetsController,
-            Executor mainExecutor, TransactionPool transactionPool) {
+            TransactionPool transactionPool,
+            Executor mainExecutor) {
         mWmService = wmService;
         mDisplayController = displayController;
         mDisplayInsetsController = displayInsetsController;
         mMainExecutor = mainExecutor;
         mTransactionPool = transactionPool;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
-    /** Starts monitor displays changes and set insets controller for each displays. */
-    public void startMonitorDisplays() {
+    /**
+     * Starts monitor displays changes and set insets controller for each displays.
+     */
+    public void onInit() {
         mDisplayController.addDisplayWindowListener(this);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayInsetsController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayInsetsController.java
index f546f110e87c..90a01f8c5295 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayInsetsController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/DisplayInsetsController.java
@@ -29,6 +29,7 @@ import android.view.InsetsVisibilities;
 import androidx.annotation.BinderThread;
 
 import com.android.wm.shell.common.annotations.ShellMainThread;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.concurrent.CopyOnWriteArrayList;
 
@@ -45,17 +46,20 @@ public class DisplayInsetsController implements DisplayController.OnDisplaysChan
     private final SparseArray<CopyOnWriteArrayList<OnInsetsChangedListener>> mListeners =
             new SparseArray<>();
 
-    public DisplayInsetsController(IWindowManager wmService, DisplayController displayController,
+    public DisplayInsetsController(IWindowManager wmService,
+            ShellInit shellInit,
+            DisplayController displayController,
             ShellExecutor mainExecutor) {
         mWmService = wmService;
         mDisplayController = displayController;
         mMainExecutor = mainExecutor;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     /**
      * Starts listening for insets for each display.
      **/
-    public void initialize() {
+    public void onInit() {
         mDisplayController.addDisplayWindowListener(this);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/DynamicOverride.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/DynamicOverride.java
index 806f795d1015..10b121bbc32c 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/DynamicOverride.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/DynamicOverride.java
@@ -92,6 +92,8 @@ import javax.inject.Qualifier;
  *
  * For example, this uses the same setup as above, but the interface provided (if bound) is used
  * otherwise the default is created:
+ *
+ * BaseModule:
  *   @BindsOptionalOf
  *   @DynamicOverride
  *   abstract Interface dynamicInterface();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTrigger.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTrigger.java
new file mode 100644
index 000000000000..482b19983850
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTrigger.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.dagger;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+import javax.inject.Qualifier;
+
+/**
+ * An annotation to specifically mark the provider that is triggering the creation of independent
+ * shell components that are not created as a part of the dependencies for interfaces passed to
+ * SysUI.
+ *
+ * TODO: This will be removed once we have a more explicit method for specifying components to start
+ *       with SysUI
+ */
+@Documented
+@Inherited
+@Qualifier
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ShellCreateTrigger {}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTriggerOverride.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTriggerOverride.java
new file mode 100644
index 000000000000..31c678968a25
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/ShellCreateTriggerOverride.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.dagger;
+
+import java.lang.annotation.Documented;
+import java.lang.annotation.Inherited;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+
+import javax.inject.Qualifier;
+
+/**
+ * An annotation for non-base modules to specifically mark the provider that is triggering the
+ * creation of independent shell components that are not created as a part of the dependencies for
+ * interfaces passed to SysUI.
+ *
+ * TODO: This will be removed once we have a more explicit method for specifying components to start
+ *       with SysUI
+ */
+@Documented
+@Inherited
+@Qualifier
+@Retention(RetentionPolicy.RUNTIME)
+public @interface ShellCreateTriggerOverride {}
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 a168cb22f6ad..3add4179ef82 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
@@ -112,16 +112,20 @@ public abstract class WMShellBaseModule {
     @WMSingleton
     @Provides
     static DisplayController provideDisplayController(Context context,
-            IWindowManager wmService, @ShellMainThread ShellExecutor mainExecutor) {
-        return new DisplayController(context, wmService, mainExecutor);
+            IWindowManager wmService,
+            ShellInit shellInit,
+            @ShellMainThread ShellExecutor mainExecutor) {
+        return new DisplayController(context, wmService, shellInit, mainExecutor);
     }
 
     @WMSingleton
     @Provides
-    static DisplayInsetsController provideDisplayInsetsController( IWindowManager wmService,
+    static DisplayInsetsController provideDisplayInsetsController(IWindowManager wmService,
+            ShellInit shellInit,
             DisplayController displayController,
             @ShellMainThread ShellExecutor mainExecutor) {
-        return new DisplayInsetsController(wmService, displayController, mainExecutor);
+        return new DisplayInsetsController(wmService, shellInit, displayController,
+                mainExecutor);
     }
 
     // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
@@ -134,16 +138,17 @@ public abstract class WMShellBaseModule {
     static DisplayImeController provideDisplayImeController(
             @DynamicOverride Optional<DisplayImeController> overrideDisplayImeController,
             IWindowManager wmService,
+            ShellInit shellInit,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
-            @ShellMainThread ShellExecutor mainExecutor,
-            TransactionPool transactionPool
+            TransactionPool transactionPool,
+            @ShellMainThread ShellExecutor mainExecutor
     ) {
         if (overrideDisplayImeController.isPresent()) {
             return overrideDisplayImeController.get();
         }
-        return new DisplayImeController(wmService, displayController, displayInsetsController,
-                mainExecutor, transactionPool);
+        return new DisplayImeController(wmService, shellInit, displayController,
+                displayInsetsController, transactionPool, mainExecutor);
     }
 
     @WMSingleton
@@ -155,42 +160,45 @@ public abstract class WMShellBaseModule {
     @WMSingleton
     @Provides
     static DragAndDropController provideDragAndDropController(Context context,
+            ShellInit shellInit,
             ShellController shellController,
             DisplayController displayController,
             UiEventLogger uiEventLogger,
             IconProvider iconProvider,
             @ShellMainThread ShellExecutor mainExecutor) {
-        return new DragAndDropController(context, shellController, displayController, uiEventLogger,
-                iconProvider, mainExecutor);
+        return new DragAndDropController(context, shellInit, shellController, displayController,
+                uiEventLogger, iconProvider, mainExecutor);
     }
 
     @WMSingleton
     @Provides
-    static ShellTaskOrganizer provideShellTaskOrganizer(@ShellMainThread ShellExecutor mainExecutor,
-            Context context,
+    static ShellTaskOrganizer provideShellTaskOrganizer(
+            ShellInit shellInit,
             CompatUIController compatUI,
             Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<RecentTasksController> recentTasksOptional
+            Optional<RecentTasksController> recentTasksOptional,
+            @ShellMainThread ShellExecutor mainExecutor
     ) {
-        return new ShellTaskOrganizer(mainExecutor, context, compatUI, unfoldAnimationController,
-                recentTasksOptional);
+        return new ShellTaskOrganizer(shellInit, compatUI, unfoldAnimationController,
+                recentTasksOptional, mainExecutor);
     }
 
     @WMSingleton
     @Provides
     static KidsModeTaskOrganizer provideKidsModeTaskOrganizer(
-            @ShellMainThread ShellExecutor mainExecutor,
-            @ShellMainThread Handler mainHandler,
             Context context,
+            ShellInit shellInit,
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
             Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<RecentTasksController> recentTasksOptional
+            Optional<RecentTasksController> recentTasksOptional,
+            @ShellMainThread ShellExecutor mainExecutor,
+            @ShellMainThread Handler mainHandler
     ) {
-        return new KidsModeTaskOrganizer(mainExecutor, mainHandler, context, syncTransactionQueue,
+        return new KidsModeTaskOrganizer(context, shellInit, syncTransactionQueue,
                 displayController, displayInsetsController, unfoldAnimationController,
-                recentTasksOptional);
+                recentTasksOptional, mainExecutor, mainHandler);
     }
 
     @WMSingleton
@@ -256,6 +264,20 @@ public abstract class WMShellBaseModule {
         return backAnimationController.map(BackAnimationController::getBackAnimationImpl);
     }
 
+    @WMSingleton
+    @Provides
+    static Optional<BackAnimationController> provideBackAnimationController(
+            Context context,
+            @ShellMainThread ShellExecutor shellExecutor,
+            @ShellBackgroundThread Handler backgroundHandler
+    ) {
+        if (BackAnimationController.IS_ENABLED) {
+            return Optional.of(
+                    new BackAnimationController(shellExecutor, backgroundHandler, context));
+        }
+        return Optional.empty();
+    }
+
     //
     // Bubbles (optional feature)
     //
@@ -282,12 +304,15 @@ public abstract class WMShellBaseModule {
     @Provides
     static FullscreenTaskListener provideFullscreenTaskListener(
             @DynamicOverride Optional<FullscreenTaskListener> fullscreenTaskListener,
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
             SyncTransactionQueue syncQueue,
             Optional<RecentTasksController> recentTasksOptional) {
         if (fullscreenTaskListener.isPresent()) {
             return fullscreenTaskListener.get();
         } else {
-            return new FullscreenTaskListener(syncQueue, recentTasksOptional);
+            return new FullscreenTaskListener(shellInit, shellTaskOrganizer, syncQueue,
+                    recentTasksOptional);
         }
     }
 
@@ -343,7 +368,7 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
-    static Optional<FreeformComponents> provideFreeformTaskListener(
+    static Optional<FreeformComponents> provideFreeformComponents(
             @DynamicOverride Optional<FreeformComponents> freeformComponents,
             Context context) {
         if (FreeformComponents.isFreeformEnabled(context)) {
@@ -440,11 +465,13 @@ public abstract class WMShellBaseModule {
     @Provides
     static Optional<RecentTasksController> provideRecentTasksController(
             Context context,
+            ShellInit shellInit,
             TaskStackListenerImpl taskStackListener,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
         return Optional.ofNullable(
-                RecentTasksController.create(context, taskStackListener, mainExecutor));
+                RecentTasksController.create(context, shellInit, taskStackListener,
+                        mainExecutor));
     }
 
     //
@@ -459,12 +486,15 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
-    static Transitions provideTransitions(ShellTaskOrganizer organizer, TransactionPool pool,
-            DisplayController displayController, Context context,
+    static Transitions provideTransitions(Context context,
+            ShellInit shellInit,
+            ShellTaskOrganizer organizer,
+            TransactionPool pool,
+            DisplayController displayController,
             @ShellMainThread ShellExecutor mainExecutor,
             @ShellMainThread Handler mainHandler,
             @ShellAnimationThread ShellExecutor animExecutor) {
-        return new Transitions(organizer, pool, displayController, context, mainExecutor,
+        return new Transitions(context, shellInit, organizer, pool, displayController, mainExecutor,
                 mainHandler, animExecutor);
     }
 
@@ -542,11 +572,13 @@ public abstract class WMShellBaseModule {
     @WMSingleton
     @Provides
     static StartingWindowController provideStartingWindowController(Context context,
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
             @ShellSplashscreenThread ShellExecutor splashScreenExecutor,
             StartingWindowTypeAlgorithm startingWindowTypeAlgorithm, IconProvider iconProvider,
             TransactionPool pool) {
-        return new StartingWindowController(context, splashScreenExecutor,
-                startingWindowTypeAlgorithm, iconProvider, pool);
+        return new StartingWindowController(context, shellInit, shellTaskOrganizer,
+                splashScreenExecutor, startingWindowTypeAlgorithm, iconProvider, pool);
     }
 
     // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
@@ -595,9 +627,11 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
-    static Optional<ActivityEmbeddingController> provideActivityEmbeddingController(
-            Context context, Transitions transitions) {
-        return Optional.of(new ActivityEmbeddingController(context, transitions));
+    static ActivityEmbeddingController provideActivityEmbeddingController(
+            Context context,
+            ShellInit shellInit,
+            Transitions transitions) {
+        return new ActivityEmbeddingController(context, shellInit, transitions);
     }
 
     //
@@ -606,24 +640,34 @@ public abstract class WMShellBaseModule {
 
     @WMSingleton
     @Provides
-    static ShellInterface provideShellSysuiCallbacks(ShellController shellController) {
+    static ShellInterface provideShellSysuiCallbacks(
+            @ShellCreateTrigger Object createTrigger,
+            ShellController shellController) {
         return shellController.asShell();
     }
 
     @WMSingleton
     @Provides
-    static ShellController provideShellController(@ShellMainThread ShellExecutor mainExecutor) {
-        return new ShellController(mainExecutor);
+    static ShellController provideShellController(ShellInit shellInit,
+            @ShellMainThread ShellExecutor mainExecutor) {
+        return new ShellController(shellInit, mainExecutor);
     }
 
     //
     // Misc
     //
 
+    // Workaround for dynamic overriding with a default implementation, see {@link DynamicOverride}
+    @BindsOptionalOf
+    @ShellCreateTriggerOverride
+    abstract Object provideIndependentShellComponentsToCreateOverride();
+
+    // TODO: Temporarily move dependencies to this instead of ShellInit since that is needed to add
+    // the callback. We will be moving to a different explicit startup mechanism in a follow- up CL.
     @WMSingleton
+    @ShellCreateTrigger
     @Provides
-    static ShellInit provideShellInitImpl(
-            ShellController shellController,
+    static Object provideIndependentShellComponentsToCreate(
             DisplayController displayController,
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController,
@@ -638,29 +682,17 @@ public abstract class WMShellBaseModule {
             Optional<UnfoldTransitionHandler> unfoldTransitionHandler,
             Optional<FreeformComponents> freeformComponents,
             Optional<RecentTasksController> recentTasksOptional,
-            Optional<ActivityEmbeddingController> activityEmbeddingOptional,
+            ActivityEmbeddingController activityEmbeddingOptional,
             Transitions transitions,
             StartingWindowController startingWindow,
-            @ShellMainThread ShellExecutor mainExecutor) {
-        return new ShellInit(shellController,
-                displayController,
-                displayImeController,
-                displayInsetsController,
-                dragAndDropController,
-                shellTaskOrganizer,
-                kidsModeTaskOrganizer,
-                bubblesOptional,
-                splitScreenOptional,
-                pipTouchHandlerOptional,
-                fullscreenTaskListener,
-                unfoldAnimationController,
-                unfoldTransitionHandler,
-                freeformComponents,
-                recentTasksOptional,
-                activityEmbeddingOptional,
-                transitions,
-                startingWindow,
-                mainExecutor);
+            @ShellCreateTriggerOverride Optional<Object> overriddenCreateTrigger) {
+        return new Object();
+    }
+
+    @WMSingleton
+    @Provides
+    static ShellInit provideShellInit(@ShellMainThread ShellExecutor mainExecutor) {
+        return new ShellInit(mainExecutor);
     }
 
     @WMSingleton
@@ -679,18 +711,4 @@ public abstract class WMShellBaseModule {
                 kidsModeTaskOrganizer, splitScreenOptional, pipOptional, oneHandedOptional,
                 hideDisplayCutout, recentTasksOptional, mainExecutor);
     }
-
-    @WMSingleton
-    @Provides
-    static Optional<BackAnimationController> provideBackAnimationController(
-            Context context,
-            @ShellMainThread ShellExecutor shellExecutor,
-            @ShellBackgroundThread Handler backgroundHandler
-    ) {
-        if (BackAnimationController.IS_ENABLED) {
-            return Optional.of(
-                    new BackAnimationController(shellExecutor, backgroundHandler, context));
-        }
-        return Optional.empty();
-    }
 }
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 588a1aa57a3a..e2bf7678f4b4 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
@@ -75,6 +75,8 @@ import com.android.wm.shell.pip.phone.PipTouchHandler;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.SplitscreenPipMixedHandler;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider;
 import com.android.wm.shell.unfold.UnfoldAnimationController;
@@ -138,6 +140,7 @@ public abstract class WMShellModule {
     @WMSingleton
     @Provides
     static BubbleController provideBubbleController(Context context,
+            ShellInit shellInit,
             ShellController shellController,
             BubbleData data,
             FloatingContentCoordinator floatingContentCoordinator,
@@ -158,8 +161,8 @@ public abstract class WMShellModule {
             @ShellBackgroundThread ShellExecutor bgExecutor,
             TaskViewTransitions taskViewTransitions,
             SyncTransactionQueue syncQueue) {
-        return new BubbleController(context, shellController, data, null /* synchronizer */,
-                floatingContentCoordinator,
+        return new BubbleController(context, shellInit, shellController, data,
+                null /* synchronizer */, floatingContentCoordinator,
                 new BubbleDataRepository(context, launcherApps, mainExecutor),
                 statusBarService, windowManager, windowManagerShellWrapper, userManager,
                 launcherApps, logger, taskStackListener, organizer, positioner, displayController,
@@ -205,18 +208,34 @@ public abstract class WMShellModule {
     @WMSingleton
     @Provides
     static FreeformTaskListener<?> provideFreeformTaskListener(
+            Context context,
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
             WindowDecorViewModel<?> windowDecorViewModel) {
-        return new FreeformTaskListener<>(windowDecorViewModel);
+        // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
+        //                    override for this controller from the base module
+        ShellInit init = FreeformComponents.isFreeformEnabled(context)
+                ? shellInit
+                : null;
+        return new FreeformTaskListener<>(init, shellTaskOrganizer,
+                windowDecorViewModel);
     }
 
     @WMSingleton
     @Provides
-    static FreeformTaskTransitionHandler provideTaskTransitionHandler(
+    static FreeformTaskTransitionHandler provideFreeformTaskTransitionHandler(
+            Context context,
+            ShellInit shellInit,
             Transitions transitions,
             WindowDecorViewModel<?> windowDecorViewModel,
             FreeformTaskListener<?> freeformTaskListener) {
-        return new FreeformTaskTransitionHandler(transitions, windowDecorViewModel,
-                freeformTaskListener);
+        // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
+        //                    override for this controller from the base module
+        ShellInit init = FreeformComponents.isFreeformEnabled(context)
+                ? shellInit
+                : null;
+        return new FreeformTaskTransitionHandler(init, transitions,
+                windowDecorViewModel, freeformTaskListener);
     }
 
     //
@@ -247,20 +266,25 @@ public abstract class WMShellModule {
     @Provides
     @DynamicOverride
     static SplitScreenController provideSplitScreenController(
+            Context context,
+            ShellInit shellInit,
             ShellController shellController,
             ShellTaskOrganizer shellTaskOrganizer,
-            SyncTransactionQueue syncQueue, Context context,
+            SyncTransactionQueue syncQueue,
             RootTaskDisplayAreaOrganizer rootTaskDisplayAreaOrganizer,
-            @ShellMainThread ShellExecutor mainExecutor,
             DisplayController displayController,
             DisplayImeController displayImeController,
-            DisplayInsetsController displayInsetsController, Transitions transitions,
-            TransactionPool transactionPool, IconProvider iconProvider,
-            Optional<RecentTasksController> recentTasks) {
-        return new SplitScreenController(shellController, shellTaskOrganizer, syncQueue, context,
-                rootTaskDisplayAreaOrganizer, mainExecutor, displayController, displayImeController,
-                displayInsetsController, transitions, transactionPool, iconProvider,
-                recentTasks);
+            DisplayInsetsController displayInsetsController,
+            DragAndDropController dragAndDropController,
+            Transitions transitions,
+            TransactionPool transactionPool,
+            IconProvider iconProvider,
+            Optional<RecentTasksController> recentTasks,
+            @ShellMainThread ShellExecutor mainExecutor) {
+        return new SplitScreenController(context, shellInit, shellController, shellTaskOrganizer,
+                syncQueue, rootTaskDisplayAreaOrganizer, displayController, displayImeController,
+                displayInsetsController, dragAndDropController, transitions, transactionPool,
+                iconProvider, recentTasks, mainExecutor);
     }
 
     //
@@ -332,14 +356,16 @@ public abstract class WMShellModule {
     @WMSingleton
     @Provides
     static PipTouchHandler providePipTouchHandler(Context context,
-            PhonePipMenuController menuPhoneController, PipBoundsAlgorithm pipBoundsAlgorithm,
+            ShellInit shellInit,
+            PhonePipMenuController menuPhoneController,
+            PipBoundsAlgorithm pipBoundsAlgorithm,
             PipBoundsState pipBoundsState,
             PipTaskOrganizer pipTaskOrganizer,
             PipMotionHelper pipMotionHelper,
             FloatingContentCoordinator floatingContentCoordinator,
             PipUiEventLogger pipUiEventLogger,
             @ShellMainThread ShellExecutor mainExecutor) {
-        return new PipTouchHandler(context, menuPhoneController, pipBoundsAlgorithm,
+        return new PipTouchHandler(context, shellInit, menuPhoneController, pipBoundsAlgorithm,
                 pipBoundsState, pipTaskOrganizer, pipMotionHelper,
                 floatingContentCoordinator, pipUiEventLogger, mainExecutor);
     }
@@ -414,9 +440,31 @@ public abstract class WMShellModule {
                 floatingContentCoordinator);
     }
 
+    @WMSingleton
+    @Provides
+    static PipParamsChangedForwarder providePipParamsChangedForwarder() {
+        return new PipParamsChangedForwarder();
+    }
+
+    //
+    // Transitions
+    //
+
+    @WMSingleton
+    @Provides
+    static SplitscreenPipMixedHandler provideSplitscreenPipMixedHandler(
+            ShellInit shellInit,
+            Optional<SplitScreenController> splitScreenOptional,
+            Optional<PipTouchHandler> pipTouchHandlerOptional,
+            Transitions transitions) {
+        return new SplitscreenPipMixedHandler(shellInit, splitScreenOptional,
+                pipTouchHandlerOptional, transitions);
+    }
+
     //
     // Unfold transition
     //
+
     @WMSingleton
     @Provides
     @DynamicOverride
@@ -426,6 +474,7 @@ public abstract class WMShellModule {
             @UnfoldTransition SplitTaskUnfoldAnimator splitAnimator,
             FullscreenUnfoldTaskAnimator fullscreenAnimator,
             Lazy<Optional<UnfoldTransitionHandler>> unfoldTransitionHandler,
+            ShellInit shellInit,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
         final List<UnfoldTaskAnimator> animators = new ArrayList<>();
@@ -433,6 +482,7 @@ public abstract class WMShellModule {
         animators.add(fullscreenAnimator);
 
         return new UnfoldAnimationController(
+                        shellInit,
                         transactionPool,
                         progressProvider.get(),
                         animators,
@@ -441,7 +491,6 @@ public abstract class WMShellModule {
                 );
     }
 
-
     @Provides
     static FullscreenUnfoldTaskAnimator provideFullscreenUnfoldTaskAnimator(
             Context context,
@@ -460,6 +509,10 @@ public abstract class WMShellModule {
             Lazy<Optional<SplitScreenController>> splitScreenOptional,
             DisplayInsetsController displayInsetsController
     ) {
+        // TODO(b/238217847): The lazy reference here causes some dependency issues since it
+        // immediately registers a listener on that controller on init.  We should reference the
+        // controller directly once we refactor ShellTaskOrganizer to not depend on the unfold
+        // animation controller directly.
         return new SplitTaskUnfoldAnimator(context, executor, splitScreenOptional,
                 backgroundController, displayInsetsController);
     }
@@ -485,8 +538,9 @@ public abstract class WMShellModule {
             @UnfoldShellTransition SplitTaskUnfoldAnimator unfoldAnimator,
             TransactionPool transactionPool,
             Transitions transitions,
-            @ShellMainThread ShellExecutor executor) {
-        return new UnfoldTransitionHandler(progressProvider.get(), animator,
+            @ShellMainThread ShellExecutor executor,
+            ShellInit shellInit) {
+        return new UnfoldTransitionHandler(shellInit, progressProvider.get(), animator,
                 unfoldAnimator, transactionPool, executor, transitions);
     }
 
@@ -502,9 +556,17 @@ public abstract class WMShellModule {
         );
     }
 
+    //
+    // Misc
+    //
+
+    // TODO: Temporarily move dependencies to this instead of ShellInit since that is needed to add
+    // the callback. We will be moving to a different explicit startup mechanism in a follow- up CL.
     @WMSingleton
+    @ShellCreateTriggerOverride
     @Provides
-    static PipParamsChangedForwarder providePipParamsChangedForwarder() {
-        return new PipParamsChangedForwarder();
+    static Object provideIndependentShellComponentsToCreate(
+            SplitscreenPipMixedHandler splitscreenPipMixedHandler) {
+        return new Object();
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
index f4e2f20f4637..2aa933d641fa 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/docs/changes.md
@@ -44,10 +44,24 @@ to SysUI where it posts the work to the main Shell thread.
 
 ### Component initialization
 To initialize the component:
-- On the Shell side, update `ShellInitImpl` to get a signal to initialize when the SysUI is started
+- On the Shell side, you potentially need to do two things to initialize the component:
+  - Inject `ShellInit` into your component and add an init callback
+  - Ensure that your component is a part of the dagger dependency graph, either by:
+    - Making this component a dependency of an existing component already exposed to SystemUI
+    - Explicitly add this component to the WMShellBaseModule @ShellCreateTrigger provider or
+      the @ShellCreateTriggerOverride provider for your product module to expose it explicitly 
+      if it is a completely independent component
 - On the SysUI side, update `WMShell` to setup any bindings for the component that depend on
   SysUI code
 
+To verify that your component is being initialized at startup, you can enable the `WM_SHELL_INIT` 
+protolog group and restart the SysUI process:
+```shell
+adb shell wm logging enable-text WM_SHELL_INIT
+adb shell kill `pid com.android.systemui`
+adb logcat *:S WindowManagerShell
+```
+
 ### General Do's & Dont's
 Do:
 - Do add unit tests for all new components
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
index c5df53b6dbc8..4697a0184eb4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragAndDropController.java
@@ -62,9 +62,9 @@ import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.sysui.ConfigurationChangeListener;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.ArrayList;
-import java.util.Optional;
 
 /**
  * Handles the global drag and drop handling for the Shell.
@@ -94,6 +94,7 @@ public class DragAndDropController implements DisplayController.OnDisplaysChange
     }
 
     public DragAndDropController(Context context,
+            ShellInit shellInit,
             ShellController shellController,
             DisplayController displayController,
             UiEventLogger uiEventLogger,
@@ -105,14 +106,29 @@ public class DragAndDropController implements DisplayController.OnDisplaysChange
         mLogger = new DragAndDropEventLogger(uiEventLogger);
         mIconProvider = iconProvider;
         mMainExecutor = mainExecutor;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
-    public void initialize(Optional<SplitScreenController> splitscreen) {
-        mSplitScreen = splitscreen.orElse(null);
-        mDisplayController.addDisplayWindowListener(this);
+    /**
+     * Called when the controller is initialized.
+     */
+    public void onInit() {
+        // TODO(b/238217847): The dependency from SplitscreenController on DragAndDropController is
+        // inverted, which leads to SplitscreenController not setting its instance until after
+        // onDisplayAdded.  We can remove this post once we fix that dependency.
+        mMainExecutor.executeDelayed(() -> {
+            mDisplayController.addDisplayWindowListener(this);
+        }, 0);
         mShellController.addConfigurationChangeListener(this);
     }
 
+    /**
+     * Sets the splitscreen controller to use if the feature is available.
+     */
+    public void setSplitScreenController(SplitScreenController splitscreen) {
+        mSplitScreen = splitscreen;
+    }
+
     /** Adds a listener to be notified of drag and drop events. */
     public void addListener(DragAndDropListener listener) {
         mListeners.add(listener);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index ba093a5fcc05..ab66107399c6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -16,6 +16,8 @@
 
 package com.android.wm.shell.freeform;
 
+import static com.android.wm.shell.ShellTaskOrganizer.TASK_LISTENER_TYPE_FREEFORM;
+
 import android.app.ActivityManager.RunningTaskInfo;
 import android.util.Log;
 import android.util.SparseArray;
@@ -27,6 +29,7 @@ import androidx.annotation.Nullable;
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
@@ -42,6 +45,7 @@ public class FreeformTaskListener<T extends AutoCloseable>
         implements ShellTaskOrganizer.TaskListener {
     private static final String TAG = "FreeformTaskListener";
 
+    private final ShellTaskOrganizer mShellTaskOrganizer;
     private final WindowDecorViewModel<T> mWindowDecorationViewModel;
 
     private final SparseArray<State<T>> mTasks = new SparseArray<>();
@@ -53,8 +57,19 @@ public class FreeformTaskListener<T extends AutoCloseable>
         T mWindowDecoration;
     }
 
-    public FreeformTaskListener(WindowDecorViewModel<T> windowDecorationViewModel) {
+    public FreeformTaskListener(
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
+            WindowDecorViewModel<T> windowDecorationViewModel) {
+        mShellTaskOrganizer = shellTaskOrganizer;
         mWindowDecorationViewModel = windowDecorationViewModel;
+        if (shellInit != null) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
+    }
+
+    private void onInit() {
+        mShellTaskOrganizer.addListenerForType(this, TASK_LISTENER_TYPE_FREEFORM);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
index 8731eb604cbc..20d77259406a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
@@ -32,6 +32,7 @@ import android.window.WindowContainerTransaction;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
@@ -49,17 +50,26 @@ public class FreeformTaskTransitionHandler
 
     private final Transitions mTransitions;
     private final FreeformTaskListener<?> mFreeformTaskListener;
+    private final WindowDecorViewModel<?> mWindowDecorViewModel;
 
     private final List<IBinder> mPendingTransitionTokens = new ArrayList<>();
 
     public FreeformTaskTransitionHandler(
+            ShellInit shellInit,
             Transitions transitions,
             WindowDecorViewModel<?> windowDecorViewModel,
             FreeformTaskListener<?> freeformTaskListener) {
         mTransitions = transitions;
         mFreeformTaskListener = freeformTaskListener;
+        mWindowDecorViewModel = windowDecorViewModel;
+        if (shellInit != null && Transitions.ENABLE_SHELL_TRANSITIONS) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
+    }
 
-        windowDecorViewModel.setFreeformTaskTransitionStarter(this);
+    private void onInit() {
+        mWindowDecorViewModel.setFreeformTaskTransitionStarter(this);
+        mTransitions.addHandler(this);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
index 79e363bcdb41..0ba4afc24c45 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/fullscreen/FullscreenTaskListener.java
@@ -32,6 +32,7 @@ import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentTasksController;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 
 import java.io.PrintWriter;
@@ -43,19 +44,34 @@ import java.util.Optional;
 public class FullscreenTaskListener implements ShellTaskOrganizer.TaskListener {
     private static final String TAG = "FullscreenTaskListener";
 
+    private final ShellTaskOrganizer mShellTaskOrganizer;
     private final SyncTransactionQueue mSyncQueue;
     private final Optional<RecentTasksController> mRecentTasksOptional;
 
     private final SparseArray<TaskData> mDataByTaskId = new SparseArray<>();
 
+    /**
+     * This constructor is used by downstream products.
+     */
     public FullscreenTaskListener(SyncTransactionQueue syncQueue) {
-        this(syncQueue, Optional.empty());
+        this(null /* shellInit */, null /* shellTaskOrganizer */, syncQueue, Optional.empty());
     }
 
-    public FullscreenTaskListener(SyncTransactionQueue syncQueue,
-            Optional<RecentTasksController> recentTasks) {
+    public FullscreenTaskListener(ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
+            SyncTransactionQueue syncQueue,
+            Optional<RecentTasksController> recentTasksOptional) {
+        mShellTaskOrganizer = shellTaskOrganizer;
         mSyncQueue = syncQueue;
-        mRecentTasksOptional = recentTasks;
+        mRecentTasksOptional = recentTasksOptional;
+        // Note: Some derivative FullscreenTaskListener implementations do not use ShellInit
+        if (shellInit != null) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
+    }
+
+    private void onInit() {
+        mShellTaskOrganizer.addListenerForType(this, TASK_LISTENER_TYPE_FULLSCREEN);
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
index 2c8ba0970ccc..73b9b89e6993 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizer.java
@@ -50,7 +50,7 @@ import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.recents.RecentTasksController;
-import com.android.wm.shell.startingsurface.StartingWindowController;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.unfold.UnfoldAnimationController;
 
 import java.io.PrintWriter;
@@ -140,18 +140,18 @@ public class KidsModeTaskOrganizer extends ShellTaskOrganizer {
 
     @VisibleForTesting
     KidsModeTaskOrganizer(
-            ITaskOrganizerController taskOrganizerController,
-            ShellExecutor mainExecutor,
-            Handler mainHandler,
             Context context,
+            ITaskOrganizerController taskOrganizerController,
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
             Optional<UnfoldAnimationController> unfoldAnimationController,
             Optional<RecentTasksController> recentTasks,
-            KidsModeSettingsObserver kidsModeSettingsObserver) {
-        super(taskOrganizerController, mainExecutor, context, /* compatUI= */ null,
-                unfoldAnimationController, recentTasks);
+            KidsModeSettingsObserver kidsModeSettingsObserver,
+            ShellExecutor mainExecutor,
+            Handler mainHandler) {
+        super(/* shellInit= */ null, taskOrganizerController, /* compatUI= */ null,
+                unfoldAnimationController, recentTasks, mainExecutor);
         mContext = context;
         mMainHandler = mainHandler;
         mSyncQueue = syncTransactionQueue;
@@ -161,27 +161,30 @@ public class KidsModeTaskOrganizer extends ShellTaskOrganizer {
     }
 
     public KidsModeTaskOrganizer(
-            ShellExecutor mainExecutor,
-            Handler mainHandler,
             Context context,
+            ShellInit shellInit,
             SyncTransactionQueue syncTransactionQueue,
             DisplayController displayController,
             DisplayInsetsController displayInsetsController,
             Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<RecentTasksController> recentTasks) {
-        super(mainExecutor, context, /* compatUI= */ null, unfoldAnimationController, recentTasks);
+            Optional<RecentTasksController> recentTasks,
+            ShellExecutor mainExecutor,
+            Handler mainHandler) {
+        // Note: we don't call super with the shell init because we will be initializing manually
+        super(/* shellInit= */ null, /* compatUI= */ null, unfoldAnimationController, recentTasks,
+                mainExecutor);
         mContext = context;
         mMainHandler = mainHandler;
         mSyncQueue = syncTransactionQueue;
         mDisplayController = displayController;
         mDisplayInsetsController = displayInsetsController;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     /**
      * Initializes kids mode status.
      */
-    public void initialize(StartingWindowController startingWindowController) {
-        initStartingWindow(startingWindowController);
+    public void onInit() {
         if (mKidsModeSettingsObserver == null) {
             mKidsModeSettingsObserver = new KidsModeSettingsObserver(mMainHandler, mContext);
         }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
index a2ff97247189..c86c1368b88e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/phone/PipTouchHandler.java
@@ -57,6 +57,7 @@ import com.android.wm.shell.pip.PipTaskOrganizer;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip.PipUiEventLogger;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.io.PrintWriter;
 
@@ -165,6 +166,7 @@ public class PipTouchHandler {
 
     @SuppressLint("InflateParams")
     public PipTouchHandler(Context context,
+            ShellInit shellInit,
             PhonePipMenuController menuController,
             PipBoundsAlgorithm pipBoundsAlgorithm,
             @NonNull PipBoundsState pipBoundsState,
@@ -173,7 +175,6 @@ public class PipTouchHandler {
             FloatingContentCoordinator floatingContentCoordinator,
             PipUiEventLogger pipUiEventLogger,
             ShellExecutor mainExecutor) {
-        // Initialize the Pip input consumer
         mContext = context;
         mMainExecutor = mainExecutor;
         mAccessibilityManager = context.getSystemService(AccessibilityManager.class);
@@ -213,9 +214,17 @@ public class PipTouchHandler {
                 mMotionHelper, pipTaskOrganizer, mPipBoundsAlgorithm.getSnapAlgorithm(),
                 this::onAccessibilityShowMenu, this::updateMovementBounds,
                 this::animateToUnStashedState, mainExecutor);
+
+        // TODO(b/181599115): This should really be initializes as part of the pip controller, but
+        // until all PIP implementations derive from the controller, just initialize the touch handler
+        // if it is needed
+        shellInit.addInitCallback(this::onInit, this);
     }
 
-    public void init() {
+    /**
+     * Called when the touch handler is initialized.
+     */
+    public void onInit() {
         Resources res = mContext.getResources();
         mEnableResize = res.getBoolean(R.bool.config_pipEnableResizeForMenu);
         reloadResources();
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
index 49042e663f6e..3d1a7e98e20d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
@@ -44,6 +44,7 @@ import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.common.annotations.ShellMainThread;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
 import com.android.wm.shell.util.SplitBounds;
 
@@ -85,28 +86,32 @@ public class RecentTasksController implements TaskStackListenerCallback,
     @Nullable
     public static RecentTasksController create(
             Context context,
+            ShellInit shellInit,
             TaskStackListenerImpl taskStackListener,
             @ShellMainThread ShellExecutor mainExecutor
     ) {
         if (!context.getResources().getBoolean(com.android.internal.R.bool.config_hasRecents)) {
             return null;
         }
-        return new RecentTasksController(context, taskStackListener, mainExecutor);
+        return new RecentTasksController(context, shellInit, taskStackListener, mainExecutor);
     }
 
-    RecentTasksController(Context context, TaskStackListenerImpl taskStackListener,
+    RecentTasksController(Context context,
+            ShellInit shellInit,
+            TaskStackListenerImpl taskStackListener,
             ShellExecutor mainExecutor) {
         mContext = context;
         mIsDesktopMode = mContext.getPackageManager().hasSystemFeature(FEATURE_PC);
         mTaskStackListener = taskStackListener;
         mMainExecutor = mainExecutor;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     public RecentTasks asRecentTasks() {
         return mImpl;
     }
 
-    public void init() {
+    private void onInit() {
         mTaskStackListener.addListener(this);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 2008a757839d..1be17f9bc5fd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -78,12 +78,14 @@ import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.common.split.SplitLayout;
 import com.android.wm.shell.common.split.SplitScreenConstants.SplitPosition;
+import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.draganddrop.DragAndDropPolicy;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.splitscreen.SplitScreen.StageType;
 import com.android.wm.shell.sysui.KeyguardChangeListener;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.LegacyTransitions;
 import com.android.wm.shell.transition.Transitions;
 
@@ -141,6 +143,7 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
     private final DisplayController mDisplayController;
     private final DisplayImeController mDisplayImeController;
     private final DisplayInsetsController mDisplayInsetsController;
+    private final DragAndDropController mDragAndDropController;
     private final Transitions mTransitions;
     private final TransactionPool mTransactionPool;
     private final SplitscreenEventLogger mLogger;
@@ -152,15 +155,21 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
     // outside the bounds of the roots by being reparented into a higher level fullscreen container
     private SurfaceControl mSplitTasksContainerLayer;
 
-    public SplitScreenController(ShellController shellController,
+    public SplitScreenController(Context context,
+            ShellInit shellInit,
+            ShellController shellController,
             ShellTaskOrganizer shellTaskOrganizer,
-            SyncTransactionQueue syncQueue, Context context,
+            SyncTransactionQueue syncQueue,
             RootTaskDisplayAreaOrganizer rootTDAOrganizer,
-            ShellExecutor mainExecutor, DisplayController displayController,
+            DisplayController displayController,
             DisplayImeController displayImeController,
             DisplayInsetsController displayInsetsController,
-            Transitions transitions, TransactionPool transactionPool, IconProvider iconProvider,
-            Optional<RecentTasksController> recentTasks) {
+            DragAndDropController dragAndDropController,
+            Transitions transitions,
+            TransactionPool transactionPool,
+            IconProvider iconProvider,
+            Optional<RecentTasksController> recentTasks,
+            ShellExecutor mainExecutor) {
         mShellController = shellController;
         mTaskOrganizer = shellTaskOrganizer;
         mSyncQueue = syncQueue;
@@ -170,28 +179,29 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
         mDisplayController = displayController;
         mDisplayImeController = displayImeController;
         mDisplayInsetsController = displayInsetsController;
+        mDragAndDropController = dragAndDropController;
         mTransitions = transitions;
         mTransactionPool = transactionPool;
         mLogger = new SplitscreenEventLogger();
         mIconProvider = iconProvider;
         mRecentTasksOptional = recentTasks;
+        // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
+        //                    override for this controller from the base module
+        if (ActivityTaskManager.supportsSplitScreenMultiWindow(context)) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
     }
 
     public SplitScreen asSplitScreen() {
         return mImpl;
     }
 
-    @Override
-    public Context getContext() {
-        return mContext;
-    }
-
-    @Override
-    public ShellExecutor getRemoteCallExecutor() {
-        return mMainExecutor;
-    }
-
-    public void onOrganizerRegistered() {
+    /**
+     * This will be called after ShellTaskOrganizer has initialized/registered because of the
+     * dependency order.
+     */
+    @VisibleForTesting
+    void onInit() {
         mShellController.addKeyguardChangeListener(this);
         if (mStageCoordinator == null) {
             // TODO: Multi-display
@@ -200,6 +210,17 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
                     mDisplayInsetsController, mTransitions, mTransactionPool, mLogger,
                     mIconProvider, mMainExecutor, mRecentTasksOptional);
         }
+        mDragAndDropController.setSplitScreenController(this);
+    }
+
+    @Override
+    public Context getContext() {
+        return mContext;
+    }
+
+    @Override
+    public ShellExecutor getRemoteCallExecutor() {
+        return mMainExecutor;
     }
 
     public boolean isSplitScreenVisible() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
index fbc992378e50..379af21ac956 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/startingsurface/StartingWindowController.java
@@ -42,10 +42,12 @@ import androidx.annotation.BinderThread;
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.util.function.TriConsumer;
 import com.android.launcher3.icons.IconProvider;
+import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.RemoteCallable;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SingleInstanceRemoteListener;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
 
 /**
  * Implementation to draw the starting window to an application, and remove the starting window
@@ -74,6 +76,7 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     private TriConsumer<Integer, Integer, Integer> mTaskLaunchingCallback;
     private final StartingSurfaceImpl mImpl = new StartingSurfaceImpl();
     private final Context mContext;
+    private final ShellTaskOrganizer mShellTaskOrganizer;
     private final ShellExecutor mSplashScreenExecutor;
     /**
      * Need guarded because it has exposed to StartingSurface
@@ -81,14 +84,20 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
     @GuardedBy("mTaskBackgroundColors")
     private final SparseIntArray mTaskBackgroundColors = new SparseIntArray();
 
-    public StartingWindowController(Context context, ShellExecutor splashScreenExecutor,
-            StartingWindowTypeAlgorithm startingWindowTypeAlgorithm, IconProvider iconProvider,
+    public StartingWindowController(Context context,
+            ShellInit shellInit,
+            ShellTaskOrganizer shellTaskOrganizer,
+            ShellExecutor splashScreenExecutor,
+            StartingWindowTypeAlgorithm startingWindowTypeAlgorithm,
+            IconProvider iconProvider,
             TransactionPool pool) {
         mContext = context;
+        mShellTaskOrganizer = shellTaskOrganizer;
         mStartingSurfaceDrawer = new StartingSurfaceDrawer(context, splashScreenExecutor,
                 iconProvider, pool);
         mStartingWindowTypeAlgorithm = startingWindowTypeAlgorithm;
         mSplashScreenExecutor = splashScreenExecutor;
+        shellInit.addInitCallback(this::onInit, this);
     }
 
     /**
@@ -98,6 +107,10 @@ public class StartingWindowController implements RemoteCallable<StartingWindowCo
         return mImpl;
     }
 
+    private void onInit() {
+        mShellTaskOrganizer.initStartingWindow(this);
+    }
+
     @Override
     public Context getContext() {
         return mContext;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellCommandHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellCommandHandler.java
index 0427efb30b9f..f4fc0c4c36bc 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellCommandHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellCommandHandler.java
@@ -45,7 +45,6 @@ public final class ShellCommandHandler {
     private final Optional<RecentTasksController> mRecentTasks;
     private final ShellTaskOrganizer mShellTaskOrganizer;
     private final KidsModeTaskOrganizer mKidsModeTaskOrganizer;
-    private final ShellExecutor mMainExecutor;
 
     public ShellCommandHandler(
             ShellController shellController,
@@ -64,7 +63,6 @@ public final class ShellCommandHandler {
         mPipOptional = pipOptional;
         mOneHandedOptional = oneHandedOptional;
         mHideDisplayCutout = hideDisplayCutout;
-        mMainExecutor = mainExecutor;
         // TODO(238217847): To be removed once the command handler dependencies are inverted
         shellController.setShellCommandHandler(this);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellController.java
index 618028c1544f..f1f317f65ba9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellController.java
@@ -44,10 +44,10 @@ import java.util.concurrent.CopyOnWriteArrayList;
 public class ShellController {
     private static final String TAG = ShellController.class.getSimpleName();
 
+    private final ShellInit mShellInit;
     private final ShellExecutor mMainExecutor;
     private final ShellInterfaceImpl mImpl = new ShellInterfaceImpl();
 
-    private ShellInit mShellInit;
     private ShellCommandHandler mShellCommandHandler;
 
     private final CopyOnWriteArrayList<ConfigurationChangeListener> mConfigChangeListeners =
@@ -57,7 +57,8 @@ public class ShellController {
     private Configuration mLastConfiguration;
 
 
-    public ShellController(ShellExecutor mainExecutor) {
+    public ShellController(ShellInit shellInit, ShellExecutor mainExecutor) {
+        mShellInit = shellInit;
         mMainExecutor = mainExecutor;
     }
 
@@ -69,15 +70,6 @@ public class ShellController {
     }
 
     /**
-     * Sets the init handler to call back to.
-     * TODO(238217847): This is only exposed this way until we can remove the dependencies from the
-     *                  init handler to other classes.
-     */
-    public void setShellInit(ShellInit shellInit) {
-        mShellInit = shellInit;
-    }
-
-    /**
      * Sets the command handler to call back to.
      * TODO(238217847): This is only exposed this way until we can remove the dependencies from the
      *                  command handler to other classes.
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
index dd7fab7c5f46..c250e0313255 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/sysui/ShellInit.java
@@ -16,7 +16,6 @@
 
 package com.android.wm.shell.sysui;
 
-import static com.android.wm.shell.ShellTaskOrganizer.TASK_LISTENER_TYPE_FULLSCREEN;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_INIT;
 
 import android.os.Build;
@@ -26,149 +25,27 @@ import android.util.Pair;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.internal.protolog.common.ProtoLog;
-import com.android.wm.shell.ShellTaskOrganizer;
-import com.android.wm.shell.activityembedding.ActivityEmbeddingController;
-import com.android.wm.shell.bubbles.BubbleController;
-import com.android.wm.shell.common.DisplayController;
-import com.android.wm.shell.common.DisplayImeController;
-import com.android.wm.shell.common.DisplayInsetsController;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.draganddrop.DragAndDropController;
-import com.android.wm.shell.freeform.FreeformComponents;
-import com.android.wm.shell.fullscreen.FullscreenTaskListener;
-import com.android.wm.shell.kidsmode.KidsModeTaskOrganizer;
-import com.android.wm.shell.pip.phone.PipTouchHandler;
-import com.android.wm.shell.recents.RecentTasksController;
-import com.android.wm.shell.splitscreen.SplitScreenController;
-import com.android.wm.shell.startingsurface.StartingWindowController;
-import com.android.wm.shell.transition.DefaultMixedHandler;
-import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
-import com.android.wm.shell.unfold.UnfoldTransitionHandler;
 
 import java.util.ArrayList;
-import java.util.Optional;
 
 /**
  * The entry point implementation into the shell for initializing shell internal state.  Classes
- * which need to setup on start should inject an instance of this class and add an init callback.
+ * which need to initialize on start of the host SysUI should inject an instance of this class and
+ * add an init callback.
  */
 public class ShellInit {
     private static final String TAG = ShellInit.class.getSimpleName();
 
-    private final DisplayController mDisplayController;
-    private final DisplayImeController mDisplayImeController;
-    private final DisplayInsetsController mDisplayInsetsController;
-    private final DragAndDropController mDragAndDropController;
-    private final ShellTaskOrganizer mShellTaskOrganizer;
-    private final KidsModeTaskOrganizer mKidsModeTaskOrganizer;
-    private final Optional<BubbleController> mBubblesOptional;
-    private final Optional<SplitScreenController> mSplitScreenOptional;
-    private final Optional<PipTouchHandler> mPipTouchHandlerOptional;
-    private final FullscreenTaskListener mFullscreenTaskListener;
-    private final Optional<UnfoldAnimationController> mUnfoldController;
-    private final Optional<UnfoldTransitionHandler> mUnfoldTransitionHandler;
-    private final Optional<FreeformComponents> mFreeformComponentsOptional;
     private final ShellExecutor mMainExecutor;
-    private final Transitions mTransitions;
-    private final StartingWindowController mStartingWindow;
-    private final Optional<RecentTasksController> mRecentTasks;
-    private final Optional<ActivityEmbeddingController> mActivityEmbeddingOptional;
 
     // An ordered list of init callbacks to be made once shell is first started
     private final ArrayList<Pair<String, Runnable>> mInitCallbacks = new ArrayList<>();
     private boolean mHasInitialized;
 
-    public ShellInit(
-            ShellController shellController,
-            DisplayController displayController,
-            DisplayImeController displayImeController,
-            DisplayInsetsController displayInsetsController,
-            DragAndDropController dragAndDropController,
-            ShellTaskOrganizer shellTaskOrganizer,
-            KidsModeTaskOrganizer kidsModeTaskOrganizer,
-            Optional<BubbleController> bubblesOptional,
-            Optional<SplitScreenController> splitScreenOptional,
-            Optional<PipTouchHandler> pipTouchHandlerOptional,
-            FullscreenTaskListener fullscreenTaskListener,
-            Optional<UnfoldAnimationController> unfoldAnimationController,
-            Optional<UnfoldTransitionHandler> unfoldTransitionHandler,
-            Optional<FreeformComponents> freeformComponentsOptional,
-            Optional<RecentTasksController> recentTasks,
-            Optional<ActivityEmbeddingController> activityEmbeddingOptional,
-            Transitions transitions,
-            StartingWindowController startingWindow,
-            ShellExecutor mainExecutor) {
-        mDisplayController = displayController;
-        mDisplayImeController = displayImeController;
-        mDisplayInsetsController = displayInsetsController;
-        mDragAndDropController = dragAndDropController;
-        mShellTaskOrganizer = shellTaskOrganizer;
-        mKidsModeTaskOrganizer = kidsModeTaskOrganizer;
-        mBubblesOptional = bubblesOptional;
-        mSplitScreenOptional = splitScreenOptional;
-        mFullscreenTaskListener = fullscreenTaskListener;
-        mPipTouchHandlerOptional = pipTouchHandlerOptional;
-        mUnfoldController = unfoldAnimationController;
-        mUnfoldTransitionHandler = unfoldTransitionHandler;
-        mFreeformComponentsOptional = freeformComponentsOptional;
-        mRecentTasks = recentTasks;
-        mActivityEmbeddingOptional = activityEmbeddingOptional;
-        mTransitions = transitions;
-        mMainExecutor = mainExecutor;
-        mStartingWindow = startingWindow;
-        // TODO(238217847): To be removed once the init dependencies are inverted
-        shellController.setShellInit(this);
-    }
-
-    private void legacyInit() {
-        // Start listening for display and insets changes
-        mDisplayController.initialize();
-        mDisplayInsetsController.initialize();
-        mDisplayImeController.startMonitorDisplays();
-
-        // Setup the shell organizer
-        mShellTaskOrganizer.addListenerForType(
-                mFullscreenTaskListener, TASK_LISTENER_TYPE_FULLSCREEN);
-        mShellTaskOrganizer.initStartingWindow(mStartingWindow);
-        mShellTaskOrganizer.registerOrganizer();
 
-        mSplitScreenOptional.ifPresent(SplitScreenController::onOrganizerRegistered);
-        mBubblesOptional.ifPresent(BubbleController::initialize);
-
-        // Bind the splitscreen impl to the drag drop controller
-        mDragAndDropController.initialize(mSplitScreenOptional);
-
-        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            mTransitions.register(mShellTaskOrganizer);
-            mActivityEmbeddingOptional.ifPresent(ActivityEmbeddingController::init);
-            mUnfoldTransitionHandler.ifPresent(UnfoldTransitionHandler::init);
-            mFreeformComponentsOptional.flatMap(f -> f.mTaskTransitionHandler)
-                    .ifPresent(mTransitions::addHandler);
-            if (mSplitScreenOptional.isPresent() && mPipTouchHandlerOptional.isPresent()) {
-                final DefaultMixedHandler mixedHandler = new DefaultMixedHandler(mTransitions,
-                        mPipTouchHandlerOptional.get().getTransitionHandler(),
-                        mSplitScreenOptional.get().getTransitionHandler());
-                // Added at end so that it has highest priority.
-                mTransitions.addHandler(mixedHandler);
-            }
-        }
-
-        // TODO(b/181599115): This should really be the pip controller, but until we can provide the
-        // controller instead of the feature interface, can just initialize the touch handler if
-        // needed
-        mPipTouchHandlerOptional.ifPresent((handler) -> handler.init());
-
-        // Initialize optional freeform
-        mFreeformComponentsOptional.ifPresent(f ->
-                mShellTaskOrganizer.addListenerForType(
-                        f.mTaskListener, ShellTaskOrganizer.TASK_LISTENER_TYPE_FREEFORM));
-
-        mUnfoldController.ifPresent(UnfoldAnimationController::init);
-        mRecentTasks.ifPresent(RecentTasksController::init);
-
-        // Initialize kids mode task organizer
-        mKidsModeTaskOrganizer.initialize(mStartingWindow);
+    public ShellInit(ShellExecutor mainExecutor) {
+        mMainExecutor = mainExecutor;
     }
 
     /**
@@ -201,13 +78,9 @@ public class ShellInit {
             final long t1 = SystemClock.uptimeMillis();
             info.second.run();
             final long t2 = SystemClock.uptimeMillis();
-            ProtoLog.v(WM_SHELL_INIT, "\t%s took %dms", info.first, (t2 - t1));
+            ProtoLog.v(WM_SHELL_INIT, "\t%s init took %dms", info.first, (t2 - t1));
         }
         mInitCallbacks.clear();
-
-        // TODO: To be removed
-        legacyInit();
-
         mHasInitialized = true;
     }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
index 0bec54399dd8..afc70a448aa8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultTransitionHandler.java
@@ -109,6 +109,7 @@ import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -136,6 +137,7 @@ public class DefaultTransitionHandler implements Transitions.TransitionHandler {
     private final TransactionPool mTransactionPool;
     private final DisplayController mDisplayController;
     private final Context mContext;
+    private final Handler mMainHandler;
     private final ShellExecutor mMainExecutor;
     private final ShellExecutor mAnimExecutor;
     private final TransitionAnimation mTransitionAnimation;
@@ -167,27 +169,33 @@ public class DefaultTransitionHandler implements Transitions.TransitionHandler {
         }
     };
 
-    DefaultTransitionHandler(@NonNull DisplayController displayController,
-            @NonNull TransactionPool transactionPool, Context context,
+    DefaultTransitionHandler(@NonNull Context context,
+            @NonNull ShellInit shellInit,
+            @NonNull DisplayController displayController,
+            @NonNull TransactionPool transactionPool,
             @NonNull ShellExecutor mainExecutor, @NonNull Handler mainHandler,
             @NonNull ShellExecutor animExecutor) {
         mDisplayController = displayController;
         mTransactionPool = transactionPool;
         mContext = context;
+        mMainHandler = mainHandler;
         mMainExecutor = mainExecutor;
         mAnimExecutor = animExecutor;
         mTransitionAnimation = new TransitionAnimation(context, false /* debug */, Transitions.TAG);
         mCurrentUserId = UserHandle.myUserId();
+        mDevicePolicyManager = mContext.getSystemService(DevicePolicyManager.class);
+        shellInit.addInitCallback(this::onInit, this);
+    }
 
-        mDevicePolicyManager = context.getSystemService(DevicePolicyManager.class);
+    private void onInit() {
         updateEnterpriseThumbnailDrawable();
         mContext.registerReceiver(
                 mEnterpriseResourceUpdatedReceiver,
                 new IntentFilter(ACTION_DEVICE_POLICY_RESOURCE_UPDATED),
                 /* broadcastPermission = */ null,
-                mainHandler);
+                mMainHandler);
 
-        AttributeCache.init(context);
+        AttributeCache.init(mContext);
     }
 
     private void updateEnterpriseThumbnailDrawable() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/SplitscreenPipMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/SplitscreenPipMixedHandler.java
new file mode 100644
index 000000000000..678e91fd8829
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/SplitscreenPipMixedHandler.java
@@ -0,0 +1,55 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.transition;
+
+import com.android.wm.shell.pip.phone.PipTouchHandler;
+import com.android.wm.shell.splitscreen.SplitScreenController;
+import com.android.wm.shell.sysui.ShellInit;
+
+import java.util.Optional;
+
+/**
+ * Handles transitions between the Splitscreen and PIP components.
+ */
+public class SplitscreenPipMixedHandler {
+
+    private final Optional<SplitScreenController> mSplitScreenOptional;
+    private final Optional<PipTouchHandler> mPipTouchHandlerOptional;
+    private final Transitions mTransitions;
+
+    public SplitscreenPipMixedHandler(ShellInit shellInit,
+            Optional<SplitScreenController> splitScreenControllerOptional,
+            Optional<PipTouchHandler> pipTouchHandlerOptional,
+            Transitions transitions) {
+        mSplitScreenOptional = splitScreenControllerOptional;
+        mPipTouchHandlerOptional = pipTouchHandlerOptional;
+        mTransitions = transitions;
+        if (Transitions.ENABLE_SHELL_TRANSITIONS
+                && mSplitScreenOptional.isPresent() && mPipTouchHandlerOptional.isPresent()) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
+    }
+
+    private void onInit() {
+        // Special handling for initializing based on multiple components
+        final DefaultMixedHandler mixedHandler = new DefaultMixedHandler(mTransitions,
+                mPipTouchHandlerOptional.get().getTransitionHandler(),
+                mSplitScreenOptional.get().getTransitionHandler());
+        // Added at end so that it has highest priority.
+        mTransitions.addHandler(mixedHandler);
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
index 881b7a1699f6..e6006c4fc6d0 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/Transitions.java
@@ -59,13 +59,13 @@ import androidx.annotation.BinderThread;
 import com.android.internal.R;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.protolog.common.ProtoLog;
-import com.android.wm.shell.ShellTaskOrganizer;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.RemoteCallable;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.annotations.ExternalThread;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -111,6 +111,7 @@ public class Transitions implements RemoteCallable<Transitions> {
     private final ShellExecutor mMainExecutor;
     private final ShellExecutor mAnimExecutor;
     private final TransitionPlayerImpl mPlayerImpl;
+    private final DefaultTransitionHandler mDefaultTransitionHandler;
     private final RemoteTransitionHandler mRemoteTransitionHandler;
     private final DisplayController mDisplayController;
     private final ShellTransitionImpl mImpl = new ShellTransitionImpl();
@@ -136,8 +137,11 @@ public class Transitions implements RemoteCallable<Transitions> {
     /** Keeps track of currently playing transitions in the order of receipt. */
     private final ArrayList<ActiveTransition> mActiveTransitions = new ArrayList<>();
 
-    public Transitions(@NonNull WindowOrganizer organizer, @NonNull TransactionPool pool,
-            @NonNull DisplayController displayController, @NonNull Context context,
+    public Transitions(@NonNull Context context,
+            @NonNull ShellInit shellInit,
+            @NonNull WindowOrganizer organizer,
+            @NonNull TransactionPool pool,
+            @NonNull DisplayController displayController,
             @NonNull ShellExecutor mainExecutor, @NonNull Handler mainHandler,
             @NonNull ShellExecutor animExecutor) {
         mOrganizer = organizer;
@@ -146,33 +150,35 @@ public class Transitions implements RemoteCallable<Transitions> {
         mAnimExecutor = animExecutor;
         mDisplayController = displayController;
         mPlayerImpl = new TransitionPlayerImpl();
+        mDefaultTransitionHandler = new DefaultTransitionHandler(context, shellInit,
+                displayController, pool, mainExecutor, mainHandler, animExecutor);
+        mRemoteTransitionHandler = new RemoteTransitionHandler(mMainExecutor);
+        shellInit.addInitCallback(this::onInit, this);
+    }
+
+    private void onInit() {
         // The very last handler (0 in the list) should be the default one.
-        mHandlers.add(new DefaultTransitionHandler(displayController, pool, context, mainExecutor,
-                mainHandler, animExecutor));
+        mHandlers.add(mDefaultTransitionHandler);
         // Next lowest priority is remote transitions.
-        mRemoteTransitionHandler = new RemoteTransitionHandler(mainExecutor);
         mHandlers.add(mRemoteTransitionHandler);
 
-        ContentResolver resolver = context.getContentResolver();
+        ContentResolver resolver = mContext.getContentResolver();
         mTransitionAnimationScaleSetting = Settings.Global.getFloat(resolver,
                 Settings.Global.TRANSITION_ANIMATION_SCALE,
-                context.getResources().getFloat(
+                mContext.getResources().getFloat(
                         R.dimen.config_appTransitionAnimationDurationScaleDefault));
         dispatchAnimScaleSetting(mTransitionAnimationScaleSetting);
 
         resolver.registerContentObserver(
                 Settings.Global.getUriFor(Settings.Global.TRANSITION_ANIMATION_SCALE), false,
                 new SettingsObserver());
-    }
 
-    private Transitions() {
-        mOrganizer = null;
-        mContext = null;
-        mMainExecutor = null;
-        mAnimExecutor = null;
-        mDisplayController = null;
-        mPlayerImpl = null;
-        mRemoteTransitionHandler = null;
+        if (Transitions.ENABLE_SHELL_TRANSITIONS) {
+            // Register this transition handler with Core
+            mOrganizer.registerTransitionPlayer(mPlayerImpl);
+            // Pre-load the instance.
+            TransitionMetrics.getInstance();
+        }
     }
 
     public ShellTransitions asRemoteTransitions() {
@@ -195,14 +201,6 @@ public class Transitions implements RemoteCallable<Transitions> {
         }
     }
 
-    /** Register this transition handler with Core */
-    public void register(ShellTaskOrganizer taskOrganizer) {
-        if (mPlayerImpl == null) return;
-        taskOrganizer.registerTransitionPlayer(mPlayerImpl);
-        // Pre-load the instance.
-        TransitionMetrics.getInstance();
-    }
-
     /**
      * Adds a handler candidate.
      * @see TransitionHandler
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java
index 05a024a0eb12..6b59e313b01b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/UnfoldAnimationController.java
@@ -24,13 +24,14 @@ import android.app.TaskInfo;
 import android.util.SparseArray;
 import android.view.SurfaceControl;
 
+import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider.UnfoldListener;
 import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
 
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.Executor;
 
 import dagger.Lazy;
 
@@ -47,7 +48,7 @@ import dagger.Lazy;
 public class UnfoldAnimationController implements UnfoldListener {
 
     private final ShellUnfoldProgressProvider mUnfoldProgressProvider;
-    private final Executor mExecutor;
+    private final ShellExecutor mExecutor;
     private final TransactionPool mTransactionPool;
     private final List<UnfoldTaskAnimator> mAnimators;
     private final Lazy<Optional<UnfoldTransitionHandler>> mUnfoldTransitionHandler;
@@ -55,28 +56,36 @@ public class UnfoldAnimationController implements UnfoldListener {
     private final SparseArray<SurfaceControl> mTaskSurfaces = new SparseArray<>();
     private final SparseArray<UnfoldTaskAnimator> mAnimatorsByTaskId = new SparseArray<>();
 
-    public UnfoldAnimationController(@NonNull TransactionPool transactionPool,
+    public UnfoldAnimationController(
+            @NonNull ShellInit shellInit,
+            @NonNull TransactionPool transactionPool,
             @NonNull ShellUnfoldProgressProvider unfoldProgressProvider,
             @NonNull List<UnfoldTaskAnimator> animators,
             @NonNull Lazy<Optional<UnfoldTransitionHandler>> unfoldTransitionHandler,
-            @NonNull Executor executor) {
+            @NonNull ShellExecutor executor) {
         mUnfoldProgressProvider = unfoldProgressProvider;
         mUnfoldTransitionHandler = unfoldTransitionHandler;
         mTransactionPool = transactionPool;
         mExecutor = executor;
         mAnimators = animators;
+        // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
+        //                    override for this controller from the base module
+        if (unfoldProgressProvider != ShellUnfoldProgressProvider.NO_PROVIDER) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
     }
 
     /**
      * Initializes the controller, starts listening for the external events
      */
-    public void init() {
+    public void onInit() {
         mUnfoldProgressProvider.addListener(mExecutor, this);
 
         for (int i = 0; i < mAnimators.size(); i++) {
             final UnfoldTaskAnimator animator = mAnimators.get(i);
             animator.init();
-            animator.start();
+            // TODO(b/238217847): See #provideSplitTaskUnfoldAnimatorBase
+            mExecutor.executeDelayed(animator::start, 0);
         }
     }
 
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 9bf32faa12bd..5d7b62905d3b 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
@@ -28,6 +28,7 @@ import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 import com.android.wm.shell.transition.Transitions.TransitionFinishCallback;
 import com.android.wm.shell.transition.Transitions.TransitionHandler;
@@ -59,11 +60,13 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
 
     private final List<UnfoldTaskAnimator> mAnimators = new ArrayList<>();
 
-    public UnfoldTransitionHandler(ShellUnfoldProgressProvider unfoldProgressProvider,
+    public UnfoldTransitionHandler(ShellInit shellInit,
+            ShellUnfoldProgressProvider unfoldProgressProvider,
             FullscreenUnfoldTaskAnimator fullscreenUnfoldAnimator,
             SplitTaskUnfoldAnimator splitUnfoldTaskAnimator,
             TransactionPool transactionPool,
-            Executor executor, Transitions transitions) {
+            Executor executor,
+            Transitions transitions) {
         mUnfoldProgressProvider = unfoldProgressProvider;
         mTransactionPool = transactionPool;
         mExecutor = executor;
@@ -71,9 +74,18 @@ public class UnfoldTransitionHandler implements TransitionHandler, UnfoldListene
 
         mAnimators.add(splitUnfoldTaskAnimator);
         mAnimators.add(fullscreenUnfoldAnimator);
+        // TODO(b/238217847): Temporarily add this check here until we can remove the dynamic
+        //                    override for this controller from the base module
+        if (unfoldProgressProvider != ShellUnfoldProgressProvider.NO_PROVIDER
+                && Transitions.ENABLE_SHELL_TRANSITIONS) {
+            shellInit.addInitCallback(this::onInit, this);
+        }
     }
 
-    public void init() {
+    /**
+     * Called when the transition handler is initialized.
+     */
+    public void onInit() {
         for (int i = 0; i < mAnimators.size(); i++) {
             mAnimators.get(i).init();
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellInitTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellInitTest.java
index 6cd5677db2c3..4bcdcaae230b 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellInitTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellInitTest.java
@@ -24,25 +24,8 @@ import android.testing.TestableLooper;
 
 import androidx.test.filters.SmallTest;
 
-import com.android.wm.shell.activityembedding.ActivityEmbeddingController;
-import com.android.wm.shell.bubbles.BubbleController;
-import com.android.wm.shell.common.DisplayController;
-import com.android.wm.shell.common.DisplayImeController;
-import com.android.wm.shell.common.DisplayInsetsController;
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.draganddrop.DragAndDropController;
-import com.android.wm.shell.freeform.FreeformComponents;
-import com.android.wm.shell.fullscreen.FullscreenTaskListener;
-import com.android.wm.shell.kidsmode.KidsModeTaskOrganizer;
-import com.android.wm.shell.pip.phone.PipTouchHandler;
-import com.android.wm.shell.recents.RecentTasksController;
-import com.android.wm.shell.splitscreen.SplitScreenController;
-import com.android.wm.shell.startingsurface.StartingWindowController;
-import com.android.wm.shell.sysui.ShellController;
 import com.android.wm.shell.sysui.ShellInit;
-import com.android.wm.shell.transition.Transitions;
-import com.android.wm.shell.unfold.UnfoldAnimationController;
-import com.android.wm.shell.unfold.UnfoldTransitionHandler;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -51,31 +34,12 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
 import java.util.ArrayList;
-import java.util.Optional;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
 @TestableLooper.RunWithLooper(setAsMainLooper = true)
 public class ShellInitTest extends ShellTestCase {
 
-    @Mock private ShellController mShellController;
-    @Mock private DisplayController mDisplayController;
-    @Mock private DisplayImeController mDisplayImeController;
-    @Mock private DisplayInsetsController mDisplayInsetsController;
-    @Mock private DragAndDropController mDragAndDropController;
-    @Mock private ShellTaskOrganizer mShellTaskOrganizer;
-    @Mock private KidsModeTaskOrganizer mKidsModeTaskOrganizer;
-    @Mock private Optional<BubbleController> mBubblesOptional;
-    @Mock private Optional<SplitScreenController> mSplitScreenOptional;
-    @Mock private Optional<PipTouchHandler> mPipTouchHandlerOptional;
-    @Mock private FullscreenTaskListener mFullscreenTaskListener;
-    @Mock private Optional<UnfoldAnimationController> mUnfoldAnimationController;
-    @Mock private Optional<UnfoldTransitionHandler> mUnfoldTransitionHandler;
-    @Mock private Optional<FreeformComponents> mFreeformComponentsOptional;
-    @Mock private Optional<RecentTasksController> mRecentTasks;
-    @Mock private Optional<ActivityEmbeddingController> mActivityEmbeddingController;
-    @Mock private Transitions mTransitions;
-    @Mock private StartingWindowController mStartingWindow;
     @Mock private ShellExecutor mMainExecutor;
 
     private ShellInit mImpl;
@@ -83,12 +47,7 @@ public class ShellInitTest extends ShellTestCase {
     @Before
     public void setUp() {
         MockitoAnnotations.initMocks(this);
-        mImpl = new ShellInit(mShellController, mDisplayController, mDisplayImeController,
-                mDisplayInsetsController, mDragAndDropController, mShellTaskOrganizer,
-                mKidsModeTaskOrganizer, mBubblesOptional, mSplitScreenOptional,
-                mPipTouchHandlerOptional, mFullscreenTaskListener, mUnfoldAnimationController,
-                mUnfoldTransitionHandler, mFreeformComponentsOptional, mRecentTasks,
-                mActivityEmbeddingController, mTransitions, mStartingWindow, mMainExecutor);
+        mImpl = new ShellInit(mMainExecutor);
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
index 3dd00329253c..7517e8ab6826 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/ShellTaskOrganizerTests.java
@@ -36,11 +36,11 @@ import static org.junit.Assume.assumeFalse;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
-import android.content.Context;
 import android.content.LocusId;
 import android.content.pm.ParceledListSlice;
 import android.os.Binder;
@@ -58,9 +58,8 @@ import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.common.ShellExecutor;
-import com.android.wm.shell.common.SyncTransactionQueue;
-import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.compatui.CompatUIController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -84,13 +83,11 @@ public class ShellTaskOrganizerTests extends ShellTestCase {
     @Mock
     private ITaskOrganizerController mTaskOrganizerController;
     @Mock
-    private Context mContext;
-    @Mock
     private CompatUIController mCompatUI;
+    @Mock
+    private ShellInit mShellInit;
 
     ShellTaskOrganizer mOrganizer;
-    private final SyncTransactionQueue mSyncTransactionQueue = mock(SyncTransactionQueue.class);
-    private final TransactionPool mTransactionPool = mock(TransactionPool.class);
     private final ShellExecutor mTestExecutor = mock(ShellExecutor.class);
 
     private class TrackingTaskListener implements ShellTaskOrganizer.TaskListener {
@@ -135,8 +132,13 @@ public class ShellTaskOrganizerTests extends ShellTestCase {
             doReturn(ParceledListSlice.<TaskAppearedInfo>emptyList())
                     .when(mTaskOrganizerController).registerTaskOrganizer(any());
         } catch (RemoteException e) {}
-        mOrganizer = spy(new ShellTaskOrganizer(mTaskOrganizerController, mTestExecutor, mContext,
-                mCompatUI, Optional.empty(), Optional.empty()));
+        mOrganizer = spy(new ShellTaskOrganizer(mShellInit, mTaskOrganizerController,
+                mCompatUI, Optional.empty(), Optional.empty(), mTestExecutor));
+    }
+
+    @Test
+    public void instantiate_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/activityembedding/ActivityEmbeddingControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/activityembedding/ActivityEmbeddingControllerTests.java
new file mode 100644
index 000000000000..bfe3b5468085
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/activityembedding/ActivityEmbeddingControllerTests.java
@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.activityembedding;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
+
+import static org.junit.Assume.assumeTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.sysui.ShellInit;
+import com.android.wm.shell.transition.Transitions;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for the activity embedding controller.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:ActivityEmbeddingControllerTests
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class ActivityEmbeddingControllerTests extends ShellTestCase {
+
+    private @Mock Context mContext;
+    private @Mock ShellInit mShellInit;
+    private @Mock Transitions mTransitions;
+    private ActivityEmbeddingController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mController = spy(new ActivityEmbeddingController(mContext, mShellInit, mTransitions));
+    }
+
+    @Test
+    public void instantiate_addInitCallback() {
+        assumeTrue(Transitions.ENABLE_SHELL_TRANSITIONS);
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayChangeControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayChangeControllerTests.java
new file mode 100644
index 000000000000..b8aa8e7cbc48
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayChangeControllerTests.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.common;
+
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spy;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.view.IWindowManager;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.sysui.ShellInit;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for the display change controller.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:DisplayChangeControllerTests
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DisplayChangeControllerTests extends ShellTestCase {
+
+    private @Mock IWindowManager mWM;
+    private @Mock ShellInit mShellInit;
+    private @Mock ShellExecutor mMainExecutor;
+    private DisplayChangeController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mController = spy(new DisplayChangeController(mWM, mShellInit, mMainExecutor));
+    }
+
+    @Test
+    public void instantiate_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayControllerTests.java
new file mode 100644
index 000000000000..1e5e153fdfe1
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayControllerTests.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.common;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+import android.view.IWindowManager;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.sysui.ShellInit;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for the display controller.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:DisplayControllerTests
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class DisplayControllerTests extends ShellTestCase {
+
+    private @Mock Context mContext;
+    private @Mock IWindowManager mWM;
+    private @Mock ShellInit mShellInit;
+    private @Mock ShellExecutor mMainExecutor;
+    private DisplayController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        mController = new DisplayController(mContext, mWM, mShellInit, mMainExecutor);
+    }
+
+    @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), eq(mController));
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
index 5b691f231d85..9967e5f47752 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayImeControllerTest.java
@@ -26,6 +26,7 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyZeroInteractions;
 
@@ -40,26 +41,32 @@ import androidx.test.filters.SmallTest;
 
 import com.android.internal.view.IInputMethodManager;
 import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 import java.util.concurrent.Executor;
 
 @SmallTest
 public class DisplayImeControllerTest extends ShellTestCase {
 
+    @Mock
     private SurfaceControl.Transaction mT;
-    private DisplayImeController.PerDisplay mPerDisplay;
+    @Mock
     private IInputMethodManager mMock;
+    @Mock
+    private ShellInit mShellInit;
+    private DisplayImeController.PerDisplay mPerDisplay;
     private Executor mExecutor;
 
     @Before
     public void setUp() throws Exception {
-        mT = mock(SurfaceControl.Transaction.class);
-        mMock = mock(IInputMethodManager.class);
+        MockitoAnnotations.initMocks(this);
         mExecutor = spy(Runnable::run);
-        mPerDisplay = new DisplayImeController(null, null, null, mExecutor, new TransactionPool() {
+        mPerDisplay = new DisplayImeController(null, mShellInit, null, null, new TransactionPool() {
             @Override
             public SurfaceControl.Transaction acquire() {
                 return mT;
@@ -68,7 +75,7 @@ public class DisplayImeControllerTest extends ShellTestCase {
             @Override
             public void release(SurfaceControl.Transaction t) {
             }
-        }) {
+        }, mExecutor) {
             @Override
             public IInputMethodManager getImms() {
                 return mMock;
@@ -79,6 +86,11 @@ public class DisplayImeControllerTest extends ShellTestCase {
     }
 
     @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+
+    @Test
     public void insetsControlChanged_schedulesNoWorkOnExecutor() {
         mPerDisplay.insetsControlChanged(insetsStateWithIme(false), insetsSourceControl());
         verifyZeroInteractions(mExecutor);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayInsetsControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayInsetsControllerTest.java
index 4a7fd3d259da..5f5a3c584ee0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayInsetsControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/DisplayInsetsControllerTest.java
@@ -19,6 +19,7 @@ package com.android.wm.shell.common;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.ArgumentMatchers.notNull;
 import static org.mockito.Mockito.times;
@@ -37,6 +38,7 @@ import androidx.test.filters.SmallTest;
 
 import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -55,6 +57,8 @@ public class DisplayInsetsControllerTest extends ShellTestCase {
     private IWindowManager mWm;
     @Mock
     private DisplayController mDisplayController;
+    @Mock
+    private ShellInit mShellInit;
     private DisplayInsetsController mController;
     private SparseArray<IDisplayWindowInsetsController> mInsetsControllersByDisplayId;
     private TestShellExecutor mExecutor;
@@ -69,11 +73,16 @@ public class DisplayInsetsControllerTest extends ShellTestCase {
         mInsetsControllersByDisplayId = new SparseArray<>();
         mDisplayIdCaptor =  ArgumentCaptor.forClass(Integer.class);
         mInsetsControllerCaptor = ArgumentCaptor.forClass(IDisplayWindowInsetsController.class);
-        mController = new DisplayInsetsController(mWm, mDisplayController, mExecutor);
+        mController = new DisplayInsetsController(mWm, mShellInit, mDisplayController, mExecutor);
         addDisplay(DEFAULT_DISPLAY);
     }
 
     @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+
+    @Test
     public void testOnDisplayAdded_setsDisplayWindowInsetsControllerOnWMService()
             throws RemoteException {
         addDisplay(SECOND_DISPLAY);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropControllerTest.java
index e209971998c8..b6dbcf204364 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/draganddrop/DragAndDropControllerTest.java
@@ -50,6 +50,7 @@ import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -57,8 +58,6 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
-import java.util.Optional;
-
 /**
  * Tests for the drag and drop controller.
  */
@@ -69,6 +68,8 @@ public class DragAndDropControllerTest extends ShellTestCase {
     @Mock
     private Context mContext;
     @Mock
+    private ShellInit mShellInit;
+    @Mock
     private ShellController mShellController;
     @Mock
     private DisplayController mDisplayController;
@@ -88,9 +89,14 @@ public class DragAndDropControllerTest extends ShellTestCase {
     @Before
     public void setUp() throws RemoteException {
         MockitoAnnotations.initMocks(this);
-        mController = new DragAndDropController(mContext, mShellController, mDisplayController,
-                mUiEventLogger, mIconProvider, mMainExecutor);
-        mController.initialize(Optional.of(mSplitScreenController));
+        mController = new DragAndDropController(mContext, mShellInit, mShellController,
+                mDisplayController, mUiEventLogger, mIconProvider, mMainExecutor);
+        mController.onInit();
+    }
+
+    @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
index 184a8dfecff9..a919ad0b4765 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/kidsmode/KidsModeTaskOrganizerTest.java
@@ -49,7 +49,7 @@ import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.DisplayInsetsController;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
-import com.android.wm.shell.startingsurface.StartingWindowController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -73,7 +73,7 @@ public class KidsModeTaskOrganizerTest extends ShellTestCase {
     @Mock private WindowContainerToken mToken;
     @Mock private WindowContainerTransaction mTransaction;
     @Mock private KidsModeSettingsObserver mObserver;
-    @Mock private StartingWindowController mStartingWindowController;
+    @Mock private ShellInit mShellInit;
     @Mock private DisplayInsetsController mDisplayInsetsController;
 
     KidsModeTaskOrganizer mOrganizer;
@@ -87,10 +87,9 @@ public class KidsModeTaskOrganizerTest extends ShellTestCase {
         } catch (RemoteException e) {
         }
         // NOTE: KidsModeTaskOrganizer should have a null CompatUIController.
-        mOrganizer = spy(new KidsModeTaskOrganizer(mTaskOrganizerController, mTestExecutor,
-                mHandler, mContext, mSyncTransactionQueue, mDisplayController,
-                mDisplayInsetsController, Optional.empty(), Optional.empty(), mObserver));
-        mOrganizer.initialize(mStartingWindowController);
+        mOrganizer = spy(new KidsModeTaskOrganizer(mContext, mTaskOrganizerController,
+                mSyncTransactionQueue, mDisplayController, mDisplayInsetsController,
+                Optional.empty(), Optional.empty(), mObserver, mTestExecutor, mHandler));
         doReturn(mTransaction).when(mOrganizer).getWindowContainerTransaction();
         doReturn(new InsetsState()).when(mDisplayController).getInsetsState(DEFAULT_DISPLAY);
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipTouchHandlerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipTouchHandlerTest.java
index 74519eaf3ebf..ecefd89d8778 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipTouchHandlerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/pip/phone/PipTouchHandlerTest.java
@@ -38,6 +38,7 @@ import com.android.wm.shell.pip.PipSnapAlgorithm;
 import com.android.wm.shell.pip.PipTaskOrganizer;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.pip.PipUiEventLogger;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -78,6 +79,9 @@ public class PipTouchHandlerTest extends ShellTestCase {
     private PipUiEventLogger mPipUiEventLogger;
 
     @Mock
+    private ShellInit mShellInit;
+
+    @Mock
     private ShellExecutor mMainExecutor;
 
     private PipBoundsState mPipBoundsState;
@@ -104,11 +108,11 @@ public class PipTouchHandlerTest extends ShellTestCase {
         PipMotionHelper pipMotionHelper = new PipMotionHelper(mContext, mPipBoundsState,
                 mPipTaskOrganizer, mPhonePipMenuController, mPipSnapAlgorithm,
                 mMockPipTransitionController, mFloatingContentCoordinator);
-        mPipTouchHandler = new PipTouchHandler(mContext, mPhonePipMenuController,
-                mPipBoundsAlgorithm, mPipBoundsState, mPipTaskOrganizer,
-                pipMotionHelper, mFloatingContentCoordinator, mPipUiEventLogger,
-                mMainExecutor);
-        mPipTouchHandler.init();
+        mPipTouchHandler = new PipTouchHandler(mContext, mShellInit, mPhonePipMenuController,
+                mPipBoundsAlgorithm, mPipBoundsState, mPipTaskOrganizer, pipMotionHelper,
+                mFloatingContentCoordinator, mPipUiEventLogger, mMainExecutor);
+        // We aren't actually using ShellInit, so just call init directly
+        mPipTouchHandler.onInit();
         mMotionHelper = Mockito.spy(mPipTouchHandler.getMotionHelper());
         mPipResizeGestureHandler = Mockito.spy(mPipTouchHandler.getPipResizeGestureHandler());
         mPipTouchHandler.setPipMotionHelper(mMotionHelper);
@@ -133,6 +137,11 @@ public class PipTouchHandlerTest extends ShellTestCase {
     }
 
     @Test
+    public void instantiate_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+
+    @Test
     public void updateMovementBounds_minMaxBounds() {
         final int shorterLength = Math.min(mPipBoundsState.getDisplayBounds().width(),
                 mPipBoundsState.getDisplayBounds().height());
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
index b1e0911bd0a8..d406a4ed3fd7 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
@@ -48,6 +48,7 @@ import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TaskStackListenerImpl;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
 import com.android.wm.shell.util.SplitBounds;
 
@@ -71,6 +72,8 @@ public class RecentTasksControllerTest extends ShellTestCase {
     private Context mContext;
     @Mock
     private TaskStackListenerImpl mTaskStackListener;
+    @Mock
+    private ShellInit mShellInit;
 
     private ShellTaskOrganizer mShellTaskOrganizer;
     private RecentTasksController mRecentTasksController;
@@ -80,10 +83,11 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void setUp() {
         mMainExecutor = new TestShellExecutor();
         when(mContext.getPackageManager()).thenReturn(mock(PackageManager.class));
-        mRecentTasksController = spy(new RecentTasksController(mContext, mTaskStackListener,
-                mMainExecutor));
-        mShellTaskOrganizer = new ShellTaskOrganizer(mMainExecutor, mContext,
-                null /* sizeCompatUI */, Optional.empty(), Optional.of(mRecentTasksController));
+        mRecentTasksController = spy(new RecentTasksController(mContext, mShellInit,
+                mTaskStackListener, mMainExecutor));
+        mShellTaskOrganizer = new ShellTaskOrganizer(mShellInit,
+                null /* sizeCompatUI */, Optional.empty(), Optional.of(mRecentTasksController),
+                mMainExecutor);
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
index c7a261f32e43..10788f9df32f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/SplitScreenControllerTests.java
@@ -51,8 +51,10 @@ import com.android.wm.shell.common.DisplayLayout;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.transition.Transitions;
 
 import org.junit.Before;
@@ -71,6 +73,7 @@ import java.util.Optional;
 public class SplitScreenControllerTests extends ShellTestCase {
 
     @Mock ShellController mShellController;
+    @Mock ShellInit mShellInit;
     @Mock ShellTaskOrganizer mTaskOrganizer;
     @Mock SyncTransactionQueue mSyncQueue;
     @Mock RootTaskDisplayAreaOrganizer mRootTDAOrganizer;
@@ -78,6 +81,7 @@ public class SplitScreenControllerTests extends ShellTestCase {
     @Mock DisplayController mDisplayController;
     @Mock DisplayImeController mDisplayImeController;
     @Mock DisplayInsetsController mDisplayInsetsController;
+    @Mock DragAndDropController mDragAndDropController;
     @Mock Transitions mTransitions;
     @Mock TransactionPool mTransactionPool;
     @Mock IconProvider mIconProvider;
@@ -88,16 +92,21 @@ public class SplitScreenControllerTests extends ShellTestCase {
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mSplitScreenController = spy(new SplitScreenController(mShellController, mTaskOrganizer,
-                mSyncQueue, mContext, mRootTDAOrganizer, mMainExecutor, mDisplayController,
-                mDisplayImeController, mDisplayInsetsController, mTransitions, mTransactionPool,
-                mIconProvider, mRecentTasks));
+        mSplitScreenController = spy(new SplitScreenController(mContext, mShellInit,
+                mShellController, mTaskOrganizer, mSyncQueue, mRootTDAOrganizer, mDisplayController,
+                mDisplayImeController, mDisplayInsetsController, mDragAndDropController,
+                mTransitions, mTransactionPool, mIconProvider, mRecentTasks, mMainExecutor));
+    }
+
+    @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
     }
 
     @Test
     public void testControllerRegistersKeyguardChangeListener() {
         when(mDisplayController.getDisplayLayout(anyInt())).thenReturn(new DisplayLayout());
-        mSplitScreenController.onOrganizerRegistered();
+        mSplitScreenController.onInit();
         verify(mShellController, times(1)).addKeyguardChangeListener(any());
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingWindowControllerTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingWindowControllerTests.java
new file mode 100644
index 000000000000..35515e3bb6e8
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/startingsurface/StartingWindowControllerTests.java
@@ -0,0 +1,80 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.startingsurface;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import android.content.Context;
+import android.hardware.display.DisplayManager;
+import android.view.Display;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.icons.IconProvider;
+import com.android.wm.shell.ShellTaskOrganizer;
+import com.android.wm.shell.ShellTestCase;
+import com.android.wm.shell.common.ShellExecutor;
+import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+/**
+ * Tests for the starting window controller.
+ *
+ * Build/Install/Run:
+ *  atest WMShellUnitTests:StartingWindowControllerTests
+ */
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class StartingWindowControllerTests extends ShellTestCase {
+
+    private @Mock Context mContext;
+    private @Mock DisplayManager mDisplayManager;
+    private @Mock ShellInit mShellInit;
+    private @Mock ShellTaskOrganizer mTaskOrganizer;
+    private @Mock ShellExecutor mMainExecutor;
+    private @Mock StartingWindowTypeAlgorithm mTypeAlgorithm;
+    private @Mock IconProvider mIconProvider;
+    private @Mock TransactionPool mTransactionPool;
+    private StartingWindowController mController;
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        doReturn(mock(Display.class)).when(mDisplayManager).getDisplay(anyInt());
+        doReturn(mDisplayManager).when(mContext).getSystemService(eq(DisplayManager.class));
+        mController = new StartingWindowController(mContext, mShellInit, mTaskOrganizer,
+                mMainExecutor, mTypeAlgorithm, mIconProvider, mTransactionPool);
+    }
+
+    @Test
+    public void instantiate_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+}
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sysui/ShellControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sysui/ShellControllerTest.java
index 1c0e46f7264e..02311bab2e4d 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sysui/ShellControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/sysui/ShellControllerTest.java
@@ -43,6 +43,8 @@ import java.util.Locale;
 public class ShellControllerTest extends ShellTestCase {
 
     @Mock
+    private ShellInit mShellInit;
+    @Mock
     private ShellExecutor mExecutor;
 
     private ShellController mController;
@@ -54,7 +56,7 @@ public class ShellControllerTest extends ShellTestCase {
         MockitoAnnotations.initMocks(this);
         mKeyguardChangeListener = new TestKeyguardChangeListener();
         mConfigChangeListener = new TestConfigurationChangeListener();
-        mController = new ShellController(mExecutor);
+        mController = new ShellController(mShellInit, mExecutor);
         mController.onConfigurationChanged(getConfigurationCopy());
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
index e2f2b71cea04..388792b63db3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/transition/ShellTransitionTests.java
@@ -84,6 +84,7 @@ import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.DisplayController;
 import com.android.wm.shell.common.ShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -117,6 +118,14 @@ public class ShellTransitionTests extends ShellTestCase {
     }
 
     @Test
+    public void instantiate_addInitCallback() {
+        ShellInit shellInit = mock(ShellInit.class);
+        final Transitions t = new Transitions(mContext, shellInit, mOrganizer, mTransactionPool,
+                createTestDisplayController(), mMainExecutor, mMainHandler, mAnimExecutor);
+        verify(shellInit, times(1)).addInitCallback(any(), eq(t));
+    }
+
+    @Test
     public void testBasicTransitionFlow() {
         Transitions transitions = createTestTransitions();
         transitions.replaceDefaultHandlerForTest(mDefaultHandler);
@@ -832,14 +841,18 @@ public class ShellTransitionTests extends ShellTestCase {
         } catch (RemoteException e) {
             // No remote stuff happening, so this can't be hit
         }
-        DisplayController out = new DisplayController(mContext, mockWM, mMainExecutor);
-        out.initialize();
+        ShellInit shellInit = new ShellInit(mMainExecutor);
+        DisplayController out = new DisplayController(mContext, mockWM, shellInit, mMainExecutor);
+        shellInit.init();
         return out;
     }
 
     private Transitions createTestTransitions() {
-        return new Transitions(mOrganizer, mTransactionPool, createTestDisplayController(),
-                mContext, mMainExecutor, mMainHandler, mAnimExecutor);
+        ShellInit shellInit = new ShellInit(mMainExecutor);
+        final Transitions t = new Transitions(mContext, shellInit, mOrganizer, mTransactionPool,
+                createTestDisplayController(), mMainExecutor, mMainHandler, mAnimExecutor);
+        shellInit.init();
+        return t;
     }
 //
 //    private class TestDisplayController extends DisplayController {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java
index 46de60772766..81eefe25704e 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/unfold/UnfoldAnimationControllerTest.java
@@ -20,7 +20,10 @@ import static com.android.wm.shell.unfold.UnfoldAnimationControllerTest.TestUnfo
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.ActivityManager.RunningTaskInfo;
@@ -33,6 +36,7 @@ import com.android.wm.shell.ShellTestCase;
 import com.android.wm.shell.TestRunningTaskInfoBuilder;
 import com.android.wm.shell.TestShellExecutor;
 import com.android.wm.shell.common.TransactionPool;
+import com.android.wm.shell.sysui.ShellInit;
 import com.android.wm.shell.unfold.animation.UnfoldTaskAnimator;
 
 import org.junit.Before;
@@ -65,6 +69,8 @@ public class UnfoldAnimationControllerTest extends ShellTestCase {
     @Mock
     private UnfoldTransitionHandler mUnfoldTransitionHandler;
     @Mock
+    private ShellInit mShellInit;
+    @Mock
     private SurfaceControl mLeash;
 
     private UnfoldAnimationController mUnfoldAnimationController;
@@ -85,6 +91,7 @@ public class UnfoldAnimationControllerTest extends ShellTestCase {
         animators.add(mTaskAnimator1);
         animators.add(mTaskAnimator2);
         mUnfoldAnimationController = new UnfoldAnimationController(
+                mShellInit,
                 mTransactionPool,
                 mProgressProvider,
                 animators,
@@ -94,6 +101,11 @@ public class UnfoldAnimationControllerTest extends ShellTestCase {
     }
 
     @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+
+    @Test
     public void testAppearedMatchingTask_appliesUnfoldProgress() {
         mTaskAnimator1.setTaskMatcher((info) -> info.getWindowingMode() == 2);
         RunningTaskInfo taskInfo = new TestRunningTaskInfoBuilder()
@@ -244,7 +256,8 @@ public class UnfoldAnimationControllerTest extends ShellTestCase {
 
     @Test
     public void testInit_initsAndStartsAnimators() {
-        mUnfoldAnimationController.init();
+        mUnfoldAnimationController.onInit();
+        mShellExecutor.flushAll();
 
         assertThat(mTaskAnimator1.mInitialized).isTrue();
         assertThat(mTaskAnimator1.mStarted).isTrue();
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java b/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
index e9ca0fdbb929..5f586c927ef7 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIInitializer.java
@@ -79,11 +79,6 @@ public abstract class SystemUIInitializer {
 
         // Stand up WMComponent
         setupWmComponent(mContext);
-        if (initializeComponents) {
-            // Only initialize when not starting from tests since this currently initializes some
-            // components that shouldn't be run in the test environment
-            mWMComponent.init();
-        }
 
         // And finally, retrieve whatever SysUI needs from WMShell and build SysUI.
         SysUIComponent.Builder builder = mRootComponent.getSysUIComponent();
@@ -102,6 +97,10 @@ public abstract class SystemUIInitializer {
                     .setDisplayAreaHelper(mWMComponent.getDisplayAreaHelper())
                     .setRecentTasks(mWMComponent.getRecentTasks())
                     .setBackAnimation(mWMComponent.getBackAnimation());
+
+            // Only initialize when not starting from tests since this currently initializes some
+            // components that shouldn't be run in the test environment
+            mWMComponent.init();
         } else {
             // TODO: Call on prepareSysUIComponentBuilder but not with real components. Other option
             // is separating this logic into newly creating SystemUITestsFactory.
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 d009280345de..fee17c785ed2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/BubblesTest.java
@@ -134,6 +134,7 @@ import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import org.junit.Before;
 import org.junit.Ignore;
@@ -218,6 +219,8 @@ public class BubblesTest extends SysuiTestCase {
     private BubbleEntry mBubbleEntry2User11;
 
     @Mock
+    private ShellInit mShellInit;
+    @Mock
     private ShellController mShellController;
     @Mock
     private Bubbles.BubbleExpandListener mBubbleExpandListener;
@@ -339,6 +342,7 @@ public class BubblesTest extends SysuiTestCase {
         when(mShellTaskOrganizer.getExecutor()).thenReturn(syncExecutor);
         mBubbleController = new TestableBubbleController(
                 mContext,
+                mShellInit,
                 mShellController,
                 mBubbleData,
                 mFloatingContentCoordinator,
@@ -389,6 +393,11 @@ public class BubblesTest extends SysuiTestCase {
     }
 
     @Test
+    public void instantiateController_addInitCallback() {
+        verify(mShellInit, times(1)).addInitCallback(any(), any());
+    }
+
+    @Test
     public void instantiateController_registerConfigChangeListener() {
         verify(mShellController, times(1)).addConfigurationChangeListener(any());
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/wmshell/TestableBubbleController.java b/packages/SystemUI/tests/src/com/android/systemui/wmshell/TestableBubbleController.java
index f901c327b76e..880ad187f910 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/wmshell/TestableBubbleController.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/wmshell/TestableBubbleController.java
@@ -39,6 +39,7 @@ import com.android.wm.shell.common.TaskStackListenerImpl;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.onehanded.OneHandedController;
 import com.android.wm.shell.sysui.ShellController;
+import com.android.wm.shell.sysui.ShellInit;
 
 import java.util.Optional;
 
@@ -49,6 +50,7 @@ public class TestableBubbleController extends BubbleController {
 
     // Let's assume surfaces can be synchronized immediately.
     TestableBubbleController(Context context,
+            ShellInit shellInit,
             ShellController shellController,
             BubbleData data,
             FloatingContentCoordinator floatingContentCoordinator,
@@ -69,13 +71,13 @@ public class TestableBubbleController extends BubbleController {
             Handler shellMainHandler,
             TaskViewTransitions taskViewTransitions,
             SyncTransactionQueue syncQueue) {
-        super(context, shellController, data, Runnable::run, floatingContentCoordinator,
+        super(context, shellInit, shellController, data, Runnable::run, floatingContentCoordinator,
                 dataRepository, statusBarService, windowManager, windowManagerShellWrapper,
                 userManager, launcherApps, bubbleLogger, taskStackListener, shellTaskOrganizer,
                 positioner, displayController, oneHandedOptional, dragAndDropController,
                 shellMainExecutor, shellMainHandler, new SyncExecutor(), taskViewTransitions,
                 syncQueue);
         setInflateSynchronously(true);
-        initialize();
+        onInit();
     }
 }