diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index 4720f55..c13225a 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -171,7 +171,7 @@
     }
 
     @Override
-    protected void onUiChangedWhileSleeping() {
+    public void onUiChangedWhileSleeping() {
         // Remove the snapshot because the content view may have obvious changes.
         UI_HELPER_EXECUTOR.execute(
                 () -> ActivityManagerWrapper.getInstance().invalidateHomeTaskSnapshot(this));
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 7ef6a4a..d040904 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -40,6 +40,8 @@
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.RemoteTransitionCompat;
+import com.android.systemui.shared.system.smartspace.ISmartspaceCallback;
+import com.android.systemui.shared.system.smartspace.ISmartspaceTransitionController;
 import com.android.wm.shell.onehanded.IOneHanded;
 import com.android.wm.shell.pip.IPip;
 import com.android.wm.shell.pip.IPipAnimationListener;
@@ -61,6 +63,7 @@
 
     private ISystemUiProxy mSystemUiProxy;
     private IPip mPip;
+    private ISmartspaceTransitionController mSmartspaceTransitionController;
     private ISplitScreen mSplitScreen;
     private IOneHanded mOneHanded;
     private IShellTransitions mShellTransitions;
@@ -74,6 +77,7 @@
     private IPipAnimationListener mPendingPipAnimationListener;
     private ISplitScreenListener mPendingSplitScreenListener;
     private IStartingWindowListener mPendingStartingWindowListener;
+    private ISmartspaceCallback mPendingSmartspaceCallback;
 
     // Used to dedupe calls to SystemUI
     private int mLastShelfHeight;
@@ -125,7 +129,8 @@
 
     public void setProxy(ISystemUiProxy proxy, IPip pip, ISplitScreen splitScreen,
             IOneHanded oneHanded, IShellTransitions shellTransitions,
-            IStartingWindow startingWindow) {
+            IStartingWindow startingWindow,
+            ISmartspaceTransitionController smartSpaceTransitionController) {
         unlinkToDeath();
         mSystemUiProxy = proxy;
         mPip = pip;
@@ -133,6 +138,7 @@
         mOneHanded = oneHanded;
         mShellTransitions = shellTransitions;
         mStartingWindow = startingWindow;
+        mSmartspaceTransitionController = smartSpaceTransitionController;
         linkToDeath();
         // re-attach the listeners once missing due to setProxy has not been initialized yet.
         if (mPendingPipAnimationListener != null && mPip != null) {
@@ -147,10 +153,14 @@
             setStartingWindowListener(mPendingStartingWindowListener);
             mPendingStartingWindowListener = null;
         }
+        if (mPendingSmartspaceCallback != null && mSmartspaceTransitionController != null) {
+            setSmartspaceCallback(mPendingSmartspaceCallback);
+            mPendingSmartspaceCallback = null;
+        }
     }
 
     public void clearProxy() {
-        setProxy(null, null, null, null, null, null);
+        setProxy(null, null, null, null, null, null, null);
     }
 
     // TODO(141886704): Find a way to remove this
@@ -642,4 +652,21 @@
             mPendingStartingWindowListener = listener;
         }
     }
+
+
+    //
+    // SmartSpace transitions
+    //
+
+    public void setSmartspaceCallback(ISmartspaceCallback callback) {
+        if (mSmartspaceTransitionController != null) {
+            try {
+                mSmartspaceTransitionController.setSmartspace(callback);
+            } catch (RemoteException e) {
+                Log.w(TAG, "Failed call setStartingWindowListener", e);
+            }
+        } else {
+            mPendingSmartspaceCallback = callback;
+        }
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 7956fcc..e52405b 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -30,6 +30,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_SHELL_TRANSITIONS;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_SPLIT_SCREEN;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SHELL_STARTING_WINDOW;
+import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SMARTSPACE_TRANSITION_CONTROLLER;
 import static com.android.systemui.shared.system.QuickStepContract.KEY_EXTRA_SYSUI_PROXY;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_TRACING_ENABLED;
 
@@ -107,6 +108,7 @@
 import com.android.systemui.shared.system.InputChannelCompat.InputEventReceiver;
 import com.android.systemui.shared.system.InputConsumerController;
 import com.android.systemui.shared.system.InputMonitorCompat;
+import com.android.systemui.shared.system.smartspace.ISmartspaceTransitionController;
 import com.android.systemui.shared.tracing.ProtoTraceable;
 import com.android.wm.shell.onehanded.IOneHanded;
 import com.android.wm.shell.pip.IPip;
@@ -164,9 +166,13 @@
                     bundle.getBinder(KEY_EXTRA_SHELL_SHELL_TRANSITIONS));
             IStartingWindow startingWindow = IStartingWindow.Stub.asInterface(
                     bundle.getBinder(KEY_EXTRA_SHELL_STARTING_WINDOW));
+            ISmartspaceTransitionController smartspaceTransitionController =
+                    ISmartspaceTransitionController.Stub.asInterface(
+                            bundle.getBinder(KEY_EXTRA_SMARTSPACE_TRANSITION_CONTROLLER));
             MAIN_EXECUTOR.execute(() -> {
                 SystemUiProxy.INSTANCE.get(TouchInteractionService.this).setProxy(proxy, pip,
-                        splitscreen, onehanded, shellTransitions, startingWindow);
+                        splitscreen, onehanded, shellTransitions, startingWindow,
+                        smartspaceTransitionController);
                 TouchInteractionService.this.initInputMonitor();
                 preloadOverview(true /* fromInit */);
                 mDeviceState.runOnUserUnlocked(() -> {